The Parser Already Knows: Lightweight Bias Correction in Constrained Decoding

Language Models (LMs) generate human-like text and code but consistently struggle to satisfy rigid syntactic and semantic constraints. Prompt engineering and fine-tuning fall short of guaranteeing strict conformance to context-free grammars (CFGs), creating a bottleneck when integrating LMs into real-world software pipelines where even minor syntax violations yield malformed outputs causing downstream failures in parsers, input validators, structured data processing engines, and auto-formalization frameworks.

Prior work on constrained generation fundamentally trades generation quality for computational cost:

1. Rejection sampling repeatedly draws completions until a valid candidate emerges within a fixed budget; however, this is computationally prohibitive and provides no formal compliance guarantee.

2. Grammar Constrained Decoding (GCD) methods (e.g., Syncode, XGrammar, Outlines, DOMINO) address this by masking non-conforming logits at each decoding step. While masking guarantees structural compliance with minimal overhead, it only evaluates whether a token is valid locally, ignoring whether it forecloses high-probability completions down the line. Consequently, masking distorts the model's output distribution, often steering generation toward unnatural or degraded outputs.

3. Online sampling techniques (e.g., ASAp, CARS, MCMC, AWRS, SMC Steering) restore lost probability mass by dynamically reweighting candidates, but require expensive online resampling at every step, severely inflating inference latency.

4. Tractable probabilistic model augmentation (e.g., GeLaTo, Ctrl-G, P-GCD) avoids runtime sampling by pairing the LM with an offline surrogate model (typically an HMM) that is exactly computable. However, learning a tractable approximation of the full LM distribution is notoriously difficult and introduces significant modeling errors.

The paper states: Because each paradigm compromises on either latency, fidelity, or scalability, none of these options is ideal for practical deployment.

The paper formally characterizes the problem. Given an LLM p, a prompt x≤s, and a grammar constraint α, the goal is to sample from the conditional distribution p(x>sx≤s, α). GCD tools provide access to a masked distribution qα defined as:

qα(xix<i) = p(xix<i)I(x<i xi ∈ Pref(L(α))) / Σt∈V p(tx<i)

where V is the set of grammatically valid next tokens. This approximation is not exact: qα only masks based on whether the prefix conforms to the grammar so far (past grammaticality), so it does not incorporate the future grammaticality of the tokens it selects.

The paper demonstrates this with an example grammar Gsk: 00000 1(01)4, where the grammar accepts binary strings of length five, and a string can only continue with 0 if it started with 0. With a uniform LLM distribution, each of the 17 grammar-satisfying strings should have probability 1/17. However, masked decoding would generate 00000 with probability 1/2 (if the first token is 0), instead of the correct 1/17.

The root cause: in each decoding step, the algorithm should take into account the future, or the probability of the next token yielding a valid sequence in the future.

SHIM is a lightweight, offline-trained probability controller that predicts a correction factor γ for each candidate next token and applies it to the masked distribution qα at every decoding step:

qα⋆(xix<i) ∝ qα(xix<i) · γ(xi, x<i)

The paper derives that the ideal correction factor simplifies to γ(xi, x<i) ≈ p(αxi, x<i), i.e., the probability that the LLM generates something that satisfies the constraint given the candidate token and prefix.

The paper's central insight is that the masking tool's own parser and lexer state already encode the grammatical structure this correction needs. Instead of training a separate model to learn that structure from scratch, SHIM reads it out directly, together with the candidate next token, at essentially no added cost.

The three features used are:

1. Parser state: captures structure across completed lexemes but stalls mid-lexeme

2. Lexer state: captures within-lexeme progress but is blind to structure beyond it

3. Candidate next token: provides an early signal of what comes next, much like an LR(1) parser uses a single lookahead token to resolve shift/reduce ambiguity

These features are complementary: parser state captures structure across completed lexemes but stalls mid-lexeme, while lexer state captures within-lexeme progress but is blind to structure beyond it.

