Enforcing Structured AI Outputs with Grammars and FSMs

This article explains how finite-state machines (FSMs) and context-free grammars (CFGs) constrain token logits during runtime to guarantee that large language models (LLMs) generate valid structured outputs like JSON or regular expressions. By intercepting the model's prediction probabilities at each decoding step, these formal language systems mask out invalid tokens before sampling occurs. This process eliminates syntax errors and hallucinations in structured data generation while maintaining the model's underlying semantic capabilities.

The Mechanics of Logit Generation and Masking

During inference, an LLM processes input text and outputs a vector of raw, unnormalized predictions called logits. This vector contains a score for every token in the model's vocabulary. Normally, a softmax function converts these logits into probabilities, and a sampling method selects the next token.

To guarantee a specific output structure, an external parser intercepts the logits before the softmax calculation. The parser evaluates the current state of the generated text and determines which tokens in the vocabulary can legally follow. It then applies a logit mask:

  • Valid tokens retain their original logit scores.
  • Invalid tokens have their logits set to negative infinity (\(-\infty\)).

When the softmax function runs on the modified logits, the tokens set to \(-\infty\) receive a probability of exactly zero. The model is mathematically forced to select a token that adheres to the defined rules.

Finite-State Machines for Regular Expressions

A regular expression defines a regular language, which can be directly represented by a finite-state machine. An FSM consists of a finite number of states, transitions between those states based on character inputs, an initial state, and one or more accepting (final) states.

To constrain an LLM using an FSM, the system must bridge the gap between character-level state transitions and multi-character subword tokens.

  1. Compilation: The target regular expression is compiled into a deterministic finite automaton (DFA).
  2. Vocabulary Mapping: Because LLM tokens often contain multiple characters, the system pre-computes how every token in the model's vocabulary interacts with the DFA. If a token starts at State A, the system traces its characters through the DFA to see if it lands on a valid State B. If any character in the token triggers an invalid transition, that token is marked as invalid for State A.
  3. Runtime Tracking: Generation begins at the DFA's initial state. At each step, the system looks up the pre-computed valid tokens for the current state, masks all other token logits to \(-\infty\), and samples a token.
  4. State Transition: Once a token is selected, the FSM advances to the new state dictated by that token, and the process repeats until the model selects an end-of-sequence token from an accepting state.

Context-Free Grammars for JSON

While FSMs handle flat patterns like regular expressions, they cannot track arbitrary nesting, such as matching opening and closing brackets in JSON. JSON requires a context-free grammar (CFG), which uses production rules to define recursive structures.

A CFG is processed using a Pushdown Automaton (PDA)—essentially an FSM equipped with a stack memory to keep track of nested contexts.

  1. Grammar Definition: The target format (e.g., a specific JSON schema) is translated into a formal CFG, often written in Backus-Naur Form (BNF).
  2. Incremental Parsing: At runtime, a parser (such as an Earley parser or a generalized LR parser) maintains a stack representing the current depth and context of the JSON structure (e.g., inside an object, expecting a key, or inside an array).
  3. Token Validation: Before the LLM samples the next token, the parser checks the vocabulary against the current stack state. For example, if the generated text is {"age":, the grammar rules dictate that the next token must start a valid JSON value (a number, string, boolean, array, or object). Tokens like } or ] are rejected.
  4. Stack Manipulation: When a token containing { or [ is selected, a new state is pushed onto the stack. When the corresponding } or ] is selected, the state is popped from the stack.

Overcoming Subword Tokenization Challenges

A major hurdle in grammar-constrained decoding is tokenization boundary misalignment. A single LLM token might span across multiple grammar rules. For example, the token ": " in JSON contains the closing quote of a key, the colon separator, and the whitespace before the value.

To resolve this, constrained decoding engines treat tokens as sequences of bytes or characters. During the validation phase, the parser feeds the characters of a candidate token through the FSM or PDA one by one. A token is considered valid only if every character in the token successfully transitions the state machine without hitting an error state.

To avoid the massive computational overhead of testing tens of thousands of tokens character-by-character at every step, modern frameworks organize the model's vocabulary into a prefix tree (trie). The parser traverses this trie alongside the grammar states, allowing it to eliminates entire branches of invalid tokens simultaneously, reducing latency to just a few milliseconds per token.