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

Abstract: Constrained decoding is essential for serving LLMs, ensuring that generated outputs follow specific structures such as JSON schema-formatted function calls. Existing systems are designed for autoregressive models and assume left-to-right generation, masking out invalid next tokens at each step. Diffusion LLMs, however, break this assumption: they sample multiple positions simultaneously from a fully-factorized mean-field distribution at each denoising step. In this paper, we present an exact and tractable algorithm for sampling from the constrained mean-field posterior under any constraint expressible as a finite automaton. Viewing finite automata as graphical models, we obtain tractable representations of the constrained distribution that enable efficient inference. The approach guarantees constraint satisfaction by construction, supports both greedy and sampling-based decoding, and is compatible with parallel and block-wise decoding under arbitrary remasking schedules. Applying depth-reduction techniques from arithmetic circuit theory, we further reduce sampling depth from linear to logarithmic in the sequence length. Empirical evaluations on Dream-7B and LLaDA-8B show substantial accuracy gains across various tasks including function calling (xLAM, BFCL), planning (Sudoku, Countdown), text-to-SQL (Spider), and math reasoning (GSM-Symbolic), with little inference overhead relative to unconstrained decoding. For example, on BFCL-Live, our approach improves Dream-7B's greedy decoding accuracy from 63.9% to 71.5%, and stochastic sampling accuracy from 22.3% to 69.0%, where the unconstrained baseline collapses, with under 5% wall-clock overhead.

• The paper introduces an exact and parallel constrained decoding method for diffusion language models that guarantees 100% constraint satisfaction.
• It leverages a graphical model interpretation of finite automata to reduce decoding depth from O(L) to O(log L), enhancing parallel efficiency.
• Experiments on tasks like function calling and Sudoku reveal significant accuracy gains with minimal computational overhead.

Constrained Decoding for Diffusion LLMs: Inference over Finite Automata

Motivation and Problem Statement

Diffusion LLMs (dLLMs) have recently emerged as a compelling alternative to autoregressive LLMs, supporting efficient, parallel, and bidirectional text generation. A significant bottleneck in deploying LLM-based systems for real-world applications—such as function calling, program synthesis, and structured database queries—is the necessity for generated outputs to conform to complex, user- or system-specified formal constraints. Existing constrained decoding methods are tailored for autoregressive, left-to-right generation by filtering out invalid next-token choices based on context; however, these approaches do not directly extend to the mean-field, simultaneous multi-token sampling paradigm of dLLMs, where the independence assumption at each step precludes position-wise token masking for constraint enforcement.

The paper "Constrained Decoding for Diffusion LLMs via Efficient Inference over Finite Automata" (2607.07026) introduces an exact, tractable, and parallelizable solution for constrained decoding in dLLMs for any constraint expressible by a finite automaton—including both deterministic (DFA) and nondeterministic (NFA) forms. The approach leverages a graphical model interpretation of automata, enabling efficient methods—rooted in arithmetic circuit theory—to overcome the standard linear-time sequential bottlenecks.

Methodology

Graphical Model Representation of Constraints

Any regular language constraint $\mathcal{C}$ can be encoded by a finite automaton $\mathcal{M}$ with states $\mathcal{S}$, edges $\mathcal{E}$, and a vocabulary $\mathcal{V}$. By viewing $\mathcal{M}$ as a hidden Markov model (HMM), the paper constructs a structured graphical model where a length-$L$ sequence $x_{1:L}$ corresponds to a path through the automaton with global satisfaction of $\mathcal{C}$ as support.

In this representation, sampling from $p_\theta(x^0 \mid x^t, \mathcal{C})$ at each denoising step in the diffusion process involves generating sequences from the product of the model's factorized mean-field distribution and the automaton's support indicator. This constructs a new chain-structured graphical model amenable to exact inference.

Tractable Sampling with Efficient Parallelization

Naïve ancestral sampling over the automaton-constrained graphical model requires sequential forward–backward message passing, scaling linearly in sequence length $\mathcal{M}$0. The major algorithmic contribution is the adaptation of depth-reduction techniques from arithmetic circuit theory, transforming the chain into a binary recursion tree over segment boundaries. This modification reduces the decoding depth from $\mathcal{M}$1 to $\mathcal{M}$2, greatly enhancing parallelism on modern hardware. Global normalization is maintained by postponing factor normalization until the full joint distribution is established, rather than local (conditional) normalization, and the induced recursion ensures statistical correctness.

