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

Wenhu Chen; Hexiang Hu; Chitwan Saharia; William W. Cohen

Source-linked AI summary

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

Wenhu Chen, Hexiang Hu, Chitwan Saharia, William W. Cohen

arXiv:2209.14491v3Sep 29, 2022cs.CVcs.AIcs.LG

TL;DR

Text-to-image models often struggle with uncommon entities despite strong performance on common ones. Re-Imagen retrieves multimodal references and integrates them into diffusion generation, achieving improved benchmark FID and higher entity faithfulness, especially for less frequent entities.

Problem
Existing text-to-image models can generate high-quality common entities but often struggle to reproduce uncommon or rare entities faithfully.
Method
Re-Imagen retrieves relevant image-text pairs from an external multimodal database and conditions a cascaded diffusion model on them alongside the text prompt.
Results
Re-Imagen improves FID on COCO and WikiImages and raises faithfulness by around 40% over Imagen on EntityDrawBench.
Takeaways & Limitations
Retrieval augmentation makes Re-Imagen particularly effective at generating realistic and faithful images for rare entities.
Takeaways & Limitations
Re-Imagen can be negatively affected by low-quality retrieved pairs, may fail on highly compositional prompts, and can distort low-level retrieved details.

Abstract

from arXiv · show

Research on text-to-image generation has witnessed significant progress in generating diverse and photo-realistic images, driven by diffusion and auto-regressive models trained on large-scale image-text data. Though state-of-the-art models can generate high-quality images of common entities, they often have difficulty generating images of uncommon entities, such as `Chortai (dog)' or `Picarones (food)'. 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-fidelity and faithful images, even for rare or unseen entities. 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. 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. 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. Furthermore, we develop a new sampling strategy to interleave the classifier-free guidance for text and retrieval conditions to balance the text and retrieval alignment. Re-Imagen achieves significant gain on FID score over COCO and WikiImage. 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. Human evaluation on EntityDrawBench shows that Re-Imagen can significantly improve the fidelity of generated images, especially on less frequent entities.

1 INTRODUCTION

Existing text-to-image models generate realistic images but are sensitive to entity frequency, motivating Re-Imagen’s retrieval-based approach for improving rare-entity fidelity. Re-Imagen retrieves multimodal references, trains on retrieval-grounded examples, and reports gains on standard and entity-focused evaluations.

State-of-the-art text-to-image models produce realistic, diverse images but their outputs are sensitive to entity frequency.
Re-Imagen retrieves relevant image-text pairs from an external multimodal knowledge base instead of relying on memorized appearances of rare entities.The references provide semantic attributes and concrete visual appearances to guide synthesis.
Re-Imagen is trained on retrieval-associated examples and uses interleaved guidance to balance text alignment with entity fidelity.Its cascaded diffusion backbone progressively generates 1024×1024 images.
Re-Imagen achieves state-of-the-art FID performance among non-fine-tuned models on COCO and WikiImages.On COCO, the gain is associated with retrieved in-domain image styles; on WikiImages, it is associated with grounding on retrieved entity information.
EntityDrawBench evaluates entities across categories and frequencies, while Figure 1 qualitatively shows better faithfulness than Imagen on less frequent entities.

2 RELATED WORK

Related work situates Re-Imagen within diffusion-based text-to-image generation and concurrent retrieval-augmented approaches. These neighboring methods differ in how retrieval examples are selected and used.

Diffusion models have achieved broad success in text-to-image generation, surpassing GANs and autoregressive Transformers in photorealism and diversity under similar model sizes.
Concurrent retrieval-augmented methods also use retrieved examples or neighbors to improve diffusion-based generation.
RDM selects neighbors using image features and replaces retrievals with user-chosen exemplars at inference, transferring artistic style to generated images.

3 MODEL

Re-Imagen extends cascaded diffusion with retrieved multimodal neighbors alongside text conditioning. Its pipeline encodes these references, incorporates them through attention, and balances text and neighbor guidance during denoising.

