PaperlayerView exact revision on arXiv

Source-linked AI summary

KNN-Diffusion: Image Generation via Large-Scale Retrieval

Shelly Sheynin, Oron Ashual, Adam Polyak, Uriel Singer, Oran Gafni, Eliya Nachmani, Yaniv Taigman

arXiv:2204.02849v2cs.CVcs.AIcs.CLcs.GRcs.LG

TL;DR

Text-to-image models typically require large paired datasets, limiting training in new domains with scarce or unlabeled data. kNN-Diffusion uses retrieval-conditioned diffusion trained without text, and reports state-of-the-art results while enabling out-of-distribution generation and mask-free semantic manipulation.

  • Problem

    Large-scale text-to-image models rely on paired text-image datasets, limiting application to new or customized domains with only unannotated images.

  • Method

    kNN-Diffusion conditions diffusion on CLIP embeddings and retrieved nearest-neighbor image embeddings, using images during training and text at inference.

  • Results

    The method achieves state-of-the-art human-evaluation and FID results across datasets and diffusion backbones, with 400 million parameters and 7-second inference.

  • Takeaways & Limitations

    Retrieval supplies conditioning information that supports textless training, out-of-distribution generation, and mask-free local semantic manipulation while preserving object identity.

Abstract

from arXiv · show

Recent text-to-image models have achieved impressive results. However, since they require large-scale datasets of text-image pairs, it is impractical to train them on new domains where data is scarce or not labeled. In this work, we propose using large-scale retrieval methods, in particular, efficient k-Nearest-Neighbors (kNN), which offers novel capabilities: (1) training a substantially small and efficient text-to-image diffusion model without any text, (2) generating out-of-distribution images by simply swapping the retrieval database at inference time, and (3) performing text-driven local semantic manipulations while preserving object identity. To demonstrate the robustness of our method, we apply our kNN approach on two state-of-the-art diffusion backbones, and show results on several different datasets. As evaluated by human studies and automatic metrics, our method achieves state-of-the-art results compared to existing approaches that train text-to-image generation models using images only (without paired text data)

1 INTRODUCTION

Existing text-to-image models require large paired datasets and substantial computation, limiting use in unlabeled or customized domains. kNN-Diffusion addresses these constraints through retrieval-conditioned diffusion, enabling textless training, distribution changes, and mask-free semantic manipulation.

  • Large-scale text-to-image models depend on paired text-image data, which limits application to new domains containing only unannotated images.
  • kNN-Diffusion conditions diffusion on CLIP embeddings and k nearest-neighbor image embeddings, allowing training without text data.At training time it uses image embeddings; at inference it uses text embeddings and retrieved neighbors.
  • Retrieval embeddings bridge training and inference distributions, support target-distribution generation and out-of-distribution samples, and reduce the information the model must store.
  • Across two diffusion backbones and multiple datasets, the method achieves state-of-the-art human-evaluation and FID results with 400 million parameters and 7-second inference.
  • The method introduces mask-free local semantic manipulation that automatically locates edits while preserving object identity and aligning changes with target text.
  • The paper evaluates textless training, retrieval-database substitution for out-of-distribution generation, and mask-free manipulation across discrete and continuous diffusion backbones.

2 RELATED WORK

Prior text-to-image systems largely rely on supervised paired data, while related textless approaches face optimization costs or distribution gaps. The paper builds on multimodal representations and retrieval to use visual data during training and text at inference.

  • Text-to-image models: Earlier text-to-image systems use GANs, VQ-VAE Transformers, or diffusion models but are generally trained with paired text-image datasets.
  • Text-to-image models: Textless alternatives include per-image CLIP optimization and models trained with perturbed CLIP image embeddings, but they remain time-consuming or require annotated auxiliary data.
  • Multi-modal feature learning: The paper identifies a text-image embedding gap: CLIP image embeddings alone do not accurately bridge the distributions, motivating retrieval-based conditioning.
  • Retrieval for generation: Retrieval methods efficiently select nearest neighbors from large databases using precomputed embeddings, providing a compact information-access mechanism.
  • Multi-modal feature learning: CLIP creates an aligned vision-language feature space that supports image generation and manipulation, and this paper uses it for textless training and retrieval.

3 METHOD

