Re-Imagen: Retrieval-Augmented
Text-to-Image Generator

Diffusion Models Diffusion models (Sohl-Dickstein et al., 2015) are latent variable models, parameterized by $\theta$, in the form of $p_{\theta}(\bm{x}_{0}):=\int p_{\theta}(\bm{x}_{0:T})d\bm{x}_{1:T}$, where $\bm{x}_{1},\cdots,\bm{x}_{T}$ are “noised” latent versions of the input image $\bm{x}_{0}\sim q(\bm{x}_{0})$. Note that the dimensionality of both latents and the image is the same throughout the entire process, with $\bm{x}_{0:T}\in\mathbb{R}^{d}$ and $d$ equals the product of <height, width, # of channels>. The process that computes the posterior distribution $q(\bm{x}_{1:T}|\bm{x}_{0})$ is also called the forward (or diffusion) process, and is implemented as a predefined Markov chain that gradually adds Gaussian noise to the data according to a schedule $\beta_{t}$:

$$q(\bm{x}_{1:T}|\bm{x}_{0})=\prod_{t=1}^{T}q(\bm{x}_{t}|\bm{x}_{t-1})\quad\quad q(\bm{x}_{t}|\bm{x}_{t-1}):=\mathcal{N}(\bm{x}_{t};\sqrt{1-\beta_{t}}\bm{x}_{t-1},\beta_{t}\bm{I})$$

(1)

Diffusion models are trained to learn the image distribution by reversing the diffusion Markov chain. Theoretically, this reduces to learning to denoise $\bm{x}_{t}\sim q(\bm{x}_{t}|\bm{x}_{0})$ into $\bm{x}_{0}$, with a time re-weighted square error loss—see Ho et al. (2020) for the complete proof:

$$\mathbb{E}_{\bm{x}_{0},\bm{\epsilon},t}[w_{t}\cdot||\hat{\bm{x}}_{\theta}(\bm{x}_{t},\bm{c})-\bm{x}_{0}||_{2}^{2}]$$

(2)

Here, the noised image is denoted as $\bm{x}_{t}:=\sqrt{\bar{\alpha}_{t}}\bm{x}_{0}+\sqrt{1-\bar{\alpha}_{t}}\bm{\epsilon}$, $\bm{x}_{0}$ is the ground-truth image, $\bm{c}$ is the condition, $\bm{\epsilon}\sim\mathcal{N}(\bf{0},I)$ is the noise term, $\alpha_{t}:=1-\beta_{t}$ and $\bar{\alpha}_{t}:=\prod_{s=1}^{t}\alpha_{s}$. To simplify notation, we will allow the condition $\bm{c}$ to include multiple conditioning signals, such as text prompts $\bm{c}_{p}$, a low-resolution image input $\bm{c}_{x}$ (which is used in super-resolution), or retrieved neighboring images $\bm{c}_{n}$ (which are used in Re-Imagen). Imagen (Saharia et al., 2022) uses a U-Net (Ronneberger et al., 2015) to implement $\bm{\epsilon}_{\theta}(\bm{x}_{t},\bm{c},t)$. The U-Net represents the reversed noise generator as follows:

$$\hat{\bm{x}}_{\theta}(\bm{x}_{t},\bm{c}):=(\bm{x}_{t}-\sqrt{1-\bar{\alpha}_{t}}\bm{\epsilon}_{\theta}(\bm{x}_{t},\bm{c},t))/\sqrt{\bar{\alpha}_{t}}$$

(3)

During the training, we randomly sample $t\sim\mathcal{U}([0,1])$ and image $\bm{x}_{0}$ from the dataset $\mathcal{D}$, and minimize the difference between $\hat{\bm{x}}_{\theta}(\bm{x}_{t},\bm{c})$ and $\bm{x}_{0}$ according to Equation 2. At the inference time, the diffusion model uses DDPM (Ho et al., 2020) to sample recursively as follows:

$$\bm{x}_{t-1}=\frac{\sqrt{\bar{\alpha}_{t-1}}\beta_{t}}{1-\bar{\alpha}_{t}}\hat{\bm{x}}_{\theta}(\bm{x}_{t},\bm{c})+\frac{\sqrt{\alpha_{t}}(1-\bar{\alpha}_{t-1})}{1-\bar{\alpha}_{t}}\bm{x}_{t}+\frac{\sqrt{(1-\bar{\alpha}_{t-1})\beta_{t}}}{\sqrt{1-\bar{\alpha}_{t}}}\bm{\epsilon}$$

(4)

The model sets $\bm{x}_{T}$ as a Gaussian noise with $T$ denoting the total number of diffusion steps, and then keeps sampling in reverse until step $T=0$, i.e. $\bm{x}_{T}\rightarrow\bm{x}_{T-1}\rightarrow\cdots$, to reach the final image $\hat{\bm{x}}_{0}$.

