🤖 AI Summary
Large language models (LLMs) struggle to simultaneously ensure syntactic correctness and distributional fidelity when generating highly structured outputs (e.g., code, mathematical expressions, markup). While grammar-constrained decoding (GCD) guarantees syntactic compliance, it severely distorts the model’s original conditional distribution. This work formally defines the *grammar alignment* problem and proposes ASAP—a novel decoding framework that strictly enforces context-free grammar constraints while provably preserving the LLM’s conditional output distribution. ASAP integrates adaptive sampling, approximate expected future state estimation, grammar-guided over-approximation of prefix feasibility, and context-aware constraint propagation. Experiments on code generation and structured NLP tasks demonstrate that ASAP achieves significantly higher likelihood under the original LLM distribution compared to baselines, while maintaining 100% syntactic validity.
📝 Abstract
Large Language Models (LLMs) struggle with reliably generating highly structured outputs, such as program code, mathematical formulas, or well-formed markup. Constrained decoding approaches mitigate this problem by greedily restricting what tokens an LLM can output at each step to guarantee that the output matches a given constraint. Specifically, in grammar-constrained decoding (GCD), the LLM's output must follow a given grammar. In this paper, we demonstrate that GCD techniques (and in general constrained decoding techniques) can distort the LLM's distribution, leading to outputs that are grammatical but appear with likelihoods that are not proportional to the ones given by the LLM, and so ultimately are low-quality. We call the problem of aligning sampling with a grammar constraint, grammar-aligned decoding (GAD), and propose adaptive sampling with approximate expected futures (ASAp), a decoding algorithm that guarantees the output to be grammatical while provably producing outputs that match the conditional probability of the LLM's distribution conditioned on the given grammar constraint. Our algorithm uses prior sample outputs to soundly overapproximate the future grammaticality of different output prefixes. Our evaluation on code generation and structured NLP tasks shows how ASAp often produces outputs with higher likelihood (according to the LLM's distribution) than existing GCD techniques, while still enforcing the desired grammatical constraints.