Original: R E -I MAGEN : R ETRIEVAL -A UGMENTED T EXT - TO -I MAGE G ENERATOR
Revised: R E -I MAGEN : R ETRIEVAL -A UGMENTED T EXT - TO -I MAGE G ENERATOR

Original: Anonymous authors
Revised: Anonymous authors

Original: Paper under double-blind review
Revised: Paper under double-blind review

Original: A BSTRACT
Revised: A BSTRACT

Original: Research on text-to-image generation has witnessed signiﬁcant progress in generating diverse and photo-realistic images, driven by diffusion and auto-regressive models trained on large-scale image-text data.
Revised: Research on text-to-image generation has witnessed signiﬁcant progress in generating diverse and photo-realistic images, driven by diffusion and auto-regressive models trained on large-scale image-text data.

Original: Though state-of-the-art models can generate high-quality images of common entities, they often have difﬁculty generating images of uncommon entities, such as ‘Chortai (dog)’ or ‘Picarones (food)’.
Revised: Though state-of-the-art models can generate high-quality images of common entities, they often have difﬁculty generating images of uncommon entities, such as ‘Chortai (dog)’ or ‘Picarones (food)’.

Original: To tackle this issue, we present the Retrieval-Augmented Text-to-Image Generator (Re-Imagen), a generative model that uses retrieved information to produce highﬁdelity and faithful images, even for rare or unseen entities.
Revised: To tackle this issue, we present the Retrieval-Augmented Text-to-Image Generator (Re-Imagen), a generative model that uses retrieved information to produce high-ﬁdelity and faithful images, even for rare or unseen entities.

Original: Given a text prompt, Re-Imagen accesses an external multi-modal knowledge base to retrieve relevant (image, text) pairs, and uses them as references to generate the image.
Revised: Given a text prompt, Re-Imagen accesses an external multi-modal knowledge base to retrieve relevant (image, text) pairs and uses them as references to generate the image.

Original: With this retrieval step, Re-Imagen is augmented with the knowledge of high-level semantics and low-level visual details of the mentioned entities, and thus improves its accuracy in generating the entities’ visual appearances.
Revised: With this retrieval step, Re-Imagen is augmented with the knowledge of highlevel semantics and low-level visual details of the mentioned entities, and thus improves its accuracy in generating the entities’ visual appearances.

Original: We train Re-Imagen on a constructed dataset containing (image, text, retrieval) triples to teach the model to ground on both text prompt and retrieval.
Revised: We train ReImagen on a constructed dataset containing (image, text, retrieval) triples to teach the model to ground on both text prompt and retrieval.

Original: Furthermore, we develop a new sampling strategy to interleave the classiﬁer-free guidance for text and retrieval condition to balance the text and retrieval alignment.
Revised: Furthermore, we develop a new sampling strategy to interleave the classiﬁer-free guidance for text and retrieval conditions to balance the text and retrieval alignment.

Original: Re-Imagen achieves new SoTA FID results ontwo image generation benchmarks, such as COCO ( i.e. , FID = 5.25) and WikiImage ( i.e. , FID = 5.82) without ﬁne-tuning.
Revised: Re-Imagen achieves signiﬁcant gain on FID score over COCO and WikiImage.

Original: To further evaluate the capabilities of the model, we introduce EntityDrawBench, a new benchmark that evaluates image generation for diverse entities, from frequent to rare, across multiple visual domains.
Revised: To further evaluate the capabilities of the model, we introduce EntityDrawBench, a new benchmark that evaluates image generation for diverse entities, from frequent to rare, across multiple object categories including dogs, foods, landmarks, birds, and characters.

Original: Human evaluation on EntityDrawBench shows that ReImagen performs on par with the best prior models in photo-realism, but with signiﬁcantly better real-world faithfulness, especially on less frequent entities.
Revised: Human evaluation on EntityDrawBench shows that Re-Imagen can signiﬁcantly improve the ﬁdelity of generated images, especially on less frequent entities.

Original:  I NTRODUCTION
Revised:  I NTRODUCTION

Original: Recent research efforts on conditional generative modeling, such as Imagen (Saharia et al., 2022),
Revised: Recent research efforts on conditional generative modeling, such as Imagen (Saharia et al., 2022),

Original: DALL · E 2 (Ramesh et al., 2022), and Parti (Yu et al., 2022), have advanced text-to-image generation to an unprecedented level, producing accurate, diverse, and even creative images from text prompts.
Revised: DALL · E 2 (Ramesh et al., 2022), and Parti (Yu et al., 2022), have advanced text-to-image generation to an unprecedented level, producing accurate, diverse, and even create images from text prompts.

Original: These models leverage paired image-text data at Web scale (with hundreds of millions of training examples), and powerful backbone generative models, i.e. , autoregressive models (Van Den Oord et al., 2017; Ramesh et al., 2021; Yu et al., 2022), diffusion models (Ho et al., 2020; Dhariwal & Nichol, 2021), etc. , and generate highly realistic images.
Revised: These models leverage paired image-text data at Web scale (with hundreds of millions of training examples), and powerful backbone generative models, i.e. , autoregressive models (Van Den Oord et al., 2017; Ramesh et al., 2021; Yu et al., 2022), diffusion models (Ho et al., 2020; Dhariwal & Nichol, 2021), etc, and generate highly realistic images.

Original: Studying these models’ generation results, we discovered their outputs are surprisingly sensitive to the frequency of the entities (or objects) in the text prompts.
Revised: Studying these models’ generation results, we discovered their outputs are surprisingly sensitive to the frequency of the entities (or objects) in the text prompts.

Original: In particular, when generating text prompts about frequent entities (see Appendix B), these models often generate realistic images, with faithful grounding to the entities’ visual appearance.
Revised: In particular, when generating text prompts about frequent entities, these models often generate realistic images, with faithful grounding to the entities’ visual appearance.

Original: However, when generating from text prompts with less frequent entities, those models either hallucinate non-existent entities, or output related frequent entities (see Figure 1), failing to establish a connection between the generated image and the visual appearance of the mentioned entity.
Revised: However, when generating from text prompts with less frequent entities, those models either hallucinate nonexistent entities, or output related frequent entities (see Figure 1), failing to establish a connection between the generated image and the visual appearance of the mentioned entity.

Original: This key limitation can greatly harm the trustworthiness of text-to-image models in real-world applications and even raise ethical concerns.
Revised: This key limitation can greatly harm the trustworthiness of text-to-image models in real-world applications and even raise ethical concerns.

Revised: In our studies, we found these models suffer from signiﬁcant quality degradation in generating visual objects associated with under-represented groups.

Original: In this paper, we propose a Re trieval-augmented Text-to-
Revised: In this paper, we propose a Re trieval-augmented Text-to-

Original: Ima ge Gen erator (Re-Imagen), which alleviates such limitations by searching for entity information in a multi-modal knowledge base, rather than attempting to memorize the appearance of rare entities.
Revised: Ima ge Gen erator (Re-Imagen), which alleviates such limitations by searching for entity information in a multi-modal knowledge base, rather than attempting to memorize the appearance of rare entities.

Original: Speciﬁcally, we deﬁne our multi-modal knowledge base encodes the visual appearances and descriptions of entities with a collection of refer- ence < image, text > pairs’.
Revised: Speciﬁcally, we deﬁne our multi-modal knowledge base encodes the visual appearances and descriptions of entities with a collection of refer- ence < image, text > pairs’.

Original: To use this resource, Re-Imagen ﬁrst uses the input text prompt to retrieve the most relevant < image, text > pairs from the external multi-modal knowledge base, then uses the retrieved knowledge as model additional inputs to synthesize the target images.
Revised: To use this resource, Re-Imagen ﬁrst uses the input text prompt to retrieve the most relevant < image, text > pairs from the external multi-modal knowledge base, then uses the retrieved knowledge as model additional inputs to synthesize the target images.

