🤖 machine learning

From Global to Factor-Wise Expert Composition in Discrete Diffusion Models

This paper introduces FactorDiff, a novel framework for discrete diffusion models that improves compositional generation by dynamically routing individual sample factors to specialized experts rather than applying global mixing weights, thereby achieving superior performance on spatially complex reasoning tasks like ARC-AGI.

Original authors: Haozhe Huang, Yudong Xu, Abhijoy Mandal, Alán Aspuru-Guzik

Published 2026-07-14

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Haozhe Huang, Yudong Xu, Abhijoy Mandal, Alán Aspuru-Guzik

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're trying to solve a giant, tricky puzzle, like the ones in the ARC-AGI benchmark, where you have to figure out the hidden rules of a grid of colored squares. To do this, you've hired a team of "expert" AI models. But here's the catch: some experts are great at drawing the shapes (the "Occupancy" experts), while others are amazing at picking the right colors (the "Color" experts), but neither is perfect at doing both at the same time.

For a long time, the standard way to combine these experts was like a team vote. Every single expert would shout out their opinion for the entire puzzle at once, and the team would pick a single "score" to decide whose voice was loudest for the whole picture. The paper calls this "per-sample composition."

The Problem with the Team Vote
The authors argue that this "one voice for the whole room" approach is flawed. Imagine a puzzle where one expert is a master at drawing borders but terrible at filling in the middle, while another is a color wizard who gets the shapes wrong. If you give the border expert a single "vote" for the whole puzzle, their bad guesses about the middle colors drag down the whole team. It's like asking a master chef to also fix the plumbing in the kitchen just because they are the best cook; when they try to fix the pipe, they mess up the soup.

The New Idea: FactorDiff (The "Specialist Switch")
The paper introduces a new method called FactorDiff. Instead of voting on the whole puzzle at once, FactorDiff acts like a smart manager who looks at every single square (or pixel) on the grid individually.

Think of it like a construction crew building a house:

When the crew needs to lay the foundation and build the walls, they listen only to the Structure Expert.
When they need to paint the walls and pick the curtains, they switch over and listen only to the Color Expert.

The paper shows that by dynamically routing each tiny piece of the puzzle to the expert who is most confident about that specific piece, you get a much better result. They call this "per-pixel routing."

What the Numbers Say
The authors tested this on 120 different tasks from the ARC-AGI dataset. Here is what they found:

When experts are specialists: If you take a "Color" expert (who gets the full puzzle right only 3.3% of the time) and an "Occupancy" expert (who gets it right only 0.2% of the time), the old "team vote" methods struggled, getting the full puzzle right only about 78.4% of the time at best. But with FactorDiff's "Specialist Switch," the team got it right 90.5% of the time. That's a huge jump!
When experts are generalists: Even when the experts are both very good at everything (getting it right 93.0% and 89.3% of the time), FactorDiff didn't make things worse. It matched the best results, getting 95.2% right, proving that this method is safe to use even when you don't need to switch experts.

What the Paper Rules Out
The paper explicitly argues against the idea that a single, global "weight" or "vote" for an expert is enough. They show that trying to fix the "team vote" method with complex math (like the "Feynman-Kac correctors" or "SuperDiff" mentioned in the paper) still fails to capture the nuance of needing different experts for different parts of the image. You can't just tweak the volume of the experts; you have to change who is speaking for where.

How Sure Are They?
The authors are very confident in these results because they ran actual experiments, not just simulations. They trained models on 120,000 example pairs and tested them on a held-out set of 200 examples per task. The results were measured directly: FactorDiff consistently outperformed the old methods, especially when the experts had different strengths.

However, the paper does admit one thing they haven't fully solved yet: their "manager" (the routing system) currently decides who to listen to based on how confident an expert looks. The authors suggest that while this works well for the tests they ran, it might not work for every possible pair of experts in the universe. They propose that in the future, we might need to learn a better way to decide who to listen to, rather than just guessing based on confidence.

In short, the paper suggests that for complex reasoning tasks, we shouldn't treat AI experts like a monolithic block. Instead, we should treat them like a team of specialists, letting the best person handle the specific part of the job they are best at, one tiny piece at a time.

Technical Summary: From Global to Factor-Wise Expert Composition in Discrete Diffusion Models

Problem Statement

Discrete diffusion models have emerged as a powerful framework for complex reasoning tasks, particularly through compositional generation, which combines multiple pre-trained "expert" models to generalize beyond their individual training distributions. Recent theoretical advancements (e.g., SuperDiff, RNE, Feynman-Kac Correctors) have improved compositional sampling by introducing time-dependent mixing weights to align composed dynamics with target distributions.

However, these state-of-the-art methods share a fundamental limitation: they operate on a per-sample basis, assigning a single global scalar weight to each expert at every timestep. This "monolithic" assumption implicitly treats every expert as equally knowledgeable across the entire state space. In complex generative tasks—such as spatial reasoning where one expert may specialize in boundaries while another handles textures—this global weighting forces a compromise. It dilutes the signal of the correct expert in their specific region of competence, failing to capture the spatial and functional independence of different experts.

