Automata-based constraints for language model decoding

Automata-based Constraints for Language Model Decoding

The paper "Automata-based constraints for language model decoding," authored by Terry Koo, Frederick Liu, and Luheng He from Google DeepMind, presents a novel approach to ensuring that the outputs of language models (LMs) conform to specified formal languages using principles of automata theory. This method is particularly advantageous for generating structured data, API calls, and code snippets, which require strict adherence to syntactical constraints.

Overview and Contributions

Language models, particularly smaller ones adapted for large-scale deployment, may occasionally generate outputs that deviate from expected formal language structures. Traditional fine-tuning methods to enforce conformance are resource-intensive and impractical for uncommon or highly specific formats. This paper proposes leveraging finite-state automata (FSA) and finite-state transducers (FST) to apply hard constraints on LMs, ensuring the generation of valid outputs efficiently.

The authors have identified key contributions in their research:

1. Detokenization as Transduction: Reformulating detokenization as an FST, which maps token sequences back to text.
2. Adaptation of Regular Expressions to Tokens: Converting FSAs that accept character sequences into FSAs that accept token sequences, enabling token-level constraints on formal languages.
3. Extensions for Practical Applications: Introducing special capturing groups and FST extensions to handle high branching factors and enhance usability.

Methodology

Finite-state Constraints: The method involves forming constraints for an LM by masking decoding logits. This involves:

• Building a mask of valid next tokens based on the current state of constraints.
• Penalizing the logits for invalid tokens to ensure only valid tokens are chosen.
• Updating the state of constraints accordingly after each token is selected.

Addressing Tokenization: A significant challenge in this approach is the tokenization used by popular LMs, which may not align with the formal grammar. This paper finds solutions within automata theory by constituting a vocabulary of token detokenization using FSTs.

By constructing an FST (T_V) for any given token vocabulary, the paper ensures that any sequence accepted by the adapted FSA (A’) corresponds to a valid detokenized sequence accepted by the original FSA (A). This decomposition separates the vocabulary-specific component from the grammar-specific component, fostering efficiency and modularity.

Extensions: Special capturing groups (terminal labels) handle cases where matching tokens would otherwise lead to a large number of outbound edges, making operations computationally expensive. By using these extensions, constraints on tokens are applied more efficiently.

Push-down Constraints

The paper extends its automata-theoretic approach to deterministic context-free languages using push-down automata (PDAs). PDAs are essentially FSAs equipped with a stack, capable of handling more complex syntactical structures. Although deterministic PDAs are constrained to deterministic context-free languages, the vast majority of structured data formats fall within this category.

Employing PDAs allows the same methodology used for FSAs to be applied to more complex grammars by creating a correspondence between character-based PDAs and token-based PDAs through FST composition.

Practical Implications and Applications

The approach's primary advantage is its efficiency and simplicity, making it applicable for real-time, large-scale language model deployments. By adopting an automata-theoretic perspective, the method is broadly applicable across different constraints and language models.

JSON Generation: One application described is generating JSON outputs that conform to specified schemas. The set of JSON expressions matching a schema is shown to be a regular language, simplifying the constraint application. Tools are developed to automatically translate JSON schemas into regular expressions, further easing implementation.

Python Dataclasses: The method can ensure that generated data conforms to Python dataclass definitions. This has practical implications for generating structured programmatic content that must adhere to specific class schemas.

Speculative Decoding: The technique enhances speculative decoding, wherein a smaller, faster model generates token sequences that a larger model then validates. By constraining the smaller model, the acceptance rate of sampled tokens is increased, significantly improving the overall efficiency of the decoding process.

Conclusion and Future Work

The reformulation of detokenization as an FST marks this paper's primary contribution, providing a framework that unifies constraint application under automata theory. The clean and efficient solutions offered for tokenization ambiguities establish this method as a practical tool for structured language generation.

The paper hints at further exploration into PDAs, considering their more compact nature and expressive power for more demanding tasks. Future research may delve into optimizing grammar specifications to avoid non-deterministic PDAs and expanding the utility of context-free constraints for longer, more complex decoding tasks.

This work lays a robust foundation for integrating automata theory with language model decoding, promising more reliable and syntactically precise AI-generated content.