Source-linked AI summary
Efficient Guided Generation for Large Language Models
Brandon T. Willard, Rémi Louf
TL;DR
Guided LLM generation must produce text that conforms to regular expressions or CFGs without the scaling costs of checking the full vocabulary at every step. The paper reformulates generation as finite-state transitions and builds vocabulary indices for efficient masking, extending the approach to CFGs and LALR(1) parsers. The resulting framework supports structured generation with O(1) average indexing cost and reports indices of around 50 MB in a Python-grammar test, while trading processing time for memory.
Problem
Guided generation must enforce regular-expression or CFG constraints, but existing approaches can have scaling costs and require O(N) vocabulary checks per generated token.
Method
The paper uses FSM formulations and vocabulary indices to determine valid tokens efficiently, extending the approach to CFGs and LALR(1) parsers.
Results
Indices constructed for an augmented Python grammar were around 50 MB, while the indexing algorithm costs O(1) on average.
Takeaways & Limitations
The approach enables structured guided generation for regular expressions, CFGs, popular data formats, and programming languages, and may also assist LLM training or fine-tuning when structured outputs are required.
Takeaways & Limitations
Vocabulary indexing trades processing for memory, and its memory costs may require conventional reductions when they are relatively high.
Abstract
from arXiv · showhide
In this article we show how the problem of neural text generation can be constructively reformulated in terms of transitions between the states of a finite-state machine. This framework leads to an efficient approach to guiding text generation with regular expressions and context-free grammars by allowing the construction of an index over a language model's vocabulary. The approach is model agnostic, allows one to enforce domain-specific knowledge and constraints, and enables the construction of reliable interfaces by guaranteeing the structure of the generated text. It adds little overhead to the token sequence generation process and significantly outperforms existing solutions. An implementation is provided in the open source Python library Outlines
1 Introduction
Guided LLM generation targets outputs conforming to regular expressions or CFGs, but existing methods can scale poorly. The paper introduces vocabulary indexing from FSMs to reduce per-token constraint checking and extend efficient guidance to CFGs and LALR(1) parsers.
- Guided generation makes LLM outputs usable under rigid formatting requirements that fine-tuning alone can make difficult or costly to capture.
- Repeatedly checking the entire vocabulary imposes a fixed O(N) cost for every generated token.N denotes the LLM vocabulary size.
- FSM-based vocabulary indexing retrieves nonzero-probability tokens with O(1) average cost while allowing guided generation to start and stop arbitrarily.
- The regular-expression approach avoids requiring a complete transducer abstraction and can extend efficient regex libraries without modifying their underlying automatons.
- The indexing approach extends to CFGs and LALR(1) parsers for guided generation over formats and languages including JSON, Python, and SQL.
2 LLM Sampling and Guided Generation
LLM sampling recursively selects tokens from next-token distributions, while guided generation masks tokens so sequences satisfy regular expressions or specified grammars. Computing these masks is costly because valid tokens depend on the previously sampled prefix, motivating efficient incomplete-string matching and parsing.
- The LLM maps a token sequence and trained parameters to a next-token probability distribution over its vocabulary.The method extends to any function that takes token sequences and returns a next-token probability distribution.
- Multinomial sampling recursively draws tokens from the categorical next-token distribution until EOS is found.
- A boolean mask restricts the support of the next-token distribution, producing conditional samples constrained by valid continuations.
- Masked generation can constrain sequences to match a regular expression or parse according to a specified grammar such as Python or SQL.
- Computing the mask is expensive because valid tokens depend on previously sampled tokens, making each step an iterative matching or parsing problem.
- The central question is how to efficiently match or parse incomplete strings and determine the mask at every generation step.
3 Iterative FSM Processing and Indexing
The paper reformulates regular-expression-guided generation as FSM state transitions, enabling vocabulary masks to be indexed and retrieved without scanning the full vocabulary at runtime. This supports continued generation from arbitrary FSM states with average O(1) masking cost.
- FSM formulation: FSM states track where generation stands after each sampled vocabulary token, allowing the process to continue without rereading the entire sequence.The formulation determines the FSM states reached after one token and carries them through subsequent sampling steps.
- FSM matching: In the floating-point example, invalid vocabulary strings are masked according to the current FSM state, and sampling a token advances that state.For example, sampling ".2" moves to state 3, where only "42" and "1" remain valid completions.
- FSM matching: The method handles tokens matching arbitrary portions of a regular expression by finding FSM sub-sequences that can start at any state.Algorithm 3 records state paths that accept a vocabulary string, including failed paths that stop when a transition is unavailable.
- Vocabulary indexing: The approach indexes vocabulary strings by the FSM states from which they can be accepted, mapping each state to valid vocabulary subsets.The map σ connects FSM states to vocabulary elements accepted from those states.
- Efficiency: O(1) average cost replaces the usual O(N) per-token vocabulary scan, where N is the language model vocabulary size.A hash map makes mask computation constant-time on average, while preprocessing occurs outside token sampling.
- Evaluation: The implementation examples use Outlines with GPT2-medium, while the efficiency comparison measures token-generation timings against Guidance.The comparison uses a vocabulary of N = 50,257 and varies max_tokens, recording one timing sample for each setting.
- Evaluation: The reported timing behavior shows a striking increase with the maximum number of sampled tokens for the compared approach.The authors describe this scaling as indicative of the growing computational problem implied by the approach.
4 Extensions to Iterative Parsing
The FSM indexing framework extends to CFG-guided generation by combining lexical FSMs with parser configurations represented by pushdown automata. This handles partially observed token sequences whose correct lexical structure may only become clear from later tokens.
- Motivation: Traditional parsing can misidentify a partially observed token, such as treating "f" as a complete NAME even though later tokens show it begins "foo".After observing "def f", the correct NAME token is determined by the continuation of the sequence.
- Motivation: The parser must allow vocabulary strings that either continue the partially matched NAME or begin a valid LPAR construct.Both continuation and grammar-driven transition possibilities are considered for the next sampled string.
- FSM representation: FSM states formalize regex sub-patterns, with the NAME example using states 0, 1, and 2, of which states 1 and 2 are accepting.The states correspond to the initial state, [^\W\d], and \w*.
- FSM representation: For the vocabulary string "f", matching can start in states 0, 1, or 2 and end in states 1 or 2, enabling continuation after either accepting state.The resulting sequences are (0, 1), (1, 2), and (2, 2).
- Vocabulary indexing: Algorithm 4 maps FSM states 0, 1, and 2 to vocabulary strings that validly expand NAME or advance to a state accepting LPAR.The illustrated valid subset is "d", "ef", "pass", " ", and "oo(".
- PDA extension: Pushdown automata extend the FSM approach by incorporating parser states, stack behavior, and grammar-allowed transitions.The indexing construction uses PDA configurations and stack values to identify paths supporting complete parses.
- PDA extension: For each parse state, combined FSMs represent the terminal symbols permitted by PDA transitions, while scanning identifies possible terminals from recent characters.For example, scanning "de" can yield DEF or NAME depending on how the continuation completes the lexical patterns.
- PDA extension: Applying FSM matching to these combined machines yields parser configurations containing PDA states, FSM states, and possible terminal symbols.The PDA transition map’s pre-image is then used to determine stack values that can read those configurations.
5 Discussion
The indexing approach trades memory for broader computational efficiency, with small naive indices observed in Python-grammar tests. The discussion also identifies extensions to training, model evaluation, and internal computation reduction.
- Vocabulary indexing trades processing for memory while removing a prohibitive run-time scaling barrier in guided generation.The authors state that memory costs are relatively low on average and can be reduced through conventional means when necessary.
- Naively constructed indices for a slightly augmented Python grammar were around 50 MB, despite containing unused and redundant configurations.The indices used unreduced DFAs, leaving additional opportunities to reduce their size.
- The indexing approach could assist with training or fine-tuning LLMs when structured outputs are required.The authors also speculate that assisted generation during training may reduce the need for models to learn syntactic details.
- Comparing masked logits from the method with raw model logits could provide an alternative way to evaluate current models and inform training objectives.The proposed comparison concerns the discrepancy between constrained masked logits and the model's unconstrained raw logits.
- Lifting computed masks into higher levels of the model architecture may reduce computation by avoiding unnecessary parameter-slice operations.The current formulation applies masks at the lowest level; lifting them further is presented as a possibility with potential computational savings.