For better generation efficiency, cascaded diffusion models (Ho et al., 2022; Ramesh et al., 2022; Saharia et al., 2022) use three separate diffusion models to generate high-resolution images gradually, going from low resolution to high resolution. The three models 64$\times$ model, 256$\times$ super-resolution model, and 1024$\times$ super-resolution model gradually increase the model resolution to $1024\times 1024$.

Classifier-free Guidance Ho & Salimans (2021) first proposed classifier-free guidance to trade off diversity and sample quality. This sampling strategy has been widely used due to its simplicity. In particular, Imagen (Saharia et al., 2022) adopts an adjusted $\epsilon$-prediction as follows:

$$\hat{\bm{\epsilon}}=w\cdot\bm{\epsilon}_{\theta}(\bm{x}_{t},\bm{c},t)-(w-1)\cdot\bm{\epsilon}_{\theta}(\bm{x}_{t},t)$$

(5)

where $w$ is the guidance weight. The unconditional $\epsilon$-prediction $\bm{\epsilon}_{\theta}(\bm{x}_{t},t)$ is calculated by dropping the condition, i.e. the text prompt.

Similar to Imagen (Saharia et al., 2022), Re-Imagen is a cascaded diffusion model, consisting of 64$\times$, 256$\times$, and 1024$\times$ diffusion models. However, Re-Imagen augments the diffusion model with the new capability of leveraging multimodal ‘knowledge’ from the external database, thus freeing the model from memorizing the appearance of rare entities. For brevity (and concreteness) we present below a high-level overview of the 64$\times$ model: the others are similar.

Main Idea As shown in Figure 2, during the denoising process, Re-Imagen conditions its generation result not only on the text prompt $\bm{c}_{p}$ (and also with $\bm{c}_{x}$ for super-resolution), but on the neighbors $\bm{c}_{n}$ that were retrieved from the external knowledge base. Here, the text prompt $\bm{c}_{p}\in\mathbb{R}^{n\times d}$ is represented using a T5 embedding (Raffel et al., 2020), with $n$ being the text length and $d$ being the embedding dimension. Meanwhile, the top-k neighbors $\bm{c}_{n}:=[\texttt{<}\text{image, text}\texttt{>}_{1},\cdots,\texttt{<}\text{image, text}\texttt{>}_{k}]$ are retrieved from external knowledge base $\mathcal{B}$, using the input prompt $p$ as the query and a retrieval similarity function $\gamma(p,\mathcal{B})$. We experimented with two different choices for the similarity function: maximum inner product scores for BM25 (Robertson et al., 2009) and CLIP (Radford et al., 2021).
Model Architecture We show the architecture of our model in Figure 3, where we decompose the UNet into the downsampling encoder (DStack) and the upsampling decoder (UStack). Specifically, the DStack takes an image, a text, and a time step as the input, and generates a feature map, which is denoted as $f_{\theta}(\bm{x}_{t},\bm{c}_{p},t)\in\mathbb{R}^{F\times F\times d}$, with $F$ denoting the feature map width and $d$ denoting the hidden dimension. We share the same DStack encoder when we encode the retrieved <image, text> pairs (with $t$ set to zero) which produce a set of feature maps $f_{\theta}(\bm{c}_{n},0)\in\mathbb{R}^{K\times F\times F\times d}$. We then use a multi-head attention module (Vaswani et al., 2017) to extract the most relevant information to produce a new feature map $f^{\prime}_{\theta}(\bm{x}_{t},\bm{c}_{p},\bm{c}_{n},t)=Attn(f_{\theta}(\bm{x}_{t},\bm{c}_{p},t),f_{\theta}(\bm{c}_{n},0))$. The upsampling stack decoder then predicts the noise term $\bm{\epsilon}_{\theta}(\bm{x}_{t},\bm{c}_{p},\bm{c}_{n},t)$ and uses it to compute $\hat{\bm{x}}_{\theta}$ with Equation 3, which is either used for regression during training or DDPM sampling.
Model Training In order to train Re-Imagen, we construct a new dataset KNN-ImageText based on the 50M ImageText-dataset used in Imagen. There are two motivations for selecting this dataset. (1) the dataset contains many similar photos regarding specific entities, which is extremely helpful for obtaining similar neighbors, and (2) the dataset is highly sanitized with fewer unethical or harmful images. For each instance in the 50M ImageText-dataset, we search over the same dataset with text-to-text BM25 similarity to find the top-2 neighbors as $\bm{c}_{n}$ (excluding the query instance). We experimented with both CLIP and BM25 similarity scores, and retrieval was implemented with ScaNN (Guo et al., 2020). We train Re-Imagen on the KNN-ImageText by minimizing the loss function of Equation 2. During training, we also randomly drop the text and neighbor conditions independently with 10% chance. Such random dropping will help the model learn the marginalized noise term $\bm{\epsilon}_{\theta}(\bm{x}_{t},\bm{c}_{p},t)$ and $\bm{\epsilon}_{\theta}(\bm{x}_{t},\bm{c}_{n},t)$, which will be used for the classifier-free guidance.
Interleaved Classifier-free Guidance Different from existing diffusion models, our model needs to deal with more than one condition, i.e., text prompts $\bm{c}_{t}$ and retrieved neighbors $\bm{c}_{n}$, which allows new options for incorporating guidance. In particular, Re-Imagen could use classifier-free guidance by subtracting the unconditioned $\epsilon$-predictions, or either of the two partially conditioned $\epsilon$-predictions. Empirically, we observed that subtracting unconditioned $\epsilon$-predictions (the standard classifier-free guidance of Figure 2) often leads to an undesired imbalance, where the outputs are either dominated by the text condition or the neighbor condition. Hence, we designed an interleaved guidance schedule that balances the two conditions. Formally, we define the two adjusted $\epsilon$-predictions as:

