Constrained decoding

Overview

A language model picks one token at a time. Constrained decoding changes that choice so the output meets a requirement. The two families below differ in how strictly the requirement is enforced.

Hard masking. A parser reads the output so far against a grammar and returns the tokens that still lead to a valid output. Every other token is removed from the choice, and the model keeps its own preferences among the allowed tokens 2305.13971 §2.2. The invariant is that the output stays inside the grammar, so structure is enforced at decoding time and not learned from examples. No finetuning is needed 2305.13971 §2.3. A trie of allowed label strings is the simplest case of the same mask: the allowed set is listed in advance rather than derived from a grammar 2202.03052 §3.4. A third mask needs neither a grammar nor a list: the allowed set is the tokens of the input, so the output can only reuse them 2104.05919 §2.1.

Soft scoring. Nothing is blocked. The requirement is a score that the search trades off against the model's own preferences, and it comes in two shapes. COLD holds the whole draft as one vector of logits per position, combines differentiable scores into an energy, and samples the draft from that energy with Langevin dynamics. Each position is then mapped to a token from the same LM's top-k candidates 2202.11705 §3.1 2202.11705 §3.3. ALICE keeps ordinary token-by-token beam search and adds a classifier's log-probability to each candidate's language-model score [arxiv:2203.09509 §3.2]. In both, the requirements are preferences, not guarantees 2202.11705 tab. 3.

A common use of hard masking is to make the reply parse as JSON for a program that reads it. Words2Contact does this with a grammar for its open models and with the provider's JSON mode for GPT-3.5 2407.14229 §3.1 fig. 3.

This page draws on six paper cards. Every number below is author-reported for that paper's setup. The papers evaluate different tasks, so their numbers are not comparable to each other.

Variants

Hard grammar mask (GCD)

Grammar-constrained decoding is the variant studied in Grammar-Constrained Decoding for Structured NLP Tasks without Finetuning. The main choice is how the grammar relates to the input.

Input-independent grammar (IIG)

One grammar per task. It does not change with the input. For entity disambiguation it allows any entity in a catalogue of about 470K entities 2305.13971 §3.2. On constituency parsing the paper reports 65.9 to 80.3% validity for IIG, and the grammar can produce trees with the wrong number of leaves, so it cannot enforce completeness 2305.13971 tab. 3 2305.13971 §3.3.

Input-dependent grammar (IDG)

The grammar is built for each input. For entity disambiguation the output is limited to the candidate entities given with the mention [arxiv:2305.13971 §3.2]. For constituency parsing the leaves must be the input words, which also needs the input-dependent form 2305.13971 §3.3. With LLaMA-33B, the paper reports 80.3 average accuracy on entity disambiguation against 68.2 for IIG 2305.13971 tab. 2, and 100.0% validity on constituency parsing against 72.3% for IIG 2305.13971 tab. 3. Card: Grammar-Constrained Decoding for Structured NLP Tasks without Finetuning.

Trie mask over a closed label set (OFA)

OFA: Unifying Architectures, Tasks, and Modalities Through a Simple Sequence-to-Sequence Learning Framework uses the same kind of mask for tasks with a fixed answer set, such as VQA answers, entailment labels and ImageNet classes. The candidate strings are stored in a trie. During beam search, any token that would leave the trie is masked [arxiv:2202.03052 §3.4] 2202.03052 app. B. The VQA trie covers the 3,129 most frequent answers 2202.03052 app. A.3. The stated reason is that the model may otherwise generate labels outside the closed set 2202.03052 §3.4.

The paper's ablation finds small gains. On Base, VQA scores 76.03 with the trie against 75.86 without it, and the gains on the five tasks tested are between 0.17 and 0.5 points 2202.03052 tab. 13. This variant mostly guards against off-list output, and the paper does not show a large accuracy effect.

Input-token restriction (GenArg)

Document-Level Event Argument Extraction by Conditional Generation restricts a BART decoder to tokens that occur in its encoder input, which holds the event template and the document. At each step every input token is scored by its dot product with the decoder state, and a softmax runs over those scores only. Every other vocabulary item gets zero probability 2104.05919 §2.1 2104.05919 eq. 2. The stated aim is to stop the model from inventing arguments that are not in the text 2104.05919 §2.1. No parser is needed, because the allowed set comes from the input itself.

