🤖 machine learning

Test-Time Compositional Generalization in Diffusion Models via Concept Discovery

This paper proposes a test-time compositional generalization method for pretrained diffusion models that discovers query-specific concepts by analyzing time-indexed score geometry to construct a product-of-experts teacher model, enabling the generation of novel configurations without relying on predefined concept libraries.

Original authors: Zekun Wang, Anant Gupta, Tianyi Zhu, Christopher J. MacLellan

Published 2026-05-11

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Zekun Wang, Anant Gupta, Tianyi Zhu, Christopher J. MacLellan

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you have a master chef who has spent years cooking thousands of meals. This chef knows how to make a "Spicy Beef Stew" and a "Sweet Fruit Salad" perfectly. But one day, you ask them to make something they've never seen before: "Spicy Fruit Stew."

In the world of AI, this is called Compositional Generalization. The challenge is: Can the chef combine the concept of "spicy beef" and the concept of "fruit salad" to create a new dish, even though they've never cooked that specific combination before?

Usually, AI models need a recipe book (a library of concepts) to do this. If "Spicy Fruit Stew" isn't in the book, the chef gets confused.

This paper introduces a new way for the AI chef to figure it out on the fly, without a recipe book. Here is how it works, using simple analogies:

1. The "Foggy Room" Analogy (The Diffusion Model)

Think of a trained AI image generator as a room filled with fog.

The Fog: The AI starts with a completely foggy, noisy image.
The Process: The AI slowly clears the fog, step by step, to reveal a clear picture.
The Secret: As the fog clears, the AI learns "landmarks." At a thick fog level, it might see a vague shape of a "face." As the fog lifts further, it sees "eyes," then "blue eyes," then "a specific person's blue eyes."

The authors realized that these "landmarks" (called modes) exist at different levels of clarity. Some are blurry (general concepts like "face"), and some are sharp (specific details like "smiling").

2. The Detective Work (Concept Discovery)

When you give the AI a weird request (like "Spicy Fruit Stew"), it doesn't panic. Instead, it acts like a detective in that foggy room.

The Clue: You show the AI a single, slightly blurry example of what you want (the "Query").
The Search: The AI runs a special search (called Gradient Ascent) through the fog. It looks for the "peaks" or high points in the fog where the data is densest.
The Discovery: It finds several distinct "peaks" that match parts of your request. Maybe it finds a peak that looks like "fruit" and another that looks like "spicy." It doesn't need a label for these; it just finds the shapes in the fog that fit.

3. The "Team of Experts" (Product-of-Experts)

Now the AI has found these "peaks" (concepts). But how do you combine them?

Imagine you have a team of experts. One expert knows about "fruit," another knows about "spicy."
The AI creates a Product-of-Experts (PoE) model. This is like a meeting where all these experts vote on what the final image should look like.
They don't just mash the images together; they mathematically combine their "opinions" (probabilities) to create a new, clear blueprint for "Spicy Fruit Stew."

4. The Two Ways to Cook (Sampling & Distillation)

Once the AI has this new blueprint, it can generate the image in two ways:

Method A: The Instant Guide (Analytic Sampling)
The AI uses the blueprint to guide the fog-clearing process directly. It's like the chef looking at the blueprint and saying, "Okay, I know exactly how to clear the fog to make this specific dish right now."
Method B: The Training Session (LoRA Distillation)
The AI takes the images it generated using the blueprint and uses them to teach itself a new "trick." It adds a small, lightweight update (called LoRA) to its brain.
- Analogy: It's like the chef tasting the new "Spicy Fruit Stew," writing down a new note in their personal notebook, and then memorizing it. Next time, they can make it instantly without needing the detective work again.

Why This Matters

The paper tested this on two things:

Colored Digits: Making a "Green Number 7" when the AI was only trained on "Red Number 7" and "Green Number 3."
Faces: Making a face with "Blond Hair" and "Big Lips" when it had never seen that exact combination before.

The Results:

Old Methods: Tried to find the closest thing it knew (e.g., "Oh, you want Green? Here is a Green 3, but it's not a 7"). The result was often wrong.
The New Method: Successfully discovered the "Green" concept and the "7" concept separately, combined them, and created a perfect "Green 7." It was better at keeping the details right (faithfulness) and making it look like a real image (generalization).

The Bottom Line

This paper shows that AI models already contain a hidden library of "building blocks" inside their foggy learning process. You don't need to give them the blocks; you just need to teach them how to find the blocks and snap them together when they see a new, weird request. It turns the AI from a rigid recipe follower into a flexible, creative problem solver.

Technical Summary: Test-Time Compositional Generalization in Diffusion Models via Concept Discovery

Problem Statement

Compositional generalization—the ability to recombine familiar parts into novel configurations—is a hallmark of human intelligence but remains a significant challenge for deep learning models. While neural networks often interpolate well within training distributions, they struggle when familiar primitives appear in unfamiliar combinations. Existing approaches to compositional generation in diffusion models typically rely on a predefined library of concepts, conditioning signals, or external labels to guide the generation of novel compositions. This creates a gap between concept discovery and compositional generalization: prior methods assume the relevant primitives are already known and available, failing to address scenarios where a model must discover query-specific concepts from scratch at test time without a fixed vocabulary or external supervision.

