Constrained Decoding for Diffusion Language Models via Efficient Inference over Finite Automata

TL;DR: Diffusion language models sample many token positions simultaneously from an independent (fully factorized) distribution at each denoising step, which breaks every existing constrained-decoding system built for left-to-right generation. This paper shows that if you view a finite automaton constraint as a hidden Markov model, the constrained per-step distribution becomes an exactly samplable chain-structured model — and a divide-and-conquer trick borrowed from arithmetic circuit theory cuts the sampling depth from O(L) to O(\log L), making the whole thing nearly free in wall-clock terms. The payoff is large: on BFCL-Live function calling, Dream-7B's stochastic sampling accuracy goes from 22.3% (baseline collapse) to 69.0%, at under 5% overhead.

Why constrained decoding breaks for diffusion LMs

Structured output — JSON tool calls, SQL, formal expressions — is one of the most economically important things LLMs do. For autoregressive models, the solution is mature and boring: at each step, intersect the model's next-token distribution with the set of tokens that keep the output valid, sample, advance. Systems like Outlines, XGrammar, and Guidance do exactly this, and it works because generation is a walk down a prefix tree. At every step, there is a well-defined "current state" of the constraint, and validity of the next token depends only on that state.

Diffusion language models (dLLMs) like Dream-7B and LLaDA-8B destroy this structure. At each denoising step, the model looks at a partially masked sequence x^t and produces a mean-field prediction — a fully factorized distribution p_\theta(x^0 \mid x^t) = \prod_i p_\theta(x_i^0 \mid x^t) over all positions at once. It samples every masked position independently, commits some of them, and re-masks the rest. There is no left-to-right order, no prefix, no single automaton state you can condition on.

The failure mode is not subtle. Suppose the constraint accepts valid real numbers of length 2, so both 1. and .1 are legal. Position 1's marginal puts mass on . (consistent with .1), and position 2's marginal puts mass on . (consistent with 1.). Sample each position independently and you can get .. — every token was locally plausible, and the sequence is garbage. Per-position token masking, the entire toolkit of autoregressive constrained decoding, cannot fix this: validity at position i depends on what is being sampled simultaneously at every other position.

Prior work on this problem either restricts the setting (DINGO handles greedy, block-wise MAP decoding over deterministic automata) or resorts to rejection sampling against context-free grammars, which gives no guarantee of producing a valid output within any finite budget — acceptance rates can be vanishingly small.

The insight: the constraint is a hidden Markov model

The goal at each denoising step is to sample from the constrained mean-field posterior:

where \mathcal{C} is the set of valid length-L sequences, expressed as a finite automaton (FA) with states \mathcal{S}, vocabulary \mathcal{V}, and edges \mathcal{E}. The indicator is a global coupling across all L positions — apparently intractable.

The move that unlocks everything: a finite automaton, viewed the right way, is a chain-structured graphical model — a hidden Markov model. Introduce latent variables z_1, \dots, z_L, where z_i is the automaton edge traversed at step i. The transition factor p(z_i \mid z_{i-1}) is an indicator that consecutive edges chain end-to-end (\mathrm{dst}(z_{i-1}) = \mathrm{src}(z_i)); the emission factor p(x_i \mid z_i) is an indicator that token x_i lies in edge z_i's label set; boundary conditions pin the start state and require the final state to be accepting. Marginalizing out the path z_{1:L} gives a distribution p_\mathcal{M}(x_{1:L}) whose support is exactly \mathcal{C}. If the automaton is deterministic (DFA), each accepted string has exactly one accepting path, so p_\mathcal{M} is uniform over \mathcal{C}.

Now the punchline. The constrained posterior is a product of two distributions:

Products of arbitrary distributions are intractable, but a factorized distribution times a chain-structured one is again chain-structured — you simply reweight each emission factor of the HMM by the model's per-position probability, leaving the transitions untouched. The result is a standard (unnormalized) HMM, and everything you'd want — exact sampling, per-token marginals, normalization — falls out of forward–backward message passing in O(L(|\mathcal{S}|^2 + |\mathcal{E}||\mathcal{V}|)) FLOPs. No rejection, no approximation, guaranteed constraint satisfaction at every denoising step by construction.