The comparison in the paper is of the whole generator, not of the restriction alone. On WikiEvents argument classification (Head F1), it reaches 64.57 against 56.16 for BERT-QA 2104.05919 tab. 5. The paper reports no run with the restriction removed, so these numbers do not show what the mask adds 2104.05919 §4.

JSON-shaped reply in a robot pipeline (Words2Contact)

Words2Contact: Identifying Support Contacts from Verbal Instructions Using Foundation Models applies the same kind of hard mask to replies with a fixed JSON shape. Each LLM call is forced to JSON, with grammar-constrained sampling for open models 2407.14229 §3.1 fig. 3. The final call, which picks the hand and the task type, is one of these constrained replies 2407.14229 §3.5. The results do not isolate what the constraint adds. The grammar is a format guard for the program that reads the reply, and it does not check whether the content is right.

Soft energy-based decoding (COLD)

COLD (COLD Decoding: Energy-based Constrained Text Generation with Langevin Dynamics) does not block tokens. It defines an energy from weighted terms: fluency under a forward and a reverse LM, future-token prediction for fixed right-hand text, and a differentiable n-gram match to a reference 2202.11705 eq. 3 2202.11705 eq. 4 2202.11705 eq. 5. Each task combines these terms differently 2202.11705 §4.1.1 2202.11705 §4.2 2202.11705 §4.3. It uses a frozen GPT2-XL and no finetuning 2202.11705 §4.

Author-reported results against the baselines in each table:

• Keyword coverage on CommonGen is 94.50% against 91.00% for NeuroLogic, a beam-search method, but perplexity is higher (54.98 against 28.61) 2202.11705 tab. 3.
• On counterfactual rewriting, the min-edit overlap with the original ending is 56.84 against 52.90 for DeLorean. Mix-and-Match has a higher overlap (about 85) but lower coherence 2202.11705 tab. 2.
• On abductive reasoning, the overall human coherence is 2.96 against 2.87 for DeLorean, while DeLorean is higher on left coherence [arxiv:2202.11705 tab. 1]. The paper reports low annotator agreement, Krippendorff's alpha of 0.36, 0.52 and 0.29 across its three human evaluations [arxiv:2202.11705 §4.1.2] 2202.11705 §4.2 2202.11705 §4.3.

Classifier-guided beam search (ALICE)

ALICE is the decoding step in ToxiGen: A Large-Scale Machine-Generated Dataset for Adversarial and Implicit Hate Speech Detection. It is a beam search over GPT-3 output. Each candidate's score is a weighted sum of its language-model log-probability and a toxicity classifier's log-probability for a target class, with both weights set to 0.5 2203.09509 §3.2 2203.09509 §3.3 2203.09509 eq. 1. The target is the opposite of the prompt's label, so the search pushes against the classifier: a toxic prompt is steered towards benign-scoring text, and a benign prompt towards toxic-scoring text 2203.09509 §3.2. Tokens copied from the prompt are blocked 2203.09509 §3.3. The method produced 14,174 of the paper's 274,186 statements, about 5% of the dataset 2203.09509 tab. 2.

The paper's 125-prompt comparison with plain top-k sampling is mixed. ALICE fools HateBERT on 26.4% of statements against 16.8% for top-k, and its output is less toxic to human raters on toxic prompts (mean 2.97 against 3.75 on a 1-to-5 scale). Its output also agrees less often with the prompt label on toxic prompts, at 40.3% against 67.7% 2203.09509 §4.3. A larger study of about 9,000 statements shows the same pattern for toxic statements, with 58.97% fooling HateBERT for ALICE against 26.88% for top-k [arxiv:2203.09509 app. E]. Card: ToxiGen.

When to use which

Shared caveats

