Grammar-Aligned Decoding

TL;DR

This paper introduces Grammar-Aligned Decoding (GAD), an adaptive sampling method that ensures outputs are both grammatical and probabilistically aligned with the LLM's distribution, improving structured output quality.

Contribution

The paper proposes the ASA p algorithm for grammar-aligned decoding, addressing distribution distortion in constrained decoding methods and ensuring outputs match the LLM's likelihoods under grammatical constraints.

Findings

ASAp often produces higher likelihood outputs than existing methods.

ASAp guarantees grammatical outputs while aligning with the LLM's distribution.

Evaluation on code and NLP tasks shows improved output quality.

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.