Scheduling Thoughts: Learning the Order of Thought in Diffusion Language Models

2026-06-22

AI Summary

Existing diffusion language models rely on heuristic decoding orders and lack systematic optimization of the "order of thought," severely compromising generation quality. This work proposes a Self-Aware Scheduling (SAS) framework that, for the first time, formalizes decoding order selection as a policy optimization problem. By deriving an upper bound on the KL divergence induced by sequential decoding mismatch, the method constructs a path-aware dense reward signal to train a lightweight policy network—while keeping the denoiser frozen—to learn optimal scheduling strategies. Evaluated on Sudoku, SAS improves accuracy from 82.0% to 91.8% (reaching 97.5% after fine-tuning) and achieves state-of-the-art results on mathematical reasoning benchmarks, attaining 76% on GSM8K and 41% on MBPP, thereby consistently outperforming existing heuristic approaches.

Abstract

Masked diffusion language models decode by iteratively unmasking tokens, where the unmasking order defines an "order of thought" that strongly influences generation quality yet is typically chosen heuristically. We derive a tractable upper bound on the sequential decoding mismatch, measured by the Kullback-Leibler divergence and expressed in terms of the model's pathwise log-likelihood, with tightness under sufficient model expressivity. This bound induces a dense self-aware reward over ordered trajectories, casting order selection as a principled policy optimization problem with a frozen denoiser. We instantiate this idea as Self-Aware Scheduling (SAS), which learns a lightweight order policy using Group Relative Policy Optimization and applies seamlessly to both any-order and semi-autoregressive decoding. On Sudoku with 1B MDM, SAS improves puzzle accuracy from 82.0% (best heuristic schedule) to 91.8%, and reaches 97.5% with second-stage fine-tuning along learned trajectories. On mathematical reasoning with LLaDA-8B, SAS improves pass@1 on GSM8K from 64% to 76% and on MBPP from 39.5% to 41%, consistently matching or exceeding heuristic schedules across generation lengths and block sizes. Project page: https://jimmyxu123.github.io/SAS

Problem

• diffusion language models
• unmasking order
• order of thought
• decoding schedule
• generation quality

Innovation

• diffusion language models
• order of thought
• self-aware scheduling
• policy optimization
• sequential decoding