• Token-level grammars are tokenizer-specific (GCD). The paper writes the grammar over characters and compiles it to the tokenizer's tokens. The same string can be split several ways by BPE, so the compiled grammar must be rebuilt for each tokenizer 2305.13971 §2.1.
• Empty output can win (GCD). If the grammar allows an empty output, the most probable complete output may be empty. The paper reports this for LLaMA-13B on closed information extraction at length-normalisation exponents from 1.0 to 2.0, and that exponents of 2.5 or higher removed it 2305.13971 §5 2305.13971 tab. 5. Choose the exponent on your own validation data.
• Validity is not accuracy (GCD). The paper's constituency-parsing outputs are all valid, yet F1 stays well below supervised parsers [arxiv:2305.13971 §4.3].
• Trie gains are small and single-run (OFA). The ablation shows one number per setting, with no spread across runs 2202.03052 tab. 13. Differences of a few tenths of a point may be noise, so check the effect on your own validation data before relying on it.
• The grammar and its consumer must agree (Words2Contact). A JSON grammar only guarantees the shape the grammar allows. In the released code, the grammar for the correction step forces keys that the code does not read, and the arithmetic grammar may reject parenthesised expressions that the prompt examples use. Both come from reading the code and grammar files, not from running them. See the Words2Contact card, Pitfalls. Test the grammar on your own prompts before relying on it.
• The input-token mask blocks words the input lacks (GenArg). The output can contain only words from the encoder input. A joining word such as "and" is therefore available only if the text contains it. The released code has a helper that would always allow it, but the helper is not called. The released code also scores a token that repeats in the input by whichever write lands last, and gives non-input tokens a large negative logit rather than minus infinity. Both differ from the paper's formula. See the GenArg card, Pitfalls.
• Soft constraints are not guaranteed (COLD). Keyword coverage is 94.50%, not 100% 2202.11705 tab. 3. Do not treat COLD as a validator.
• The constraints are checked on the soft draft (COLD). The final tokens are chosen in a separate step after sampling, from the LM's top-k candidates 2202.11705 §3.3 2202.11705 eq. 6.
• Prompt labels are a weak proxy (ALICE). The paper trains on the prompt's label 2203.09509 §5. For ALICE output that proxy is the weakest, since human ratings agree with it less often 2203.09509 §4.3. Check labels on your own data before training on them.
• The released code departs from the paper's weights (ALICE). The paper has one weighted sum 2203.09509 eq. 1. The released code applies the classifier term twice and keeps only 5 vocabulary candidates per step, not the 100 the paper states. Both change the effective balance between model and classifier. The ToxiGen card lists these under "Paper vs code".
• Several samples are usually drawn (COLD). The paper selects one of 16 or 32 samples per input, so the reported runtime per sample understates the total cost 2202.11705 app. B 2202.11705 tab. 5.

Implementation sketch

All variants share an outer loop that asks the model for a choice and applies a rule. They differ in what the rule touches: the next token (hard masking), a beam of candidates (ALICE), or a whole draft (COLD).

Hard masking

The core step is the same for all grammar variants. The paper uses constrained beam search with a beam size of 2 2305.13971 app. F. The greedy version below shows the mask and leaves out the beam bookkeeping.

# Grammar-constrained greedy decoding  [arxiv:2305.13971 §2.2]
def constrained_greedy(lm, parser, prompt_tokens, max_new_tokens):
    # The parser state is the grammar position reached so far  [§2.2]
    state = parser.start()
    out = []
    for step in range(max_new_tokens):
        # Next-token log-probs; an API that hides them cannot be masked  [§2.2]
        log_probs = lm.next_token_log_probs(prompt_tokens + out)
        # Tokens that keep the output inside the grammar  [§2.2]
        allowed = parser.allowed_next_tokens(state)
        # Every token outside the allowed set is removed from the choice  [§2.2]
        best = max(allowed, key=lambda t: log_probs[t])
        out.append(best)
        state = parser.advance(state, best)
        # Stop once the output is a complete member of the grammar  [app. F]
        if parser.is_complete(state):
            break
    return out

For input-dependent grammars, build parser from the candidates or tree constraints for this input before the loop 2305.13971 §3.2 2305.13971 §3.3. For a closed label set, parser is a trie over the label token sequences, and allowed_next_tokens returns the children of the current trie node 2202.03052 app. B.

For the input-token mask, the allowed set is the input itself. The scores come from the paper's dot product, and the embedding table is the one the paper calls Emb(w), which the paper does not name 2104.05919 eq. 2.