3.1 PRELIMINARIES: Diffusion models reverse a forward Markov process that gradually adds Gaussian noise, learning to denoise noisy images toward the data distribution.At inference, DDPM sampling starts from Gaussian noise and recursively produces the final image.
3.1 PRELIMINARIES: Cascaded diffusion uses separate 64×, 256×, and 1024× models to progressively generate 1024 × 1024 images.The stages move from low resolution to high resolution for generation efficiency.
3.2 GENERATING IMAGE WITH MULTI-MODAL KNOWLEDGE: Re-Imagen conditions denoising on the text prompt and top-k image-text neighbors retrieved from an external knowledge base.The retrieved references supply semantic attributes and concrete visual appearance for the mentioned entities.
3.2 GENERATING IMAGE WITH MULTI-MODAL KNOWLEDGE: The DStack encodes the noisy image, prompt, and retrieved pairs; cross-attention augments the image representation before UStack predicts noise.Retrieved pairs are encoded with the shared DStack using t set to zero, producing feature maps for attention.
3.2 GENERATING IMAGE WITH MULTI-MODAL KNOWLEDGE: Training constructs KNN-ImageText from Imagen’s 50M ImageText dataset, associating each instance with text-similar nearest neighbors as retrieval references.The dataset is selected for entity-similar photos and sanitization, while neighbors are found using text-to-text search.
3.2 GENERATING IMAGE WITH MULTI-MODAL KNOWLEDGE: Interleaved classifier-free guidance alternates text-enhanced and neighbor-enhanced noise predictions to balance text alignment and retrieved-image fidelity.Standard unconditional guidance can make outputs dominated by either condition; a predefined ratio η controls the alternation.

4 EXPERIMENTS

Experiments evaluate Re-Imagen on general text-to-image benchmarks and the rare-entity EntityDrawBench, then analyze retrieval count and the text–entity faithfulness trade-off. Retrieval improves FID and entity faithfulness, particularly for less-frequent entities, while sampling controls the balance between text alignment and retrieved-entity fidelity.

Evaluation setup: Re-Imagen is evaluated with automatic FID measurements on COCO and WikiImages and human ratings on EntityDrawBench for long-tail entities.The model uses separate benchmark settings to measure general photorealistic generation and rare-entity capability.
COCO and WikiImages: Roughly 2.0 absolute FID improvement over Imagen is achieved on COCO through external-database retrieval.Retrieval from an out-of-domain database still improves FID by 0.4 over Imagen, while in-domain retrieval performs better on COCO.
COCO and WikiImages: On WikiImages, the out-of-domain LAION database outperforms the in-domain database because entity-focused images make relevant entity retrieval more important.The larger LAION-400M database also provides a higher chance of retrieving related entities.
EntityDrawBench: Human evaluation shows much higher faithfulness than existing models while maintaining similar photorealism, with faithfulness improving around 40% over Imagen.The result indicates that retrieved knowledge is incorporated into generation.
EntityDrawBench: Re-Imagen is less sensitive to entity frequency than other models, showing only minor faithfulness drops for infrequent entities.EntityDrawBench partitions entities into frequent and infrequent groups using the top 50% frequency threshold.
Analysis: Increasing retrieval count from 1 to 4 has little FID effect on COCO and WikiImages but dramatically improves EntityDrawBench faithfulness.Multiple images are especially important for grounding visual entities.
Analysis: The sampling parameter η exposes a trade-off between text-prompt alignment and retrieved-entity faithfulness, with η around 0.5 usually balancing both.Lower η increases entity faithfulness but decreases text alignment.

5 CONCLUSIONS

Re-Imagen is a retrieval-augmented diffusion model designed to generate realistic, faithful images, with particular effectiveness for prompts mentioning rare entities. Its retrieval-based modeling also introduces several limitations.

Re-Imagen demonstrates effectiveness in generating realistic and faithful images through automatic FID measures and human evaluation.The evaluations use COCO, WikiImage, and the newly introduced EntityDrawBench.
Re-Imagen is particularly effective at generating images from text prompts that mention rare entities.
Low-quality retrieved image-text pairs can negatively influence the generated image.The model is sensitive to the quality of the references on which it is conditioned.
Re-Imagen sometimes fails to generate high-quality images for highly compositional prompts involving multiple entities.
Its super-resolution model may fail to capture low-level details of retrieved entities, causing visual distortion.

ETHICS STATEMENT

Re-Imagen shares the ethical challenges of web-scale text-to-image models while offering additional control that may improve faithfulness and support tail entities. The authors also identify risks of misuse and plan cautious release practices.

