Search papers, labs, and topics across Lattice.
This paper addresses the limitations of Grammar Constrained Decoding (GCD) in Language Models (LMs) by introducing a lightweight logit correction method that leverages internal parser and lexer states to restore the model's true probability distribution. By doing so, it mitigates the bias introduced by rigid token masking, which often leads to suboptimal outputs, without incurring the computational costs associated with online sampling. The proposed method consistently outperforms traditional masking and online sampling across various grammars, demonstrating that even minimal contextual information can significantly enhance output quality while maintaining inference efficiency.
Internal parser states can effectively restore a language model's true distribution, closing the gap left by rigid token masking.
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.
