arXiv preprint ·

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

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

Published
Source
arXiv preprint, posted before peer review
Original
Read it at arXiv
Topics
Machine LearningArtificial IntelligenceComputation & Language / NLP

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.

The authors’ abstract, as arXiv lists it. arXiv’s listings, abstracts included, are free for anyone to reuse.

Read it on ReadSearch

Each opens after you create an account, and brings you back to this paper.

Already a reader? Sign in

All new papers in Machine Learning

Get the few papers worth your time in Machine Learning every morning.

ReadSearch reads the new papers in your field each day and brings you the few worth reading, each with its key finding in a sentence.

Papers belong to their authors and publishers; every page here links to the original. All research topics