Constrained Adaptive Rejection Sampling
Abstract
Language Models (LMs) are increasingly used in applications where generated outputs must satisfy strict semantic or syntactic constraints. Existing approaches to constrained generation fall along a spectrum: greedy constrained decoding methods enforce validity during decoding but distort the LM’s distribution, while rejection sampling (RS) preserves fidelity but wastes computation by discarding invalid outputs. Both extremes are problematic in domains such as program fuzzing, where both validity and diversity of samples are essential. We present Constrained Adaptive Rejection Sampling (CARS), an approach that strictly improves the sample-efficiency of RS without distributional distortion. CARS begins with unconstrained LM sampling and adaptively rules out constraint-violating continuations by recording them in a trie and subtracting their probability mass from future draws. This adaptive pruning ensures that prefixes proven invalid are never revisited, acceptance rates improve monotonically, and the resulting samples exactly follow the constrained distribution. In experiments on a variety of domains---e.g., program fuzzing and molecular generation---CARS consistently achieves higher efficiency---measured in the number of LM forward passes per valid sample---while also producing stronger sample diversity than both Greedy Constrained Decoding (GCD) and methods that approximate the LM's distribution.
Lay Summary
AI language models are increasingly asked to produce outputs that must obey strict rules---valid computer programs, well-formed data records, or chemically synthesizable molecules. Existing methods face a painful trade-off. One family steers the model away from invalid choices as it writes, but this subtly biases which valid outputs come out. The other lets the model write freely and discards invalid attempts, which wastes enormous amounts of computation: in some of our experiments, fewer than one in a hundred attempts were valid! We developed Constrained Adaptive Rejection Sampling (CARS), which captures the strengths of both approaches. Each time the model produces an invalid output, CARS remembers not just that specific mistake but every nearby mistake the rules would have caught, and excludes all of them from future attempts. As more samples are drawn, wasted effort shrinks rapidly, while the valid outputs the model produces remain mathematically identical to what an infinitely patient version of plain "try-and-discard" would have produced. This matters wherever both correctness and variety count---automated software testing, where diverse valid programs help expose bugs, and molecular discovery, where researchers need many chemically plausible candidates. CARS makes such applications dramatically faster without sacrificing the statistical guarantees that make their results trustworthy.