Source-linked AI summary
Constrained Sampling for Language Models Should Be Easy: An MCMC Perspective
Emmanuel Anaya Gonzalez, Sairam Vaidya, Kanghee Park, Ruyi Ji, Taylor Berg-Kirkpatrick, Loris D'Antoni
TL;DR
Existing constrained decoding can enforce validity but may distort the LM distribution, which matters for applications such as program fuzzing that require diverse valid inputs. The paper uses valid-output proposals with likelihood-aware Metropolis-Hastings transitions to address this problem. Its MCMC samplers improve KL divergence and fuzzing code coverage while satisfying the stated convergence properties.
Problem
Existing constrained-decoding methods often distort the LM distribution, limiting faithful generation of diverse valid inputs for program fuzzing.
Method
The framework proposes only constraint-satisfying outputs and applies a Metropolis-Hastings acceptance criterion guided by the LM likelihood.
Results
MCMC variants reduce KL divergence more than GCD and ASAp after 10 steps, with geomean reductions up to 5.07× versus GCD and 8.70× versus ASAp.
Takeaways & Limitations
The framework provides constrained samples that preserve the target distribution asymptotically and improve practical fuzzing outcomes where diverse inputs are important.
Takeaways & Limitations
Existing constrained decoding and ASAp provide important comparison limitations: GCD distorts the target distribution, while ASAp can require thousands of samples to converge [34].
Abstract
from arXiv · showhide
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.
1 Introduction
Existing constrained decoding can enforce hard constraints but often distorts the LM’s distribution, undermining applications that need diverse valid samples. The paper introduces MCMC-based constrained sampling designed to satisfy constraints, converge monotonically to the conditional distribution, and work efficiently.
- Hard-constrained decoding can distort the LM’s learned distribution, degrading downstream performance.
- Program fuzzing requires syntactically valid and distributionally diverse seeds that explore different execution paths.
- The paper asks whether constrained samplers can satisfy constraints, converge monotonically to the true conditional distribution, and estimate it efficiently in finite time.
- The proposed MCMC framework generates valid proposals and uses LM-likelihood-guided Metropolis-Hastings acceptance to achieve constraint satisfaction and monotonic convergence.
- Across synthetic and real-world fuzzing tasks, the method yields better KL divergence and higher code coverage than existing approaches.
2 The Ideal Properties of Constrained Sampling
Constrained sampling seeks valid sequences from a grammar while preserving the LM’s conditional distribution, but rejection sampling and constrained decoding have important efficiency or faithfulness limitations. The paper formalizes three desired properties and proposes an MCMC approach that provably guarantees the first two and is practically efficient.
- Grammar-Aligned Sampling: Grammar-aligned sampling targets the LM distribution restricted to sequences in the grammar’s language, including CFG-defined valid programs.
- Rejection Sampling: Rejection sampling preserves the conditioned distribution but may require unboundedly many rejected outputs when valid patterns are unlikely under the LM.
- Constrained Decoding: Constrained decoding masks invalid next tokens but does not preserve the LM’s underlying distribution [34].
- Limitations of Existing Approaches: GCD’s empirical KL divergence from the target distribution does not decrease with more samples on the SQLite fuzzing problem.
- Limitations of Existing Approaches: ASAp can eventually reach the target distribution, but its intermediate approximations are non-monotonic and may require thousands of samples to converge [34].
- Desired Properties of a Grammar-Aligned Sampler: A good sampler should satisfy constraint satisfaction, monotonic convergence, and efficiency; the proposed MCMC approach provably guarantees the first two and is practically efficient.
3 Grammar-Aligned MCMC Sampling
The framework uses grammar-constrained proposals within a Metropolis-Hastings chain, so every candidate is valid while likelihood-based acceptance steers sampling toward the constrained distribution. Its guarantees depend on proposal conditions, and different prefix-selection strategies trade off local refinement, diversity, and convergence speed.
- Constrained generation: Metropolis-Hastings accepts or rejects each valid proposal using target and proposal likelihood ratios, preserving the desired stationary distribution.The shared normalization factor in the constrained target cancels, enabling efficient acceptance-probability computation.
- Proposal distributions: Proposal choices trade off simplicity, local structure, diversity, and mixing speed, with uniform, priority, and restart strategies offering distinct behaviors.Priority resamples higher-uncertainty positions, uniform sampling preserves partial structure, and restart generates a new grammar-constrained sample from scratch.
- Proposal distributions: The proposal mechanism samples a prefix and regenerates its suffix with grammar-constrained decoding, enabling structured local edits while retaining constraint satisfaction.The prefix position is drawn from a configurable distribution over sequence positions.
- Constrained generation: Every proposal is generated by grammar-constrained decoding, ensuring that all sampled candidates satisfy the constraint by construction.The chain starts from a valid sample and repeatedly proposes valid continuations from prefixes of the current sequence.
- Proposal distributions: The sampler satisfies constraint satisfaction and monotonic convergence when the truncation distribution assigns nonzero probability to the empty prefix.This condition allows the chain to meet the stated theoretical guarantees for grammar-aligned sampling.
4 Experiments
Experiments evaluate MCMC against GCD and ASAp on synthetic grammar-constrained benchmarks and XML/SQL fuzzing. MCMC monotonically approaches the target distribution and produces seeds that improve branch coverage, with MCMC-Priority achieving the strongest reported fuzzing results.
- 4.1 Benchmarks by Park et al. [34]: The evaluation compares MCMC with GCD and ASAp on SLIA, BV4, and constituency-parsing grammar benchmarks, using up to 10 sampling steps.The benchmark suite contains 15 SLIA, 15 BV4, and 6 constituency-parsing problems.
- 4.1 Benchmarks by Park et al. [34]: MCMC variants reduce KL divergence monotonically with sampling steps, whereas ASAp converges only in the limit and lacks monotonic convergence.The trend holds across variants, with randomness causing some fluctuation.
- 4.1 Benchmarks by Park et al. [34]: At 10 steps, MCMC reduces KL divergence versus GCD by 2.11×–5.07× and versus ASAp by 3.99×–8.70× across its variants.The respective geomean reductions are 2.11×, 2.42×, and 5.07× versus GCD, and 3.99×, 5.08×, and 8.70× versus ASAp.
- 4.2 Fuzzing experiments: In SQL fuzzing, KL divergence decreases and converges for all MCMC variants, while coverage rankings inversely mirror KL divergence rankings, with overall differences relatively minor.The coverage comparison reports MCMC at 10 steps and measures branch coverage over time.
- 4.2 Fuzzing experiments: MCMC-Priority with 10 steps achieves 33.70% SQL and 12.81% XML coverage, exceeding GCD by 1.2× and 1.15× and ASAp by 1.12× and 1.15×, respectively.MCMC-based seeds significantly outperform Grammarinator, GCD, and ASAp; 5-step variants are lower than 10-step variants but generally above GCD.
- 4.2 Fuzzing experiments: The results support that MCMC convergence to the desired distribution is associated with higher fuzzing coverage, while MCMC-Priority may better explore variants of the same output.This interpretation connects distributional convergence with branch exploration in the reported fuzzing experiments.
5 Related Work
Related work includes constrained decoding, gradient-based sampling, constraint-aligned decoding, and controlled generation. These approaches vary in constraint representation and distributional fidelity, while ASAp targets the constrained distribution but converges slowly in practice.
- Constrained Decoding: Constrained decoding enforces syntactic or lexical restrictions using regular languages or CFGs but can distort the LM distribution.The cited literature modifies token-by-token decoding to enforce constraints.
- Gradient-based sampling: Gradient-based methods guide generation with relaxed differentiable surrogates and suit soft or semantic constraints, but remain inefficient and do not guarantee valid outputs.
- Constraint-Aligned Decoding: Top-k, nucleus, and beam search are unreliable estimators of the constrained distribution when hard constraints are introduced.This motivates methods designed to sample multiple outputs while approximating the LM distribution.
- Constraint-Aligned Decoding: ASAp approximates the true constrained distribution but may produce skewed intermediate estimates and can require thousands of samples before converging.Its limiting behavior does not provide monotonic convergence during sampling.
- Controlled generation: GeLaTo and Ctrl-G combine autoregressive LMs with HMMs for constraints representable as deterministic finite automata.Given a prefix, an HMM estimates the probability that a generated suffix will satisfy the DFA constraint.
6 Conclusion
The paper concludes that its MCMC framework samples directly within the constrained space while asymptotically preserving the LM-defined target distribution. It reports practical fuzzing improvements but acknowledges narrow experimental coverage and opportunities for richer proposal mechanisms.
- MCMC directly samples constrained outputs and asymptotically preserves the LM-defined target distribution, unlike approaches limited by slow convergence or rejection inefficiency.
- The framework improves applications requiring diverse valid inputs, especially program fuzzing where higher-quality diverse samples translate into higher code coverage.
- Limitations: The experiments use one model, two fuzzed projects, one prompt choice, 1-hour trials, and only three proposal distributions.The authors identify richer proposal mechanisms and combinations with beam or nucleus sampling as open directions.
D.1 Benchmarks
The benchmarks evaluate constrained sampling on grammar-intensive XML and SQL targets, using standardized prompts, grammar-derived inputs, and controlled fuzzing measurements.
- Each benchmark uses a fixed in-context prompt pairing a specification with a solution before generating new constrained solutions.
- Fuzzing experiments use five independent one-hour AFL++ runs for each benchmark and method, with other parameters held at defaults.
- Coverage is measured with LLVM instrumentation, which adds no more than 2% runtime overhead.
- Rejection sampling produces valid outputs at rates below 1% under the benchmark grammar constraints and is therefore excluded from comparison.
D.8 Ablation Study
The ablation study examines how MCMC proposal families and chain length affect convergence and fuzzing coverage across XML and SQL benchmarks.
- Experimental design: The ablation isolates proposal family and step count by plotting Priority, Restart, and Uniform variants separately against GCD.
- Number of steps: Coverage grows monotonically with the number of MCMC steps, but k=5 already captures at least 95% of the gain achieved by k=10 on both benchmarks.
- Proposal families: Even the weakest MCMC setting, k=2, surpasses GCD’s final coverage by 4–5%.
- XML benchmark: The XML ablation includes KL divergence and branch-coverage analyses, with results summarized over repeated runs and confidence intervals.
- Additional metrics: Additional analyses measure line and function coverage using the same LLVM instrumentation.
D.10 Overall Coverage Results
The MCMC sampler’s branch-coverage gains extend consistently to function and line coverage across both SQL and XML benchmarks.
- Improvements in branch coverage consistently translate to higher function and line coverage across SQL and XML.
E Properties and Proofs
The framework is proven to produce grammar-valid samples and to converge monotonically toward the target constrained distribution under stated conditions.
- Constraint Satisfying: Algorithm 1 always outputs a sequence in the constrained language L(G), for any language model, grammar, chain length, and truncation distribution.
- Proof strategy: The proof applies a general Markov-chain convergence theorem requiring irreducibility, aperiodicity, and a stationary distribution.
- Monotonically Converging: Under the stated positive-proposal condition, the output distribution converges monotonically to P_G as chain length k approaches infinity.
- Monotonically Converging: The convergence proof verifies countable states, irreducibility, aperiodicity, and the target distribution’s stationarity under Metropolis-Hastings.
- Monotonically Converging: Metropolis-Hastings makes the constrained language-model distribution P_G stationary for the constructed Markov chain.
F Benchmarks by Park et al. [34]
The benchmark evaluates how closely Grammar-Aligned MCMC Sampling approaches the target distribution across tasks and proposal variants. It compares MCMC with GCD and ASAp using KL divergence, convergence over increasing steps, and XML coverage.
- Table 3 reports XML coverage means with 95% confidence intervals over five trials and relative branch-coverage gains versus GCD for k ∈{2, 5, 10}.
- Figures 13–14 compare KL(P_G∥MCMC-T(10)) with GCD and ASAp(10) across benchmark tasks, where below-diagonal points indicate better MCMC approximation.The comparisons cover Uniform, Priority, and Restart proposal distributions.
- Figures 15–17 track KL divergence for ASAp(k) and MCMC-T(k) as the number of steps increases across SLIA, BV4, and CP subsets.Lower KL divergence indicates a better approximation to P_G.
- KL estimates use 100 samples for GCD, MCMC, and ASAp, while P_G uses all samples collected during the corresponding runs; convergence plots show 95% bootstrap confidence bands.