💬 NLP

When Diffusion Breaks Constraints: Sequential Autoregressive Generation with RL and MCTS

This paper demonstrates that diffusion models fundamentally struggle with constrained generation tasks due to their inability to sample from low-dimensional feasible regions, and proposes a sequential autoregressive approach enhanced by reinforcement learning and Monte Carlo tree search as a more effective alternative for satisfying strict geometric and physical constraints.

Original authors: Zirui Zhao, Boye Niu, Harold Soh, David Hsu, Wee Sun Lee

Published 2026-05-14

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Zirui Zhao, Boye Niu, Harold Soh, David Hsu, Wee Sun Lee

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

The Big Problem: The "Perfect Puzzle" Dilemma

Imagine you are trying to build a specific shape out of seven wooden puzzle pieces (a Tangram). You are given a description like "a bird sitting on a branch."

You have two ways to try to solve this:

The "Spray and Pray" Method (Diffusion Models): Imagine you have a magical machine that shoots out random arrangements of the seven pieces all at once. It tries to guess the right shape by learning from thousands of pictures. The problem? The machine is great at making things look mostly right, but it often fails at the strict rules. It might make the bird's wing overlap with its body, or leave a gap so the pieces aren't connected. In the real world, these "rules" (no overlapping, must be connected) are hard constraints. If you break one, the whole solution is trash.
The "Step-by-Step" Method (Autoregressive Models): Instead of shooting out the whole picture at once, you place one piece, then another, then another. You check the rules after every single move.

The Paper's Discovery: The authors found that the "Spray and Pray" method (Diffusion) is terrible at these strict puzzle tasks. Even if you tell the machine, "Hey, don't let pieces overlap," it still fails almost 100% of the time on hard puzzles. It's like trying to thread a needle while blindfolded and spinning in a circle; the target is just too small and specific for the machine to hit by accident.

Why Does the "Spray and Pray" Method Fail?

The paper uses a mathematical concept called "Feasible Mass."

Think of the entire universe of possible puzzle arrangements as a giant, empty warehouse.

The "Good" Arrangements: The arrangements that actually fit the rules (no overlaps, connected, look like a bird) are like a few tiny, invisible specks of dust floating in that warehouse.
The "Bad" Arrangements: Everything else (overlapping pieces, disconnected parts) fills up the rest of the warehouse.

