Lookahead Path Likelihood Optimization for Diffusion LLMs
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.
Lay Summary
Large language models usually generate answers from left to right, but diffusion language models can fill in different parts of an answer in many possible orders. This flexibility creates a new challenge: some generation orders lead to much better answers than others, but it is not obvious how to choose them. In this work, we study how to evaluate and guide these generation paths. We introduce a way to score whether a partially generated answer is likely to lead to a good final answer, and we use this score to guide the model toward more promising paths while still exploring alternatives. This helps diffusion language models produce better answers on tasks involving math, coding, and planning. Beyond improving performance, our results also show that the order in which a model fills in an answer is an important factor in how well it solves a problem.