Original: Consequently, the retrieved references provide knowledge regarding the semantic attributes and the concrete visual appearance of mentioned entities to guide Re-Imagen to paint the entities in the target images.
Revised: Consequently, the retrieved references provide knowledge regarding the semantic attributes and the concrete visual appearance of mentioned entities to guide Re-Imagen to paint the entities in the target images.

Original: The backbone of Re-Imagen is a cascaded diffusion model (Ho et al., 2022), which contains three independent generation stages (implemented as U-Nets (Ronneberger et al., 2015)) to gradually produce high-resolution ( i.e. , 1024 ⇥ 1024) images.
Revised: The backbone of Re-Imagen is a cascaded diffusion model (Ho et al., 2022), which contains three independent generation stages (implemented as U-Nets (Ronneberger et al., 2015)) to gradually produce high-resolution ( i.e. , 1024 × 1024) images.

Original: In particular, we train Re-Imagen on a dataset constructed from the image-text dataset used by Imagen (Saharia et al., 2022), where each data instance is associated with the top-k nearest neighbors within the dataset, based on text-only BM score.
Revised: In particular, we train Re-Imagen on a dataset constructed from the image-text dataset used by Imagen (Saharia et al., 2022), where each data instance is associated with the top-k nearest neighbors within the dataset, based on text-only BM score.

Original: The retrieved top-k < image, text > pairs will be used as a reference for the model attend to.
Revised: The retrieved top-k < image, text > pairs will be used as a reference for the model attend to.

Original: During inference, we design an interleaved guidance schedule that switches between text guidance and retrieval guidance, which ensures both text alignment and entity alignment.
Revised: During inference, we design an interleaved guidance schedule that switches between text guidance and retrieval guidance, which ensures both text alignment and entity alignment.

Original: We show some examples generated by Re-Imagen, and compare them against Imagen in Figure 1.
Revised: We show some examples generated by Re-Imagen, and compare them against Imagen in Figure 1.

Original: We can qualitatively observe thatthe our images are more faithful to the appearance of the reference entity.
Revised: We can qualitatively observe that our images are more faithful to the appearance of the reference entity.

Original: To further quantitatively evaluate Re-Imagen, we present zero-shot text-to-image generation results on two challenging datasets: COCO (Lin et al., 2014) and WikiImages (Chang et al., 2022) 1 .
Revised: To further quantitatively evaluate Re-Imagen, we present zero-shot text-to-image generation results on two challenging datasets: COCO (Lin et al., 2014) and WikiImages (Chang et al., 2022) 1 .

Original: ReImagen uses an external non-overlapping image-text database as the knowledge base for retrieval and then grounds on the retrieval to synthesize the target image.
Revised: ReImagen uses an external non-overlapping image-text database as the knowledge base for retrieval and then grounds on the retrieval to synthesize the target image.

Original: We show that Re-Imagen achieves the state-of-the-art performance for text-to-image generation on COCO and WikiImages, measured in FID score Heusel et al. (2017), among non-ﬁne-tuned models.
Revised: We show that Re-Imagen achieves the state-of-the-art performance for text-to-image generation on COCO and WikiImages, measured in FID score (Heusel et al., 2017), among non-ﬁne-tuned models.

Original: For  non-entity-centric dataset
Revised: For the non-entity-centric dataset

Original: COCO, the perform gain is coming from biasing the model to generate images with similar styles as the retrieved in-domain images.
Revised: COCO, the performance gain is coming from biasing the model to generate images with similar styles as the retrieved in-domain images.

Original: For the entity-centric dataset WikiImages, the performance gain comes from grounding the generation on retrieved images containing similar entities.
Revised: For the entity-centric dataset WikiImages, the performance gain comes from grounding the generation on retrieved images containing similar entities.

Original: We further evaluate Re-Imagen on a more challenging benchmark — EntityDrawBench, to test the model’s ability to generate a variety of infrequent entities (dogs, landmarks, foods) in different scenes.
Revised: We further evaluate Re-Imagen on a more challenging benchmark — EntityDrawBench, to test the model’s ability to generate a variety of infrequent entities (dogs, landmarks, foods, birds, animated characters) in different scenes.

Original: We compare Re-Imagen with Imagen (Saharia et al., 2022), DALL-E 2 (Ramesh et al., 2022) and StableDiffusion (Rombach et al., 2022) in terms of faithfulness and photorealism with human raters.
Revised: We compare Re-Imagen with Imagen (Saharia et al., 2022), DALLE 2 (Ramesh et al., 2022) and StableDiffusion (Rombach et al., 2022) in terms of faithfulness and photorealism with human raters.

Original: We demonstrate thatthat Re-Imagen can reach around 71% on  the faithfulness score, greatly improving from Imagine’s score of 27% and DALL-E 2’s score of 48% .
Revised: We demonstrate that Re-Imagen can generate faithful and realistic images on 80% over input prompts, beating the existing best models by at least 30% on EntityDrawBench.

Original: Analysis shows that the improvement mostly comes from low-frequency  entities.
Revised: Analysis shows that the improvements are mostly coming from low-frequency visual entities.

Original: To summarize, our key contributions are: (1) a novel retrieval-augmented text-to-image model ReImagen, which achieves SoTA FID scores on two dasets; (2) interleaved classiﬁer-free guidance during sampling to ensure both text alignment and entity ﬁdelity; and (3) We introduce EntityDrawBench and show that Re-Imagen can signiﬁcantly improve faithfulness on less-frequent entities.
Revised: To summarize, our key contributions are: (1) a novel retrieval-augmented text-to-image model ReImagen, which improves FID scores on two datasets; (2) interleaved classiﬁer-free guidance during sampling to ensure both text alignment and entity ﬁdelity; and (3) We introduce EntityDrawBench and show that Re-Imagen can signiﬁcantly improve faithfulness on less-frequent entities.

Original:  R ELATED W ORK
Revised:  R ELATED W ORK

Original: Text-to-Image Diffusion Models There has been a wide-spread success (Ashual et al., 2022;
Revised: Text-to-Image Diffusion Models There has been a wide-spread success (Ashual et al., 2022;

Original: Ramesh et al., 2022; Saharia et al., 2022; Nichol et al., 2021) in modeling text-to-image generation with diffusion models, which has outperformed GANs (Goodfellow et al., 2014) and auto-regressive
Revised: Ramesh et al., 2022; Saharia et al., 2022; Nichol et al., 2021) in modeling text-to-image generation with diffusion models, which has outperformed GANs (Goodfellow et al., 2014) and auto-regressive

Original: Transformers (Ramesh et al., 2021) in photorealism and diversity (under similar model size), without training instability and mode collapsing issues.
Revised: Transformers (Ramesh et al., 2021) in photorealism and diversity (under similar model size), without training instability and mode collapsing issues.

Original: Among them, some recent large text-to-image models such as Imagen (Saharia et al., 2022), GLIDE (Nichol et al., 2021), and DALL-E2 (Ramesh et al., 2022) have demonstrated excellent generation from complex prompt inputs.
Revised: Among them, some recent large text-to-image models such as Imagen (Saharia et al., 2022), GLIDE (Nichol et al., 2021), and DALL-E2 (Ramesh et al., 2022) have demonstrated excellent generation from complex prompt inputs.

Original: These models achieve highly ﬁne-grained control over the generated images with text inputs.
Revised: These models achieve highly ﬁne-grained control over the generated images with text inputs.

Original: However, they do not perform explicit grounding over external visual knowledge and are restricted to memorizing the visual appearance of every possible visual entity in their parameters.
Revised: However, they do not perform explicit grounding over external visual knowledge and are restricted to memorizing the visual appearance of every possible visual entity in their parameters.

Original: This makes it difﬁcult for them to generalize to rare or even unseen entities.
Revised: This makes it difﬁcult for them to generalize to rare or even unseen entities.

Original: In contrast, Re-Imagen is designed to free the diffusion model from memorizing, as models are encouraged to retrieve semantic neighbors from the knowledge base and use retrievals as context to paint the image.
Revised: In contrast, Re-Imagen is designed to free the diffusion model from memorizing, as models are encouraged to retrieve semantic neighbors from the knowledge base and use retrievals as context to paint the image.