$\displaystyle\begin{split}\hat{\bm{\epsilon}}_{p}&=w_{p}\cdot\bm{\epsilon}_{\theta}(\bm{x}_{t},\bm{c}_{p},\bm{c}_{n},t)-(w_{p}-1)\cdot\bm{\epsilon}_{\theta}(\bm{x}_{t},\bm{c}_{n},t)\\ \hat{\bm{\epsilon}}_{n}&=w_{n}\cdot\bm{\epsilon}_{\theta}(\bm{x}_{t},\bm{c}_{p},\bm{c}_{n},t)-(w_{n}-1)\cdot\bm{\epsilon}_{\theta}(\bm{x}_{t},\bm{c}_{p},t)\end{split}$

(6)

where $\hat{\bm{\epsilon}}_{p}$ and $\hat{\bm{\epsilon}}_{n}$ are the text-enhanced and neighbor-enhanced $\epsilon$-predictions, respectfully. Here, $w_{p}$ is the text guidance weight and $w_{n}$ is the neighbor guidance weight. We then interleave the two guidance predictions by a certain predefined ratio $\eta$. Specifically, at each guidance step, we sample a [0, 1]-uniform random number $R$, and $R<\eta$, we use $\hat{\bm{\epsilon}}_{p}$, and otherwise $\hat{\bm{\epsilon}}_{n}$. We can adjust $\eta$ to balance the faithfulness w.r.t text description or the retrieved image-text pairs.

In these two experiments, we used the standard non-interleaved classifier-free guidance (Figure 2) with $T$=1000 steps for both the 64$\times$ diffusion model and 256$\times$ super-resolution model. The guidance weight $w$ for the 64$\times$ model is swept over [1.0, 1.25, 1.5, 1.75, 2.0], while the 256$\times$256 super-resolution models’ guidance weight $w$ is swept over [1.0, 5.0, 8.0, 10.0]. We select the guidance $w$ with the best FID score, which is reported in Table 1. We also demonstrate examples in Figure 4.
COCO Results COCO is the most widely-used benchmark for text-to-image generation models. Although COCO does not contain many rare entities, it does contain unusual combinations of common entities, so it is plausible that retrieval augmentation could also help with some challenging text prompts. We adopt FID (Heusel et al., 2017) score to measure image quality. Following the previous literature, we randomly sample 30K prompts from the validation set as input to the model. The generated images are compared with the reference images from the full validation set (42K). We list the results in two columns: FID-30K denotes that the model with access to the in-domain COCO train set, while Zero-shot FID-30K does not have access to any COCO data.

Re-Imagen can achieve a significant gain on FID by retrieving from external databases: roughly a 2.0 absolute FID improvement over Imagen. Its performance is even better than fine-tuned Make-A-Scene (Gafni et al., 2022). We found that Re-Imagen retrieving from OOD database achieves less gain than IND database, but still obtains a 0.4 FID improvement over Imagen. When comparing with other retrieval-augmented models, Re-Imagen is shown to outperform KNN-Diffusion and Memory-Driven T2I models by a significant margin of 11 FID score. We also note that Re-Imagen-small is also competent in the FID, which outperforms normal-sized Imagen with fewer parameters.

WikiImages Results WikiImages is constructed based on the multimodal corpus provided in WebQA (Chang et al., 2022), which consists of <image, text> pairs crawled from Wikimedia Commons²²2https://commons.wikimedia.org/wiki/Main_Page. We filtered the original corpus to remove noisy data (see AppendixB), which leads to a total of 320K examples. We randomly sample 22K as our validation set to perform zero-shot evaluation, we further sample 20K prompts from the dataset as the input. Similar to the previous experiment, we also adopt the guidance weight schedule as before and evaluate 256$\times$256 images.