kNN-Diffusion trains an image-generation model with image-only data by combining shared text-image embeddings with nearest-neighbor retrieval. The framework also supports retrieval-conditioned generation across diffusion backbones and local semantic manipulation without user-provided masks.

  • Core framework: The method maps text and images into a shared CLIP space, enabling image-conditioned training and text-conditioned inference.Training uses image embeddings, while inference uses text embeddings together with retrieved image embeddings.
  • Core framework: A non-trainable image index supplies nearest-neighbor embeddings that condition the trainable diffusion-based image-generation network.The retrieval model contains pretrained text and image encoders plus an indexed image-embedding database.
  • Retrieval conditioning: During training, retrieval uses the input image embedding; during inference, it uses the text embedding to retrieve conditioning images from the same index.The model conditions on the input embedding and its k nearest neighbors in both stages, with the query modality changing from image to text.
  • Diffusion backbones: The approach is applied to discrete and continuous diffusion backbones, including a conditional Transformer for the discrete model.The discrete process tokenizes images with VQGAN and models reverse transitions with a conditional Transformer.
  • Text-only image manipulation: For local semantic manipulation, training replaces a random image region with an aligned nearest neighbor and trains restoration from the manipulated image plus a local CLIP embedding.At inference, the desired local modification is represented by a text embedding rather than a user-provided mask.

4 EXPERIMENTS

The experiments evaluate kNN-Diffusion across datasets, diffusion backbones, image manipulation, out-of-distribution generation, and ablations. Results indicate strong image quality, text alignment, controllable retrieval-based generation, and benefits from kNN conditioning.

  • Evaluation setup: The evaluation covers zero-shot generation, image-only sticker generation, text-only manipulation, out-of-distribution generation, and ablation studies.Experiments use objective FID and human preferences for image quality and text alignment.
  • Photo-realistic generation: kNN-Diffusion achieves the lowest FID across MS-COCO, CUB, and LN-COCO scenarios and receives the highest human ratings for image quality and text alignment.The comparison includes LAFITE, FuseDream, and a no-kNN baseline.
  • Sticker generation: On the stickers dataset, kNN-Diffusion achieves the lowest FID and outperforms baselines in human evaluations with both continuous and discrete diffusion backbones.It is significantly more favorable than the corresponding model trained without kNN for text alignment and image quality.
  • Image manipulation: The model performs local text-driven edits while preserving object identity, whereas comparison methods either lack text correspondence, alter identity, or struggle with structural changes.The method does not require user-provided masks.
  • Out-of-distribution generation: Swapping retrieval databases enables out-of-distribution generation without fine-tuning, including sticker-domain, aesthetic-filtered, and online-search settings.Filtering retrieved images by higher aesthetics scores yields generated images with higher mean aesthetics scores.
  • Ablation study: The ablation study evaluates retrieval index size and kNN conditioning approaches, while nearest-neighbor experiments compare inference settings using 1, 5, 10, 20, 100, and 1000 neighbors.The supplied passage states that the best performance is achieved using 10 neighbors.

5 CONCLUSION

The paper concludes that large-scale retrieval can train a text-free text-to-image model while shifting some burden of learning novel concepts to an external knowledge base. It also supports adaptation to test-time samples and text-driven semantic manipulation without masks, with favorable quality and text-alignment evaluations.

  • Conclusion: Large-scale retrieval enables training a text-to-image model without text data while reducing the model’s burden of learning novel concepts.The conclusion links this to using a relatively small model.
  • Conclusion: An external knowledge base allows the model to adapt to new samples observed during testing.The paper presents this as a capability enabled by retrieval.
  • Conclusion: Retrieval also supports text-driven semantic manipulations without user-provided masks.The conclusion presents this as a separate technique enabled by the retrieval method.
  • Conclusion: Human studies and automatic metrics find the method preferable to baselines in image quality and text alignment.The conclusion states this comparison without reporting a numerical value.

6 APPENDIX

The appendix describes the human-evaluation protocol and the large-scale datasets used to train and assess the model. Human preferences are collected through repeated ratings, while the modified PMD dataset aggregates multiple public text-image sources.

  • Human evaluation: Human evaluations use 600 samples per experiment, with each sample scored by five Mechanical Turk participants and decided by majority opinion.Comparisons ask separately about image quality and text matching.
  • Dataset construction: The modified PMD dataset combines publicly available sources including SBU Captions, Localized Narratives, Conceptual Captions, Visual Genome, and other web-scale datasets.The passage states that the combined dataset contains 69 million text-image examples.

6.3 ABLATION STUDY

The ablation study examines how retrieval index size and kNN conditioning affect performance. Increasing the inference index size improves FID, while multiple conditioning designs are compared.

  • Index size: FID improves as the inference index grows from 10% to 70% of PMD images, with scores of 13.92, 13.85, 13.72, and 13.65.The values correspond respectively to 10%, 30%, 50%, and 70% index sizes.
  • kNN conditioning: The study compares self-attention, cross-attention, and concatenation-based approaches for conditioning on kNN embeddings.Each approach combines the image embedding with retrieved-neighbor information differently.

6.4 RETRIEVAL MODEL