Methodology: FactorDiff

The authors propose FactorDiff, a factor-wise composition framework that decomposes the generation process from the sample level down to the factor level (e.g., spatial pixels or grid cells). Instead of applying a single weight to an entire generated state, FactorDiff dynamically routes each factor to the most relevant expert.

Core Components

Factor-Wise Routing: The state space is decomposed into user-defined factors. During the sampling process, a routing map assigns each factor (e.g., a pixel in a grid) to a specific expert. This allows heterogeneous expert contributions within a single generated state.
Confidence-Based Selection: The framework instantiates routing by dynamically selecting the expert with the highest confidence at each position during inference. The routing variable $z_v$ for a variable $v$ is determined by:
$z_v = \arg\max_{k} C(p_k(x_v | x_t))$
where $C$ is a confidence metric (the margin between the top two prediction probabilities). The composed reverse step then selects the prediction from the chosen expert for that specific position.
Native 2D Masked Diffusion Architecture: To support this, the authors introduce a 2D masked diffusion architecture that preserves grid topology. Unlike prior works that flatten grids into 1D sequences (introducing topological mismatches), this approach treats the grid as a lattice graph.
- Experts: Three complementary model types are trained: a Full model (all colors), an Occupancy model (background vs. foreground), and a Color model (foreground colors only).
- Backbone: The denoising network uses an iterative refinement architecture (either Convolutional or Transformer-based) with learnable 2D positional embeddings to ensure spatial locality.

Theoretical Justification

Existing methods often define a target distribution as a weighted product of experts ( $p(x) \propto \prod p_m(x)^{\alpha_m}$ ). The authors argue that for complementary specialists, a global exponent $\alpha_m$ is misspecified because it forces an expert's predictions (even erroneous ones outside their domain) to influence the entire state. FactorDiff replaces this with a factor-wise product:
$p_{comp}(x) \propto \prod_{f \in \mathcal{F}} \prod_{m=1}^M (p_m(x_f))^{\alpha_{m,f}}$
This formulation allows expertise to vary across factors. In the hard-routing case, $\alpha_{m,f}$ is one-hot, selecting a single expert per factor without requiring additional normalization or importance correction, as the product of valid conditionals defines a valid joint transition.

Key Contributions

FactorDiff Framework: A novel composition method that casts expert combination as routing over factors, specifically instantiated with position-level expert selection for grid-structured diffusion models.
2D Masked Diffusion Architecture: A native architecture that maintains grid topology, enabling the training of complementary experts (Occupancy, Color, Full) whose predictions can be composed at inference time.
Empirical Validation on ARC-AGI: Extensive evaluation on the ARC-AGI benchmark demonstrates that simple factor-specific routing consistently outperforms complex global scalar weighting schemes, particularly on tasks requiring logical consistency and spatial disentanglement.

Experimental Results

The authors evaluated FactorDiff on a subset of 120 ARC-AGI tasks using the RE-ARC generator for in-distribution (IID) evaluation.

Specialization Regime (RQ1): When combining highly specialized experts (e.g., Occupancy + Full, or Color + Occupancy), FactorDiff significantly outperformed per-sample baselines (SuperDiff, RNE, FKC).
- Color + Occupancy: Individual specialists scored near zero on the full task (3.3% and 0.2% exact accuracy). FactorDiff composed them to achieve 90.5% exact accuracy, a massive improvement over the best per-sample baseline (78.4%).
- Occupancy + Full: FactorDiff achieved 94.0% exact accuracy, outperforming the best single expert (93.0%) and all per-sample composition methods.
Generalist Regime (RQ2): When combining two strong generalist models (Full + Full), FactorDiff was "safe," matching the best per-sample baseline (95.2%) without degrading performance. This indicates that spatially varying selection does not introduce instability when experts broadly agree.
Ablation: Even when composing two full models with identical objectives but different architectures (Trans-IR vs. Conv-IR), FactorDiff achieved higher accuracy (94.8%) than either individual model, successfully stitching correct regions from both.

Significance and Claims

The paper claims that FactorDiff provides a simple and effective mechanism for combining arbitrary specialists within a single diffusion trajectory. Its primary significance lies in demonstrating that per-sample global weighting is insufficient for tasks where experts have spatially or functionally specialized domains. By moving to factor-wise composition, the method captures complementary expertise that global weights cannot represent, achieving superior generalization on tasks requiring strict logical consistency.

The authors note that while their current implementation relies on confidence-based routing (assuming confidence correlates with correctness), the framework naturally extends to learned or task-conditioned routing. They position this work as a step toward more granular composition in discrete diffusion, moving beyond the "monolithic" assumptions of current state-of-the-art methods.