Exact method corrects bias in constrained text generation models

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

Machine LearningArtificial IntelligenceComputation and Language

Summary

Generating text with specific rules can be tricky for some AI models because the way they build sentences step-by-step can introduce subtle biases. The authors studied these biases in a type of masked diffusion language model, which fills in missing words repeatedly to create sentences that follow certain constraints. They found that previous methods caused a bias away from the model's true probabilities of valid word sequences. To fix this, the authors developed a new method called TWISTER that corrects the bias exactly, ensuring the generated text strictly follows the rules without skewing probabilities. This approach uses advanced sampling techniques to keep the outputs both valid and unbiased.

What this means in practice

Authors

Aditya Thimmaiah, Lara Marinov, Jayanth Srinivasa, Haris Vikalo, Junyi Jessy Li, Milos Gligoric

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.
View PDFOpen arXiv