From Table 1, we found that using the OOD database (LAION) actually achieves better performance than using the IND database. Unlike COCO, WikiImages contains mostly entity-focused images, thus the importance of finding relevant entities in the database is more important than distilling the styles from the training set—and since the scale of LAION-400M is 100x larger than an in-domain database, the chance of retrieving related entities is much higher, which leads to better performance. One example is depicted in the lower part of Figure 4, where the LAION retrieval finds ‘Island of San Giorgio Maggiore’, which helps the model generate the classical Renaissance-style church.

Dataset Construction We introduce EntityDrawBench to evaluate the model’s capability to generate diverse sets of entities in different visual scenes. Specifically, we pick various types of visual entities (dog breeds, landmarks, foods, birds, and animated characters) from Wikipedia Commons, Google Landmarks and Fandom to construct our prompts. In total, we collect 250 entity-centric prompts for evaluation. These prompts are mostly very unique and we cannot find them on the Internet, let alone the model training data. The dataset construction details are in Appendix C. To evaluate the model’s capability to ground on broader types of entities, we also randomly select 20 objects like ‘sunglasses, backpack, vase, teapot, etc’ and write creative prompts for them. We compare our generation results with the results from DreamBooth (Ruiz et al., 2022) in Appendix H.

We use the constructed prompt as the input and its corresponding image-text pairs as the ‘retrieval’ for Re-Imagen, to generate four 1024$\times$1024 images. For the other models, we feed the prompts directly also to generate four images. We pick the best image of 4 random samples to rate its Photorealism and Faithfulness by human raters. For photorealism, we rate 1 if the image is moderately realistic without noticeable artifacts. For the faithfulness measure, we rate 1 if the image is faithful to both the entity appearance and the text description.

EntityDrawBench Results We use the proposed interleaved classifier-free guidance (subsection 3.2) for the 64$\times$ diffusion model, which runs for 256 diffusion steps under a strong guidance weight of $w$=30 for both text and neighbor conditions. For the 256$\times$ and 1024$\times$ resolution models, we use a constant guidance weight of 5.0 and 3.0, respectively, with 128 and 32 diffusion steps. The inference speed is 30-40 secs for 4 images on 4 TPU-v4 chips. We demonstrate our human evaluation results for faithfulness and photorealism in Table 2.

We can observe that Re-Imagen can in general achieve much higher faithfulness than the existing models while maintaining similar photorealism scores. When comparing with our backbone Imagen, we see the faithfulness score improves by around 40%, which indicates that our model is paying attention to the retrieved knowledge and assimilating it into the generation process.

We further partition the entities into ‘frequent’ and ‘infrequent’ categories based on their frequency (top 50% as ‘frequent’). We plot the faithfulness score for ‘frequent’ and ‘infrequent’ separately in Figure 5. We can see that Re-Imagen is less sensitive to the frequency of the input entities than the other models with only minor performance drops. This study reflects the effectiveness of text-to-image generation models on long-tail entities. More generation examples are shown in Appendix F.

Comparison to Other Models We demonstrate some examples from different models in Figure 6. As can be seen, the images generated from Re-Imagen strike a good balance between text alignment and entity fidelity. Unlike image editing to perform in-place modification, Re-Imagen can transform the neighbor entities both geometrically and semantically according to the text guidance. As a concrete example, Re-Imagen generates the Braque Saint-Germain (2nd row in Figure 6) on the grass, in a different viewpoint from to the reference image.

Impact of Number of Retrievals The number of retrievals $K$ is an important factor for Re-Imagen. We vary the number of $K$ for all three datasets to understand their impact on the model performance. From Figure 7, we found that on COCO and WikiImages, increasing K from 1 to 4 does not lead to many changes in the FID score. However, on EntityDrawBench, increasing K will dramatically improve the faithfulness of generated image. It indicates the importance of having multiple images to help Re-Imagen ground on the visual entity. We provide visual examples in Appendix D.

Text and Entity Faithfulness Trade-offs In our experiments, we found that there is a trade-off between faithefulness to the text prompt and faithfulness to the retrieved entity images. Based on Equation 6, by adjusting $\eta$, i.e.the proportion of $\hat{\epsilon}_{p}$ and $\hat{\epsilon}_{n}$ in the sampling schedule, we can control Re-Imagen so as to generate images that explore this tradeoff: decreasing $\eta$ will increase the entity’s entity faithfulness but decrease the text alignment. We found that having $\eta$ around 0.5 is usually a ‘sweet spot’ that balances both conditions.