Original: Re-Imagen improves the grounding of the diffusion models to real-world knowledge and is therefore capable of faithful image synthesis.
Revised: Re-Imagen improves the grounding of the diffusion models to real-world knowledge and is therefore capable of faithful image synthesis.

Original: Concurrent Work There are several concurrent works (Li et al., 2022; Blattmann et al., 2022;
Revised: Concurrent Work There are several concurrent works (Li et al., 2022; Blattmann et al., 2022;

Original: Ashual et al., 2022), that also leverage retrieval to improve diffusion models.
Revised: Ashual et al., 2022), that also leverage retrieval to improve diffusion models.

Original: RDM (Blattmann et al., 2022) is trained similarly to Re-Imagen, using examples and near neighbors, but the neighbors in RDM are selected using image features, and at inference time retrievals are replaced with user-chosen exemplars.
Revised: RDM (Blattmann et al., 2022) is trained similarly to Re-Imagen, using examples and near neighbors, but the neighbors in RDM are selected using image features, and at inference time retrievals are replaced with user-chosen exemplars.

Original: RDM was shown to effectively transfer artistic style from exemplars to generated images.
Revised: RDM was shown to effectively transfer artistic style from exemplars to generated images.

Original: In contrast, our proposed Re-Imagen conditions on both text and multi-modal neighbors to generate the image, includes retrieval at inference time, and is demonstrated to improve performance on rare images (as well as more generally).
Revised: In contrast, our proposed Re-Imagen conditions on both text and multi-modal neighbors to generate the image includes retrieval at inference time and is demonstrated to improve performance on rare images (as well as more generally).

Original: KNN-Diffusion (Ashual et al., 2022) is more closely related work to us, as it also uses retrieval to the quality of generated images.
Revised: KNN-Diffusion (Ashual et al., 2022) is more closely related work to us, as it also uses retrieval to the quality of generated images.

Original: However, KNN-Diffusion uses discrete image representations, while Re-Imagen uses the raw pixels, and
Revised: However, KNN-Diffusion uses discrete image representations, while Re-Imagen uses the raw pixels, and

Original: Re-Imagen’s retrieved neighbors can be < image, text > pairs, while KNN-Diffusion’s are only images.
Revised: Re-Imagen’s retrieved neighbors can be < image, text > pairs, while KNN-Diffusion’s are only images.

Original: Quantitatively, Re-Imagen outperforms KNN-Diffusion on the COCO dataset signiﬁcantly.
Revised: Quantitatively, Re-Imagen outperforms KNN-Diffusion on the COCO dataset signiﬁcantly.

Original: Others Due to the space limit, we provide an additional literature review inthe Appendix A.
Revised: Others Due to the space limit, we provide an additional literature review in Appendix A.

Original:  M ODEL
Revised:  M ODEL

Original: In this section, wediscuss our proposed Re-Imagen in detail. We start with background knowledge, in the form of a brief overview of the cascaded diffusion models used by Imagen (Wang et al., 2022).
Revised: In this section, we start with background knowledge, in the form of a brief overview of the cascaded diffusion models used by Imagen.

Original: Next, we describe the concrete technical details on how we incorporatemulti-modal retrieval for Re-Imagen.
Revised: Next, we describe the concrete technical details of how we incorporate retrieval for Re-Imagen.

Original: Finally, we discusstraining and the interleaved guidance sampling for Re-Imagen.
Revised: Finally, we discuss interleaved guidance sampling.

Original: 3. P RELIMINARIES
Revised: 3. P RELIMINARIES

Original: Diffusion Models Diffusion models (Sohl-Dickstein et al., 2015) are latent variable models, parameterized by ✓ , in the form of p ✓ ( x 0 ) :
Revised: Diffusion Models Diffusion models (Sohl-Dickstein et al., 2015) are latent variable models, parameterized by θ , in the form of p θ ( x 0 ) :

Original: = R p ✓ ( x 0: T ) d x 1: T , where x 1 , · · · , x T are “noised” latent versions of the input image x ⇠ q ( x 0 ) .
Revised: = (cid:82) p θ ( x 0: T ) d x 1: T , where x 1 , · · · , x T are “noised” latent versions of the input image x 0 ∼ q ( x 0 ) .

Original: Note that the dimensionality of both latents and the image are the same throughout the entire process, with x 0: T 2 R d and d equals the product of < height, width, # of channels > .
Revised: Note that the dimensionality of both latents and the image is the same throughout the entire process, with x 0: T ∈ R d and d equals the product of < height, width, # of channels > .

Original: The process that computes the posterior distribution q ( x 1: T | x 0 ) is also called the forward (or diffusion) process, and is implemented as a predeﬁned Markov chain that gradually adds Gaussian noise to the data according to a schedule  t :
Revised: The process that computes the posterior distribution q ( x 1: T | x 0 ) is also called the forward (or diffusion) process, and is implemented as a predeﬁned Markov chain that gradually adds Gaussian noise to the data according to a schedule β t :

Original: [Equation]
Revised: [Equation]

Original: Diffusion models are trained to learn the image distribution by reversing the diffusion Markov chain.
Revised: Diffusion models are trained to learn the image distribution by reversing the diffusion Markov chain.

Original: Theoretically, this reduces to learning to denoise x t ⇠ q ( x t | x 0 ) into x 0 , with a time re-weighted square error loss—see Ho et al.
Revised: Theoretically, this reduces to learning to denoise x t ∼ q ( x t | x 0 ) into x 0 , with a time re-weighted square error loss—see Ho et al.

Original: (2020) for the complete proof:
Revised: (2020) for the complete proof:

Original: [Equation]
Revised: [Equation]

Original: Here, the noised image is denoted as x t :=
Revised: Here, the noised image is denoted as x t :=

Original: p ¯ ↵ t x 0 + (1  ¯ ↵ t ) ✏ , x 0 is the ground-truth image, c is the condition, ✏ ⇠ N ( 0 , I ) is the noise term, ↵ t := 1   t and ¯ ↵ t := Q t s =1 ↵ s .
Revised: √ ¯ α t x 0 + √ 1 − ¯ α t (cid:15) , x 0 is the ground-truth image, c is the condition, (cid:15) ∼ N ( 0 , I ) is the noise term, α t := 1 − β t and ¯ α t := (cid:81) ts =1 α s .