Re-Imagen is exposed to social-bias challenges because it uses web-scale datasets similar to those used by prior text-to-image models.
Its additional control could be used to produce realistic images of specific people in misleading contexts for misinformation.
The model can reduce hallucination and increase the faithfulness of generated images to user intent.
Support for tail entities may benefit minorities and users in smaller communities by improving generation of culturally specific landmarks, foods, and artifacts.
The authors plan to balance external auditing with the risks of unrestricted access through a responsible-use framework and cautious code and API release.
Related work includes retrieval-augmented generative models and text-guided image editing.

B WIKIIMAGES DATASET

WikiImages was derived from WebQA using Wikimedia Commons images collected through Bing Visual Search and further filtered because the original dataset contained substantial noise.

WikiImages is taken from the WebQA dataset and contains images crawled from Wikimedia Commons through the Bing Visual Search API.
The collection process seeded searches with natural scenes and gradually refined the search pool to obtain more visually interesting images.
Most WikiImages images contain entities from Wikipedia or Wikidata, but the original dataset contained heavy noise and was further filtered.

C ENTITYDRAWBENCH

EntityDrawBench evaluates entity generation across several categories using sampled source images and prompt templates. Its construction pairs candidate entities with retrieved knowledge and varied textual descriptions.

EntityDrawBench construction: EntityDrawBench samples 50 candidates for each category, including dog breeds, birds, landmarks, foods, and film characters.
EntityDrawBench construction: The benchmark uses appropriately paired source images as retrieved knowledge for each entity category.
EntityDrawBench construction: For every entity category, the benchmark uses five prompt templates containing an entity-name placeholder.
EntityDrawBench construction: The construction process lists entity names, finds source images from Wikimedia, and generates related prompts.
EntityDrawBench construction: Figure 10 lists the prompt templates for the different entity categories.

D IMPACT OF RETRIEVAL NUMBER K

The section examines how retrieval count and classifier-free guidance affect Re-Imagen’s generation. Multiple retrievals help capture entity appearance, while interleaved guidance better balances text and retrieval as sampling steps increase.

D IMPACT OF RETRIEVAL NUMBER K: Increasing retrievals from K=1 to K=2 helps the model better capture the visual appearance of entities.Figure 11 provides generation examples illustrating the benefit of multiple retrievals.
D IMPACT OF RETRIEVAL NUMBER K: The study compares joint-condition, weighted, and interleaved guidance sampling for combining text and retrieval conditions.Weighted guidance sums text- and neighbor-enhanced scores, whereas interleaved guidance switches between them at a specified ratio.
D IMPACT OF RETRIEVAL NUMBER K: With 200 sampling steps, weighted sampling scores better than interleaved sampling, but interleaved sampling scores better as sampling steps increase.Both weighted and interleaved guidance balance the text and retrieval conditions better than joint decoding in the reported dog-image comparison.

F GENERATION EXAMPLES

This section presents additional Re-Imagen generation examples, including EntityDrawBench outputs and imaginary scenes. The cited material identifies the figures but does not state outcome comparisons.

F GENERATION EXAMPLES: Figures 14 and 15 provide additional uncherry-picked EntityDrawBench examples for different models.The figures are presented as extra generation examples rather than a reported quantitative comparison.
F GENERATION EXAMPLES: The section’s cited passages document example figures without specifying their visual outcomes.
F GENERATION EXAMPLES: Figure 16 provides generation results for an imaginary scene.

H COMPARISON WITH DREAMBOOTH

The section compares Re-Imagen with DreamBooth using similar input images and presents additional generation examples. It also documents failure cases involving prompt interpretation, retrieval grounding, and condition mixing.

H COMPARISON WITH DREAMBOOTH: Re-Imagen is compared with DreamBooth using almost the same input images as in the DreamBooth study.The resulting generations are displayed in Figures 17, 18, and 19.
H COMPARISON WITH DREAMBOOTH: Figures 17–19 display Re-Imagen generation results for the comparison with DreamBooth.
H COMPARISON WITH DREAMBOOTH: Reported failure modes include overly strong text priors, incorrect grounding of retrieval text on images, and mixing two conditions into one object.Examples include interpreting “Zoom” as zoom-in imagery, misreading Escudella, and merging an Australian Pinscher with a rabbit.