landmark

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

Wenhu Chen, Hexiang Hu, Chitwan Saharia, William W. Cohen · Google · 2022-09 · arXiv:2209.14491

Why it matters

Retrieves image-text pairs from an external knowledge base so the model can render rare and long-tail entities it has not memorised, and introduces a benchmark for exactly that failure. The paper frames this as an entity-memorisation problem rather than as a general theory that generation failures are knowledge failures; the atlas files it as the origin of the retrieval bet, which is an extrapolation of its motivation rather than its own stated thesis.

What this paper does

read: full text

Before this

Text-to-image diffusion and autoregressive models trained on large-scale image-text data (this paper compares against Imagen and DALL-E 2) could generate high-quality images of common entities, but the paper's own analysis found output quality was 'surprisingly sensitive to the frequency of the entities' in training data, so uncommon entities such as a specific dog breed or a regional food were rendered poorly or unfaithfully because the model had not seen enough examples to memorize their appearance.

The problem

How to generate faithful images of entities that are rare or effectively unseen in the training distribution, without requiring the model to have memorized their visual appearance from web-scale pretraining alone.

The idea

Give a text-to-image diffusion model access to an external multimodal knowledge base at generation time: retrieve image-text pairs relevant to the prompt and feed them to the model as additional conditioning, so the model can copy visual and semantic detail rather than recall it purely from parametric memory.

How it works

Re-Imagen fine-tunes Imagen's cascaded diffusion model (a 2.5B base model plus 750M and 400M super-resolution stages). A KNN-ImageText training set is built by retrieving each training pair's top-2 nearest neighbors via BM25 text similarity over a 50M-pair corpus, producing (image, text, retrieval) triples. Retrieved image-text pairs are encoded with a shared DStack encoder and fused into the denoising network via a cross-attention module that lets the diffusion model attend to neighbor features. Text and neighbor conditioning are each randomly dropped 10% of the time during training to support classifier-free guidance, and the paper introduces an interleaved sampling strategy that alternates, via a randomly sampled ratio, between a text-emphasizing and a neighbor-emphasizing noise prediction to balance prompt fidelity against retrieved-reference fidelity.

Evidence

COCO 30K FID: 5.25 in-domain, 6.88 zero-shot (retrieving from LAION) versus Imagen's 7.27 zero-shot baseline. WikiImages FID: 5.88 in-domain, 5.80 out-of-domain, versus Imagen's 6.44. On the paper's own EntityDrawBench benchmark, human-evaluated faithfulness averages 0.80 for Re-Imagen versus 0.58 for DALL-E 2 and 0.28 for Imagen, roughly a 30% improvement over the best baseline, with the largest gains concentrated on less-frequent entities.

Limitations

The paper concedes that low-quality retrieved images degrade generation quality, that highly compositional prompts involving multiple entities remain difficult, that the super-resolution stages fail to preserve low-level detail from retrieved entities and introduce visual distortion, and that the model inherits hallucination and bias issues from the underlying web-scale training data.

Why it matters

It demonstrates, with a working system and a purpose-built benchmark, that at least one class of generation failure (rare-entity fidelity) is better addressed by giving the model external reference material than by scaling the model or its training data. The atlas frames this as the origin of the more general bet that some generation failures are knowledge failures fixable by retrieval; the paper's own framing is narrower than that generalization, centered specifically on entity frequency and memorization rather than a stated general theory of knowledge failures, so the broader framing is a reasonable extrapolation from this paper's motivation rather than a claim the paper itself makes explicitly.

Abstract, in the authors' own words

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.

Research line

Agentic and search-augmented generationemerging

Method note — the shared flow-matching interpolation

Every flow-matching / rectified-flow paper in this atlas trains toward a straight-line path between a noise sample x₀ and a data sample x₁:

x_t = (1-t)\,x_0 + t\,x_1, \qquad v_\theta(x_t, t) \approx x_1 - x_0

Builds on

Nothing recorded yet.

Built on by

Nothing recorded yet.