Original: To simplifySpeciﬁcally, Re-Imagen learns a UNet to iteratively predict ✏ ( x t c n c p , t that denoises the notation, we will allow the condition c  include multiple conditioning signals, such as text prompts c p , a low-resolution image input c x (which is used in super-resolution), or retrieved neighboring images c n (which are used in Re-Imagen).
Revised: To simplify notation, we will allow the condition c to include multiple conditioning signals, such as text prompts c p , a low-resolution image input c x (which is used in super-resolution), or retrieved neighboring images c n (which are used in Re-Imagen).

Original: Imagen (Saharia et al., 2022) uses a U-Net (Ronneberger et al., 2015) to implement ✏ ✓ ( x t , c , t ) .
Revised: Imagen (Saharia et al., 2022) uses a U-Net (Ronneberger et al., 2015) to implement (cid:15) θ ( x t , c , t ) .

Original: The U-Net represents the reversed noise generator as follows:
Revised: The U-Net represents the reversed noise generator as follows:

Original: [Equation]
Revised: [Equation]

Original: During the training, we randomly sample t ⇠ U ([0 , 1]) and image x 0 from the dataset D , and minimize the difference between ˆ x ✓ ( x t , c ) and x 0 according to Equation 2.
Revised: During the training, we randomly sample t ∼ U ([0 , 1]) and image x 0 from the dataset D , and minimize the difference between ˆ x θ ( x t , c ) and x 0 according to Equation 2.

Original: At the inference time, the diffusion model uses DDPM
Revised: At the inference time, the diffusion model uses DDPM

Original: (Ho et al., 2020) to sample recursively as follows:
Revised: (Ho et al., 2020) to sample recursively as follows:

Original: [Equation]
Revised: [Equation]

Original: The model sets x T as a Gaussian noise with T denoting the total number of diffusion steps, and then keep sampling in reverse until step T = 0 , i.e. x T ! x T  1 ! · · · , to reach the ﬁnal image ˆ x 0 .
Revised: The model sets 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. x T → x T − 1 → · · · , to reach the ﬁnal image ˆ x 0 .

Original: For better generation efﬁciency, cascaded diffusion models (Ho et al., 2022; Ramesh et al., 2022;
Revised: For better generation efﬁciency, cascaded diffusion models (Ho et al., 2022; Ramesh et al., 2022;

Original: Saharia et al., 2022) use three separate diffusion models to generate high-resolution images gradually, going from low resolution to high resolution.
Revised: Saharia et al., 2022) use three separate diffusion models to generate high-resolution images gradually, going from low resolution to high resolution.

Original: The three models 64 ⇥ model, 256 ⇥ super-resolution model and 1024 ⇥ super-resolution model gradually increase the model resolution to 1024 ⇥ 1024 .
Revised: The three models 64 × model, 256 × super-resolution model, and 1024 × super-resolution model gradually increase the model resolution to 1024 × 1024 .

Original: Classiﬁer-free Guidance Ho & Salimans (2021)
Revised: Classiﬁer-free Guidance Ho & Salimans (2021)

Original: ﬁrst proposed classiﬁer-free guidance to trade off diversity and sample quality.
Revised: ﬁrst proposed classiﬁer-free guidance to trade off diversity and sample quality.

Original: This sampling strategy has been widely used due to its simplicity.
Revised: This sampling strategy has been widely used due to its simplicity.

Original: In particular, Imagen (Saharia et al., 2022) adopts an adjusted ✏ -prediction as follows:
Revised: In particular, Imagen (Saharia et al., 2022) adopts an adjusted (cid:15) -prediction as follows:

Original: [Equation]
Revised: [Equation]

Original: where w is the guidance weight.
Revised: where w is the guidance weight.

Original: The unconditional ✏ -prediction ✏ ✓ ( x t , t ) is calculated by dropping the condition, i.e. the text prompt.
Revised: The unconditional (cid:15) -prediction (cid:15) θ ( x t , t ) is calculated by dropping the condition, i.e. the text prompt.

Original: 3. G ENERATING I MAGE WITH M ULTI -M ODAL K NOWLEDGE
Revised: 3. G ENERATING I MAGE WITH M ULTI -M ODAL K NOWLEDGE

Original: Similar to Imagen (Saharia et al., 2022), Re-Imagen is a cascaded diffusion model, consisting of 64 ⇥ , 256 ⇥ , and 1024 ⇥ diffusion models.
Revised: Similar to Imagen (Saharia et al., 2022), Re-Imagen is a cascaded diffusion model, consisting of 64 × , 256 × , and 1024 × diffusion models.

Original: 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.
Revised: 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.

Original: For brevity (and concreteness) we present below a high-level overview of the 64 ⇥ model: the others are similar.
Revised: For brevity (and concreteness) we present below a high-level overview of the 64 × model: the others are similar.

Original: Main Idea As shown in Figure 2, during the denoising process, Re-Imagen conditions its generation result not only on the text prompt c p (and also with c x for super-resolution), but on the neighbors c n that were retrieved from the external knowledge base.
Revised: Main Idea As shown in Figure 2, during the denoising process, Re-Imagen conditions its generation result not only on the text prompt c p (and also with c x for super-resolution), but on the neighbors c n that were retrieved from the external knowledge base.

Original: Here, the text prompt c p 2 R n ⇥ d is represented using a T5 embedding (Raffel et al., 2020), with n being the text length and d being the embedding dimension.
Revised: Here, the text prompt c p ∈ R n × d is represented using a T5 embedding (Raffel et al., 2020), with n being the text length and d being the embedding dimension.

Original: Meanwhile, the top-k neighbors c n :=
Revised: Meanwhile, the top-k neighbors c n :=

Original: [ < image, text > 1 , · · · , < image, text > k ] are retrieved from external knowledge base B , using the input prompt p as the query and a retrieval similarity function  ( p, B ) .
Revised: [ < image, text > 1 , · · · , < image, text > k ] are retrieved from external knowledge base B , using the input prompt p as the query and a retrieval similarity function γ ( p, B ) .

Original: 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).
Revised: 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).

Original: Model Architecture We show the architecture of our model in Figure 3, where we decompose the
Revised: Model Architecture We show the architecture of our model in Figure 3, where we decompose the

Original: UNet into the downsampling encoder (DStack) and the upsampling decoder (UStack).
Revised: UNet into the downsampling encoder (DStack) and the upsampling decoder (UStack).

Original: Speciﬁcally, the DStack takes an image, a text, and a time step as the input, and generates a feature map, which is denoted as f ✓ ( x t , c p , t ) 2 R F ⇥ F ⇥ d , with F denoting the feature map width and d denoting the hidden dimension.
Revised: Speciﬁcally, the DStack takes an image, a text, and a time step as the input, and generates a feature map, which is denoted as f θ ( x t , c p , t ) ∈ R F × F × d , with F denoting the feature map width and d denoting the hidden dimension.

Original: We share the same DStack encoder when we encode the retrieved < image, text > pairs (with t set to zero) which produces a set of feature maps f ✓ ( c n , 0) 2 R K ⇥ F ⇥ F ⇥ d .
Revised: 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 θ ( c n , 0) ∈ R K × F × F × d .

Original: 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 0 ✓ ( x t , c p , c n , t ) =
Revised: 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 (cid:48) θ ( x t , c p , c n , t ) =

Original: Attn ( f ✓ ( x t , c p , t ) , f ✓ ( c n , 0)) .
Revised: Attn ( f θ ( x t , c p , t ) , f θ ( c n , 0)) .

Original: The upsampling stack decoder then predicts the noise term ✏ ✓ ( x t , c p , c n , t ) and uses it to compute ˆ x ✓ with Equation 3, which is either used for regression during training or DDPM sampling.
Revised: The upsampling stack decoder then predicts the noise term (cid:15) θ ( x t , c p , c n , t ) and uses it to compute ˆ x θ with Equation 3, which is either used for regression during training or DDPM sampling.

Original: Model Training
Revised: Model Training

Original: In order to train Re-Imagen, we construct a new dataset KNN-ImageText based on the 50M ImageText-dataset used in Imagen.
Revised: In order to train Re-Imagen, we construct a new dataset KNN-ImageText based on the 50M ImageText-dataset used in Imagen.

Original: There are two motivations for selecting this dataset.
Revised: There are two motivations for selecting this dataset.

Original: (1) the dataset contains many similar photos regarding speciﬁc entities, which is extremely helpful for obtaining similar neighbors, and (2) the dataset is highly sanitized with fewer unethical or harmful images.
Revised: (1) the dataset contains many similar photos regarding speciﬁc entities, which is extremely helpful for obtaining similar neighbors, and (2) the dataset is highly sanitized with fewer unethical or harmful images.

Original: For each instance in the 50M ImageText-dataset, we search over the same dataset with text-to-text BM25 similarity to ﬁnd the top-2 neighbors as c n (excluding the query instance).
Revised: For each instance in the 50M ImageText-dataset, we search over the same dataset with text-to-text BM25 similarity to ﬁnd the top-2 neighbors as c n (excluding the query instance).

Original: We experimented with both CLIP and BM25 similarity scores, and retrieval was implemented with
Revised: We experimented with both CLIP and BM25 similarity scores, and retrieval was implemented with

Original: ScaNN (Guo et al., 2020).
Revised: ScaNN (Guo et al., 2020).

Original: We train Re-Imagen on the KNN-ImageText by minimizing the loss function of Equation 2.
Revised: We train Re-Imagen on the KNN-ImageText by minimizing the loss function of Equation 2.

