Constrained Sampling for Language Models Should Be Easy: An MCMC Perspective

TL;DR

This paper introduces an MCMC-based constrained sampling method for language models that guarantees constraint satisfaction, converges to the true distribution, and is computationally efficient, improving diversity and validity in applications like program fuzzing.

Contribution

The paper presents a novel MCMC framework for constrained sampling in language models that ensures constraint satisfaction, convergence, and efficiency, addressing limitations of prior methods.

Findings

Outperforms existing constrained sampling methods on synthetic benchmarks.

Achieves higher diversity and validity in program fuzzing tasks.

Ensures principled exploration of constrained output space.

Abstract

Constrained decoding enables Language Models (LMs) to produce samples that provably satisfy hard constraints. However, existing constrained-decoding approaches often distort the underlying model distribution, a limitation that is especially problematic in applications like program fuzzing, where one wants to generate diverse and valid program inputs for testing purposes. We propose a new constrained sampling framework based on Markov Chain Monte Carlo (MCMC) that simultaneously satisfies three core desiderata: constraint satisfying (every sample satisfies the constraint), monotonically converging (the sampling process converges to the true conditional distribution), and efficient (high-quality samples emerge in few steps). Our method constructs a proposal distribution over valid outputs and applies a Metropolis-Hastings acceptance criterion based on the LM's likelihood, ensuring principled and efficient exploration of the constrained space. Empirically, our sampler outperforms existing methods on both synthetic benchmarks and real-world program fuzzing tasks.