Constrained Decoding for Diffusion Language Models via Efficient Inference over Finite Automata
This work solves the problem of enforcing structured constraints (e.g., JSON schemas) in diffusion language models, which previously lacked efficient constrained decoding methods, enabling their use in practical applications.
The paper introduces an exact and tractable algorithm for constrained decoding in diffusion language models under any constraint expressible as a finite automaton, achieving substantial accuracy gains across tasks like function calling, planning, text-to-SQL, and math reasoning with minimal inference overhead (e.g., improving Dream-7B's greedy decoding accuracy from 63.9% to 71.5% on BFCL-Live with under 5% wall-clock overhead).
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 language models, 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.