# One greedy step restricted to input tokens  [arxiv:2104.05919 §2.1] [eq. 2]
def input_restricted_step(dec_state, input_ids, input_emb):
    # Dot product of the decoder state with each input token's embedding
    scores = input_emb[input_ids] @ dec_state          # (n,)
    # Softmax over input positions only; other vocabulary gets probability 0
    probs = softmax(scores)
    # Greedy choice, as in the paper's worked example  [§2.1]
    return input_ids[argmax(probs)]

Soft scoring: beam step (ALICE)

One step of ALICE's classifier-weighted beam search, in the paper's form, not the released code. The released code's weights differ, as the caveats note.

# One classifier-guided beam step  [arxiv:2203.09509 §3.2] [eq. 1]
def alice_step(beams, lm, clf, target_class, beam_size):
    # Top next-token candidates per beam from the LM  [§3.3]
    candidates = lm.top_k_next(beams, k=100)
    # Drop tokens copied from the prompt; keep punctuation  [§3.3]
    candidates = drop_prompt_tokens(candidates)
    # Equal-weight sum of LM and classifier log-probs  [eq. 1] [§3.3]
    scores = 0.5 * candidates.lm_logprob
    scores += 0.5 * clf.logprob(candidates, target_class)
    # Keep the best hypotheses for the next step  [§3.3]
    return keep_best(candidates, scores, beam_size)

Soft scoring: energy and draft (COLD)

This is a sketch of the paper's equations, not of the released code. The paper's update is a plain Langevin step, while the released code uses Adam; see the COLD card before reproducing.

# Energy-based soft decoding, then top-k discretisation  [arxiv:2202.11705 §3.1]
def cold_sketch(lm, constraints, length, steps, k):
    # One logit vector per position, from the initialisation  [§3.4]
    soft = init_soft_sequence(lm, length)
    for n in range(steps):
        # Energy is minus the weighted sum of constraint scores  [eq. 1]
        energy = -sum(w * f(soft) for w, f in constraints)
        # Gradient step plus Gaussian noise; the noise shrinks over time  [eq. 2]
        soft = soft - eta * grad(energy, soft) + noise(n)
    # Pick each token from the LM's top-k candidates  [eq. 6] [§3.3]
    return discretise(lm, soft, k)

Contradictions

None between the six papers' reported numbers. They evaluate different tasks, so their figures cannot be set against each other. The OFA trie result is a small gain in a single-run ablation, and the GCD, COLD and ALICE papers do not test that setting. The GenArg paper reports no restriction ablation. Its released RAMS script uses an unrestricted BART model, and the paper does not say which model produced its RAMS numbers 2104.05919 §4.2. See the GenArg card, Paper vs code. ALICE's 125-prompt result is a trade-off rather than a contradiction: it fools the classifier more often while its labels agree less with the prompt. Within the COLD paper, its equations and its released code differ in several places, such as the optimizer and the weights; the COLD card lists them under "Paper vs code". Check those before reproducing its numbers. The ALICE paper's code also differs from its equations, as noted in the caveats above.

Papers

• Grammar-Constrained Decoding for Structured NLP Tasks without Finetuning: masks an LLM's next tokens with an incremental grammar parser, with input-independent and input-dependent grammars, and no finetuning.
• COLD Decoding: Energy-based Constrained Text Generation with Langevin Dynamics: samples a soft token sequence from a weighted energy over constraint terms, then maps it to text with a frozen GPT2-XL.
• OFA: Unifying Architectures, Tasks, and Modalities Through a Simple Sequence-to-Sequence Learning Framework: one encoder-decoder for vision, language and cross-modal tasks; uses trie-masked beam search for closed label sets, with an ablation of its effect.
• ToxiGen: A Large-Scale Machine-Generated Dataset for Adversarial and Implicit Hate Speech Detection: generates a 274k-statement dataset with GPT-3; a classifier-guided beam search (ALICE) produces a small share of it.
• Document-Level Event Argument Extraction by Conditional Generation: a BART generator fills event templates from a whole document, with output restricted to input tokens.
• Words2Contact: Identifying Support Contacts from Verbal Instructions Using Foundation Models: an LLM and VLM pipeline for robot contact points; its LLM replies are forced to JSON.