Original: During training, we also randomly drop the text and neighbor conditions independently with 10% chance.
Revised: During training, we also randomly drop the text and neighbor conditions independently with 10% chance.

Original: Such random dropping will help the model learn the marginalized noise term ✏ ✓ ( x t , c p , t ) and ✏ ✓ ( x t , c n , t ) , which will be used for the classiﬁer-free guidance.
Revised: Such random dropping will help the model learn the marginalized noise term (cid:15) θ ( x t , c p , t ) and (cid:15) θ ( x t , c n , t ) , which will be used for the classiﬁer-free guidance.

Original: Interleaved Classiﬁer-free Guidance Different from existing diffusion models, our model needs to deal with more than one condition, i.e. , text prompts c t and retrieved neighbors c n , which allows new options for incorporating guidance.
Revised: Interleaved Classiﬁer-free Guidance Different from existing diffusion models, our model needs to deal with more than one condition, i.e. , text prompts c t and retrieved neighbors c n , which allows new options for incorporating guidance.

Original: In particular, Re-Imagen could use classiﬁer-free guidance by subtracting the unconditioned ✏ -predictions, or either of the two partially conditioned ✏ -predictions.
Revised: In particular, Re-Imagen could use classiﬁer-free guidance by subtracting the unconditioned (cid:15) -predictions, or either of the two partially conditioned (cid:15) -predictions.

Original: Empirically, we observed that subtracting unconditioned ✏ -predictions (the standard classiﬁer-free guidance of Figure 3.1) often leads to an undesired imbalance, where the outputs are either dominated by the text condition or the neighbor condition.
Revised: Empirically, we observed that subtracting unconditioned (cid:15) -predictions (the standard classiﬁer-free guidance of Figure 3.1) often leads to an undesired imbalance, where the outputs are either dominated by the text condition or the neighbor condition.

Original: Hence, we designed an interleaved guidance schedule that balances the two conditions.
Revised: Hence, we designed an interleaved guidance schedule that balances the two conditions.

Original: Formally, we deﬁne the two adjusted ✏ -predictions as:
Revised: Formally, we deﬁne the two adjusted (cid:15) -predictions as:

Original: [Equation]
Revised: [Equation]

Original: (6) where ˆ ✏ p and ˆ ✏ n are the text-enhanced and neighbor-enhanced ✏ -predictions, respectfully.
Revised: where ˆ (cid:15) p and ˆ (cid:15) n are the text-enhanced and neighbor-enhanced (cid:15) -predictions, respectfully.

Original: Here, w p is the text guidance weight and w n is the neighbor guidance weight.
Revised: Here, w p is the text guidance weight and w n is the neighbor guidance weight.

Original: We then interleave the two guidance predictions by a certain predeﬁned ratio ⌘ .
Revised: We then interleave the two guidance predictions by a certain predeﬁned ratio η .

Original: Speciﬁcally, at each guidance step, we sample a [0, 1]-uniform random number R , and R < ⌘ , we use ˆ ✏ p , and otherwise ˆ ✏ n .
Revised: Speciﬁcally, at each guidance step, we sample a [0, 1]-uniform random number R , and R < η , we use ˆ (cid:15) p , and otherwise ˆ (cid:15) n .

Original: We can adjust ⌘ to balance the faithfulness w.r.t text description or the retrieved image-text pairs.
Revised: We can adjust η to balance the faithfulness w.r.t text description or the retrieved image-text pairs.

Original: In EntityDrawBench experiment, we found that ⌘ = 0 .

Original: 55 can lead to better quality.

Original:  E XPERIMENTS
Revised:  E XPERIMENTS

Original: Re-Imagen consists of three submodels: a 2.5B 64 ⇥ 64 text-to-image model, a 750M 256 ⇥super-resolution model and a 400M 1024 ⇥ 1024 super-resolution model. 
Revised: Re-Imagen consists of three submodels: a 2.5B 64 × 64 text-to-image model, a 750M 256 ×super-resolution model and a 400M 1024 × 1024 super-resolution model. We also have a Re-Imagensmall with 1.4B 64 × 64 text-to-image model to understand the impact of model size.

Original:  We ﬁnetune these models on the constructed KNN-ImageText dataset.
Revised: [Table] We ﬁnetune these models on the constructed KNN-ImageText dataset.

