# Test-Time Compositional Generalization in Diffusion Models via Concept Discovery

arXiv: https://arxiv.org/abs/2605.07078 · HTML: https://mlsanity.com/paper/2605.07078/

## Key facts

- Cited by 1 papers in the map (data through 14 Oct 2026).
- Placed in Diffusion foundations and sampling, one of 14 directions in Diffusion models.
- Builds on 33 of its 53 references that are in the map.
- First posted on arXiv 8 May 2026.

## Placed in

- Diffusion models / [Diffusion foundations and sampling](https://mlsanity.com/diffusion/diffusion-foundations-and-sampling/)

## Builds on (most cited first, 30 of 33)

- [Learning Transferable Visual Models From Natural Language Supervision](https://mlsanity.com/paper/2103.00020/) (2021)
- [Denoising Diffusion Probabilistic Models](https://mlsanity.com/paper/2006.11239/) (2020)
- [Score-Based Generative Modeling through Stochastic Differential Equations](https://mlsanity.com/paper/2011.13456/) (2020)
- [Deep Unsupervised Learning using Nonequilibrium Thermodynamics](https://mlsanity.com/paper/1503.03585/) (2015)
- [Classifier-Free Diffusion Guidance](https://mlsanity.com/paper/2207.12598/) (2022)
- [Fashion-MNIST: a Novel Image Dataset for Benchmarking Machine Learning Algorithms](https://mlsanity.com/paper/1708.07747/) (2017)
- [Generative Modeling by Estimating Gradients of the Data Distribution](https://mlsanity.com/paper/1907.05600/) (2019)
- [Deep Learning Face Attributes in the Wild](https://mlsanity.com/paper/1411.7766/) (2014)
- [DreamBooth: Fine Tuning Text-to-Image Diffusion Models for Subject-Driven Generation](https://mlsanity.com/paper/2208.12242/) (2022)
- [Diffusion Policy: Visuomotor Policy Learning via Action Diffusion](https://mlsanity.com/paper/2303.04137/) (2023)
- [An Image is Worth One Word: Personalizing Text-to-Image Generation using Textual Inversion](https://mlsanity.com/paper/2208.01618/) (2022)
- [Interpretability Beyond Feature Attribution: Quantitative Testing with Concept Activation Vectors (TCAV)](https://mlsanity.com/paper/1711.11279/) (2017)
- [Variational Diffusion Models](https://mlsanity.com/paper/2107.00630/) (2021)
- [Multi-Concept Customization of Text-to-Image Diffusion](https://mlsanity.com/paper/2212.04488/) (2022)
- [Diffusion-LM Improves Controllable Text Generation](https://mlsanity.com/paper/2205.14217/) (2022)
- [Improved Precision and Recall Metric for Assessing Generative Models](https://mlsanity.com/paper/1904.06991/) (2019)
- [Large Language Diffusion Models](https://mlsanity.com/paper/2502.09992/) (2025)
- [This Looks Like That: Deep Learning for Interpretable Image Recognition](https://mlsanity.com/paper/1806.10574/) (2018)
- [Network Dissection: Quantifying Interpretability of Deep Visual Representations](https://mlsanity.com/paper/1704.05796/) (2017)
- [NVAE: A Deep Hierarchical Variational Autoencoder](https://mlsanity.com/paper/2007.03898/) (2020)
- [Compositional Visual Generation with Composable Diffusion Models](https://mlsanity.com/paper/2206.01714/) (2022)
- [Neural Module Networks](https://mlsanity.com/paper/1511.02799/) (2015)
- [Ladder Variational Autoencoders](https://mlsanity.com/paper/1602.02282/) (2016)
- [Reduce, Reuse, Recycle: Compositional Generation with Energy-Based Diffusion Models and MCMC](https://mlsanity.com/paper/2302.11552/) (2023)
- [Good-Enough Compositional Data Augmentation](https://mlsanity.com/paper/1904.09545/) (2019)
- [A Benchmark for Systematic Generalization in Grounded Language Understanding](https://mlsanity.com/paper/2003.05161/) (2020)
- [A Phase Transition in Diffusion Models Reveals the Hierarchical Nature of Data](https://mlsanity.com/paper/2402.16991/) (2024)
- [Learning to Compose Neural Networks for Question Answering](https://mlsanity.com/paper/1601.01705/) (2016)
- [DreamTime: An Improved Optimization Strategy for Diffusion-Guided 3D Generation](https://mlsanity.com/paper/2306.12422/) (2023)
- [Unsupervised Compositional Concepts Discovery with Text-to-Image Generative Models](https://mlsanity.com/paper/2306.05357/) (2023)

## Built on by (most cited first, 1 of 1)

- [Diffusion Models and Concept Formation](https://mlsanity.com/paper/2609.13047/) (2026)

## Abstract (arXiv)

Compositional generalization requires models to produce novel configurations from familiar parts. In diffusion models, prior compositional generation methods typically assume that the relevant concepts or conditioning signals are already available. We instead ask whether a pretrained diffusion model can discover query-specific concepts from the time-indexed scores it learns for the noisy marginals pt(xt) and compose them at test time. Given a single out-of-distribution query, our method performs gradient ascent on sθ(xt,t) ≈ nablaxtlog pt(xt) at multiple noising timesteps to recover local density modes, maps these modes into clean-space Gaussians, greedily selects relevant prototypes with a submodular likelihood objective, and combines them into a product-of-experts (PoE) teacher model with an analytic score. This teacher model can be sampled directly through classifier-free guidance or used to generate a sample pool for training a new class embedding and low-rank adapter. On held-out composition benchmarks built from ColorMNIST and CelebA, both the analytic PoE sampler and the low-rank adapted model outperform query-only and nearest trained-class baselines. These results suggest that the time-indexed score geometry of the diffusion model contains reusable density-mode concepts that support test-time compositional generation without a predefined concept library.