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

Constrained decoding is a critical requirement for deploying Large Language Models (LLMs) in structured environments, such as generating valid JSON for function calls or syntactically correct SQL. While autoregressive models handle this by masking invalid tokens during left-to-right generation, diffusion language models (dLLMs) present a unique challenge: they sample tokens simultaneously from a fully-factorized distribution, meaning local token validity does not guarantee global sequence validity.

To solve this, the authors treat the constraint as a finite automaton and map it to a chain-structured graphical model. By reweighting the model's mean-field predictions with the automaton's transition factors, they create a tractable posterior distribution that ensures all generated sequences satisfy the constraint by construction. To overcome the sequential bottleneck of standard message passing, they employ depth-reduction techniques from arithmetic circuit theory, transforming the inference process into a divide-and-conquer tree structure that reduces sampling depth from linear to logarithmic in sequence length.

The proposed method provides a robust way to enforce hard constraints in dLLMs without sacrificing the benefits of parallel decoding. Empirical evaluations on models like Dream-7B and LLaDA-8B demonstrate that this approach significantly improves performance across diverse tasks, including function calling, planning (Sudoku), and math reasoning. Notably, on the BFCL-Live benchmark, the method improved greedy decoding accuracy from 63.9% to 71.5% and, more impressively, boosted stochastic sampling accuracy from 22.3% to 69.0%—a scenario where unconstrained models typically fail. The entire process incurs less than 5% wall-clock overhead, making it highly efficient for production use.

As diffusion-based language models continue to challenge the dominance of autoregressive architectures, the ability to integrate them into existing software pipelines is essential. This work bridges a major gap by providing a mathematically rigorous and computationally efficient way to force dLLMs to adhere to strict output formats. By enabling logarithmic-depth sampling, the authors ensure that these models can be deployed in latency-sensitive applications that require high structural precision.