Grammar-Constrained Decoding for Structured NLP Tasks without Finetuning

In plain words

Language models are good at writing fluent text but often break a required format: a list of facts that must use known names, a link that must point to one of a fixed set of options, or a bracketed tree that must be balanced. The usual fix is to train the model further on examples of the format. This paper asks whether the format can instead be enforced while the model writes.

The idea is to describe the allowed output as a grammar, the set of rules that every valid answer must follow. While the model picks its next word piece, a small checker looks at the text so far and lists the pieces that would still lead to a valid answer. Every other piece is removed from the choice. The model keeps its own preferences among the allowed pieces, so the answer stays plausible as well as valid.

The grammar can depend on the input. For example, when the task is to choose one entity from a short list of candidates for a given mention, the list changes with each mention, and so does the grammar. The authors show this on three tasks: pulling facts out of text, linking a mention to an entity, and writing out the syntax tree of a sentence. Without any training, the constrained model does much better than the same model left free to write, and on the first two tasks it can match or beat models that were trained for the job.

TL;DR

• Decoding-time mask: an incremental parser returns the next allowed tokens from a context-free grammar, and only those tokens are sampled or searched §2.2.
• Grammars are written over characters and compiled to the tokenizer's tokens, using Grammatical Framework §2.1 app. B.
• Input-dependent grammars (IDG) restrict the output to per-input candidates and give the best results in the paper's tables §4 tab. 2 tab. 3.
• No finetuning; works for open-weight models with access to next-token distributions, not for API-only models §2.2 §6 Limitations.

Problem

Few-shot prompting of LLaMA-family models gives output that often breaks the required structure. Unconstrained LLaMA-33B reaches 17.5 F1 on closed IE, 54.1 average accuracy on entity disambiguation, and 42.9 F1 with 64.2% valid trees on constituency parsing tab. 1 tab. 2 tab. 3. Finetuning on task data is the usual remedy, and the paper argues it is costly when training data is scarce app. A.

Method