The retrieval model uses FAISS to index CLIP image embeddings and efficiently retrieve nearest neighbors from large image databases. An inverted file index restricts searches to relevant Voronoi clusters for scalability.

  • FAISS stores image indices with corresponding CLIP image embeddings and retrieves nearest neighbors in milliseconds.The database can store billions of elements, enabling large-scale retrieval during training.
  • The photo-realistic and sticker datasets contain 69M and 400M samples, respectively.
  • The inverted file index partitions the 512-dimensional CLIP embedding space into Voronoi cells for efficient search.Queries retrieve neighbors only from the cluster containing the query and 19 additional clusters.

6.5 DISCRETE KNN MODEL

The discrete kNN model combines diffusion over VQ-VAE image tokens with CLIP-based image or text conditioning and retrieved neighbors. It supports text-conditioned generation, classifier-free guidance, and strong qualitative alignment without COCO training or retrieval data.

  • Image Tokenization: Images are represented as discrete tokens using a VQ-VAE-based tokenizer with a 32×32 latent grid and a 2887-element codebook.The adapted VQGAN encoder downsamples images to 32×32 tokens.
  • Training and Inference: During training, image embeddings and their kNN embeddings condition diffusion, while inference uses text embeddings and retrieved image embeddings.The training algorithm retrieves neighbors for each image, concatenates the conditioning embeddings, and denoises sampled tokens.
  • Transformer: The model is a 400-million-parameter decoder-only Transformer with self-attention, cross-attention, and timestep conditioning.The same architecture is used for sticker and photo-realistic generation.
  • Classifier-free guidance: Classifier-free guidance combines conditional and null-conditioned logits, using λ = 8 in all experiments.The extrapolation is applied directly to model logits rather than pixel values.
  • Results: On COCO validation, the comparison includes models trained with paired text-image data, whereas kNN-Diffusion used neither COCO training nor COCO retrieval data.
  • Results: The model generates text-aligned stickers and composes multiple concepts, while outperforming baselines on challenging prompts qualitatively.The comparison includes a version using inference-time kNN and the same model without inference-time kNN.

6.6 CONTINUOUS KNN MODEL

The continuous kNN model adapts a diffusion U-Net to condition on CLIP image embeddings and retrieved neighbors during training. Human evaluators preferred the kNN model over its no-kNN counterpart for both image quality and text alignment.

  • Architecture: The continuous model uses a diffusion U-Net that accepts multiple CLIP embeddings as conditioning and removes the text transformer during training.
  • Conditioning: The kNN-conditioned model uses CLIP image embeddings together with their nearest neighbors, unlike the no-kNN model.Both models are trained without paired text-image data.
  • Results: 66.8% of human judgments preferred the kNN model for image quality, and 69.4% preferred it for text alignment.The comparison is against the no-kNN model.
  • Resolution: The model generates 64×64 images before open-source super-resolution upsampling to 256×256, 512×512, and 1024×1024.
  • Training Objective: The continuous model is trained with mean-squared error loss to predict diffusion noise over N = 1000 steps.The noise predictor is a U-Net conditioned on y.

6.7 TEXT-ONLY IMAGE MANIPULATION

The manipulation procedure replaces a randomly selected local image area with an aligned nearest-neighbor patch during training. The model then uses the local area’s CLIP embedding to reconstruct the original image, enabling mask-free local semantic manipulation.

  • Evaluation: Additional examples and baseline comparisons for the manipulation process are provided in Figures 18–20.
  • Manipulated Image Construction: Training selects a random local area and replaces it with the corresponding patch from the first nearest neighbor.
  • Manipulated Image Construction: ECC alignment aligns the nearest-neighbor image with the original before the local patch replacement.
  • Training Objective: The model receives the manipulated image and the local area’s CLIP embedding, then predicts the original image.

6.8 ADDITIONAL SAMPLES

The additional samples illustrate kNN-Diffusion outputs across challenging text prompts, sticker generation, photorealistic images, and text-guided manipulations. Comparisons also examine kNN-based generation against non-kNN and competing editing methods.

  • Additional samples: Additional samples test the photo-realistic model on challenging text inputs.These examples are provided in Figures 14, 17, and 16.
  • Additional samples: Figure 13 compares discrete kNN-Diffusion stickers with 10-nearest-neighbor CLIP retrieval and a discrete model without kNN.
  • Additional samples: Figures 15–17 show samples from the COCO validation set and continuous and discrete kNN-Diffusion sticker models.
  • Additional samples: Figures 18–20 provide manipulation examples and qualitative comparisons with Text2LIVE and Textual Inversion.In Figure 19, the bottom row contains the paper’s outputs and the top row contains Text2LIVE outputs; Figure 20 uses the same row ordering for Textual Inversion.
Loading 2204.02849v2…