Two immediate bonuses:

Constrained remasking confidence. dLLMs decide which positions to commit each step using a confidence score (LLaDA: highest probability; Dream: lowest entropy), normally computed from the unconstrained marginals. With the HMM in hand, you can compute per-token marginals under the constrained distribution and use those instead. This changes the decoding order itself, not just the tokens — and the ablation shows it matters (more below).

Arbitrary schedules and NFAs. Because sampling is exact at every step regardless of which positions are masked, the method works with parallel decoding, block-wise decoding, and any remasking schedule. Nondeterministic automata (NFAs) also work — the machinery is identical — with one caveat: p_\mathcal{M} then weights sequences by the number of accepting paths, so the sampler becomes a path-count-weighted proxy rather than exact. This matters for Spider (SQL), where the equivalent DFA would blow up exponentially and the NFA has ~10k states and ~100k edges.

The scan trick: log-depth exact sampling

There's a serious systems problem hiding here. Forward–backward and ancestral sampling are inherently sequential: z_i depends on z_{i-1}, so you make L tiny GPU calls per denoising step — each on the order of 10^6 FLOPs, which means runtime is dominated by kernel-launch overhead, not compute. And this repeats at every denoising step. The chain-based variant more than doubles total latency (+114% in their measurements).

The fix will feel familiar if you know parallel prefix sums or associative scans (the same idea that made linear state-space models fast). Conditioned on the two boundary states z_\ell and z_r of a segment, the tokens in the left half and right half are conditionally independent given the midpoint state z_m:

So instead of sampling z_1, z_2, \dots, z_L in sequence, sample the midpoint state first, then recurse into both halves in parallel. The midpoint conditionals p(z_m \mid z_\ell, z_r) \propto p(z_m \mid z_\ell)\,p(z_r \mid z_m) require multi-step transition matrices p(z_{i+2^k} \mid z_i), which are built bottom-up by repeated squaring — 2^k-step transitions from 2^{k-1}-step ones — in \log L rounds of batched matrix multiplication. Then a top-down pass samples midpoints level by level, and finally all tokens emit in parallel.

The trade is more FLOPs (matrix products cost an extra factor of |\mathcal{S}|) for exponentially fewer sequential steps — a bargain on GPUs where the chain variant was launch-bound anyway. Appendix B proves the tree sampler's output distribution is exactly the constrained posterior; this is not an approximation.

The evidence

The evaluation covers Dream-7B and LLaDA-8B on function calling (xLAM, BFCL), planning (4×4 Sudoku, Countdown), text-to-SQL (Spider), and math (GSM-Symbolic), under both greedy (T{=}0) and stochastic (T{=}1) decoding. Three findings stand out.

Stochastic sampling collapses unconstrained dLLMs; constraints restore it. At temperature 1, unconstrained Dream-7B falls to 22.3% on BFCL-Live JSON (from 63.9% greedy) — mostly because outputs stop satisfying the format at all (31.2% constraint satisfaction rate). Constrained decoding recovers 69.0%, essentially closing the gap to greedy. This is the headline result, and it's meaningful: it says dLLMs can't currently be sampled at nonzero temperature for structured tasks at all without something like this.

The format gap is structural, not semantic. Switching output format from JSON to Python syntax craters unconstrained Dream on BFCL-Live from 63.9% to 22.4% — same functions, same arguments, different surface syntax. Constrained decoding nearly erases the gap (71.5% JSON vs. 69.7% Python). The model knows what to call; it fumbles how to write it down. Appendix D's qualitative examples are telling: baseline failures include emitting prose instead of a call, PerPage vs. perPage casing, and hallucinated method names like todo.complete — all eliminated by construction. With constraints, Dream reaches 69.7% on BFCL-Live Python vs. 70.8% for Llama-3.1-8B, largely closing the diffusion-vs-autoregressive gap on this benchmark.