The Diffusion model tries to spray paint the whole warehouse at once, hoping to hit those tiny specks of dust. Because the "good" area is so incredibly small (mathematically, it's a "low-dimensional submanifold"), the model almost never hits it. It's like trying to hit a specific grain of sand on a beach by throwing a handful of sand from a helicopter.

The Solution: The "Smart Builder" (GAG MCTS)

The authors propose a new way to solve this: Sequential Autoregressive Generation with Reinforcement Learning and Search.

Let's break down their solution, which they call GAG MCTS, using an analogy of a Master Architect and a Team of Interns:

The Step-by-Step Approach (Autoregressive): Instead of building the whole bird at once, the AI places one piece, then checks if it fits. Then it places the next piece attached to the first one. This immediately cuts out the "impossible" moves (like placing a piece inside another piece).
The Reinforcement Learning (The Reward System): The AI learns by playing the game over and over. If it builds a bird that looks good and follows the rules, it gets a "gold star" (reward). If it fails, it gets a "thumbs down." Over time, it learns which moves lead to gold stars.
The "Look-Ahead" Search (MCTS): This is the secret sauce. Imagine you are playing chess. You don't just look at the move you are making now; you think, "If I move here, what happens next? Can I win in 5 moves?"
- The AI uses Monte Carlo Tree Search (MCTS) to simulate thousands of future possibilities in its head before making a move.
- It asks: "If I place this piece here, will I get stuck later?" If the answer is yes, it avoids that move, even if the move looks okay right now.

The "Adversarial" Twist

The paper also mentions a clever trick to make the AI smarter at judging what a "bird" looks like.

The Problem: The AI's "judge" (a reward model) was getting tricked. It would give a gold star to a pile of blocks that looked sort of like a bird, even if it was garbage.
The Fix: The authors set up a game of "Fake vs. Real." The AI tries to build a fake bird to trick the judge. The judge tries to spot the fake. They play this game against each other (Adversarial Training). Eventually, the judge becomes so sharp it can spot even the tiniest mistake, and the builder becomes so skilled it can only build perfect birds.

The Results: Who Won?

The authors tested this on two puzzles:

Tangrams: Seven pieces forming shapes like "a sitting person" or "a goose."
Rectangle Packing: Fitting rectangles into a box without overlapping.

The Results:

Diffusion Models (The Spray and Pray): Failed miserably. On the hardest puzzles, they succeeded less than 5% of the time. They just couldn't hit the tiny "good" area.
Step-by-Step without Search: Did better (around 60-80% success), but often got stuck in dead ends where they couldn't finish the puzzle.
GAG MCTS (The Smart Builder): Won almost every time (95-99% success). By thinking ahead and checking the rules at every step, it navigated the "tiny specks of dust" in the warehouse perfectly.

The Takeaway

The paper concludes that for tasks with strict, hard rules (like engineering designs, molecular structures, or floor plans where things cannot overlap), the current popular "spray and pray" AI models are fundamentally broken. They are trying to solve a puzzle by guessing the whole picture at once, which is mathematically impossible when the rules are so tight.

Instead, we need to switch to step-by-step builders that can think ahead (search) and learn from their mistakes (reinforcement learning). It's the difference between throwing a dart at a board while blindfolded versus walking up to the board, aiming carefully, and placing the dart exactly where it needs to go.

Technical Summary: When Diffusion Breaks Constraints

Problem Statement

Data-driven generative models, particularly diffusion models, excel in unconstrained domains like language and image synthesis but struggle significantly when tasked with constrained planning and design. In scenarios such as engineering inverse design, molecular generation, multi-robot planning, and floorplan synthesis, feasible solutions often reside on low-dimensional submanifolds or in small, potentially disconnected regions of the output space. These regions are defined by hard constraints (e.g., non-overlap, connectivity, physical feasibility) where a single violation invalidates the entire output.

The paper investigates this failure mode through two specific case studies:

Tangram Generation: Composing seven fixed 2D shapes into a silhouette matching a text description, subject to strict non-overlap and connectivity constraints.
Rectangle Composition: Arranging rectangular pieces within a text-specified bounding box under non-overlap constraints, where difficulty is modulated by the fill ratio of the pieces relative to the container.

The authors posit that continuous density matching (the core mechanism of diffusion) is structurally ill-suited for these problems because the "relative feasible mass" (the volume of valid solutions relative to the ambient space) is vanishingly small. Consequently, even with guidance or projection, diffusion models fail to generate valid samples with high probability.

Methodology

To overcome the limitations of continuous generators, the paper proposes reformulating constrained generation as a discrete autoregressive sequential generation problem, augmented with Reinforcement Learning (RL) and search.

1. Reparameterization and Discretization

Instead of generating all variables simultaneously in a continuous space, the problem is decomposed into a sequence of local placement actions.

Local Reparameterization: Given a partial assembly, the next piece is generated by attaching it to an existing side at an anchor-aligned pose.
Discretization: The allowable positions and orientations for the next piece are discretized (e.g., 45° rotation increments, specific anchor points), transforming the problem into a standard discrete autoregressive task.
Action Masking: A pairwise action mask filters out immediate geometric violations (overlaps) between the new piece and already placed pieces, though it does not guarantee future feasibility.

2. Generative Adversarial Gumbel MCTS (GAG MCTS)

The core proposed method combines discrete autoregressive generation with Monte Carlo Tree Search (MCTS) and adversarial reward refinement:

Policy-Value Network: A transformer-based architecture fusing a pre-trained ViT (image) and a frozen BERT (text) encoder. It outputs action logits and a scalar value estimate.
Gumbel MCTS: The search algorithm uses the Gumbel-Top-k trick for root action selection and Sequential Halving to explore the tree. Transitions are executed via an exact geometric simulator rather than a learned dynamics model.
Adversarial Reward Refinement: To address the difficulty of learning a robust reward function (where standard CLIP-based rewards can be fooled by "near-miss" invalid configurations), the authors employ a generative adversarial loop:
- Positive Data ( $D^+$ ): Labeled valid configurations.
- Negative Data ( $D^-$ ): Invalid configurations generated by the current policy's MCTS rollouts.
- Training: The reward network (based on CLIP) is trained to distinguish $D^+$ from $D^-$ using a preference loss, while maintaining the original image-text alignment via a contrastive loss to prevent catastrophic forgetting. The policy is then improved via PPO or imitation learning on the MCTS trajectories.

Key Contributions

Theoretical Analysis of Feasible-Mass Barriers: The paper formalizes the failure of continuous generators using a "feasible-mass bound." It demonstrates that for constrained tasks, the relative volume of valid solutions ( $\mu_y$ ) is exponentially small. Guidance mechanisms can increase the concentration factor ( $K_{\theta, \beta}$ ) but cannot alter $\mu_y$ , explaining why diffusion models degrade sharply as constraints tighten.
Sequential Constraint-Aware Generation: The authors demonstrate that reparameterizing the problem into a discrete, sequential process bypasses the "tube-hitting" problem of continuous spaces. By building constraint satisfaction into the action space (via masking) and using search to look ahead, the model can navigate the sparse feasible regions effectively.
GAG MCTS Framework: The integration of Gumbel MCTS with adversarial reward refinement provides a robust method for constrained generation. The adversarial component is crucial for training reward models that are robust to the specific "near-miss" failures common in search-based generation.

Experimental Results

The proposed GAG MCTS method was evaluated against diffusion baselines (with energy guidance) and autoregressive baselines (with and without RL/search) on both Tangram and Rectangle Composition tasks.

Tangram Assembly:
- Diffusion Models: Achieved near-zero validity (0–4%) even with strong energy guidance, confirming the theoretical barrier.
- Autoregressive (Supervised): Improved validity to ~54–88% but struggled with long-horizon feasibility.
- PPO: Lifted validity to ~99.3% but showed lower fidelity (FID) and precision compared to search methods.
- GAG MCTS: Achieved near-saturation validity (~99.3%) with the best FID and precision/recall scores. It successfully balanced feasibility with semantic alignment.
Rectangle Composition:
- Easy Regime (Low Fill Ratio): Diffusion and autoregressive methods performed reasonably well.
- Hard Regime (High Fill Ratio): Diffusion and standard autoregressive methods collapsed (success rates dropped to 1–30%). GAG MCTS maintained high success rates (50–53% on hard splits, 95–96% on easy), significantly outperforming baselines.
Human Study: A forced-choice study with 5 participants confirmed that GAG MCTS generations were preferred over autoregressive baselines for matching text descriptions, validating the semantic quality metrics.

Significance and Claims

The paper claims that the failure of diffusion models on hard constrained-generation tasks is a structural limitation of continuous density matching on low-dimensional or disconnected feasible sets. It argues that symptomatic fixes (like projection or guidance) cannot close the gap when the feasible mass is vanishingly small.

The primary significance lies in demonstrating that sequential, constraint-aware generation combined with search (MCTS) and adversarial reward refinement offers a promising alternative. The authors suggest that for problems where feasible sets are low-dimensional (due to contact precision) or sequentially depleted (due to packing density), discrete action spaces with lookahead are superior to continuous samplers.

The paper remains modest regarding its scope, noting that while the lessons are derived from tangram and rectangle puzzles, the geometric issues apply to broader domains like 3D CAD, floorplan synthesis, and molecular generation. It acknowledges limitations, including the high inference cost of search compared to single-pass baselines and the reliance on weakly calibrated CLIP-based rewards under data scarcity.