Table of Contents
Fetching ...

Lookahead-then-Verify: Reliable Constrained Decoding for Diffusion LLMs under Context-Free Grammars

Yitong Zhang, Yongmin Li, Yuetong Liu, Jia Li, Xiaoran Jia, Zherui Li, Ge Li

TL;DR

This work tackles the challenge of enforcing context-free grammar constraints in non-autoregressive diffusion LLMs by introducing LAVE, a lookahead-then-verify constrained decoding framework. LAVE leverages parallel token distributions across masked positions to sample complete prefixes and verify their extendability with a CFG parser, ensuring that every accepted token preserves the potential to produce a valid sentence. A cache-enhanced recovery mechanism further helps the model escape stalled inference states, while ablations demonstrate that both lookahead verification and recovery meaningfully improve syntactic and functional correctness with negligible runtime overhead. Across three CFG-driven benchmarks and four dLLMs, LAVE substantially improves syntactic reliability and often functional correctness, highlighting its practical value for reliable formal-language generation with diffusion models. The approach preserves the non-autoregressive nature of dLLMs and offers robust performance across hyperparameters, making CFG-constrained generation more feasible in real-world applications.

Abstract

Diffusion Large Language Models (dLLMs) have demonstrated promising generative capabilities and are increasingly used to produce formal languages defined by context-free grammars, such as source code and chemical expressions. However, as probabilistic models, they still struggle to generate syntactically valid outputs reliably. A natural and promising direction to address this issue is to adapt constrained decoding techniques to enforce grammatical correctness during generation. However, applying these techniques faces two primary obstacles. On the one hand, the non-autoregressive nature of dLLMs renders most existing constrained decoding approaches inapplicable. On the other hand, current approaches specifically designed for dLLMs may allow intermediate outputs that are impossible to complete into valid sentences, which significantly limits their reliability in practice. To address these challenges, we present LAVE, a constrained decoding approach specifically designed for dLLMs. Our approach leverages a key property of dLLMs, namely their ability to predict token distributions for all positions in parallel during each forward pass. Whenever a new token is proposed by model, LAVE performs lookahead using these distributions to efficiently and reliably verify the validity of the proposed token. This design ensures reliable constraints by reliably preserving the potential for intermediate outputs to be extended into valid sentences. Extensive experiments across four widely used dLLMs and three representative benchmarks demonstrate that LAVE consistently outperforms existing baselines and achieves substantial improvements in syntactic correctness, while incurring negligible runtime overhead.

Lookahead-then-Verify: Reliable Constrained Decoding for Diffusion LLMs under Context-Free Grammars

TL;DR

This work tackles the challenge of enforcing context-free grammar constraints in non-autoregressive diffusion LLMs by introducing LAVE, a lookahead-then-verify constrained decoding framework. LAVE leverages parallel token distributions across masked positions to sample complete prefixes and verify their extendability with a CFG parser, ensuring that every accepted token preserves the potential to produce a valid sentence. A cache-enhanced recovery mechanism further helps the model escape stalled inference states, while ablations demonstrate that both lookahead verification and recovery meaningfully improve syntactic and functional correctness with negligible runtime overhead. Across three CFG-driven benchmarks and four dLLMs, LAVE substantially improves syntactic reliability and often functional correctness, highlighting its practical value for reliable formal-language generation with diffusion models. The approach preserves the non-autoregressive nature of dLLMs and offers robust performance across hyperparameters, making CFG-constrained generation more feasible in real-world applications.

Abstract

Diffusion Large Language Models (dLLMs) have demonstrated promising generative capabilities and are increasingly used to produce formal languages defined by context-free grammars, such as source code and chemical expressions. However, as probabilistic models, they still struggle to generate syntactically valid outputs reliably. A natural and promising direction to address this issue is to adapt constrained decoding techniques to enforce grammatical correctness during generation. However, applying these techniques faces two primary obstacles. On the one hand, the non-autoregressive nature of dLLMs renders most existing constrained decoding approaches inapplicable. On the other hand, current approaches specifically designed for dLLMs may allow intermediate outputs that are impossible to complete into valid sentences, which significantly limits their reliability in practice. To address these challenges, we present LAVE, a constrained decoding approach specifically designed for dLLMs. Our approach leverages a key property of dLLMs, namely their ability to predict token distributions for all positions in parallel during each forward pass. Whenever a new token is proposed by model, LAVE performs lookahead using these distributions to efficiently and reliably verify the validity of the proposed token. This design ensures reliable constraints by reliably preserving the potential for intermediate outputs to be extended into valid sentences. Extensive experiments across four widely used dLLMs and three representative benchmarks demonstrate that LAVE consistently outperforms existing baselines and achieves substantial improvements in syntactic correctness, while incurring negligible runtime overhead.
Paper Structure (28 sections, 9 equations, 5 figures, 5 tables, 1 algorithm)

This paper contains 28 sections, 9 equations, 5 figures, 5 tables, 1 algorithm.

Figures (5)

  • Figure 1: AR LLMs generate tokens autoregressively, so every intermediate output forms a complete prefix, whereas dLLMs generate tokens non-autoregressively, yielding intermediate outputs that are often incomplete prefixes due to remaining [MASK].
  • Figure 2: The dLLM proposes a token for one masked position, which is subsequently verified through Lookahead-Based Verification. A proposed token is accepted if at least one lookahead complete prefix is extendable under the target grammar; otherwise, it is rejected.
  • Figure 3: Experimental results from our preliminary empirical study.
  • Figure 4: Outputs generated by NO-CD, FS-CD, IG-CD, and LAVE on Dream-7B for a problem from CPP-Bench, respectively.
  • Figure 5: The performance of LAVE with different values of the lookahead size $N$ and proposal budget $\tau$. We use the gray dashed line to represent the employed hyperparameters.