Twist, Don't Tilt: Trajectory-Exact Constrained Decoding for Masked Diffusion Models

2026-09-28

AI Summary

This study addresses the trajectory bias problem in constrained decoding for masked diffusion models by proposing the TWISTER decoder. This method introduces automaton twisting into a Sequential Monte Carlo (SMC) framework for the first time, integrating finite-state automata with Feynman-Kac particle filtering to derive and correct step-level sampling tilts. It formally proves that the resulting target distribution is equivalent to an unbiased Doob h-transform path measure. The proposed approach enables exact computation and efficient sampling under regular language constraints, ensuring that generated sequences strictly satisfy syntactic structures while preserving the model’s original probability distribution. Consequently, this work significantly improves the statistical consistency of constrained generation.

Abstract

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.

Problem

• Masked Diffusion Language Models
• Constrained Decoding
• Trajectory Bias
• Sequential Monte Carlo

Innovation

• Masked Diffusion Language Models
• Constrained Decoding
• Sequential Monte Carlo
• Trajectory Bias
• Feynman-Kac Correction