Lookahead Path Likelihood Optimization for Diffusion LLMs

TL;DR

Diffusion LLMs enable arbitrary-order generation but decoding quality hinges on the unmasking trajectory. The paper introduces Path LL as a trajectory-aware objective that correlates with downstream accuracy, and POKE as an efficient lookahead estimator; these feed into a Sequential Monte Carlo search (POKE-SMC) to dynamically identify high-quality unmasking paths. Empirical results on LLaDA-8B-Instruct and LLaDA-1.5-8B across six reasoning benchmarks show consistent accuracy gains and favorable accuracy–compute Pareto frontiers compared to decoding-time scaling baselines. The approach highlights the value of trajectory-aware objectives and principled lookahead in diffusion-based decoding, offering a scalable path to improved inference-time performance for dLLMs.

Abstract

Diffusion Large Language Models (dLLMs) support arbitrary-order generation, yet their inference performance critically depends on the unmasking order. Existing strategies rely on heuristics that greedily optimize local confidence, offering limited guidance for identifying unmasking paths that are globally consistent and accurate. To bridge this gap, we introduce path log-likelihood (Path LL), a trajectory-conditioned objective that strongly correlates with downstream accuracy and enables principled selection of unmasking paths. To optimize Path LL at inference time, we propose POKE, an efficient value estimator that predicts the expected future Path LL of a partial decoding trajectory. We then integrate this lookahead signal into POKE-SMC, a Sequential Monte Carlo-based search framework for dynamically identifying optimal unmasking paths. Extensive experiments across 6 reasoning tasks show that POKE-SMC consistently improves accuracy, achieving 2%--3% average gains over strong decoding-time scaling baselines at comparable inference overhead on LLaDA models and advancing the accuracy--compute Pareto frontier.

Figures (7)

• Figure 1: What defines a good order?(Left)dLLM performance across diverse tasks using different ordering mechanisms. No single static heuristic excels universally, whereas our method achieves consistent gains by dynamically optimizing Path LL. (Right)Comparison of proxy objectives on GSM8K. When ranking samples by proxy scores, Path LL exhibits a much stronger correlation with accuracy than the ELBO or Path Entropy.
• Figure 2: High-Fidelity Path LL Estimation. We aggregate path LL estimates across all intermediate steps for a GSM8K sample. While the product LL ($\sum\log q_\theta$) severely underestimates the true LL under large $|\mathcal{A}|$ (small $K$), POKE effectively bridges this gap and achieves accurate estimates.
• Figure 3: Inference-time scaling efficiency on GSM8K, MATH500, and Countdown. Accuracy is plotted against average inference latency (s) per sample. Our method (purple) consistently establishes a superior Pareto frontier across all tasks, scaling more effectively than Majority Voting (blue) and E-SMC (green).
• Figure 4: Comparison of Accuracy Curves (via Global Pooling) across Proxy Scores on GSM8K.
• Figure 5: Comparison of incorrect and correct reasoning trajectories on a GSM8K example. Despite identical prompts, premature commitment in the incorrect trajectory leads to an erroneous final answer, whereas coherent reasoning yields the correct solution.

Theorems & Definitions (9)

• Definition 5.1: Total Correlation
• Lemma 5.2: Cumulative Total Correlation Bound
• Definition 6.1: Mutual information
• Definition 6.2: Potential Function
• Proposition 6.3
• proof
• Proposition 6.4
• proof
• Lemma 6.5: New Pinning Lemma anari2025parallel