ImageRAGTurbo: Towards One-step Text-to-Image Generation with Retrieval-Augmented Diffusion Models

TL;DR

Experimental results on fast text-to-image generation demonstrate that this novel approach to efficiently finetune few-step diffusion models via retrieval augmentation produces high-fidelity images without compromising latency compared to existing methods.

Abstract

Diffusion models have emerged as the leading approach for text-to-image generation. However, their iterative sampling process, which gradually morphs random noise into coherent images, introduces significant latency that limits their applicability. While recent few-step diffusion models reduce the number of sampling steps to as few as one to four steps, they often compromise image quality and prompt alignment, especially in one-step generation. Additionally, these models require computationally expensive training procedures. To address these limitations, we propose ImageRAGTurbo, a novel approach to efficiently finetune few-step diffusion models via retrieval augmentation. Given a text prompt, we retrieve relevant text-image pairs from a database and use them to condition the generation process. We argue that such retrieved examples provide rich contextual information to the UNet denoiser that helps reduce the number of denoising steps without compromising image quality. Indeed, our initial investigations show that using the retrieved content to edit the denoiser's latent space ($\mathcal{H}$-space) without additional finetuning already improves prompt fidelity. To further improve the quality of the generated images, we augment the UNet denoiser with a trainable adapter in the $\mathcal{H}$-space, which efficiently blends the retrieved content with the target prompt using a cross-attention mechanism. Experimental results on fast text-to-image generation demonstrate that our approach produces high-fidelity images without compromising latency compared to existing methods.

Figures (7)

• Figure 1: Illustrative examples of ImageRAGTurbo, where the visual concepts are highlighted by colored boxes. ImageRAGTurbo outperforms Stable Diffusion Turbo (adversarially distilled Stable Diffusion without retrieval) in generating accurate visual concepts. Even with a single step, ImageRAGTurbo performs comparably to Stable Diffusion with 50 steps. Please zoom in for better quality.
• Figure 2: Overview of the proposed ImageRAGTurbo framework for efficiently finetuning the few-step diffusion models with retrieval-augmented generation. The framework involves two main branches: i) a standard denoising branch (highlighted by green), and ii) a retrieval branch (highlighted by purple). For a target prompt $\text{p}^{\bm{tgt}}$, it will first be converted to embeddings with a pretrained text encoder $\tau_\phi(\cdot)$. Then we query a database based on the target text embeddings to obtain the retrieved text-latent embeddings ($\tau_\phi(\text{p}^{\bm{retr}}), \bm{z}^{retr}_t$), which is then fed into the encoder of the denoiser to obtain the retrieved $\mathcal{H}$-space feature $\bm{h}_t^{retr}$. Finally, the $\mathcal{H}$-space feature $\bm{h}_t^{retr}$ in the retrieved branch is injected into the denoising branch by the proposed $\mathcal{H}$-space adapter to guide the generation.
• Figure 3: Performance of direct $\mathcal{H}$-space injection, shown as a histogram of TIFA scores across various categories.
• Figure 4: The architecture of the discriminator used for latent adversarial training, which adds noise to the samples $\hat{\bm{z}}_0$ from the student model and $\bm{z}_0$ from the teacher model, and differentiate them.
• Figure 5: Detailed histogram of TIFA scores across various categories. Despite still lagging behind 50-step Stable Diffusion v2-1-base model, our 1-step ImageRAGTurbo achieves comparable or even slightly higher TIFA score in certain categories such as object, activity, and material.