Source-linked AI summary
Twist, Don't Tilt: Trajectory-Exact Constrained Decoding for Masked Diffusion Models
Aditya Thimmaiah, Lara Marinov, Jayanth Srinivasa, Haris Vikalo, Junyi Jessy Li, Milos Gligoric
TL;DR
MDLM constrained decoding seeks outputs that satisfy structural constraints, but exact constrained draws at each denoising step can compose into a biased trajectory distribution. The paper characterizes that tilt and introduces TWISTER, whose Feynman–Kac correction targets the unbiased constraint-conditioned path law.
Problem
Method
TWISTER uses the step-exact decoder as an SMC proposal and applies Feynman–Kac corrections based on local and clamped-partition sum ratios.
Source: §1 · p. 2, §6 · p. 10Results
Step-exact sampling is tilted by inverse re-freezing ratios, while Feynman–Kac twists correct that tilt to target the Doob path law.
Source: §3.4 · p. 7, §4.1 · p. 8Takeaways & Limitations
For regular-language constraints, the correction is exactly and efficiently computable using quantities already computed for step-exact sampling.
Source: §1 · p. 2The analysis assumes complete reveal at the terminal step and full support for denoiser predictions at masked positions.
Source: §3 · p. 4
Abstract
from arXiv · showhide
Constrained decoding for Masked Diffusion Language Models (MDLMs) aims to ensure that generated outputs satisfy a specified structure or syntax constraint. MDLMs generate outputs by repeatedly unmasking masked positions present in their current state. Recent strategies for constrained decoding constrain the model's per-step mean-field posterior (which factorizes over masked positions) by enforcing the desired constraint with an automaton. The resulting chain-structured factor graph allows exact constrained sampling via dynamic programming. However, despite each draw being exact and constraint-satisfying, we prove that their composition, in general, tilts away from the model's relative probabilities over valid trajectories, thus leading to trajectory bias. We derive an exact expression for this bias as a product of ratios measuring how valid continuation mass changes when the denoiser is reconditioned, and characterize when the bias vanishes. We then correct the bias by introducing TWISTER, the first automaton-twisted Sequential Monte Carlo decoder for MDLMs, using the step-exact decoder as the proposal. We show that for regular language constraints, the Feynman-Kac correction is exactly computable, with the twists obtained efficiently using quantities pre-computed for step-exact sampling. We prove that the resulting Feynman-Kac model targets the unbiased Doob h-transformed path law conditioned on constraint satisfaction.
1 INTRODUCTION
MDLM outputs may need to satisfy format constraints, but their denoising order differs from left-to-right decoding. Although step-exact constrained sampling is exact locally, composing those steps can bias trajectories; TWISTER corrects that bias.
1 INTRODUCTION
- MDLM outputs may require structural constraints, such as compliance with a JSON schema, despite denoising in an order unlike left-to-right decoding.Earlier approaches constrain the per-step mean-field posterior with an automaton; DINGO uses MAP decoding, while later work samples exactly at each step.Source: §1 · pp. 1–2
- Composing locally exact, constraint-satisfying MDLM steps can still produce trajectories biased relative to the globally constrained distribution.The bias arises when a scheduled subset of sampled tokens is retained and the denoiser is reconditioned for the next step.Source: §1 · p. 2
- TWISTER corrects trajectory bias with a Feynman–Kac correction using ratios of local and clamped-partition sums as incremental potentials.The needed quantities are already computed by FFBS in the step-exact sampler, making the correction exactly and efficiently computable.Source: §1 · p. 2
2 BACKGROUND
The background introduces notation, deterministic finite automata and regular languages, MDLM denoising, and twisted Sequential Monte Carlo. These concepts establish how constraints are represented and how trajectory distributions can be guided.
2 BACKGROUND
- The section establishes notation, including [N] for the natural numbers up to N and indicator notation J·K.It also uses t<T with t starting from 0 and generally denotes unnormalized quantities with a tilde.Source: §2 · p. 2
2.1 DETERMINISTIC FINITE AUTOMATA AND REGULAR LANGUAGES
- A deterministic finite automaton consists of an alphabet, states, an initial state, accepting states, and a transition function; its accepted words form a regular language.A word is accepted when the automaton ends in an accepting state after consuming it.Source: §2.1 · p. 2
2.2 MASKED DIFFUSION LANGUAGE MODELS
- MDLMs generate clean token sequences by repeatedly unmasking positions, with a denoiser predicting masked tokens and a scheduler selecting which predictions to reveal.The denoiser's predictions factorize over masked positions, and the decoder commits sampled tokens only at positions in the reveal set.Source: §2.2 · pp. 2–3
2.3 TWISTED SEQUENTIAL MONTE CARLO
- Sequential Monte Carlo approximates trajectory distributions with weighted particles, while twisted SMC moves terminal-weight information into incremental potentials to guide particles as trajectories grow.The ideal twist at a state is the expected terminal weight under remaining native transitions, or the exact continuation probability.Source: §2.3 · p. 3
3 STEP-EXACT DOES NOT MEAN TRAJECTORY-EXACT
Exact constraint satisfaction at each denoising step does not ensure unbiased sampling of complete valid trajectories. The section defines the unbiased target, characterizes step-exact decoding through frozen-chain lookaheads, and traces trajectory bias to re-freezing the denoiser.
3 STEP-EXACT DOES NOT MEAN TRAJECTORY-EXACT
- Step-exact draws can each satisfy the constraint yet compose into a biased path law because later denoiser reconditioning changes valid-continuation mass.Only a scheduled subset of sampled tokens is retained in each successor state; the rest are discarded.Source: §3 · p. 4, §3.3 · p. 6
3.1 DOOB-TRANSFORMED DECODER
- The unbiased target is the native trajectory law conditioned on a valid terminal state, which preserves relative probabilities among valid trajectories.Source: §3.1 · p. 4
3.2 AUTOMATON-CONSTRAINED STEP-EXACT DECODER
- The step-exact decoder constrains the denoiser’s factorized mean-field posterior with an automaton, yielding a chain-structured CRF that FFBS can sample exactly.FFBS also computes the local and clamped partition sums used to score valid completions.Source: §3.2 · p. 5
3.3 PARTITION SUMS ARE LOOKAHEADS OF A FROZEN DENOISER MARKOV CHAIN
- Step-exact kernels are Doob transforms of native kernels frozen at the current state, but re-freezing at each successor prevents their composition from being one such transform.Source: §3.3 · p. 6
3.4 TRAJECTORY BIAS
- The step-exact path law is tilted by the product of inverse re-freezing ratios, and it is unbiased exactly when that product is constant across valid trajectories.A ratio compares an intermediate state’s valid mass under the current-state denoiser with its mass under the predecessor-state denoiser.Source: §3.4 · pp. 6–7
- Trajectory bias vanishes for vacuous constraints and when all masked positions are revealed in one step.Source: §3.4 · p. 7
4 TWISTER: AUTOMATON-TWISTED SEQUENTIAL MONTE CARLO
TWISTER corrects the trajectory tilt of composed step-exact decoding with an automaton-twisted Sequential Monte Carlo procedure. For regular-language constraints, its Feynman–Kac correction is computed from partition sums already used for step-exact sampling.
4.1 FEYNMAN–KAC CORRECTOR FOR THE RE-FREEZING TILT
- TWISTER uses step-exact kernels as proposals and computes its twists from local partition sums already available through FFBS.The ideal lookahead requires future denoiser calls; the local partition sum is a tractable surrogate from one denoiser call.Source: §4.1 · p. 8
- The Feynman–Kac twists correct the re-freezing tilt, making the resulting model target the Doob path law.The correction is applied incrementally because the tilt factorizes across steps.Source: §4.1 · p. 8, §4 · p. 8
- Weighting xa1 by its re-freezing ratio 9/5 changes the step-exact choice 1/3:2/3 to the Doob choice 9/19:10/19.The example illustrates the correction of a step-exact choice toward the Doob choice.Source: §4.1 · p. 8
4.2 TWISTER ALGORITHM
- TWISTER initializes K weighted particles at x0, propagates them using cached FFBS messages, and reweights them with the exact Feynman–Kac potential.For nonterminal successors it queries the denoiser for the successor partition sum; when resampling is triggered, particles and caches are resampled together.Source: §4.2 · p. 8
4.3 TRAJECTORY BIAS EXPERIMENT
- The bias experiment compares constrained samples with rejection-sampled native trajectories using pairwise total variation distance and estimates a finite-sample noise floor.It evaluates six models at T ∈ {1, 2, 4, 8, 16} steps, averaging over token-class partitions and permutations.Source: §4.3 · p. 9
5 RELATED WORK
Prior MDLM constrained decoders enforce constraints at individual denoising steps, while this paper targets the globally constrained trajectory law. The work also distinguishes formal constraint enforcement from SMC methods for steering diffusion samples or improving sample quality.
Constrained Decoding for Autoregressive Models
- Autoregressive constrained decoders mask invalid next tokens, and SMC approaches add incremental reweighting or stronger automaton-based proposals to target globally constrained distributions.The paper focuses instead on constrained decoding for MDLMs, where step-exact decoding accounts for future valid mass under a denoiser frozen at the current state.Source: §5 · pp. 9–10
Constrained Decoding for Diffusion Models
- DINGO uses automaton dynamic programming for per-step MAP decoding, while Dang & Ermon (2026) samples exactly from the per-step automaton-constrained posterior.Both enforce constraint satisfaction at each step, but their local exactness does not preserve relative probabilities over multi-step trajectories.Source: §5 · p. 10
- Other diffusion SMC methods steer samples with temperature scaling or external rewards, or reweight MDLM trajectories by confidence; TWISTER uses SMC to enforce formal constraints.These approaches address different objectives from TWISTER’s constraint-enforcement goal.Source: §5 · p. 10
6 CONCLUSION
The paper shows that composing locally exact MDLM constrained steps can tilt probabilities over valid trajectories. It characterizes the tilt and introduces TWISTER to correct it incrementally for regular-language constraints.
6 CONCLUSION
- Composed local sampling from an MDLM’s automaton-constrained posterior can tilt the distribution over globally constrained trajectories despite being exact at each step.The tilt follows from rescoring valid completions under different denoiser conditionings.Source: §6 · p. 10
- The trajectory tilt factorizes across denoising steps as a product of local and clamped-partition sum ratios.These ratios arise when valid completions of the same intermediate states are rescored under different denoiser conditionings.Source: §6 · p. 10
- TWISTER uses the step-exact decoder as its proposal and exactly computes regular-language Feynman–Kac potentials from partition sums obtained during step-exact sampling.The correction targets the unbiased Doob path law conditioned on constraint satisfaction.Source: §6 · p. 10
A PROOF FOR STEP-EXACT KERNEL
The step-exact transition kernel is obtained by marginalizing the reveal set from the joint kernels, under a positivity condition on partition sums.
A PROOF FOR STEP-EXACT KERNEL
- The proposition states the step-exact kernel for native kernels when partition sums are positive on the native-kernel support.The condition is required for all states in the support of κt+1.Source: A PROOF FOR STEP-EXACT KERNEL · p. 14
- The step-exact transition kernel is obtained by marginalizing out the reveal set from the joint kernels.This identifies the marginalization used to obtain the transition kernel.Source: A PROOF FOR STEP-EXACT KERNEL · p. 14
B PROOF FOR FROZEN LOOKAHEAD
The frozen lookahead for the constraint-satisfying terminal law is expressed using categorical probabilities under the initial state’s denoiser.
B PROOF FOR FROZEN LOOKAHEAD
- The frozen lookahead is given as a product of categorical probabilities for masked positions under the initial state.The proposition frames this as the lookahead for the constraint-satisfying terminal law induced by frozen native kernels.Source: B PROOF FOR FROZEN LOOKAHEAD · p. 14
C PROOFS FOR TRAJECTORY-BIAS BOUNDARY CASES
The kernel-level exactness lemma characterizes when composing step-exact kernels yields a trajectory-exact path law.
C.1 KERNEL-LEVEL TRAJECTORY EXACTNESS
- Trajectory exactness is characterized by a per-step condition over reachable states with positive partition sums and lookahead.The lemma states the condition must hold at every time step and for each supported successor state.Source: §C.1 · p. 15
C.2 VACUOUS CONSTRAINTS
The paper identifies two boundary cases in which step-exact decoding has no trajectory bias: vacuous constraints and one-step reveal-all.
C.2 VACUOUS CONSTRAINTS
- A vacuous regular-language constraint produces no trajectory bias.The corollary states that its no-bias condition holds at every step.Source: §C.2 · p. 15
C.3 ONE-STEP REVEAL ALL
- With a one-step horizon that completely unmasks the initial realization, step-exact decoding has no trajectory bias.The corollary states that the no-bias condition holds for T = 1.Source: §C.3 · p. 15
D PROOF FOR TWIST FIXES THE TILT
The proof establishes that TWISTER’s Feynman–Kac twists make the step-exact decoder target the Doob path law. The accompanying algorithm uses step-exact proposals and re-freezing corrections, with adaptive resampling.
D PROOF FOR TWIST FIXES THE TILT
- The Feynman–Kac twists correct the re-freezing tilt, making the step-exact decoder path law target the Doob path law.The theorem states this target is the constraint-conditioned Doob path law.Source: D PROOF FOR TWIST FIXES THE TILT · p. 16
E TWISTER ALGORITHM (WITH ADAPTIVE RESAMPLING)
- TWISTER propagates particles with cached step-exact FFBS messages, then weights them using the exact potential from partition-sum ratios.It queries the denoiser at successor states to obtain successor partition sums; particles are resampled when effective sample size falls below the threshold.Source: E TWISTER ALGORITHM (WITH ADAPTIVE… · p. 16
RE-FREEZING CORRECTION
- The step-exact propagation and re-freezing correction phases are parallelizable across conditionally independent particles until weight normalization and resampling.Source: RE-FREEZING CORRECTION · p. 17
F CONSTRAINT SATISFACTION EXPERIMENT
The constraint-satisfaction experiment evaluates decoding strategies on 94 JSON-Mode-Eval datapoints and reports token-prefix DFA sizes for their schemas. It compares four strategies across six open MDLMs, measuring parse validity, schema validity, and generation time.
F CONSTRAINT SATISFACTION EXPERIMENT
- The 94 evaluated schemas produced token-prefix DFAs with a mean of 188 states, most between roughly 100 and 250, and a range of 34–625.Six of the dataset's 100 datapoints were excluded because their schemas could not be compiled into regular expressions by Outlines.Source: F CONSTRAINT SATISFACTION EXPERIMENT · p. 17
- The comparison includes DINGO† mode selection, Dang & Ermon† stepwise sampling, and TWISTER with K =1 or K =4 particles.DINGO† selects a locally optimal constraint-satisfying block; Dang & Ermon† samples the automaton-constrained mean-field posterior with FFBS.Source: F CONSTRAINT SATISFACTION EXPERIMENT · pp. 17–18
- Evaluation uses six models from the DREAM, DREAMCODER, and LLaDA families, with generation length and denoising steps set to 1.5× the gold answer length.The models are DREAM-7B-BASE, DREAM-7B-INST, DREAMCODER-7B-BASE, DREAMCODER-7B-INST, LLADA-8B-BASE, and LLADA-8B-INST.Source: F CONSTRAINT SATISFACTION EXPERIMENT · p. 18
- Table 1 reports parse validity, schema validity, and average generation time for the decoding methods on JSON-Mode-Eval.Parse validity measures whether outputs are valid JSON; schema validity measures conformity to the datapoint's expected schema, and time is reported in seconds.Source: F CONSTRAINT SATISFACTION EXPERIMENT · pp. 18–19, Table 1 · p. 19