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

Masked diffusion language models generate text by iteratively unmasking tokens, a process where the order of unmasking—the 'order of thought'—significantly impacts generation quality. Traditionally, this order is determined by fixed, heuristic rules based on local token uncertainty. The authors investigate whether this order can be treated as a learnable policy to improve reasoning performance, rather than relying on myopic heuristics.

The authors derive a tractable upper bound on the sequential decoding mismatch (measured by KL divergence) and express it through the model's pathwise log-likelihood. This theoretical derivation provides a dense, 'self-aware' reward signal that assigns credit to each unmasking decision. They implement this as Self-Aware Scheduling (SAS), a framework that uses Group Relative Policy Optimization (GRPO) to learn an optimal unmasking policy while keeping the base diffusion model frozen. This allows the model to learn which tokens to reveal first to make subsequent reasoning steps more reliable.

SAS demonstrates significant improvements over standard heuristic schedules across multiple domains. On a 1B parameter masked diffusion model, SAS increased Sudoku puzzle accuracy from 82.0% to 91.8%, reaching 97.5% with additional fine-tuning. On larger models (LLaDA-8B), SAS improved pass@1 accuracy on GSM8K from 64% to 76% and on MBPP from 39.5% to 41%. The results suggest that optimizing the order of thought is a powerful, controllable primitive for enhancing reasoning in diffusion-based architectures.

This work shifts the paradigm of diffusion model inference from static, human-designed heuristics to learned, model-specific strategies. By casting order selection as a principled policy optimization problem, the authors provide a scalable way to improve reasoning capabilities without requiring full model retraining, offering a flexible 'plug-and-play' enhancement for existing diffusion language models.