The Parser Already Knows: Lightweight Bias Correction in Constrained Decoding

Grammar Constrained Decoding (GCD) ensures that language models (LMs) produce syntactically valid outputs by masking out invalid tokens. However, this rigid masking distorts the model's original probability distribution, often forcing the model toward suboptimal or unnatural sequences. Existing solutions to this problem, such as online sampling, are computationally expensive and significantly increase inference latency. This paper asks: can we restore the LM's true probability distribution during constrained decoding using a lightweight, offline-trained correction that avoids the overhead of iterative resampling?

The authors propose SHIM, a probability controller that acts as a refinement layer on top of existing masking tools. The key insight is that the internal parser and lexer states—already maintained by incremental parsers to perform masking—contain sufficient information to predict future grammatical validity.

SHIM uses these states, along with the candidate next token, as features for a lightweight model (such as a logistic regressor or a small MLP). This model is trained offline to predict a correction factor (gamma) that adjusts the masked logits. By adding this factor to the logits during inference, the system effectively reweights the remaining valid tokens to better align with the base model's original, unconstrained distribution.

• Improved Probabilistic Integrity: SHIM significantly reduces the KL divergence between the constrained output and the base model's true distribution compared to standard masking and online sampling methods.
• Efficiency: Because the correction is a simple logit adjustment based on pre-computed parser/lexer states, it adds negligible overhead to the inference process, making it much faster than online sampling.
• Lookahead Capability: Even the simplest variant of the model, which relies only on the candidate next token, performs competitively, demonstrating that the next token itself carries significant implicit information about future grammatical constraints.