(Figure 1)

Figure 1: Chain-structured graphical model and its tree-structured equivalent, enabling $\mathcal{M}$3 parallel constrained sampling through divide-and-conquer message passing and sampling.

Marginal-Based Remasking and Confidence

In unconstrained dLLMs, denoising steps use model-assigned tokenwise confidence (entropy or probability) to decide which positions to commit. The proposed method adjusts this procedure, using marginals derived from the constrained posterior as the remasking confidence, further enhancing accuracy and stability—especially when the model’s unconstrained mean-field is inconsistent with valid outputs.

Experimental Results

The evaluation encompasses two open-weight dLLMs, Dream-7B and LLaDA-8B, across a spectrum of structured generation tasks:

• Function calling (xLAM, BFCL): Output must adhere to JSON or Python schema formats.
• Planning (Sudoku, Countdown): Global output rules, e.g., fixed cells or valid steps.
• Text-to-SQL (Spider): Query generation validated against provided DB schema.
• Math reasoning (GSM-Symbolic): Outputs must parse as symbolic mathematical expressions.

Notably, the approach always guarantees $\mathcal{M}$4 constraint satisfaction, eliminating any syntactic or schema-violating output. Strong numerical improvements are observed. For instance, on the BFCL-Live function calling task (Dream-7B):

• Greedy decoding accuracy: rises from 63.9% (unconstrained) to 71.5% (constrained).
• Stochastic decoding accuracy: surges from 22.3% (baseline, which often collapses) to 69.0%, closely matching greedy performance.

(Figure 2)

Figure 2: Accuracy vs. denoising steps on BFCL Simple split. The constrained method maintains robust performance even as the number of diffusion steps decreases, whereas the unconstrained baseline degrades sharply.

For Sudoku, as the number of prefilled cells decreases (i.e., harder puzzles), constrained decoding dramatically outperforms unconstrained baselines, particularly mitigating format or task-structure errors that would otherwise be catastrophic.

Figure 3: Sudoku 4$\mathcal{M}$54 accuracy versus the number of prefilled digits; constrained decoding remains high-performing on hard instances.

Runtime analysis demonstrates that the log-depth sampler adds little wall-clock overhead (typically $\mathcal{M}$6) compared to the unconstrained baseline, while massively improving constraint adherence and accuracy. Moreover, remasking based on constrained marginals yields additional gains over unconstrained confidence.

Theoretical and Practical Implications

The work provides a formal, unifying algorithmic mechanism for constrained generation in dLLMs, matching the expressive power and reliability of autoregressive constrained decoding for automaton-expressible constraints. Theoretically, the proposal generalizes beyond prior works like DINGO by supporting arbitrary masking strategies, blockwise/parallel decoding, sampling as well as greedy decoding, and nondeterministic automata.

Practically, the methodology is highly compatible with production workloads that demand structural guarantees (e.g., API calls, code, data pipelines, browser actions), mitigating the unacceptable rates of schema breakage otherwise observed in unconstrained sampling—especially under temperature or stochasticity for robustness and diversity. The parallelization strategy is particularly significant for scaling inference efficiency to large batch sizes and long sequences on modern accelerators.

Future Directions

While the algorithm subsumes all regular constraints, further research is warranted for context-free grammar decoding, which is critical for more expressive languages (such as full programming languages or deeply nested data formats) where finite-state representations are exponentially large or intractable. Moreover, the approach could be extended to dynamic/learned constraints beyond static automata and integrated into safety-critical workflows requiring high-precision structured outputs. There are also open questions about combining this layer of constraint enforcement with other post-hoc and pre-training alignment interventions.

Conclusion

This paper rigorously bridges a key algorithmic gap in the deployment of diffusion LLMs for constrained structured generation. By viewing constraints as graphical models and implementing globally normalized, efficient inference and sampling algorithms, the proposed method ensures exact constraint satisfaction with negligible computational cost and substantial empirical gains, delivering a step change in the reliability and applicability of dLLMs to structured-generation tasks where output validity is non-negotiable.