Tasks as formal languages §2.1. The output of an NLP task is described by a context-free grammar over characters (abstract grammar). A token-level grammar for a specific tokenizer is generated from it with Grammatical Framework (GF). The paper notes that BPE can tokenize the same string in several ways (for example, three different tokenizations of the bracket string [[[), so the token-level grammar is one of several possible mappings. Its workaround, adding spaces to force single-bracket tokens, is described as tokenizer-specific and not principled §2.1.

Grammar-constrained decoding §2.2. At each step the parser takes the partial output and the grammar and returns the set of allowed next tokens. The decoder then prunes all other tokens. The method is said to work with greedy search, beam search and sampling, and with any autoregressive model whose next-token distribution is accessible §2.2.

# Grammar-constrained beam search  [§2.2] [app. F]
# Input:  lm (next-token log-probs), parser (completion engine over the
#         token-level grammar), prompt_tokens, beam_size, max_new_tokens
# Output: finished token sequences, best first

def constrained_beam_search(
    lm,
    parser,
    prompt_tokens,
    beam_size,
    max_new_tokens,
):
    # Each beam keeps its generated tokens, score and parser state  [§2.2]
    beams = [Beam(tokens=[], score=0.0, state=parser.start())]
    finished = []
    for step in range(max_new_tokens):
        candidates = []
        for beam in beams:
            # Ask the completion engine which next tokens keep the output
            # inside the grammar  [§2.2]
            allowed = parser.allowed_next_tokens(beam.state)
            # Get the LM's log-probabilities over the vocabulary  [§2.2]
            log_probs = lm.next_token_log_probs(prompt_tokens + beam.tokens)
            # Forbidden tokens are dropped before scoring; the paper does not
            # say whether the remaining mass is renormalised  [§2.2]
            for token in allowed:
                candidates.append(extend(beam, token, log_probs[token]))
        # Keep the best beam_size candidates; beams that reached a complete
        # output move to finished  [app. F]
        beams, done = split_finished(select_top(candidates, beam_size))
        finished.extend(done)
    # Rank by length-adjusted score S / m^alpha, alpha = 1.0  [app. F] [§5]
    return sorted(finished, key=length_adjusted_score, reverse=True)

Input-dependent grammars §2.1 §3.2. Some tasks need a grammar that depends on the input. For entity disambiguation the output is limited to the candidate entities supplied with the mention; the paper calls the candidate set mention-dependent. For constituency parsing the terminal nodes must be the input words, which also needs an input-dependent grammar. For the entity task the input-independent alternative (IIG) allows any entity in the full catalog (about 470K entities, from the cited dataset of Le and Titov) §3.2.

Few-shot learning with GCD §2.3. The LLM is prompted with a short instruction and several demonstrations (the paper picks them from the training set), then decoded under the grammar. The grammar is a second source of task knowledge alongside the demonstrations §2.3. The paper notes that prompts are sensitive to wording and format and that its selection was manual app. E.

Three grammars §3 fig. 2. The paper lists 14 example grammars in Figure 2 and uses three in experiments: closed information extraction (input-independent, a context-free grammar over the 2.7M Wikidata entities and 888 relations) §3.1; entity disambiguation (regular, input-dependent) §3.2; and constituency parsing (context-free, input-dependent for the completeness constraint) §3.3. Figure 1 illustrates GCD applied to cIE fig. 1. The fully-expanded linearization is the main setting: each triplet is written as [s] subject [r] relation [o] object §3.1 [code: configs/hydra_conf/experiment/IE/inference/stable_fe.yaml#L11].

Hyperparameters

Everything not marked "code" is from the paper. "Not stated" means the paper and its appendix do not give the value.

Results

Author-reported. Setup: few-shot, no finetuning, 4 shots (cIE, ED) and 8 shots (CP) per the table captions; LLaMA and Vicuna backbones; test splits only §3 tab. 1 tab. 2 tab. 3. The tables give 90% confidence intervals, which the paper reports as widths between 0.1 and 0.3 (ED) and 4.0 and 6.0 (CP).

Closed information extraction, SynthIE-text-small, 4 shots tab. 1. Micro-averaged triplet scores; GenIE is a weakly supervised T5-base trained on REBEL, not on the SynthIE train split §4.1.

The table gives 90% confidence intervals for each cell tab. 1; the constrained LLaMA-33B row reports ± 0.9 on precision, ± 0.8 on recall and ± 0.7 on F1.

The paper highlights two comparisons: constrained LLaMA-33B (F1 36.0) is above GenIE (F1 34.8), which the paper calls weakly supervised §4.1; and on entity disambiguation, constrained LLaMA-33B (average 80.3) is above GENRE trained only on the AIDA data (average 80.0), though still below the full GENRE model (average 88.8) §4.2 tab. 2.

Entity disambiguation, micro-accuracy, 4 shots tab. 2. Columns are the six datasets and their average.

The IDG and IIG rows use different candidate sets; see Pitfalls before comparing them. Supervised rows are reproduced from the paper's table; the GENRE and ReFinED rows are cited to their own papers in the source.

Constituency parsing, PTB, 8 shots, test restricted to gold trees under 64 tokens tab. 3. The bespoke rows use the full PTB test set tab. 3 caption.

Other results.

• Subject-collapsed linearization (the subject is written once, as the object of the previous triplet) is more compact but scores lower than fully expanded, and constraints still help at every size app. G tab. 7.
• Adding all 888 relations to the cIE prompt gives only a small gain over the baseline app. G tab. 8.
• Vicuna on CP (appendix only): with IDG, Vicuna-7B reaches 35.6 precision, 31.9 recall and 41.4 tag accuracy, and Vicuna-13B 51.6 precision, 44.4 recall and 42.4 tag accuracy, both with 100.0% validity. The paper says Vicuna is even worse than LLaMA on CP. The table has no F1 column tab. 9 app. G.
• Latency: GCD adds about 69 ms per token on cIE, 1 ms on ED and 4 ms on CP, measured on a consumer CPU; the LLM forward pass is 54 to 136 ms per token on an A100 for 7B to 65B models tab. 4 §4.4.

Paper vs code

Repository: linked from the paper's abstract, so treated as official; commit d45378338131a7fc08fc06aeec5eee26238cb25d; MIT license LICENSE. Discrepancies, each with both citations:

• Beam size for CP. Paper: beam size 2 for all decoding app. F. Code: CP uses a beam width of 1 stable.yaml. The CP numbers in tab. 3 cannot be reproduced with the paper's stated setting without checking which the authors used.
• Shots for CP. Paper: 8 shots tab. 3 caption. Code: 4 demonstrations for the CP test set cp_ptb_stable.yaml.
• Output selection for cIE. Paper: "most probable non-empty generation" app. F. Code: the second beam (index 1) is returned for cIE stable_fe.yaml; the first non-empty generation is used for CP hf_gen_utils.py.
• cIE test set. Paper: 50K SynthIE test samples §3.1; table caption: "SynthIE-text-small" tab. 1. Code loads a file named test_small_ordered.jsonl ie_synthie_small_stable.yaml. The number of samples in that file is not stated anywhere in the paper.
• Max new tokens. Not in the paper; set in code: 256 (cIE), 64 (ED), 256 (CP). The CP comment says the target was 128 tokens and the decoding limit was raised to avoid truncating brackets stable.yaml.
• Grammars are not in the repository. The compiled grammar objects are downloaded separately from a Hugging Face dataset, and the GF sources are in a separate helper repository download_data.md README.md. The grammar text itself cannot be checked here.
• ED context window. Paper: left and right context limited to 10 tokens each app. D. No code in the repository performs this trimming; the datasets are preprocessed files ed_aida_stable.yaml. Not verified.
• Decoding masks. The code hands the grammar's allowed set to a prefix-restricting hook gf_constrained.py. Whether the remaining probabilities are renormalised depends on the Hugging Face generation code, which is not in the repository; the paper does not say.

No other discrepancies found in the configs for early stopping, temperature, and length penalty.

Pitfalls

• Empty top output on cIE and CP. With grammars that allow an empty output, the most likely complete output can be the empty string. The paper reports that for LLaMA-13B the top generation was always empty on cIE, and that raising the length-normalisation exponent alpha to 2.5 or higher removes the empty top output; at alpha 1.0 to 2.0 it stayed empty §5 tab. 5. Vicuna-13B produced a non-empty top output 46% of the time, and Vicuna-7B never did §5. Plan for the empty-output case, and pick alpha on your own validation data: the paper gives only one model's sweep.
• Token-boundary ambiguity. A character grammar becomes a token grammar per tokenizer, and BPE can split the same string several ways. The paper's workaround (forcing single-bracket tokens) is not principled and must be redone for each tokenizer §2.1.
• Validity guarantee depends on the grammar. With IDG, CP validity is 100% tab. 3. With IIG, validity is only 65.9 to 80.3% and F1 stays below IDG tab. 3. The IIG grammar can produce trees with the wrong number of leaves, so it cannot enforce completeness §3.3.
• Latency. Large grammars are expensive per token. The cIE grammar (2.7M entities) costs 69 ms per token on CPU, comparable to the LLM forward pass tab. 4. The appendix reports 0.05 s for a 279K-entity grammar and 0.5 s for a 5.9M-entity grammar app. H. The unit basis of those appendix numbers is not stated beyond the figures. The paper says no multi-threading was used app. H.
• API-hosted models are excluded. Decoding needs the next-token distribution. Models behind APIs that do not expose logits cannot be constrained this way §2.2 §6 Limitations.
• Input-dependent grammars are not like-for-like. For IDG the candidate set is part of the grammar and some mentions are dropped when the candidate set is empty (for example 17 of 257 on ACE2004). For IIG, mentions not in the knowledge base are dropped (for example 4 of 727 on AQUAINT) app. D. Per-dataset numbers in tab. 2 therefore come from different subsets.
• Prompt sensitivity. The paper says wording and demonstration format matter, and the chosen prompt was picked by manual inspection on validation data. The ED prompt differs by dataset (Prompt 1 for AIDA and WIKI, Prompt 2 for the others) app. E tab. 6.
• CP numbers are on a restricted test set. Only gold trees under 64 tokens are used for LLM rows, so they are not comparable to the bespoke rows, which use the full test set tab. 3.
• Supervised baseline choice. On cIE, GenIE is weakly supervised (trained on REBEL), not on the same train split as the LLM rows' test data §4.1.
• Constraints do not fix reasoning. The paper states the CP outputs are valid but still far below supervised F1 (95% and above), so validity is not accuracy §4.3.

Lineage

• Builds on: GCD and completion engines from Synchromesh (Poesia et al. 2022), and incremental parsing with Grammatical Framework (Angelov 2009) §2.2.
• Compared against: GenIE (cIE), GENRE, BLINK and ReFinED (ED), and bespoke constituency parsers (CP) §3 tab. 1 tab. 2 tab. 3.
• Superseded by: not stated.