Constrained Decoding
Constrained decoding forces a language model’s output to satisfy a formal rule, such as a regular expression, a grammar, or a JSON Schema, by changing which tokens the sampler may pick at each step. Asking the model in the prompt to “reply only with valid JSON” makes valid JSON more likely but guarantees nothing, because every invalid token keeps a nonzero probability at every step. Constrained decoding makes invalid tokens impossible.
The mechanism
At step \(t\), a constraint checker looks at the text generated so far and computes the set \(A_t\) of tokens that can still lead to an acceptable output. The logits of all tokens outside \(A_t\) are set to \(-\infty\) before softmax, so their probability is exactly zero. The distribution actually sampled from is
\[ \tilde p(v\mid y_{<t},x)=\frac{p(v\mid y_{<t},x)\,\mathbf 1[v\in A_t]}{\sum_{u\in A_t}p(u\mid y_{<t},x)}, \]
where \(\mathbf 1[\cdot]\) is 1 if the condition holds and 0 otherwise. In words: take the model’s probabilities, keep only legal tokens, and rescale so they sum to 1. The model itself is unchanged, and any decoding method (greedy, temperature, top-p) can run on the masked distribution. An end-of-sequence token is allowed only where the text so far is already a complete valid output.
The practical question is how to compute \(A_t\) fast. For regular expressions, a finite automaton (see Finite Automaton) tracks the current state, and an index from state to allowed tokens makes each step a table lookup. For formats with arbitrary nesting, such as general JSON, a context-free grammar with a stack is needed. A JSON Schema is compiled into such a grammar or automaton, with the limits described in JSON Schema. Willard and Louf describe the finite-state approach, and Geng et al. study grammar-constrained decoding.
What it guarantees
If generation finishes inside the constraint, the output matches the format. This removes the failure where a response starts with prose (“Sure, here is the JSON”) or has a missing brace. It guarantees only syntax and the structural rules encoded in the constraint. A schema-conforming reply can still have the wrong label for the input, because the constraint says which labels are allowed, not which is right. Two further cases fall outside the guarantee: the generation is cut off by the maximum-token limit before the closing brace, and the model refuses and the provider reports the refusal outside the schema. Applications still need to check completion status and validate semantics.
Distribution distortion
Masking looks only one token ahead. A token can be legal as a prefix and still lead to a state where the model puts nearly all its probability on forbidden continuations. Renormalizing then hands the leftover probability to whatever is legal, however unlikely the model found it.
A hypothetical example: the only valid outputs are the strings “yes” (tokens ye then s) and “no”. The model’s first-token probabilities are ye 0.30, no 0.10, and maybe 0.60 (not allowed). After ye it is nearly sure to continue with an illegal token: s has 0.02 and maybe 0.98.
- With masking, step one renormalizes over the two legal tokens:
ye0.75,no0.25. Step two has onlys, so it is forced. The output is “yes” with probability 0.75. - The model’s own probability for the string “yes” is 0.30 x 0.02 = 0.006, and for “no” it is 0.10. Conditioned on being valid, “yes” has probability 0.006/0.106, about 0.057.
Masked sampling answers “yes” 75% of the time while the model, asked to choose between the two valid strings, mostly prefers “no”. The constraint changed the distribution over valid outputs, not merely removed invalid ones. Park et al. analyze this effect for grammar-constrained decoding and propose a sampler that follows the model’s probabilities conditioned on validity. A related issue is tokenization: the same string can be spelled by different token sequences, and a character-level mask may admit segmentations that the tokenizer would never produce, which the model has seen rarely.
Compared with neighbors
- Prompting for a format relies on the model’s compliance; constraints enforce it.
- Validate and retry lets the model generate freely, checks the result, and re-asks on failure. It costs extra calls. Repeated independent samples from the same prompt, accepted only when valid, follow the model’s distribution conditioned on validity. Re-asking with a changed prompt can change that distribution. Token masking also changes the relative probabilities of valid strings, as the example shows.
- JSON mode (offered by some providers) guarantees syntactically valid JSON but not a particular schema. Schema-constrained output goes further, within the subset of JSON Schema the provider supports.
References
- Willard and Louf, Efficient Guided Generation for Large Language Models — https://arxiv.org/abs/2307.09702
- Park et al., Grammar-Aligned Decoding — https://arxiv.org/abs/2405.21047
Discover more from Insightful Data Lab
Subscribe to get the latest posts sent to your email.
