Test-Time Compositional Generalization in Diffusion Models via Concept Discovery

TL;DR

This paper explores how pretrained diffusion models can discover and compose concepts at test time for out-of-distribution queries without predefined concept libraries, using gradient-based methods and a product-of-experts approach.

Contribution

It introduces a novel method for test-time concept discovery and composition in diffusion models, enabling out-of-distribution generation without prior concept annotations.

Findings

The proposed method outperforms baselines on ColorMNIST and CelebA composition benchmarks.

The time-indexed score geometry contains reusable density-mode concepts.

The approach enables test-time compositional generation without a predefined concept library.

Abstract

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 $p_t(x_t)$ and compose them at test time. Given a single out-of-distribution query, our method performs gradient ascent on $s_\theta(x_t,t) \approx \nabla_{x_t}\log p_t(x_t)$ 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.