Methodology

The paper proposes a framework that repurposes a pretrained diffusion model as an implicit hierarchy of time-indexed density modes to discover and compose concepts at test time. The method operates in three primary stages:

Concept Discovery via Mode Ascent:
Instead of using the diffusion model solely for reverse sampling, the authors utilize the learned time-indexed score function $s_\theta(x_t, t) \approx \nabla_{x_t} \log p_t(x_t)$ to locate local density modes of the noisy marginals $p_t(x_t)$ . Starting from a single out-of-distribution (OOD) query $x_q$ , the method performs gradient ascent on the unconditional score at multiple noise timesteps $t$ to recover local modes.
- Theoretical Basis: Leveraging Tweedie's formula, the authors demonstrate that these modes correspond to noise-free "concept prototypes" in the clean data space ( $x_0$ ), effectively acting as discrete centroids at varying levels of abstraction depending on the noise level.
- Representation: Each discovered mode is mapped back to $x_0$ -space and approximated as a local Gaussian expert $q_j(x_0) = \mathcal{N}(x_0; m_j, \Sigma_j)$ , where the mean is derived from the mode and the covariance is estimated via Hutchinson probes.
Product-of-Experts (PoE) Composition:
The discovered Gaussian experts are combined into a single analytic teacher model using a Product-of-Experts (PoE) formulation.
- Selection: A greedy submodular optimization algorithm selects the subset of prototypes most relevant to the query, maximizing a per-dimension log-likelihood coverage objective.
- Composition: The selected experts are combined into a weighted PoE distribution $q_T$ . Because the experts are Gaussian, the resulting composition is also Gaussian with a closed-form score function. This allows for the derivation of an analytic noise prediction $\epsilon_{q_T}(x_t, t)$ without requiring a trained conditional token for the novel composition.
Compositional Sampling and Distillation:
- Analytic Sampling: The analytic score of the PoE teacher is integrated into a classifier-free guidance (CFG) framework. The model contrasts the PoE prediction against the unconditional prediction to guide sampling toward the discovered composition.
- LoRA Distillation: To improve in-manifold generation and create a reusable conditional, the method optionally distills the PoE teacher into the base model. This involves training a new class embedding and a Low-Rank Adaptation (LoRA) adapter on a pool of samples generated by the PoE teacher. This step absorbs the discovered composition back into the model's score manifold.

Key Contributions

Bridging Discovery and Composition: The work bridges concept discovery and compositional generalization by treating a pretrained diffusion model as an implicit hierarchy of density modes. It enables the recovery of discrete concept prototypes at test time without a predefined primitive pool.
Test-Time Composition Mechanism: The authors introduce a method that greedily selects relevant modes using a submodular likelihood objective and combines them into a PoE teacher model with a closed-form score, enabling direct sampling without retraining the base model on the specific composition.
Empirical Validation: The framework is evaluated on two benchmarks, ColorMNIST and CelebA, demonstrating that discovered prototypes and their PoE composition encode richer information about unseen compositions than query-only sampling or nearest trained-class retrieval.

Results

Experiments on ColorMNIST (digit identity, color, background) and CelebA (face attributes) show that the proposed method outperforms several baselines:

Baselines: The method is compared against "Query-only" (sampling around the query with fixed variance), "Top-1/Top-3" (retrieving the nearest trained classes), and standard CFG.
Performance: Both the analytic PoE sampler and the LoRA-distilled variant achieve superior FID, CLIP similarity, and Precision/Recall metrics compared to baselines.
- On ColorMNIST, the LoRA variant achieves an F1 score of 87.6% (Faithfulness) and 76.3% (Generalization), significantly outperforming the Top-1 baseline (~6.5% F1).
- On CelebA, the PoE method achieves the best Generalization FID (140.3) and CLIP scores, while LoRA further improves the balance of coverage and fidelity, achieving the highest Generalization Recall and F1.
Qualitative Analysis: Visual results indicate that while baselines often fail to preserve the specific combination of attributes (e.g., changing digit identity or drifting to different attribute combinations), the proposed method successfully reconstructs the OOD composition and generalizes to other instances of the same class.

Significance and Claims

The paper claims that the time-indexed score geometry of diffusion models contains reusable density-mode concepts that support test-time compositional generation without a predefined concept library.

Theoretical Insight: The results suggest that diffusion models implicitly encode a hierarchy of concepts where modes of noisy marginals serve as prototypes at different levels of abstraction.
Practical Implication: The method allows for the generation of novel configurations from a single OOD query, offering a pathway for compositional generalization that does not rely on external supervision or pre-specified vocabularies.
Limitations: The authors modestly note that the protocol assumes primitive factors are known (standard compositional setting) and that the method relies on the inductive biases of the pretrained backbone. While it can interpolate between known primitives, it cannot fabricate genuinely new primitives that extrapolate beyond the training support (e.g., a white background in a model trained only on dark backgrounds).

The work provides both theoretical and empirical evidence that diffusion models can be repurposed for concept discovery and composition, potentially extending to broader applications in language modeling and robotic control.