Grammar-Aligned Decoding (2405.21047v1)

Abstract: 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.

• The paper introduces Grammar-Aligned Decoding (GAD) to enforce syntactic validity while preserving LLM probability distributions.
• It proposes the Adaptive Sampling with Approximate Expected Futures (ASAp) algorithm to overcome limitations of Grammar-Constrained Decoding.
• Empirical results in code generation and parsing demonstrate improved coherence and statistical fidelity over existing methods.

Grammar-Aligned Decoding

The paper "Grammar-Aligned Decoding" presents a detailed examination of the challenges and methodologies associated with ensuring syntactical correctness and probabilistic alignment in the outputs of LLMs, particularly in generating structured data such as code and formal grammars. The authors introduce a novel concept: Grammar-Aligned Decoding (GAD), which aims to simultaneously enforce grammatical correctness and preserve the inherent probability distribution learned by an LLM.

Key Contributions

1. Problem Identification: The authors identify a fundamental issue with current constrained decoding techniques, specifically Grammar-Constrained Decoding (GCD). While GCD can ensure outputs conform to a predefined grammar by masking invalid tokens, it distorts the original probability distribution of the LLM, often resulting in low-quality outputs which adhere to grammar but are semantically incoherent or improbable according to the model's training.
2. Introduction of Grammar-Aligned Decoding (GAD): GAD is introduced as a theoretical framework that seeks to decouple grammatical correctness from the probabilistic fidelity of LLM outputs. The authors formalize GAD as the task of sampling from a distribution that is proportional to the LLM's distribution but restricted to grammatically valid outputs. They emphasize that simply preserving the grammar is insufficient without maintaining the statistical properties of the underlying model.
3. Proposed Solution - Adaptive Sampling with Approximate Expected Futures (ASAp): The paper proposes ASAp as a new decoding algorithm to address the limitations of GCD. ASAp incrementally approximates the expected probability of future grammaticality, using samples to refine this approximation iteratively. Unlike GCD, ASAp aims to provide unbiased sampling with respect to the LLM's learned distribution, gradually converging closer to the actual LLM-driven probabilities through iterative sampling improvements.
4. Empirical Evaluation: The efficacy of ASAp is demonstrated through experiments on tasks such as code generation and constituency parsing. The results indicate that ASAp can improve the likelihood of generated outputs more effectively than existing GCD techniques, aligning closer to the LLM's probability distribution while still satisfying syntactic constraints.

Implications and Future Directions

The paper highlights significant implications for both theoretical and practical applications in AI:

• Theoretical Implications: By addressing the distribution distortion problem of GCD, GAD provides a framework for future research to enhance the fidelity of LLM outputs in structured tasks, particularly where formal semantics are crucial.
• Practical Implications: ASAp could improve applications in automated code synthesis, mathematical modeling, and any domain requiring syntactically correct but statistically probable text generation. It enhances the ability of LLMs to produce usable and semantically coherent outputs in structured environments, expanding their applicability in industry and academia.
• Future Research: The introduction of GAD and ASAp opens avenues for further research into improving convergence speed and efficiency. Optimization techniques such as beam search and more advanced sampling strategies might be explored to enhance ASAp’s practicality, especially in computationally intensive tasks.

In conclusion, the paper makes a substantial contribution to the field of natural language processing and structured text generation by addressing a core limitation in LLM decoding methods. By proposing a framework that ensures both grammatical correctness and distributional accuracy, the authors set the stage for developing more reliable and effective LLM applications.