The tree sampler makes it nearly free. On 200 BFCL examples (A6000, Dream): unconstrained 67.5% accuracy at 24.3 s/sample; chain-based constrained 78.5% at +114% latency; tree-based constrained 79.2% at +4% latency. That last number is what makes this deployable rather than academic.

The ablation (Table 2) separates two mechanisms: constraining the samples ("Mf", keeping unconstrained confidence for remasking) versus also using constrained marginals as the remasking signal ("Mar"). Mar wins consistently, with the biggest gaps where the base model is weakest — Dream on xLAM-Python goes 56.1% (base) → 68.4% (Mf) → 76.4% (Mar). This is a nice piece of understanding: pure filtering can't rescue a model whose mean-field prediction is misaligned with the constraint; you also need the constraint to reorder which positions get committed first. There's also a robustness result worth noting: as denoising steps drop from 256 to 16 (more parallel commits per step), the unconstrained baseline deteriorates sharply while the constrained method holds up — constraints and aggressive parallelism are complementary, which matters since fast parallel decoding is the whole selling point of dLLMs.

What to be skeptical about

"Exact" has a precise, limited meaning. The method exactly samples the constrained mean-field posterior at each step — not p_{\mathrm{dLLM}}(x^0 \mid x^0 \in \mathcal{C}), the model's full generative distribution conditioned on the constraint. That target is intractable, and this per-step projection is the natural analogue of what autoregressive constrained decoding does (which is likewise a greedy approximation with known distortions). Fine, but don't read "exact" as "unbiased conditional sampling from the model."

NFA path-counting bias. For nondeterministic automata (used for Spider), sequences are weighted by their number of accepting paths, so the sampler tilts toward strings with more parses. Spider gains are correspondingly modest under greedy decoding (52.7 → 53.0 for Dream), though the sampling-mode recovery (15.5 → 52.1) is still large. Relatedly, GSM-Symbolic shows the format-guarantee ceiling: constraint satisfaction was already ~80% at greedy, and greedy accuracy barely moves (39.9 → 40.2). Constraints fix syntax, not reasoning — Appendix D's failure table (schema-valid but semantically wrong argument values) makes this explicit.

Finite automata only. No context-free grammars, so no full recursive JSON or unrestricted SQL — the Spider automaton encodes an FA-approximable restriction of SQL grammar over the given schema. The DFA-blowup problem is real (it's why Spider needs an NFA), and the O(|\mathcal{S}|^3)-flavored cost of the tree sampler's matrix squaring could bite for much larger automata; the biggest tested is ~19.5k states. Whether the "+4% overhead" figure survives at longer sequence lengths (tested at L \le 256) and larger constraint machines is exactly the right thing to watch.

Scale and scope. Two 7–8B open models, short generations, and Sudoku's DFA only enforces format and prefilled cells (not row/column validity — that would need an exponentially large automaton, another reminder of the expressivity ceiling).

Why it matters anyway

If diffusion LMs are going to serve real workloads — and the Mercury/Gemini-Diffusion/Seed-Diffusion push suggests vendors believe they will — they need the structured-output tooling autoregressive models take for granted. This paper supplies the first version of that tooling with hard guarantees, working sampling (not just greedy MAP), compatibility with arbitrary remasking schedules, and negligible overhead. The two ideas that make it work are individually classical — FA-as-HMM product constructions, and parallel-scan depth reduction — but their combination in the diffusion decoding loop is genuinely new, and the observation that constrained marginals should drive the remasking order is a contribution in its own right. It also quietly resolves a puzzle: dLLMs weren't bad at function calling because they don't understand tools; they were bad because nobody could keep their parallel sampler inside a schema.

Read Section 4.3 and Appendix B first — the log-depth sampler and its exactness proof are the technical heart, and the conditional-independence recursion is elegant enough to reconstruct from Figure 2 alone. Then Table 2 for the filtering-vs-reordering ablation, which is where the mechanistic understanding lives.