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

2026-07-08

AI Summary

This work addresses the lack of constraint decoding methods tailored to the parallel denoising nature of diffusion language models, as existing autoregressive techniques are ill-suited for direct application. The authors propose the first exact constrained decoding algorithm for such models, which encodes structural constraints—e.g., JSON format—as finite automata and performs posterior-constrained sampling under the mean-field distribution at each denoising step. By leveraging deep compression via arithmetic circuits, the method reduces sampling depth from linear to logarithmic, enabling compatibility with both parallel and block-wise decoding. Supporting any structure expressible by a finite automaton, the approach significantly improves performance on Dream-7B and LLaDA-8B: on BFCL-Live, greedy decoding accuracy rises from 63.9% to 71.5%, and random sampling accuracy jumps from 22.3% to 69.0%, with less than 5% additional inference overhead.

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 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.

Problem

• Constrained Decoding
• Diffusion Language Models
• Finite Automata
• Structured Generation
• Non-autoregressive Generation

Innovation

• constrained decoding
• diffusion language models
• finite automata
• mean-field posterior
• depth reduction