SHIM is implemented at three levels of feature availability and model capacity:

1. MLP: A Multi-Layer Perceptron with two hidden layers (64 and 32 units, ReLU activations) using all three features

2. LR FULL: A lightweight Logistic Regressor trained on the same full feature set

3. LR NT: A Logistic Regressor restricted to the candidate next token alone

Training data is generated offline: for each of 1000 generated sequences per grammar, the masking tool's incremental parser is run over the sequence, recording at every decoding step the parser state, lexer state, and candidate next token as feature rows, labeled according to whether the resulting sequence is grammatically valid.

The method functions as a refinement layer on top of existing masking tools, including Syncode, XGrammar, and Outlines – together, over 17,000 GitHub stars.

The paper evaluates on:

• INV-BV: SyGuS's bit-vector invertibility-condition split, using Mistral-7B-Instruct-v0.2 as the base LLM

• Spider: text-to-SQL benchmark with larger, open vocabulary

Baselines: Syncode (masking-based GCD) and ASAp (online sampling, capped at 100 samples per grammar due to computational infeasibility of running to convergence).

Metrics: KL divergence to the target distribution and average inference time per sample.

Full-feature MLP correction: Reduces KL divergence over Syncode on 11 of 14 grammars and over ASAp on all 14 grammars, running roughly 4.6× faster than ASAp on average (3.64s vs. 16.62s).

Lightweight logistic regression (LR FULL): Matches or beats Syncode on 12 of 14 grammars, running in 2.52s on average – faster than Syncode's own 2.97s and roughly 7× faster than ASAp's 16.62s. The paper explains: "Computing γ adds a small per-step cost, but by correcting for future grammaticality, LR FULL avoids the forced, longer completions that plain masking's short-sighted choices can produce; the resulting savings from needing fewer steps outweigh that added cost."

Next-token-only correction (LR NT): Beats Syncode on exactly half the grammars but is never far behind even when it loses. Runs in 2.45s on average, marginally faster than Syncode.

The strongest variant achieves up to 113× lower KL divergence than masking and up to 235× lower divergence than online sampling, while all variants add negligible additional inference cost over standard masking.

Figure 5 shows KL divergence as a function of k (number of unique valid vanilla-LLM samples used to estimate the target): "MLP and LR FULL stay low and largely flat... LR NT stabilizes rather than climbs but is not always low. ASAp instead climbs steadily and ends the worst of all methods on three of the four grammars shown, indicating its early competitiveness reflected an under-sampled evaluation rather than genuine closeness to the target distribution."

On four query pairs from Spider, lightweight correctors match or improve on Syncode's masked distribution on most rows (store product, entertainment awards, cre theme park), though Syncode remains best on election.

The LR NT results are conditional: on store product, LR NT improves on the full-feature variants, consistent with the LR(1) lookahead intuition; but on entertainment awards, it is substantially worse than every full-feature variant. The paper explains this through query characteristics: store product's query matches a substring via LIKE %Scanner% where the next token is nearly the whole signal, while entertainment awards has near-duplicate column names (e.g., Num of Audience, NumOfAudience, Total Audience) that only diverge many tokens in.

ASAp scales poorly on complex SQL rows: "On store product, Syncode and our correctors generate a valid sample in 1.7–4.0s on average, while ASAp takes 18.87s, roughly an order of magnitude slower. This gap narrows on the other three rows... but ASAp's cost grows more steeply, reaching 93–108s on election and cre theme park." On cre theme park, ASAp's vocabulary had to be simplified to obtain a timing result.

1. SHIM resolves the trilemma of latency, fidelity, and scalability by learning only the probability adjustment based on internal parser state rather than masking

2. It exploits parser and lexer states that are already computed as a necessary part of incremental parsing for masking – extracting them adds negligible overhead while leaving the base LM's weights completely untouched

3. Unlike HMM-based augmentation, SHIM only needs to learn the narrow correction induced by masking rather than distill the LLM's entire output distribution

