Test-Time Compositional Generalization in Diffusion Models via Prototype Discovery

Compositional generalization asks a model to produce what it has never seen by recombining what it knows. Diffusion methods for this problem typically assume that the reusable concepts are known in advance, provided through labels, prompts, or a learned concept library. We ask whether a pretrained diffusion model can instead discover its own reusable prototypes. A prototype is a mode of a noisy marginal $p_t(x_t)$. We perform gradient ascent on the learned score $s_{\theta }(x_t,t)\approx {\nabla }_{x_t}\log p_t(x_t)$ from a datum to recover a mode and fit a local Gaussian around it, and the noise level $t$ sets the level of abstraction. On pretrained CIFAR-10 and CelebA-HQ DDPMs, the prototypes recovered from different data are distinct at fine noise levels and become less separated as noise increases, while a datum's own prototype changes sharply across levels. A single datum therefore induces several distinct, reusable prototypes rather than one prototype under increasing blur. With these discovered prototypes, at test time we select those that explain a single out-of-distribution query and compose them into a product of experts whose closed-form score guides the composable diffusion sampler. On held-out compositions of ColorMNIST and CelebA, the samples generated from one query generalize to other instances of the unseen class and outperform the same sampler conditioned on trained concept tokens. A pretrained diffusion model thus already contains prototypes at many levels of abstraction, and they can be recovered and composed without supervision.