The Parser Already Knows: Lightweight Bias Correction in Constrained Decoding

2026-08-10

AI Summary

This work addresses the distortion of a language model’s original probability distribution caused by rigid masking in grammar-constrained decoding, which often yields valid but suboptimal outputs, while existing distribution-recovery methods incur substantial computational overhead. The authors propose a lightweight, offline-trained logit correction approach that leverages lexical and parser internal states—along with candidate next tokens—from the incremental parsing process as priors to recover the true distribution without modifying model weights. This method achieves effective correction at zero additional inference cost and implicitly captures lookahead effects using only candidate tokens. Experiments across multiple grammars demonstrate that the proposed technique significantly reduces the divergence between masked and true distributions, consistently outperforming both pure masking and online resampling baselines, with even its lightest variant matching or surpassing their performance.

Abstract

Grammar Constrained Decoding (GCD) forces Language Models (LMs) to produce syntactically valid outputs by masking out non-conforming tokens at each step. However, rigid masking distorts the model's underlying probability distribution, often biasing generation toward valid but suboptimal outputs. While online sampling restores this distribution, it requires computationally expensive iterative resampling. As a result, existing methods force a compromise between output quality and inference latency. Our key insight is that the internal parser and lexer states inherently maintained during incremental parsing already encode future grammatical validity -- exactly the information required to restore the LM's true distribution. We propose a lightweight, offline-trained logit correction conditioned on this syntactic and lexical state together with candidate next tokens. Because these states 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. Across several grammars, this correction substantially closes the gap between the masked distribution and the LM's true distribution, consistently outperforming both masking and online sampling. Even its lightest variant, which relies on the candidate next token alone, still matches or exceeds both baselines: the next token itself carries an implicit lookahead, much like how parsers commonly use a lookahead token to resolve ambiguous decisions. By restoring the probability mass that masking removes, it reconciles the LM's probabilistic integrity with grammar conformance.

Problem

• Grammar Constrained Decoding
• bias correction
• language models
• probability distribution
• constrained generation

Innovation

• Grammar Constrained Decoding
• Bias Correction
• Incremental Parsing
• Logit Correction
• Lightweight Inference