Improving Grammar Constraints Generation Alignment by Sampling Highly Probable Playouts

Large Language Models (LLMs) can generate structured outputs such as program code, mathematical formulas, or well-formed markup language, but some of their outputs may end up non-grammatical, especially when considering very precise and obscure structures. Grammar Constrained Decoding (GCD) is used to force each subsequent token generated to follow a particular grammar. However, simple techniques use a distorted distribution of the LLMs. Grammar Aligned Decoding (GAD) is the problem of aligning sampling with a grammar constraint. Adaptive sampling with approximate expected futures (ASAp) has been proposed to solve GAD but requires storing the change in the probabilities of the LLM for each token in a datastructure. Furthermore, ASAp samples randomly from the aligned distribution as it is computed, which may not be optimal to maximize the alignment speed. We propose Greedy Best First Search with Greedy Sampling (GBFSGS) as a method to obtain high likelihood samples, which is more likely to maximize the alignment speed of the distribution. GBFSGS is able to solve GAD while storing the new probabilities in a datastructure with one element per sample drawn instead of per token changed. We show how the method finds the most probable samples of a distribution, including under a grammar constraint. We show how to implement the algorithm to avoid duplicate samples. We empirically evaluate our method on three sets of problems to evaluate the behavior of the methods. We show GBFSGS is able to select better samples than sampling to improve the alignment of the distribution while reducing the memory needed to store the aligned distribution.