Lookahead Path Likelihood Optimization for Diffusion LLMs

2026-02-03

AI Summary

This work addresses the limitation of diffusion-based large language models (dLLMs), whose reasoning performance is constrained by the demasking order, as existing heuristic strategies—focused solely on local confidence—fail to ensure global consistency. To overcome this, the authors propose Path Log-Likelihood (Path LL) as a path-level objective function strongly correlated with downstream accuracy. They introduce the POKE estimator to prospectively predict the future Path LL of partial trajectories and integrate it within a Sequential Monte Carlo (SMC) framework, forming POKE-SMC for dynamic optimal path search. This approach transcends the limitations of conventional greedy strategies, achieving average accuracy gains of 2%–3% across six reasoning tasks and significantly advancing the accuracy-computation Pareto frontier on the LLaDA model with comparable computational cost.

Technology Category

• Search and Optimization: Learning to Search
• Planning, Routing, and Scheduling: Planning with Language Models
• Reasoning under Uncertainty: Stochastic Optimization

Application Category

• Search and Retrieval-Augmented AI: Search Tool Learning with LLM: Teaching LLMs to invoke search and make use of retrieved information
• Semantics and Knowledge: Data modeling to support human-machine intelligence, including LLMs agents, intelligent system behavior, explanations, and user-friendly interactions
• User Modeling, Personalization and Recommendation: Large Language Models (LLM) for user modeling and recommendation

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.

Problem

• Diffusion LLMs
• unmasking order
• inference performance
• decoding trajectory
• global consistency

Innovation

• Path Log-Likelihood
• Diffusion LLMs
• POKE
• Sequential Monte Carlo
• Unmasking Order