Original: We evaluate the model under two settings: (1) automatic evaluation on COCO and WikiImages dataset, to measure the model’s general performance[Table] to generate photorealistic images, and ( evaluation on  to tail
Revised: We evaluate the model under two settings: (1) automatic evaluation on COCO and WikiImages dataset, to measure the model’s general performance to generate photorealistic images, and (2) human evaluation on the newly introduced EntityDrawBench, to measure the model’s capability to generate long-tail entities.

Original: Training and Evaluation details The ﬁne-tuning was run for 200K steps on 64 TPU-v4 chips and completed within two days.
Revised: Training and Evaluation details The ﬁne-tuning was run for 200K steps on 64 TPU-v4 chips and completed within two days.

Original: We use Adafactor for the 64 ⇥ model and Adam for the 256 ⇥ super-resolution model with a learning rate of 1e-4.
Revised: We use Adafactor for the 64 × model and Adam for the 256 × superresolution model with a learning rate of 1e-4.

Original: We set the number of neighbors k
Revised: We set the number of neighbors k

Original: [-=
-] and set  = BM25 during training.
Revised: [-=
-]=2 and set γ =BM during training.

Original: For the image-text database B , we consider three different variants: (1) the COCO/WikiImages training set, which contains non-overlapping small-scale in-domain imagetext pairs, (2) the ImageText dataset containing 50M < image, caption > pairs, and (3) the LAION dataset (Schuhmann et al., 2021) containing 400M < image, text > crawled pairs.
Revised: For the image-text database B , we consider three different variants: (1) the indomain training set, which contains non-overlapping small-scale in-domain image-text pairs from COCO or WikiImages, (2) the out-of-domain LAION dataset (Schuhmann et al., 2021) containing 400M < image, text > crawled pairs.

Original: Since indexing ImageText and LAION with CLIP encodings is expensive, we only considered the BM25 retriever for these databases.
Revised: Since indexing ImageText and LAION with CLIP encodings is expensive, we only considered the BM25 retriever for these databases.

Original: For the COCO/WikiImages training set, we used both BM25 and CLIP.

Original: 4. E VALUATION ON COCO AND W IKI I MAGES
Revised: 4. E VALUATION ON COCO AND W IKI I MAGES

Original: In these two experiments, we used the standard non-interleaved classiﬁer-free guidance (Figure 3.1) with T =1000 steps for both the 64 ⇥ diffusion model and 256 ⇥ super-resolution model.
Revised: In these two experiments, we used the standard non-interleaved classiﬁer-free guidance (Figure 3.1) with T =1000 steps for both the 64 × diffusion model and 256 × super-resolution model.

Original: The guidance weight w for the 64 ⇥ model is swept over [1.0, 1.25, 1.5, 1.75, 2.0], while the 256 ⇥ 256 superresolution models’ guidance weight w is swept over [1.0, 5.0, 8.0, 10.0].
Revised: The guidance weight w for the 64 × model is swept over [1.0, 1.25, 1.5, 1.75, 2.0], while the 256 × 256 superresolution models’ guidance weight w is swept over [1.0, 5.0, 8.0, 10.0].

Original: We select the guidance w with the best FID score, which is reported in Table 1.
Revised: We select the guidance w with the best FID score, which is reported in Table 1.

Original: We also demonstrate examples in Figure 4.
Revised: We also demonstrate examples in Figure 4.

Original: COCO Results COCO is the most widely-used benchmark for text-to-image generation models.
Revised: COCO Results COCO is the most widely-used benchmark for text-to-image generation models.

Original: 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 for some challenging text prompts.
Revised: 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.

Original: We adopt FID (Heusel et al., 2017) score to measure image quality.
Revised: We adopt FID (Heusel et al., 2017) score to measure image quality.

Original: Following the previous literature, we randomly sample 30K prompts from the validation set as input to the model.
Revised: Following the previous literature, we randomly sample 30K prompts from the validation set as input to the model.

Original: The generated images are compared with the reference images from the full validation set (42K).
Revised: The generated images are compared with the reference images from the full validation set (42K).

Original: We list the results in two columns: FID-30K denotes  the model with access to the  COCO train set (either to ﬁne-tune or retrieve from), while Zero-shot FID-30K does not have access to any COCO data.
Revised: 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.

Original: Re-Imagen(with the COCO database) can achieve a signiﬁcant gain on FID-30K without ﬁnetuning: roughly a 2.0 absolute FID improvement over Imagen.
Revised: Re-Imagen can achieve a signiﬁcant gain on FID by retrieving from external databases: roughly a 2.0 absolute FID improvement over Imagen.

Original: The performance is even better than ﬁne-tuned Make-A-Scene (Gafni et al., 2022), but slightly worse than ﬁne-tuned 20B Parti.
Revised: Its performance is even better than ﬁne-tuned MakeA-Scene (Gafni et al., 2022).

Original: In contrast, Re-Imagen retrieving from out-of-domain databases (LAION) achieves less gain, but still obtains a 0.4 FID improvement over Imagen.
Revised: 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.

Original: Re-Imagen outperforms KNN-Diffusion, another retrieval-augmented diffusion model, by a large margin.
Revised: When comparing with other retrieval-augmented models, Re-Imagen is shown to outperform KNN-Diffusion and MemoryDriven T2I models by a signiﬁcant 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.

Original: Since COCO does not contain infrequent entities, ‘entity knowledge’ is not important. In contrast, retrieving from the training set can provide useful ‘style knowledge’ for the model to ground on.
Revised: As COCO does not contain infrequent entities, retrievals from the in-domain database mainly provide useful ‘style knowledge’ for the model to ground on.

Original: Re-Imagen is able to adaptthe generated images tothe same style of the COCO distribution, it can achieve amuch better FID score.
Revised: Re-Imagen can better adapt to COCO distribution, thus achieving a better FID score.

Original: As can be seen in the upper part of from Figure 4,Re-Imagen with
Revised: As can be seen in the upper part of from Figure 4,

Original: [Figure] [Table]
Revised: Re-Imagen with retrieval generates images of the same style as COCO, while without retrieval, the output is still high quality, but the style is less similar to COCO.

Original: 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 2 .
Revised: 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 2 .

Original: We ﬁltered the original corpus to remove noisy data (see the AppendixC), which leads to a total of 320K examples.
Revised: We ﬁltered the original corpus to remove noisy data (see AppendixB), which leads to a total of 320K examples.

Original: We randomly sample 22K as our validation set to perform zero-shot evaluation, we further sample 20K prompts from the dataset as the input.
Revised: We randomly sample 22K as our validation set to perform zero-shot evaluation, we further sample 20K prompts from the dataset as the input.

Original: Similar to the previous experiment, we also adopt the guidance weight schedule as before and evaluate 256 ⇥ 256 images.
Revised: Similar to the previous experiment, we also adopt the guidance weight schedule as before and evaluate 256 × 256 images.

Original: We report our experimental results in Table 2 and mainly compare with Imagen and Stable-Diffusion.

Original: From Table 2, we found that usingout-of-domain LAION-400M as the  database  actually achieves better performance than usingin-domain WikiImages as the  database.
Revised: From Table 1, we found that using the OOD database (LAION) actually achieves better performance than using the IND database.

Original: Unlike COCO, WikiImages contains mostly entity-focused images, thus the importance of ﬁnding relevant entities is the database is more important than distilling the styles from the training set—and since the scale of LAION-400M is 100x larger than WikiImages-300K, the chance of retrieving related entities is much higher, which leads to better performance.
Revised: Unlike COCO, WikiImages contains mostly entity-focused images, thus the importance of ﬁnding 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.

Original: One example is depicted in the lower part of Figure 4, where the LAION retrieval ﬁnds ‘Island of San Giorgio Maggiore’, which helps the model generate the classical Renaissance-style church.
Revised: One example is depicted in the lower part of Figure 4, where the LAION retrieval ﬁnds ‘Island of San Giorgio Maggiore’, which helps the model generate the classical Renaissance-style church.

Original: When generating without retrieval, the model is not able to generate the speciﬁc church.

Original: This indicates the importance of having relevant entities in the retrievals for WikiImages dataset and also explains why LAION database achieves the best results.

Original: We also present more examples from WikiImages in the Appendix C.

Original: 4. E NTITY F OCUSED E VALUATION ON E NTITY D RAW B ENCH
Revised: 4. E NTITY F OCUSED E VALUATION ON E NTITY D RAW B ENCH

Original: Dataset Construction
Revised: Dataset Construction

Original: We introduce EntityDrawBench to evaluate the model’s capability to generate diverse sets of entities in different visual scenes.
Revised: We introduce EntityDrawBench to evaluate the model’s capability to generate diverse sets of entities in different visual scenes.

Original: Speciﬁcally, we pick three types of visual entities (dog breeds, landmarks,  and foods) from Wikipedia Commons and Google Landmarks to [Table]
Revised: Speciﬁcally, we pick various types of visual entities (dog breeds, landmarks, foods, birds, and animated characters) from Wikipedia Commons,

Original:  construct our prompts.
Revised: Google Landmarks and Fandom to construct our prompts.

Original: In total, we collect 150 entity-centric prompts for evaluation.
Revised: In total, we collect 250 entity-centric prompts for evaluation.

Original: These prompts are mostly  unique and we cannot ﬁnd corresponding images with Google Image Search.
Revised: These prompts are mostly very unique and we cannot ﬁnd them on the Internet, let alone the model training data.

Original: More construction details are in Appendix D.
Revised: The dataset construction details are in Appendix C.

Revised: 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.

Revised: We compare our generation results with the results from DreamBooth (Ruiz et al., 2022) in Appendix H.

Original: We use the  prompt as the input and its corresponding image-text pairs as the ‘retrieval’ for ReImagen, to generate four 1024 ⇥ 1024 images.
Revised: We use the constructed prompt as the input and its corresponding image-text pairs as the ‘retrieval’ for Re-Imagen, to generate four 1024 × 1024 images.

Original: For the other models, we feed the prompts directly also to generate four images.
Revised: For the other models, we feed the prompts directly also to generate four images.

Original: Wewill pick the best image of these four samples to rate its Photorealism and Faithfulness.
Revised: We pick the best image of 4 random samples to rate its Photorealism and Faithfulness by human raters.

Original: For photorealism, we assign 1 if the image is moderately realistic without noticeable artifacts, otherwise, we assign a score of 0.
Revised: For photorealism, we rate 1 if the image is moderately realistic without noticeable artifacts.

Original: For the faithfulness measure, we assign 1 if the image is faithful to both the entity source and the text description, otherwise, we assign 0.
Revised: For the faithfulness measure, we rate 1 if the image is faithful to both the entity appearance and the text description.

Original: Experimental Results We use the proposed interleaved classiﬁer-free guidance (subsection 3.2) for the 64 ⇥ diffusion model, which runs for 256 diffusion steps under a strong guidance weight of w =30 for both text and neighbor conditions.
Revised: EntityDrawBench Results We use the proposed interleaved classiﬁer-free guidance (subsection 3.2) for the 64 × diffusion model, which runs for 256 diffusion steps under a strong guidance weight of w =30 for both text and neighbor conditions.

Original: For the 256 ⇥ and 1024 ⇥ resolution models, we use a constant guidance weight of 5.0 and 3.0, respectively, with 128 and 32 diffusion steps.
Revised: For the 256 × and 1024 × resolution models, we use a constant guidance weight of 5.0 and 3.0, respectively, with 128 and 32 diffusion steps.

Original:  The inference speed is 30-40 secs for 4 images on 4 TPU-v4 chips.
Revised: [Table] The inference speed is 30-40 secs for 4 images on 4 TPU-v4 chips.

Original: We demonstrate our human evaluation results for faithfulness and photorealism in Table 3.
Revised: We demonstrate our human evaluation results for faithfulness and photorealism in Table 2.

Original: We can observe that Re-Imagen can in general achieve much higher faithfulness than the existing models while maintaining similar photorealism scores.
Revised: We can observe that Re-Imagen can in general achieve much higher faithfulness than the existing models while maintaining similar photorealism scores.

Original: When comparing with our backbone Imagen, we see the faithfulness score jumps from 27% to 71%, which indicates that our model is paying attention to the retrieved knowledge and assimilating it into the generation process.
Revised: 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.

Original: We further partition the entities into ‘frequent’ and ‘infrequent’ categories based on their frequency (top 50% as ‘frequent’) in Imagen’s training corpus.
Revised: We further partition the entities into ‘frequent’ and ‘infrequent’ categories based on their frequency (top 50% as ‘frequent’).

Original: We plot  faithfulness score for ‘frequent’ and ‘infrequent’ separately in Figure 5 .
Revised: We plot the faithfulness score for ‘frequent’ and ‘infrequent’ separately in Figure 5.

Original: We can see that our model is less sensitive to the frequency of the input entities than the other models, with only a 10-20% drop on infrequent entities. In contrast, both Imagen and DALL-E 2 drop by 40%-50% on infrequent entities.
Revised: 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.

Original: This study reﬂects the effectiveness of text-to-image generation models on long-tail entities.
Revised: This study reﬂects the effectiveness of text-toimage generation models on long-tail entities.

Revised: More generation examples are shown in Appendix F.

Original: Comparison to Other Models We demonstrate some examples from different models in Figure 6.
Revised: Comparison to Other Models We demonstrate some examples from different models in Figure 6.

Original: As can be seen, the images generated from Re-Imagen strike a good balance between text alignment and entity ﬁdelity.
Revised: As can be seen, the images generated from Re-Imagen strike a good balance between text alignment and entity ﬁdelity.

Original: Unlike image editing to perform in-place modiﬁcation, Re-Imagen can transform the neighbor entities both geometrically and semantically according to the text guidance.
Revised: Unlike image editing to perform in-place modiﬁcation, Re-Imagen can transform the neighbor entities both geometrically and semantically according to the text guidance.

Original: 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.
Revised: 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.

Original: 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 onEquation 6, by adjusting ⌘ , i.e. the proportion of ˆ ✏ p and ˆ ✏ n in the sampling schedule, we can control Re-Imagen so as to generate images that explore this tradeoff: decreasing ⌘ will increase the entity’s image faithfulness but decrease the text alignment.
Revised: 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.

Original: In contrast, increasing the value ⌘ will increase the text alignment and decrease the similarity to retrieved image.
Revised: 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.

Original: We demonstrate this in Figure 7.
Revised: We provide visual examples in Appendix D.

Original: With small ⌘ , the model ignores the text description and simply copies the retrieved image, while with large ⌘ ,  the model reverts to standard Imagen, without using to  the input image much.
Revised: 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 η , i.e. the proportion of ˆ (cid:15) p and ˆ (cid:15) n in the sampling schedule, we can control Re-Imagen so as to generate images that explore this tradeoff: decreasing η will increase the entity’s entity faithfulness but decrease the text alignment.

Original: We found that having ⌘ around 0.5 is usually a ‘sweet spot’ that balances both conditions.
Revised: We found that having η around 0.5 is usually a ‘sweet spot’ that balances both conditions.

Original:  C ONCLUSIONSAND L IMITATIONS
Revised:  C ONCLUSIONS

Original: We present Re-Imagen, a retrieval-augmented diffusion model, and demonstrate its effectiveness in generating realistic and faithful images.
Revised: We present Re-Imagen, a retrieval-augmented diffusion model, and demonstrate its effectiveness in generating realistic and faithful images.

Original: We exhibit such advantages not only through automatic FID measures on standard benchmarks ( i.e. , COCO and WikiImage) but also through human evaluation on the newly introduced EntityDrawBench.
Revised: We exhibit such advantages not only through automatic FID measures on standard benchmarks ( i.e. , COCO and WikiImage) but also through human evaluation of the newly introduced EntityDrawBench.

Original: We further demonstrate that our model is particularly effective in generating an image from text that mentions rare entities.
Revised: We further demonstrate that our model is particularly effective in generating an image from text that mentions rare entities.

Original: Re-Imagen still suffers from well-known issues in text-to-image generation, which we review below in section 5.
Revised: Re-Imagen still suffers from well-known issues in text-to-image generation, which we review below in section 5.

Original: In addition, Re-Imagen also has some unique limitations due to the retrieval-augmented modeling.
Revised: In addition, Re-Imagen also has some unique limitations due to the retrieval-augmented modeling.

Original: First, because Re-Imagen is sensitivethe to retrieved image-text pairs it is conditioned on, when the retrieved image is of low-quality, there will be a negative inﬂuence on the generated image.
Revised: First, because Re-Imagen is sensitive to retrieved image-text pairs it is conditioned on when the retrieved image is of low quality, there will be a negative inﬂuence on the generated image.

Original: Second, Re-Imagen sometimes still fail to ground on the retrieved entities when the entity’s visual appearance is out of the generation space.
Revised: Second, Re-Imagen sometimes still fails to generate high-quality images with highly compositional prompts, where multiple entities are involved.

Original: Third, we noticed that the super-resolution model is less effective, and frequently misses low-leveltexture details of the visual entities.
Revised: Thirdly, the super-resolution model is still not competent at capturing low-level details of retrieved entities leading to visual distortion.

Original: In future work, we plan to further investigate the above limitations and address them.
Revised: In future work, we plan to further investigate the above limitations and address them.

Original: E THICS S TATEMENT
Revised: E THICS S TATEMENT

Original: Strong text-to-image generation models, i.e. , Imagen (Saharia et al., 2022) and Parti (Yu et al., 2022), raise ethical challenges along dimensions such as the social bias .
Revised: Strong text-to-image generation models, i.e. , Imagen (Saharia et al., 2022) and Parti (Yu et al., 2022), raise ethical challenges along dimensions such as the social bias .

Original: Re-Imagen is exposed to the same challenges, as we employed Web-scale datasets that are similar to the prior models.
Revised: Re-Imagen is exposed to the same challenges, as we employed Web-scale datasets that are similar to the prior models.

Original: The retrieval-augmented modeling techniques of Re-Imagen has substantially improved the controllability and attribution of the generated image.
Revised: The retrieval-augmented modeling techniques of Re-Imagen have substantially improved the controllability and attribution of the generated image.

Original: Like many basic research topics, this additional control could be used for beneﬁcial or harmful purposes.
Revised: Like many basic research topics, this additional control could be used for beneﬁcial or harmful purposes.

Original: One obvious danger is that Re-Imagen (or similar models) could be used for malicious purposes like spreading misinformation, e.g., by producing realistic images of speciﬁc people in misleading visual contexts.
Revised: One obvious danger is that Re-Imagen (or similar models) could be used for malicious purposes like spreading misinformation, e.g., by producing realistic images of speciﬁc people in misleading visual contexts.

Original: On the other side, additional control has many potential beneﬁts.
Revised: On the other side, additional control has many potential beneﬁts.

Original: One general beneﬁt is that Re-Imagen can reduce hallucination and increase the faithfulness of the generated image to the users’ intent.
Revised: One general beneﬁt is that Re-Imagen can reduce hallucination and increase the faithfulness of the generated image to the user’s intent.

Original: Another beneﬁt is that the ability to work with tail entities makes the model more useful for minorities and other users in smaller communities: for example, Re-Imagen is more effective at generating images of landmarks famous in smaller communities or cultures, and generating images of indigenous foods and cultural artifacts.
Revised: Another beneﬁt is that the ability to work with tail entities makes the model more useful for minorities and other users in smaller communities: for example, Re-Imagen is more effective at generating images of landmarks famous in smaller communities or cultures and generating images of indigenous foods and cultural artifacts.

Original: We argue that this model can help decrease the frequency-caused bias in current neural network based AI systems.
Revised: We argue that this model can help decrease the frequency-caused bias in current neural network-based AI systems.

Original: Considering such potential threats to the public, we have currently decide not to release the code or a public demo.
Revised: Considering such potential threats to the public, we will be cautious about code and API release.

Original: In future work, we will explore a framework for responsible use that balances the value of external auditing of research with the risks of unrestricted open access, allowing this work to be used in a safe and beneﬁcial way. 
Revised: In future work, we will explore a framework for responsible use that balances the value of external auditing of research with the risks of unrestricted open access, allowing this work to be used in a safe and beneﬁcial way. A E XTENDED LITERATURE REVIEW

Revised: As aforementioned in the main text, in this section, we provide an additional review of related works, on (1) Retrieval-augmented Generative Models; and (2) Text-guided Image Editing.

Revised: Retrieval-Augmented Generative Models Knowledge grounding has also drawn signiﬁcant attention in the natural language processing (NLP) community.

Revised: Different semi-parametric models like KNN-LM (Khandelwal et al., 2019), RAG (Lewis et al., 2020), REALM (Guu et al., 2020),

Revised: RETRO (Borgeaud et al., 2021) have been proposed to leverage external textual knowledge into the transformer language models.

Revised: These models have demonstrated great advantages in increasing the language model’s faithfulness and reducing the computation/memory cost.

Revised: Such attempts have also been made in visual tasks like image recognition (Long et al., 2022), 2-D scene reconstruction (Siddiqui et al., 2021), and image inpainting (Xu et al., 2021).

Revised: Our proposed method follows the same theme to incorporate visual knowledge into a pre-trained text-to-image generation model to help the model generalize to long-tail entities or even unseen entities without scaling up the parameters.

Revised: Text-Guided Image

Revised: Editing The work of text-guided image editing aims at preserving the object’s appearance while changing certain contexts in the image.

Revised: Previously, GANs (Goodfellow et al., 2014) have been used to achieve signiﬁcant performance on image editing (Zhu et al., 2016;

Revised: Abdal et al., 2019;

Revised: Zhu et al., 2020; Roich et al., 2021;

Revised: Tov et al., 2021;

Revised: Wang et al., 2022; Alaluf et al., 2022).

Revised: The problem is also known as inversion as it normally requires ﬁnding the initial noise vector added in the generation process.

Revised: More recently, Prompt-to-Prompt (Hertz et al., 2022) propose to use pre-trained text-image models for image editing.

Revised: Image editing is focused on performing in-place modiﬁcations to the input image, either changing the global styles or editing a local region speciﬁcally without modifying the object’s appearance.

Revised: However, we treat the retrieved image as ‘knowledge’ and ground on it to synthesize new images.

Revised: Thus, we are not restricted to in-place modiﬁcations and are able to perform more sophisticated transformations over the objects.

Revised: B W IKI I MAGES D ATASET

Revised: The WikiImages dataset is taken from WebQA (Chang et al., 2022).

Revised: Images were crawled from

Revised: Wikimedia Commons via the Bing Visual Search API.

Revised: Since lots of Wikimedia’s topics are not visually interesting, the authors seeded with natural scenes and gradually reﬁne the search pool to obtain more interesting images.

Revised: The images are mostly containing entities from Wikipedia or WikiData.

Revised: However, the original dataset still contains heavy noises.

Revised: Therefore, we apply further ﬁltering to obtain the more plausible ones for image generation.

Revised: Speciﬁcally, we remove all the image-text pairs with text lengths larger than 15 tokens and all the text with a date or wiki-id information.

Revised: C E NTITY D RAW B ENCH

Revised: For dog breeds and birds, we sample 50 from Wikipedia Commons 3 as our candidates.

Revised: For landmarks, we sample 50 from Google Landmarks (Weyand et al., 2020) as our candidate.

Revised: For foods, we sample 50 from Wikipedia 4 as our candidates.

Revised: For ﬁlm characters, we collected 50 images from

Revised: Starwars from Fandom 5 .

Revised: We use appropriately paired source images as the retrieved ‘knowledge’.

Revised: For each entity category, we write 5 prompt templates with an entity name placeholder, which describes the entity in different scenes.

Revised: Each entity will sample a template and replace the placeholder with the entity’s name to generate a prompt, which is used as input to the text-to-image generation model.

Revised: 3 https://commons.wikimedia.org/wiki/List_of_dog_breedshttps://en.wikipedia.org/wiki/List_of_cuisineshttps://starwars.fandom.com/wiki/Main_Page

Revised: D I MPACT OF R ETRIEVAL N UMBER K

Revised: We change the retrieval number K from 1 to 2 to see its impact on the model output.

Revised: We show some examples in Figure 11 to demonstrate the advantage of having multiple retrievals to help the model better capture the visual appearance of the given entities.

Revised: E C LASSIFIER GUIDANCE - FREE SAMPLING STRATEGY

Revised: We demonstrate different types of sampling strategy to leverage two conditions: standard joint condition guidance sampling, weighted guidance sampling, and our proposed interleaved guidance sampling.

Revised: The standard joint condition guidance only considers the joint diffusion score (cid:15) ( x t , c n , c p ) to meet both conditions.

Revised: In contrast, weighted guidance sampling uses the weighted sum of text-enhanced epsilon ˆ (cid:15) p and neighbor-enhanced epsilon ˆ (cid:15) n .

Revised: Our interleaved classiﬁer guidance switches between ˆ (cid:15) p and ˆ (cid:15) n , with a ratio of η : 1 − η .

Revised: We plot their conceptual difference in Figure 12.

Revised: Essentially, (cid:15) n and (cid:15) p do not have dependency in weighted sampling, however, they are dependent in interleaved sampling.

Revised: In an extreme case where (cid:15) n and (cid:15) p are contradictory to each other, the model will get stuck in a local region.

Revised: In contrast, Interleaved sampling can alleviate this issue.

Revised: We compare 20 dog images generated from these three sampling strategies in EntityDrawBench.

Revised: We vary the number of diffusion step to observe their human evaluation score curve and show case some generated outputs in Figure 13.

Revised: As can be seen, the joint decoding is either dominated by the retrieval image or by the text prompt.

Revised: Weighted and Interleave can help balance the two conditions to generate better images.

Revised: We also found that with less sampling steps K=200, “weighted” sampling actually achieves better results than “interleaved” sampling.

Revised: However, as the sampling steps increase, our proposed “interleaved” sampling achieves better human evaluation score.

Revised: F G ENERATION E XAMPLES

Revised: We provide more generation examples in Figure 14 and Figure 15.

Revised: [Figure]

Revised: G I MAGINARY E XAMPLES

Revised: We provide generation results for imaginary scene in Figure 16.

Revised: H C OMPARISON WITH D REAM B OOTH

Revised: We also add comparison to DreamBooth (Ruiz et al., 2022).

Revised: We adopt almost the same input images from DreamBooth and display our generation results in Figure 17, Figure 18 and Figure 19.

Revised: I F AILURE E XAMPLES

Revised: We found that Re-Imagen can also fail in a lot of cases.

Revised: We demonstrate a few examples in Figure 20.

Revised: As can be seen, the model sometimes has a few failure modes: (1) the text input prior is too strong like ‘Zoom’ will be interpreted as a ‘Zoom-in’ picture by the model.

Revised: (2) the model cannot ground the retrieval text on the retrieval image, for example, the model believes that only the ‘beef tenderloin inside the bowl’ is ‘Escudella’ rather than the whole stew, therefore generating ‘beef tenderloin on the grass’.

Revised: (3) the model can sometimes mess up two conditions, for example, the reference ‘Australian Pinscher’ and the ‘rabbit’ in the prompt gets mixed into a single object.