Step-Aware Policy Optimization for Reasoning in Diffusion Large Language Models

ICLR 2026 Conference Submission14870 Authors

Keywords: Reasoning, LLM, Diffusion language models

TL;DR: We propose a diffusion-step aware policy optimization for improving reasoning ability of diffusion large language models

Abstract: Diffusion language models (dLLMs) offer a promising, non-autoregressive paradigm for text generation, yet training them for complex reasoning remains a key challenge. Current reinforcement learning approaches often rely on sparse, outcome-based rewards, which can reinforce flawed reasoning paths that lead to coincidentally correct answers. We argue that this stems from a fundamental mismatch with the natural structure of reasoning. We first propose a theoretical framework that formalizes complex problem-solving as a hierarchical selection process, where an intractable global constraint is decomposed into a series of simpler, localized logical steps. This framework provides a principled foundation for algorithm design, including theoretical insights into the identifiability of this latent reasoning structure. Motivated by this theory, we identify \textit{unstructured refinement}—a failure mode where a model's iterative steps do not contribute meaningfully to the solution—as a core deficiency in existing methods. We then introduce Step-Aware Policy Optimization (SAPO), a novel RL algorithm that aligns the dLLM's denoising process with the latent reasoning hierarchy. By using a process-based reward function that encourages incremental progress, SAPO guides the model to learn structured, coherent reasoning paths. Our empirical results show that this principled approach significantly improves performance on challenging reasoning benchmarks and enhances the interpretability of the generation process.

Primary Area: foundation or frontier models, including LLMs

Submission Number: 14870