4. The method is interpretable and cheap to retarget to new grammars

The paper concludes: These results demonstrate that lightweight, offline-trained probability control effectively closes the gap between speed and distributional fidelity in grammar-constrained decoding. On BV4, the full-feature MLP achieves up to 113× lower KL divergence than Syncode (mean 23×, median 7×) and up to 235× lower divergence than online sampling (mean 56×, median 21×). On Spider, SHIM matches or surpasses baseline distributions while maintaining fast inference speeds, whereas online sampling scales poorly on complex SQL rows, failing to complete within practical time budgets.

Improvements to AI Systems:

1. Grammar-Constrained Decoding with Distributional Fidelity: Integrate SHIM as a post-processing layer on top of existing GCD tools (Syncode, XGrammar, Outlines). The AI system now generates outputs that satisfy grammar constraints while preserving the original LLM's probability distribution, avoiding the unnatural or degraded outputs caused by naive masking. This is critical for production systems where both syntactic validity and semantic quality matter (e.g., code generation, SQL queries, structured data extraction).

2. Zero-Overhead Parser-State Exploitation: The system reads parser and lexer states that are already computed during incremental parsing for masking, using them as features to predict correction factors. This adds negligible computational cost (often faster than plain masking due to fewer forced long completions) while improving fidelity. The AI can now run constrained generation in real-time applications without sacrificing latency.

3. Lightweight, Retargetable Probability Control: Deploy the LR FULL variant (logistic regression on parser state, lexer state, and candidate token) as a default for new grammars. It requires only offline training data generation (running the masking tool over 1000 sequences per grammar) and trains in minutes. The system can rapidly adapt to new domain-specific grammars (e.g., new SQL dialects, custom configuration languages) without retraining the base LLM or expensive HMM distillation.

4. Next-Token-Only Correction for High-Speed Scenarios: Use the LR NT variant (candidate token alone) when inference speed is paramount. It provides a meaningful fidelity improvement over masking on many grammars with near-zero added cost, making it suitable for high-throughput serving environments (e.g., API endpoints, batch processing).

5. Adaptive Variant Selection Based on Grammar Complexity: The system can automatically choose between MLP, LR FULL, and LR NT based on grammar characteristics. For grammars with high structural ambiguity (e.g., near-duplicate column names in SQL), full-feature variants are preferred; for grammars where the next token is highly predictive (e.g., substring matching), LR NT suffices. This enables optimal latency-fidelity trade-offs per use case.

6. Offline Training Data Generation Pipeline: The system includes an automated pipeline that generates training data by running the masking tool's incremental parser over random sequences, recording parser/lexer states and candidate tokens with validity labels. This enables continuous improvement as new grammars are added, without manual feature engineering.

7. Replacement for Expensive Online Sampling: The system can replace ASAp and similar online resampling methods in applications where they are currently used for higher fidelity. SHIM achieves comparable or better distributional fidelity (up to 235× lower KL divergence than online sampling) while running 4.6–7× faster, making it feasible for real-time constrained generation where online sampling was previously impractical.

8. Improved Structured Output Generation for LLM Agents: The system enables LLM agents to reliably produce structured outputs (e.g., JSON, YAML, domain-specific languages) that are both grammatically valid and semantically natural. This is particularly valuable for autonomous agents that must interact with APIs, databases, or formal verification tools where malformed outputs cause cascading failures.

9. Interpretable Correction Factors: The system provides interpretable γ values at each decoding step, allowing developers to understand why certain tokens were favored or suppressed. This aids debugging and enables fine-grained control over generation behavior (e.g., adjusting the correction strength for specific grammar constructs).

10. Scalable Constrained Generation for Large Vocabularies: The system handles grammars with large open vocabularies (as demonstrated on Spider's text-to-SQL) without the exponential cost of online sampling. This enables constrained generation for complex, real-world schemas and domain-specific languages where vocabulary size previously limited applicability.