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

Unknown

Conceptio › Archive › arXiv CS

arXiv CSopen access

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

Unknown · 2026 · arxiv_cs

arXiv CS · Papers · License: Open Access · 2026

Open Source ↗Direct PDF ↓

machine learning, deep learning, neural networks

Scheduling Thoughts: Learning the Order of Thought in Diffusion Language Models Jiawei Xu 1 Minghui Liu 1 Aakriti Agrawal 1 Yifan Chen 2 † Furong Huang 1 †

arXiv:2606.23567v1 [cs.LG] 22 Jun 2026

Abstract

fixed left-to-right factorization of autoregressive language models (Radford et al., 2018). Crucially, diffusion decoding is not tied to a single generation order: inference depends on an unmasking schedule that decides which positions are revealed at each step. This schedule is more than an implementation detail: it is an implicit “order of thought”, and it determines which commitments the model makes early, which constraints it propagates, and which ambiguities it strategically postpones. Because different schedules trace different conditional trajectories, they induce different output distributions, and the schedule choice can substantially affect generation quality (Kim et al., 2025; Zheng et al., 2023).

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 loglikelihood, with tightness under sufficient model expressivity. This bound induces a dense selfaware 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 secondstage 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.

Heuristic schedules are fast, but they can think myopically. Most existing schedules are fixed greedy heuristics based on per-position uncertainty (e.g., confidence, margin, entropy) (Zheng et al., 2023; Kim et al., 2025; Ben-Hamu et al., 2025). While efficient, such rules optimize the next step rather than the whole decoding trajectory: they often pick “easy” tokens that become locally certain under the current partial context, even when a different reveal would better shape future conditionals. Empirically, we find that these myopic choices underperform on complex reasoning tasks (see Table 1 and Figure 2 in Section 6.1). Even worse, expert-designed logical orders can be suboptimal for an imperfect pretrained diffusion model: what is logically natural for humans is not necessarily the path along which the model’s conditional predictions are most reliable (see Section 6.1). These observations motivate a principled approach that learns the order, instead of hard-coding it.

1. Introduction

Key idea and theory: learn the schedule by optimizing the model’s trajectory likelihood. We treat diffusion decoding as a latent-order generative process (Huang et al., 2025; Wang et al., 2025d) and ask a direct question: which reveal order makes the model most faithful to the data distribution? Concretely, we formalize scheduling as minimizing the KL divergence between the data distribution and the distribution induced by a given decoding procedure. Our main theoretical contribution is to show that this intractable objective admits a tractable bound expressed in terms of the denoiser’s pathwise (trajectory) log-likelihood under teacher forcing. This yields a dense, model-self-aware reward that assigns credit to each reveal decision, rather than only the

Diffusion decoding has a hidden degree of freedom: the order of thought. Masked diffusion language models (Austin et al., 2021; Lou et al., 2023; Hoogeboom et al., 2021b; Shi et al., 2024) generate discrete sequences by iteratively unmasking tokens, offering a flexible alternative to the † Equal advising. 1 University of Maryland, College Park University of California, Los Angeles. Correspondence to: Yifan Chen <[email protected]>, Furong Huang <[email protected]>. 2

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

final sample. Under this reward, a good schedule reveals the token maximizing the expected future log-likelihood of the remaining trajectory, i.e., it chooses the next “thought” to make downstream reasoning easiest for the model, aligning the decoding path with the model’s predictive strengths rather than with local uncertainty alone.

math and coding reasoning, and show additional gains from second-stage fine-tuning along learned trajectories under the same self-aware objective.

2. Preliminaries and Problem Setup 2.1. Masked Diffusion Models

Method: Self-Aware Scheduling (SAS) for plug-and-play order learning. Guided by the theory, we propose SelfAware Scheduling (SAS), a lightweight framework that learns an order policy while keeping the diffusion model fixed. At each decoding step, the policy outputs a distribution over currently masked positions and selects (or samples) one position to reveal next. SAS is plug-and-play (no retraining of the base denoiser), and it can be optimized efficiently with Group Relative Policy Optimization (GRPO) (Shao et al., 2024). The same learned policy applies seamlessly to both any-order sequential decoding and semi-autoregressive (block) decoding (Arriola et al., 2025), turning “how to unmask” into a trainable component of inference.

Masked diffusion models generate discrete sequences by iteratively demasking tokens, analogous to the denoising process in continuous diffusion models (Song et al., 2020; Ho et al., 2020). Let X denote a finite vocabulary and n the sequence length. A sequence is x = (x1 , . . . , xn ) ∈ X n . For a mask pattern M ⊆ [n] := {1, . . . , n}, define the masked sequence xM ∈ (X ∪ {[MASK]})n by ( [MASK] if i ∈ M, M (x )i := (1) xi if i ∈ / M. A masked diffusion model (or denoiser) parameterized by θ defines conditional distributions over masked tokens given observed context. Specifically, the model pθ provides position-wise predictions piθ (· | xM ), i ∈ M, representing the distribution over token xi conditioned on the partial observation xM . Here θ denotes the learned parameters of the base model.

Results: optimizing the order of thought is a strong lever for reasoning. Across diverse reasoning benchmarks, SAS yields consistent and substantial improvements over heuristic and expert-designed schedules. On large-scale Sudoku with a 1B masked diffusion model, SAS improves puzzle accuracy from 82.0% (best heuristic) to 91.8%. On mathematical reasoning (GSM8K) with LLaDA-8B (Nie et al., 2025b), it improves pass@1 from 64% to 76%. We further show that a simple second-stage fine-tuning procedure along the learned trajectories using our self-aware objective provides additional gains. Together, these results position decoding order as a controllable and optimizable reasoning primitive: beyond learning what to generate, we can learn when to commit to each piece of information—in effect, learning an order of thinking that the model itself finds most reliable.

2.2. Unmasking Schedules: the Order of Thought Generation proceeds by iteratively unmasking positions, a process known as decoding. The unmasking schedule can be interpreted as an “order of thought” for generation. The flexibility of the masked diffusion model is that it allows any-order decoding in principle, and thus the unmasking order, or schedule, matters. The goal of our paper is to introduce a methodology to learn the unmasking order. We define the schedule vϕ , parameterized by ϕ, as the “order policy”, which takes a partial sequence as input and outputs a distribution over which position(s) to unmask next.

Summary of Contributions. • Theory (self-aware objective for order learning): We cast diffusion decoding as a latent-order generative process and derive a tractable bound linking decoding mismatch (KL divergence) to the denoiser’s cumulative pathwise log-likelihood, yielding a dense, modelself-aware reward for credit assignment over reveal decisions.

We denote by x̃(t) the partial sequence at decoding step t, where some positions contain tokens from X (revealed) and others contain [MASK] (not yet revealed). Starting from the fully masked sequence, at each step the algorithm selects which position(s) to unmask and samples their values from the model. We distinguish between two decoding regimes: sequential decoding, which unmasks a single position per step, and parallel decoding, which unmasks multiple positions simultaneously. In this work, we focus on the sequential setting to isolate the effects of ordering, while the parallel regime is detailed in Appendix A.1.

• Algorithm (SAS): We introduce Self-Aware Scheduling, a plug-and-play framework that learns an unmasking-order policy for a frozen diffusion LM via GRPO, applicable to both any-order and semiautoregressive decoding.

Unmasking procedure. Let St ⊆ [n] denote the set of revealed positions at step t, and Mt = [n]\St denote the set of masked positions. Starting from x̃(0) = x[n] (fully masked) with S0 = ∅, at each step t = 1, . . . , n: (1) Select

• Empirical validation (reasoning gains + second stage): We demonstrate large, consistent improvements over heuristic and expert schedules on Sudoku, 2

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

next position: it ∼ vϕ (· | x̃(t−1) ) where vϕ (· | x̃(t−1) ) outputs a distribution over [n] \ St−1 . (2) Sample token: (t) x̃it ∼ piθt (· | x̃(t−1) ). (3) Update: St ← St−1 ∪ {it }; set (t)

3. Unmasking Order Optimization: From Error Bounds to Reward Design

(t−1)

x̃j = x̃j for all j ̸= it . The complete run produces an unmasking order σ = (i1 , . . . , in ) and final sequence x̃(n) ∈ X n .

The Intractability Problem seq Directly minimizing KL(π(x)∥Pθ,ϕ (x)) requires averaging over intractable n! decoding orders.

Induced distributions. The decoding procedure induces distributions over output sequences. Characterizing these distributions is essential to quantify the discrepancy between the generative process and the true data distribution, thereby providing a principled objective for optimizing the order policy. To achieve so, we define notations for the decoding algorithm under teacher forcing (Williams & Zipser, 1989), where the policy vϕ and model pθ are evaluated on partially revealed versions of a fixed target sequence x, as will be explained below. This allows us to compute the probability the algorithm assigns to generating x under a particular unmasking order.

Bridge with theory to derive reward The Self-Aware Framework 1. Derive tractable ELBO on likelihood (Theorem 3.1). 2. Use ELBO as a pathwise reward Rθ for policy learning. 3. Optimizing Rθ theoretically bounds generation error (Theorem 3.2). Figure 1. High-level overview of Section 3: We overcome the intractability of latent ordering by deriving a theoretical framework that allows us to optimize a tractable, self-aware reward.

This section formalizes how the unmasking order affects the induced decoding distribution. Our primary goal is to derive a principled objective for optimizing the unmasking policy in sequential decoding. We also characterize the additional discrepancy introduced by parallel decoding via an information-theoretic quantity (Section B.1).

For a target sequence x and unmasking order σ = (i1 , . . . , in ), define the teacher-forced trajectory x̃(0) (σ), . . . , x̃(n) (σ) where x̃(t) (σ) reveals positions {i1 , . . . , it } from x: ( xj if j ∈ {i1 , . . . , it }, (t) x̃ (σ)j := (2) [MASK] otherwise.

seq Let π(x) be the target distribution on X N and let Pθ,ϕ (x) be the output law induced by sequential decoding (Section 2.1). We aim to fit the induced decoder distribution to the data by maximizing its expected log-likelihood, seq max Ex∼π log Pθ,ϕ (x) , (7)

In other words, x̃(t) (σ) is the partial sequence obtained by revealing the first t positions in order σ using their values from target x. Under teacher-forcing with target x and order σ, the order policy (parameterized by ϕ) assigns probability vϕ (σ | x) :=

n Y

vϕ (it | x̃(t−1) (σ))

θ,ϕ

seq which is equivalent to minimizing KL(π∥Pθ,ϕ ) up to the constant H(π).

(3)

t=1

3.1. Self-aware Reward as Pathwise Likelihood

to the unmasking order σ, where each factor evaluates the policy vϕ on the teacher-forced state x̃(t−1) (σ). The model (parameterized by θ) assigns path likelihood pθ (x | σ) :=

n Y

piθt (xit | x̃(t−1) (σ)),

For a data sample x and an unmasking order σ = (i1 , . . . , in ), let pθ (x | σ) denote the pathwise likelihood assigned by the diffusion model when the tokens of x are revealed according to σ under teacher forcing. We define the self-aware reward as

(4)

t=1

Rθ (x, σ) := log pθ (x | σ).

where each factor evaluates the model pθ at position it given teacher-forced state x̃(t−1) (σ). The sequential algorithm induces the joint distribution over sequences and orders seq Pθ,ϕ (x, σ) := vϕ (σ | x) · pθ (x | σ),

(5)

with marginal distribution over outputs X seq Pθ,ϕ (x) = vϕ (σ | x) · pθ (x | σ).

(6)

(8)

Thus, the reward is simply the frozen model’s own pathwise log-likelihood evaluated along a specific order. It scores an order by how well the diffusion model can explain the data when forced to reveal tokens in that order. We train the order policy by minimizing the negative expected self-aware reward:

minθ,ϕ LSAS (θ, ϕ) := minθ,ϕ −Ex∼π Eσ∼vϕ (·|x) Rθ (x, σ) . (9)

Remark 2.1. Heuristic schedules can be formalized as a special case of our order policy framework where the policy vϕ is deterministic and non-trainable. Details are provided in Appendix A.3.

In this work we primarily learn ϕ with θ frozen, and optionally apply a second-stage fine-tuning of θ under a fixed learned order (Section 6.4). 3

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

We next record two simple facts connecting this pathwise objective to the marginal likelihood over orders. These results are not needed to motivate the reward itself; rather, they clarify how the reward relates to the likelihood-based objective.

easier generation paths for the frozen denoiser, while improving θ improves the pathwise predictions. Both reduce the joint mismatch, as we observe empirically in Section 6.4. The proof is in Appendix B.4. Remark 3.5. Although the main text focuses on sequential decoding, the same self-aware loss also provides control over parallel decoding mismatch. In Appendix B, we show that the additional error caused by block updates is upperbounded by a term governed by within-block conditional dependencies, which, together with the self-aware loss, upper bounds the total variation error of the generated distribution.

Theorem 3.1 (Pathwise likelihood lower bound). For any π and order policy vϕ , the marginal log-likelihood satisfies seq Ex∼π log Pθ,ϕ (x) ≥ Ex∼π Eσ∼vϕ (·|x) log pθ (x | σ) .

(10) seq Proof. Since Pθ,ϕ (x) = Eσ∼vϕ (·|x) [pθ (x | σ)], Jensen’s seq inequality gives log Pθ,ϕ (x) ≥ Eσ∼vϕ (·|x) [log pθ (x | σ)]. Taking expectation over x ∼ π proves the claim.

4. Methodology and Policy Optimization We focus on sequential (one-by-one) unmasking, in which the decoding procedure reveals exactly one position per step. Concretely, we freeze the diffusion model parameters θ (and its denoising head) and optimize only the order-policy parameters ϕ using reinforcement learning.

3.2. KL Interpretation for Sequential Decoding We also relate the self-aware loss to the marginal mismatch seq KL(π∥Pθ,ϕ ). This gives a compact interpretation of the pathwise objective as a joint distribution matching objective over data and orders.

4.1. Learning with Self-aware Reward

Theorem 3.2 (Joint KL identity and marginal bound). Let Qϕ (x, σ) := π(x)vϕ (σ|x) and Pθ,ϕ (x, σ) := pθ (x|σ)vϕ (σ|x) be the data-policy and model-policy joint distributions, respectively. The self-aware loss LSAS satisfies seq KL(π(x)∥Pθ,ϕ (x)) ≤ LSAS (θ, ϕ) − H(π)

= KL(Qϕ (x, σ)∥Pθ,ϕ (x, σ)).

Monte Carlo approximation of the objective. For a target sequence x ∼ π and order σ, we define the return as the pathwise log-likelihood of the target tokens under the frozen denoiser: Rθ (x, σ) =

(11)

N X

log piθt (xit | x̃(t−1) (σ)).

(12)

t=1

This serves as a dense, self-aware signal for credit assignment.

where H(π) is the constant data entropy. The proof is provided in Appendix B.2. Theorem 3.2 shows that minimizing the self-aware loss is equivalent, up to the constant H(π), to minimizing a joint KL mismatch over (x, σ). This joint mismatch also upper bounds the marginal seq generation mismatch KL(π(x)∥Pθ,ϕ (x)).

Order-policy objective (optimize ϕ only). We learn the order policy vϕ (σ | x) by maximizing this expected return while keeping the diffusion model θ fixed: h i ϕ⋆ ∈ arg max Ex∼π Eσ∼vϕ (·|x) Rθ (x, σ) . (13)

Corollary 3.3 (Joint realizability and tightness). If the model family is expressive enough to allow for joint realizability—i.e., there exist parameters (θ⋆ , ϕ⋆ ) such that KL(Qϕ⋆ (x, σ)∥Pθ⋆ ,ϕ⋆ (x, σ)) = 0—then the marginal misseq match KL(π∥Pθ,ϕ ) is zero, and the pathwise lower bound of Theorem 3.1 is tight, satisfying LSAS (θ⋆ , ϕ⋆ ) = H(π).

4.2. Parameterization of the Order Policy State-dependent categorical policy over remaining indices. At step t, the policy observes the current partial state st−1 := x̃(t−1) and chooses the next index it from the masked set Mt−1 . We parameterize the order policy as a categorical distribution over indices:

The proof is provided in Appendix B.3. Corollary 3.3 states that under joint realizability the bound seq is achievable, yielding KL(π(x)∥Pθ,ϕ (x)) = 0. More generally, if approximate realizability holds with KL(Qϕ⋆ (x, σ)∥Pθ⋆ ,ϕ⋆ (x, σ)) ≤ ϵ, then the marginal misseq match is also bounded as KL(π(x)∥Pθ,ϕ (x)) ≤ ϵ. Remark 3.4. In the regime of a perfect diffusion model pθ , the upper bound in Theorem 3.2 is tight with LSAS (θ⋆ , ϕ) = H(π) for any order policy ϕ. Thus, order learning is most useful when the denoiser is imperfect: improving ϕ finds

vϕ (σ | x) =

N Y

vϕ it | st−1 ,

t=1

(14)

vϕ i | st = 0 if i ∈ / Mt . Scoring form (masked softmax). Given a partial state st , we compute a scalar score uϕ,i (st ) ∈ R for each position 4

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

i ∈ [N ] and define

et al., 2025; Kim et al., 2025; Li et al., 2025; Yu et al., 2025), with variants incorporating spatial-temporal structure (Huang et al., 2026; Wang et al., 2025b) or token dependencies (Azangulov et al., 2025). We instead learn a sequential policy via RL to optimize for global trajectory coherence.

vϕ (i | st ) =

exp uϕ,i (st )/τ 1{i ∈ Mt } , P j∈Mt exp uϕ,j (st )/τ

(15)

where τ > 0 is a temperature. In our implementation, uϕ,i (st ) is produced by a lightweight policy head on top of frozen diffusion-model features. We refer the reader to Appendix C for details of order policy design.

Learning orders. Finding effective decoding orders has been explored in non-autoregressive generation (Wu et al., 2018; Zhu et al., 2019; Gu et al., 2019; Welleck et al., 2019; Stern et al., 2019). For diffusion LMs, recent work optimizes schedules using task-specific verifiable rewards: Huang et al. (2025); Zhao et al. (2025a) jointly train the policy and MDM, while Hong et al. (2025); Jazbec et al. (2025) learn policies for frozen models using sparse terminal rewards. We propose minimizing KL divergence—a general objective that yields dense rewards and applies even when verifiable rewards are unavailable. Since discrete diffusion equivalently represents any-order autoregressive models (Hoogeboom et al., 2021a; Ou et al., 2024), related approaches derive similar ELBO-based objectives for diffusion molecular generation (Wang et al., 2025d), though requiring two additional networks and more complex formulation. We also note related work on optimizing orders for efficient parallel decoding to accelerate generation (Chen et al., 2024; Shih et al., 2023; Park et al., 2024), complementary to the statistical efficiency we pursue. Our theoretical bounds extend to parallel decoding (Appendix B.1).

4.3. GRPO training MDP formulation. Order learning can be cast as a finitehorizon Markov Decision Process (MDP): • State st−1 : the partially revealed sequence x̃(t−1) • Action at ∈ Mt−1 : choose the next index it to unmask. • Transition: deterministic teacher-forced reveal, st = T (st−1 , at ) obtained by setting position it to xit . • Reward: rt = log piθt (xit | st−1 ). The trajectory corresponds to an order σ, and the return equals the self-aware reward Rθ (x, σ). Group Relative Policy Optimization (GRPO). We optimize vϕ using GRPO (Shao et al., 2024), a critic-free method that employs group-relative normalization. For each input x, we sample a group of G outputs {σ (i) }G i=1 from the bi is computed by standardold policy vϕold . An advantage A izing the cumulative reward Ri against the group statistics: bi = (Ri − µG )/(σG + ϵadv ). The GRPO objective maxiA mizes a clipped surrogate function: " G N 1 X 1 X min gi,t (ϕ)Âi,t , JGRPO (ϕ) = E G i=1 N t=1 # (16) clip gi,t (ϕ), 1 − ϵ, 1 + ϵ Âi,t (i)

with the per-step probability ratio gi,t =

Diffusion LLMs post-training. Recent RL-based posttraining focuses on refining the diffusion backbone itself. For instance, D1 (Zhao et al., 2025b) introduces a GRPO variant tailored for discrete diffusion, while other works concentrate on improving policy gradient estimation (Tang et al., 2025; Wang et al., 2025c; Lin et al., 2025; Zhu et al., 2025; Wang et al., 2025a) to enhance reasoning capabilities. These approaches typically use fixed confidence-based schedules during optimization. In contrast, we propose a two-stage framework prioritizing generation order: first optimizing the unmasking policy with a frozen diffusion head, then fine-tuning the head along learned trajectories. This approach is computationally efficient without expensive RL updates on the full model.

(i)

vϕ (σt |st−1 ) (i) (i) vϕold (σt |st−1 )

Each training example pairs a prompt c with a target completion x. The policy operates exclusively on the completion positions, treating the prompt as fixed context. Thus, the process begins at s0 = concat(c, x̃(0) ), where the prompt is fully visible and the target is entirely masked.

6. Experiments In this section, we empirically validate that our self-aware reward formulation is both effective and scalable. Our experiments span diverse reasoning benchmarks—including logic puzzles, mathematical reasoning, and code generation—and cover model scales ranging from 1B to 8B parameters.

5. Related Work Heuristic schedules for diffusion LLMs. The unmasking schedule critically determines information accrual in masked diffusion models (Nie et al., 2025b; Ye et al., 2025; Shi et al., 2024). Prior work predominantly uses trainingfree heuristics that prioritize tokens by confidence or local certainty (Ben-Hamu et al., 2025; Wei et al., 2025; Hong

Experimental Objectives. Our evaluation is guided by three primary research questions:(i) Performance: Can optimizing the unmasking order alone yield consistent improvements in reasoning tasks? (ii) Emergent Strategy:

Scheduling Thoughts: Learning the Order of Thought in Diffusion Language Models Table 1. Sudoku performance under different unmasking-order strategies. We report puzzle accuracy (fraction of boards solved) and cell accuracy (fraction of correctly filled cells).

(a) Accuracy

O RDER S TRATEGY

P UZZLE ACC . (%)

C ELL ACC . (%)

R ANDOM C ONFIDENCE (T OP -1) M ARGIN E NTROPY H UMAN ORDER

62.7 82.0 80.0 81.0 76.6

89.80 98.56 98.38 98.53 96.18

O RDER P OLICY (O URS )

91.8

99.21

(b) Order Visualization

Figure 2. Left: Our method (Order Policy) achieves significantly higher success accuracy (91.8%) on Sudoku puzzles compared to all baselines. Right: Visualization of the learned unmasking order. White cells represent fixed initial hints. Blue cells indicate generated tokens, where color intensity corresponds to the decoding step: lighter blue denotes early unmasking (highconfidence/foundational moves), while darker blue denotes later unmasking. For example, we red-boxed the first move located, which corresponds to the naked single of the initial state—the easiest position to solve.

as the [MASK] token and model the completion of masked cells as conditional generation. We instantiate with a 1B-scale masked diffusion model: SMDM-1B (Nie et al., 2025a). We first fine-tune the diffusion model on the 1M training puzzles for one epoch; we then train a lightweight MLP policy on 1,000 puzzles (sampled from the training dataset) and report results on the remaining test split. Baselines. At inference time, we compare the learned policy against heuristic decoding orders: (i) R ANDOM (Zheng et al., 2024); (ii) C ONFIDENCE (Zheng et al., 2023); (iii) M ARGIN (Kim et al., 2025); (iv) E NTROPY (Zheng et al., 2023); and (v) a human E XPERT order (Shah et al., 2024) based on Sudoku solving logic. Results. Table 1 reports puzzle and cell accuracy across all order baselines. Our learned order policy achieves the best performance by a large margin, confirming that optimizing unmasking order alone can substantially improve structured reasoning. Interestingly, the human E XPERT order underperforms: it is approximately 5% worse than C ONFIDENCE and about 15% worse than our learned policy in puzzle accuracy. This motivates a deeper analysis of the order strategies in Sudoku. To understand the decision logic of our Self-Aware Schedule, we visualize the decoding order of a specific puzzle in Figure 2b. The resulting heatmap demonstrates a clear “easy-to-hard” curriculum. The policy prioritizes deterministic unmasking cells given the current context (e.g., naked singles). By resolving these high-certainty tokens first (light blue), the model propagates constraints to the more ambiguous regions (dark blue).

Figure 3. Self-aware reward vs. 0/1 terminal reward for learning the order policy on GSM8K (5-shot). Notation L--B: total length L and block size B. Bars represent the accuracy difference in percentage points.

What decoding strategies emerge from the learned policy, and how do they correlate with the underlying task structure? (iii) Generalization: Does the learned order policy generalize across varying sequence lengths and decoding modalities (e.g., full diffusion vs. semi-autoregression)? We investigate (i) across all benchmarks, restrict our structural analysis (ii) to the interpretable domain of Sudoku, and assess (iii) using mathematical and code generation tasks.

6.2. Order Analysis of Sudoku and Kendall’s τ To analyze how different schedules align with human expert ordering—while accounting for the fact that many Sudoku moves are interchangeable—we introduce an equivalenceclass variant of Kendall’s τ (Kendall, 1948) which is agnostic to the relative ordering of moves that are strategyequivalent. We provide details in Appendix D.1.

6.1. Sudoku: Controlled Study of Unmasking Order Setup. We construct a large-scale Sudoku corpus from (Rao, 2019). From the provided training split, we subsample 1M puzzles for training. Each Sudoku board is represented as a length-81 discrete sequence where each token corresponds to one cell value and 0 denotes the blank cell. We treat 0

Results. Table 2 reports the mean τeq , the fraction of puzzles with τeq > 0, and the fraction with τeq > 0.5. While 6

Scheduling Thoughts: Learning the Order of Thought in Diffusion Language Models Table 2. Comparison with human expert order on Sudoku via equivalence-class Kendall’s τ (2000 puzzles). O RDER S TRATEGY

M EAN τeq

F RAC . τeq > 0

F RAC . τeq > 0.5

R ANDOM C ONFIDENCE (T OP -1) M ARGIN E NTROPY

-0.002 0.514 0.516 0.513

0.49 1.00 0.99 1.00

0.00 0.58 0.59 0.58

O RDER P OLICY (O URS )

0.735

0.59

0.47

Table 3. Second-stage fine-tuning of the diffusion model with a fixed unmasking order on Sudoku. Fine-tuning yields significant gains across schedules: Human Order: Accuracy improves from 76.6% → 97.8%. Order Policy (Ours): Accuracy improves from 91.8% → 97.5%. Notably, our self-discovered curriculum achieves parity with human expertise without needing ground-truth traces. F IXED ORDER FOR FINE - TUNING

B EFORE

A FTER

76.6 91.8

97.8 97.5

H UMAN ORDER O RDER P OLICY (O URS )

our learned order policy attains the highest agreement with expert ordering, the alignment remains imperfect (τeq < 1), suggesting that the model’s optimal path differs from human logical progression. When viewed alongside the performance gains in Table 1, this observation suggests a critical insight: for masked diffusion LMs, strict adherence to human solving orders is not a prerequisite for optimal performance, and the model may benefit from alternative reasoning pathways. In other words, expert logic provides a strong prior over precedence constraints, but the optimal schedule for a pretrained diffusion decoder is shaped from modeling uncertainty and statistical correlations, which our self-aware objective explicitly exploits.

achieves superior or competitive performance across the entire spectrum of decoding regimes. Notably, we show while Left-to-Right scheduled decoding is competitive at shorter lengths (L = 128), our order policy achieves superior performance even in full diffusion generation at longer lengths. This confirms that our scheduling strategy captures fundamental reasoning structures and is highly robust to shifts in generation horizon, effectively transferring well beyond its training configuration. 6.4. Second-Stage Fine-Tuning with the Self-Aware Objective

6.3. Generalization to Math and Code Reasoning

Corollary 3.3 shows that, in the realizable setting, jointly optimizing the unmasking order policy and the diffusion model can in principle drive the sequential decoding error to zero. Motivated by this, we propose a simple second-stage optimization: after learning an order policy, we freeze the policy and further fine-tune the diffusion model using the same self-aware objective.

Setup. We then move to study whether an order policy learned with the self-aware objective transfers to largerscale reasoning tasks beyond Sudoku. We use LL A DA-8BI NSTRUCT (Nie et al., 2025b) as the frozen base masked diffusion model and parameterize the order policy as a lightweight one-layer Transformer. We train the policy under the full diffusion (any-order) sequential decoding paradigm from the training splits of GSM8K (Cobbe et al., 2021) and MBPP (Austin et al., 2021). The policy is trained with teacher-forced trajectories at a maximum generation length of 512. Training details and inference costs are provided in Appendix D.2.

Fixed trajectory optimization. The self-aware objective (Eq. (9)) can be used to optimize both θ and ϕ. After learning an order policy v̄, we fix it and replace the stochastic order distribution by a deterministic rollout determined by v̄. This reduces the training signal from a corpus of possible orders to a single order trajectory, yielding a simple supervised fine-tuning loss for the diffusion model.

Evaluation protocol. We evaluate (i) generalization across generation lengths and (ii) transfer to semi-autoregressive (block) inference. For length generalization, we report results for target lengths in {128, 256, 512}. For semiautoregressive decoding, we vary the number of decoding blocks K ∈ {1, 4, 8} and apply the same learned policy to schedule these blocks at inference time. We denote each setting as promptlen/total–block, e.g., 256--64 indicates total length 256 with block size 64 (semi-autoregressive), while 256--256 corresponds to full diffusion decoding. We evaluate the results using lm-eval-harness.

Second-stage objective. With the fixed order σ̂, the optimization over θ becomes h i max Ex∼π Rθ x, σ̂ ≡ min LSFT (θ), (17) θ

which is the negative log-likelihood of ground-truth tokens along the fixed unmasking schedule. We validate this second-stage optimization on Sudoku using two fixed orders: (i) the learned order policy v̄ (greedy rollout), and (ii) a human expert order. Table 3 reports puzzle accuracy before and after second-stage fine-tuning. With the learned order, fine-tuning the diffusion head substantially improves puzzle accuracy, from 91.8% to 97.5%. We also observe that fine-tuning under the human expert order leads to a large improvement from 76.6% to 97.8%.

Results. In Figure 4, it demonstrates the robust generalization of our approach on GSM8K and MBPP. While distinct heuristic baselines exhibit inconsistent performance across different settings, our learned order policy uniquely 7

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

Left-to-Right

Margin

Entropy

Ours

GSM8K

L = 128

0.8

Accuracy

Confidence

L = 256

L = 512

0.7 0.6 1

Number of Decoding Blocks (K)

MBPP

L = 128

0.425

Accuracy

L = 256

L = 512

0.400 0.375 0.350 0.325

Number of Decoding Blocks (K)

Figure 4. Comparison across generation length and semi-autoregressive decoding. We evaluate our proposed order policy against the other baselines across varying sequence lengths (L) and block numbers. The red dashed line represents standard Left-to-Right decoding heuristic of LLaDA-8B. Our method consistently outperforms the baseline. Notably, while Left-to-Right decoding is competitive at shorter lengths (L = 128), our order policy still achieves superior performance even in full diffusion generation at longer lengths. Table 4. Cross-domain transfer of learned order policies. “Math → Code” denotes training the order policy on math and evaluating it on code, while “Code → Math” denotes the reverse. Numbers in parentheses indicate gains over the strongest heuristic baseline under the same decoding setting.

Notably, our order learning matches the performance of supervised learning under expert ordering, demonstrating that the learned policy discovers a curriculum as effective as human expertise without requiring ground-truth ordering traces. Ultimately, these results underscore the importance of scheduling thoughts effectively before refining the thoughts themselves. While this work primarily focuses on learning the order schedule, it demonstrates that our self-aware objective serves as an effective and efficient post-training method on diffusion model weights, achieving significant gains without complicated RL on the model.

Transfer setting

256-256

256-64

256-32

Math → Code Code → Math

0.41 (+1.5%) 0.713 (+8.5%)

0.40 (+1.4%) 0.7967 (+0.27%)

0.39 (+0.33%) 0.8033 (+0.8%)

merely memorize task-specific orders, but learns transferable scheduling principles about which positions are more informative to reveal early. The gains are especially pronounced for Code → Math in the 256-256 setting, while remaining positive across smaller block sizes.

6.5. Cross-domain Transfer of Learned Order Policies To evaluate whether the learned order policy captures reusable scheduling behavior rather than only task-specific patterns, we test cross-domain transfer between math and code reasoning tasks. In particular, we train the order policy on one domain and evaluate it on the other without further task-specific adaptation. Table 4 reports the transferred policy’s performance under different decoding settings, where L-B denotes generation length and block size. The value in parentheses is the gain over the strongest heuristic baseline under the same setting.

6.6. Comparison with Other Order Learning Methods Our framework learns an unmasking-order policy using a dense self-aware reward derived from the diffusion model’s pathwise likelihood (Section 4). An alternative is to apply RLVR-style policy optimization (Guo et al., 2025) directly to the order policy using a sparse terminal reward, e.g., r ∈ {0, 1} indicating whether the final answer is correct (Jazbec et al., 2025; Hong et al., 2025). We compare our dense self-aware reward against a 0/1 terminal reward for learning the order policy on GSM8K under our settings.

The transferred policy outperforms the best heuristic baseline in all tested settings. This suggests that SAS does not

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

For the terminal-reward baseline, we assign reward 1 if the generated final answer matches the ground-truth answer and 0 otherwise, and optimize the same GRPO objective. Figure 3 reports 5-shot GSM8K accuracy comparison across generation lengths and decoding regimes. Overall, our dense self-aware reward is consistently comparable and often better across settings, suggesting stability from dense, modelconsistent rewards. We also report our comparison results with other concurrent order learning methods (Hong et al., 2025; Jazbec et al., 2025) in Appendix D.4.

References

7. Conclusion

Azangulov, I., Pandeva, T., Prasad, N., Zazo, J., and Karmalkar, S. Parallel sampling from masked diffusion models via conditional independence testing. arXiv preprint arXiv:2510.21961, 2025.

Arriola, M., Gokaslan, A., Chiu, J. T., Yang, Z., Qi, Z., Han, J., Sahoo, S. S., and Kuleshov, V. Block diffusion: Interpolating between autoregressive and diffusion language models. arXiv preprint arXiv:2503.09573, 2025. Austin, J., Odena, A., Nye, M., Bosma, M., Michalewski, H., Dohan, D., Jiang, E., Cai, C., Terry, M., Le, Q., et al. Program synthesis with large language models. arXiv preprint arXiv:2108.07732, 2021.

We introduced Self-Aware Scheduling (SAS), a plug-andplay framework for optimizing the unmasking order of pretrained masked diffusion language models. Our theoretical analysis guarantees that optimizing our objective reduces an upper bound on the target mismatch, and achieves exactness under realizability. Combined with reinforcement learning for policy optimization, SAS improves both structured and open-ended reasoning across domains and scales— from Sudoku with a 1B model to math and code reasoning with LLaDA-8B—and remains robust across generation lengths and semi-autoregressive block decoding. Overall, SAS provides a general recipe for aligning any-order diffusion inference with probabilistic objectives, and suggests a practical path toward self-improving diffusion LLMs via iterative scheduling and model refinement. A natural future direction, in light of Remark 3.5, is to combine SAS and dependency-awareness to accelerate the method further.

Ben-Hamu, H., Gat, I., Severo, D., Nolte, N., and Karrer, B. Accelerated sampling from masked diffusion models via entropy bounded unmasking. arXiv preprint arXiv:2505.24857, 2025. Chen, H., Ren, Y., Ying, L., and Rotskoff, G. Accelerating diffusion models with parallel sampling: Inference at sublinear time complexity. Advances in Neural Information Processing Systems, 37:133661–133709, 2024. Chen, Y., Vanden-Eijnden, E., and Xu, J. Lipschitz-guided design of interpolation schedules in generative models. arXiv preprint arXiv:2509.01629, 2025. Cobbe, K., Kosaraju, V., Bavarian, M., Chen, M., Jun, H., Kaiser, L., Plappert, M., Tworek, J., Hilton, J., Nakano, R., et al. Training verifiers to solve math word problems. arXiv preprint arXiv:2110.14168, 2021.

Limitations Several limitations remain. First, SAS learns from a teacherforced self-aware reward, i.e., the frozen denoiser’s pathwise likelihood on ground-truth trajectories. This creates a potential gap between training and free-running inference. Second, our current policies are trained largely on a per-task basis, so cross-task and cross-domain transfer are promising but not yet fully characterized. Third, while SAS is plugand-play with respect to the frozen backbone at inference time, learning the order policy still requires additional training compute, and the policy introduces a small but non-zero inference overhead. Addressing these issues through better transfer, reduced teacher-forcing mismatch, and hybrid selfaware/task-level rewards is an important direction for future work.

Gu, J., Liu, Q., and Cho, K. Insertion-based decoding with automatically inferred generation order. Transactions of the Association for Computational Linguistics, 7:661– 676, 2019. Guo, D., Yang, D., Zhang, H., Song, J., Zhang, R., Xu, R., Zhu, Q., Ma, S., Wang, P., Bi, X., et al. Deepseek-r1: Incentivizing reasoning capability in llms via reinforcement learning. arXiv preprint arXiv:2501.12948, 2025. Ho, J., Jain, A., and Abbeel, P. Denoising diffusion probabilistic models. Advances in neural information processing systems, 33:6840–6851, 2020. Hong, C., An, S., Kim, M.-S., and Ye, J. C. Improving discrete diffusion unmasking policies beyond explicit reference policies. arXiv preprint arXiv:2510.05725, 2025.

Impact Statement This paper presents work whose goal is to advance the field of Machine Learning. There are many potential societal consequences of our work, none which we feel must be specifically highlighted here.

Hoogeboom, E., Gritsenko, A. A., Bastings, J., Poole, B., Berg, R. v. d., and Salimans, T. Autoregressive diffusion models. arXiv preprint arXiv:2110.02037, 2021a. 9

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

Hoogeboom, E., Nielsen, D., Jaini, P., Forré, P., and Welling, M. Argmax flows and multinomial diffusion: Learning categorical distributions. Advances in neural information processing systems, 34:12454–12465, 2021b.

Park, Y.-H., Lai, C.-H., Hayakawa, S., Takida, Y., and Mitsufuji, Y. Jump your steps: Optimizing sampling schedule of discrete diffusion models. In The Thirteenth International Conference on Learning Representations, 2024.

Huang, P., Liu, T., Liu, Z., Yan, Y., Wang, S., Xiao, T., Chen, Z., and Sun, M. Empirical analysis of decoding biases in masked diffusion models, 2026. URL https: //arxiv.org/abs/2508.13021.

Radford, A., Narasimhan, K., Salimans, T., Sutskever, I., et al. Improving language understanding by generative pre-training, 2018. Rao, R. 9 million sudoku puzzles and solutions. https://www.kaggle.com/datasets/ rohanrao/sudoku, 2019. Accessed: 2025-09-24.

Huang, Z., Chen, Z., Wang, Z., Li, T., and Qi, G.-J. Reinforcing the diffusion chain of lateral thought with diffusion language models. arXiv preprint arXiv:2505.10446, 2025.

Shah, K., Dikkala, N., Wang, X., and Panigrahy, R. Causal language modeling can elicit search and reasoning capabilities on logic puzzles. Advances in Neural Information Processing Systems, 37:56674–56702, 2024.

Jazbec, M., Olausson, T. X., Béthune, L., Ablin, P., Kirchhof, M., Monterio, J., Turrisi, V., Ramapuram, J., and Cuturi, M. Learning unmasking policies for diffusion language models. arXiv preprint arXiv:2512.09106, 2025.

Shao, Z., Wang, P., Zhu, Q., Xu, R., Song, J., Bi, X., Zhang, H., Zhang, M., Li, Y., Wu, Y., et al. Deepseekmath: Pushing the limits of mathematical reasoning in open language models. arXiv preprint arXiv:2402.03300, 2024.

Kendall, M. G. Rank correlation methods. Griffin, 1948. Kim, J., Shah, K., Kontonis, V., Kakade, S. M., and Chen, S. Train for the worst, plan for the best: Understanding token ordering in masked diffusions. In Forty-second International Conference on Machine Learning, 2025.

Shi, J., Han, K., Wang, Z., Doucet, A., and Titsias, M. Simplified and generalized masked diffusion for discrete data. Advances in neural information processing systems, 37:103131–103167, 2024.

Kingma, D., Salimans, T., Poole, B., and Ho, J. Variational diffusion models. Advances in neural information processing systems, 34:21696–21707, 2021.

Shih, A., Belkhale, S., Ermon, S., Sadigh, D., and Anari, N. Parallel sampling of diffusion models. Advances in Neural Information Processing Systems, 36:4263–4276, 2023.

Li, P., Zhou, Y., Muhtar, D., Yin, L., Yan, S., Shen, L., Liang, Y., Vosoughi, S., and Liu, S. Diffusion language models know the answer before decoding. arXiv preprint arXiv:2508.19982, 2025.

Song, Y., Sohl-Dickstein, J., Kingma, D. P., Kumar, A., Ermon, S., and Poole, B. Score-based generative modeling through stochastic differential equations. arXiv preprint arXiv:2011.13456, 2020.

Lin, N., Zhang, J., Hou, L., and Li, J. Boundary-guided policy optimization for memory-efficient rl of diffusion large language models. arXiv preprint arXiv:2510.11683, 2025.

Stern, M., Chan, W., Kiros, J., and Uszkoreit, J. Insertion transformer: Flexible sequence generation via insertion operations. In International Conference on Machine Learning, pp. 5976–5985. PMLR, 2019.

Lou, A., Meng, C., and Ermon, S. Discrete diffusion modeling by estimating the ratios of the data distribution. arXiv preprint arXiv:2310.16834, 2023.

Tang, X., Dolga, R., Yoon, S., and Bogunovic, I. wd1: Weighted policy optimization for reasoning in diffusion language models. arXiv preprint arXiv:2507.08838, 2025.

Nie, S., Zhu, F., Du, C., Pang, T., Liu, Q., Zeng, G., Lin, M., and Li, C. Scaling up masked diffusion models on text. In ICLR, 2025a. Nie, S., Zhu, F., You, Z., Zhang, X., Ou, J., Hu, J., Zhou, J., Lin, Y., Wen, J.-R., and Li, C. Large language diffusion models. arXiv preprint arXiv:2502.09992, 2025b.

Wang, C., Rashidinejad, P., Su, D., Jiang, S., Wang, S., Zhao, S., Zhou, C., Shen, S. Z., Chen, F., Jaakkola, T., et al. Spg: Sandwiched policy gradient for masked diffusion language models. arXiv preprint arXiv:2510.09541, 2025a.

Ou, J., Nie, S., Xue, K., Zhu, F., Sun, J., Li, Z., and Li, C. Your absorbing discrete diffusion secretly models the conditional distributions of clean data. arXiv preprint arXiv:2406.03736, 2024.

Wang, W., Fang, B., Jing, C., Shen, Y., Shen, Y., Wang, Q., Ouyang, H., Chen, H., and Shen, C. Time is a feature: Exploiting temporal dynamics in diffusion language models. arXiv preprint arXiv:2508.09138, 2025b. 10

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

Wang, Y., Yang, L., Li, B., Tian, Y., Shen, K., and Wang, M. Revolutionizing reinforcement learning framework for diffusion large language models. arXiv preprint arXiv:2509.06949, 2025c.

Zhu, W., Hu, Z., and Xing, E. Text infilling. arXiv preprint arXiv:1901.00158, 2019. Zhu, Y., Guo, W., Choi, J., Molodyk, P., Yuan, B., Tao, M., and Chen, Y. Enhancing reasoning for diffusion llms via distribution matching policy optimization. arXiv preprint arXiv:2510.08233, 2025.

Wang, Z., Shi, J., Heess, N., Gretton, A., and Titsias, M. Learning-order autoregressive models with application to molecular graph generation. In Forty-second International Conference on Machine Learning, 2025d. Wei, Q., Zhang, Y., Liu, Z., Liu, D., and Zhang, L. Accelerating diffusion large language models with slowfast: The three golden principles. arXiv preprint arXiv:2506.10848, 2025. Welleck, S., Brantley, K., Iii, H. D., and Cho, K. Nonmonotonic sequential text generation. In International Conference on Machine Learning, pp. 6716–6726. PMLR, 2019. Williams, R. J. and Zipser, D. A learning algorithm for continually running fully recurrent neural networks. Neural computation, 1(2):270–280, 1989. Wu, L., Tan, X., He, D., Tian, F., Qin, T., Lai, J., and Liu, T.-Y. Beyond error propagation in neural machine translation: Characteristics of language also matter. arXiv preprint arXiv:1809.00120, 2018. Ye, J., Xie, Z., Zheng, L., Gao, J., Wu, Z., Jiang, X., Li, Z., and Kong, L. Dream 7b: Diffusion large language models. arXiv preprint arXiv:2508.15487, 2025. Yu, R., Ma, X., and Wang, X. Dimple: Discrete diffusion multimodal large language model with parallel decoding. arXiv preprint arXiv:2505.16990, 2025. Zhao, H., Liang, D., Tang, W., Yao, D., and Kallus, N. Diffpo: Training diffusion llms to reason fast and furious via reinforcement learning. arXiv preprint arXiv:2510.02212, 2025a. Zhao, S., Gupta, D., Zheng, Q., and Grover, A. d1: Scaling reasoning in diffusion large language models via reinforcement learning. arXiv preprint arXiv:2504.12216, 2025b. Zheng, K., Chen, Y., Mao, H., Liu, M.-Y., Zhu, J., and Zhang, Q. Masked diffusion models are secretly timeagnostic masked models and exploit inaccurate categorical sampling. arXiv preprint arXiv:2409.02908, 2024. Zheng, L., Yuan, J., Yu, L., and Kong, L. A reparameterized discrete diffusion model for text generation. arXiv preprint arXiv:2302.05737, 2023. Zhou, J., Lu, T., Mishra, S., Brahma, S., Basu, S., Luan, Y., Zhou, D., and Hou, L. Instruction-following evaluation for large language models, 2023. URL https: //arxiv.org/abs/2311.07911. 11

Appendix Table of Contents 1. Background (Appendix A, p. 13) • Masked diffusion models (Section A.1) • Unmasking schedules: sequential and parallel decoding (Section A.2) • Heuristic Schedules (Section A.3) • Unmasking-order policy learning algorithm (Alg. 1) 2. Theoretical Guarantees (Appendix B, p. 16) • Parallelization error and conditional total correlation (Section B.1) • Proof of Theorem 3.2 (Section B.2) • Proof of Corollary 3.3 (Section B.3) • Proof of Remark 3.4 (Section B.4) 3. Implementation Details (Appendix C, p. 21) • Order policy design goals and constraints (Section C.1) • Policy inputs, architecture, and action selection (Sections C.2–C.5) • Training (Section C.6) and inference (Section C.7) details 4. Training and Evaluation Results (Appendix D, p. 23) • Sudoku: additional analyses (Section D.1) • Math and code reasoning with LLaDA-8B (Section D.2) • Comparison with other order learning methods (Section D.4) 5. Comparison with Hidden-State Policy (Appendix E, p. 28) • Inputs, architecture, and action selection (Sections E.1–E.3) • Empirical comparison and discussion (Section E.4)

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

INFERENCE

TRAINING

Masked Input Masked Input Trained Order Policy Order Policy

Aggregated Rewards

Sampled Orders

GRPO Optimization

od liho tion Like imiza x Ma

Self-Aware Reward

Decoded Sequence

Figure 5. Overview of the Self-Aware Scheduling (SAS) framework. Left (Training): The order policy vϕ samples multiple unmasking permutations (σ1 , σ2 , . . . ) for a frozen diffusion model. These orders are evaluated using our Self-Aware Reward, which measures the cumulative log-likelihood of the generation under that specific order. The policy is then optimized via GRPO to favor orders that maximize the model’s own confidence. Right (Inference): The trained policy dictates the optimal unmasking sequence, guiding the diffusion model to generate high-quality outputs efficiently.

A. Background This section provides detailed background material supplementing the main article. A.1. Masked Diffusion Models Masked diffusion language models generate discrete sequences by iteratively demasking tokens, analogous to the denoising process in continuous diffusion models. We review the modeling framework and sampling procedures. Notation. Let X denote a finite vocabulary and n the sequence length. A sequence is x = (x1 , . . . , xn ) ∈ X n . For a mask pattern M ⊆ [n] := {1, . . . , n}, define the masked sequence xM ∈ (X ∪ {[MASK]})n by ( [MASK] (x )i := xi M

if i ∈ M, if i ∈ / M.

(18)

Model. A masked diffusion model (or denoiser) parameterized by θ defines conditional distributions over masked tokens given observed context. Specifically, the model pθ provides position-wise predictions piθ (· | xM ),

i ∈ M,

(19)

representing the distribution over token xi conditioned on the partial observation xM . Here θ denotes the learned parameters of the base model. 13

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

Training. Let π denote the data distribution over X n and q(M ) a distribution over mask patterns. The model parameters θ are trained via cross-entropy: " # X i M L(θ) := −Ex∼π, M ∼q log pθ (xi | x ) . (20) i∈M

Sampling. Generation proceeds by iteratively unmasking positions, a process known as decoding. We denote by x̃(t) the partial sequence at decoding step t, where some positions contain tokens from X (revealed) and others contain [MASK] (not yet revealed). Starting from the fully masked sequence, at each step the algorithm selects which position(s) to unmask and samples their values from the model. The flexibility of the masked diffusion model is that it allows any order decoding in principle, and thus the unmasking order, or schedule, matters. A.2. Unmasking schedules The unmasking schedule can be interpreted as an “order of thought” for generation. We introduce a learned policy vϕ parameterized by ϕ, which takes a partial sequence as input and outputs a distribution over which position(s) to unmask next. We distinguish two decoding regimes based on how many positions are unmasked per step: Sequential decoding. Unmask one position per step over n steps. Let St ⊆ [n] denote the set of revealed positions at step t, and Mt = [n]\St denote the set of masked positions. Starting from x̃(0) = x[n] (fully masked) with S0 = ∅, at each step t = 1, . . . , n: 1. Select next position: it ∼ vϕ (· | x̃(t−1) ) where vϕ (· | x̃(t−1) ) outputs a distribution over [n] \ St−1 (t)

2. Sample token: x̃it ∼ piθt (· | x̃(t−1) ) (t)

(t−1)

3. Update: St ← St−1 ∪ {it }; set x̃j = x̃j

for all j ̸= it

The complete run produces an unmasking order σ = (i1 , . . . , in ) and final sequence x̃(n) ∈ X n . Parallel decoding. To accelerate PK generation, unmask multiple positions simultaneously. Fix a step budget K < n and block sizes (b1 , . . . , bK ) with k=1 bk = n. Let Sk denote the set of revealed positions after k steps. Starting from x̃(0) = x[n] with S0 = ∅, at each step k = 1, . . . , K: 1. Select block Bk ⊆ [n] \ Sk−1 of size bk (typically by sampling bk positions from vϕ (· | x̃(k−1) ) without replacement) (k)

2. Sample all tokens in block: x̃i

∼ piθ (· | x̃(k−1) ) for each i ∈ Bk (k)

3. Update: Sk ← Sk−1 ∪ Bk ; set x̃j

(k−1)

= x̃j

for all j ∈ / Bk

Crucially, all positions within block Bk are sampled using the same partial context x̃(k−1) , enabling parallelization but ignoring correlations between positions unmasked in the same block. Induced distributions. Both decoding procedures induce distributions over output sequences. To characterize these distributions, we define notations for the decoding algorithm under teacher forcing, where the policy vϕ and model pθ are evaluated on partially revealed versions of a fixed target sequence x. This allows us to compute the probability the algorithm assigns to generating x under a particular unmasking order. For a target sequence x and unmasking order σ = (i1 , . . . , in ), define the teacher-forced trajectory x̃(0) (σ), . . . , x̃(n) (σ) where x̃(t) (σ) reveals positions {i1 , . . . , it } from x: ( xj if j ∈ {i1 , . . . , it }, (t) x̃ (σ)j := (21) [MASK] otherwise. In other words, x̃(t) (σ) is the partial sequence obtained by revealing the first t positions in order σ using their values from target x. 14

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

Sequential case. Under teacher-forcing with target x and order σ, the policy (parameterized by ϕ) assigns probability vϕ (σ | x) :=

n Y

vϕ (it | x̃(t−1) (σ))

(22)

t=1

to the unmasking order σ, where each factor evaluates the policy vϕ on the teacher-forced state x̃(t−1) (σ). The model (parameterized by θ) assigns path likelihood pθ (x | σ) :=

n Y

piθt (xit | x̃(t−1) (σ)),

(23)

t=1

(24)

with marginal distribution over outputs X

seq Pθ,ϕ (x) =

vϕ (σ | x) · pθ (x | σ).

(25)

σ∈Sn

Parallel case. For block decomposition β = (B1 , . . . , BK ), define the schedule likelihood under policy parameters ϕ: vϕ (β | x) :=

K Y

vϕ (Bk | x̃(k−1) (β)),

(26)

k=1

where x̃(k−1) (β) denotes the teacher-forced state after revealing blocks B1 , . . . , Bk−1 from target x. The model assigns path likelihood K Y Y pθ (x | β) := piθ (xi | x̃(k−1) (β)). (27) k=1 i∈Bk

The parallel algorithm induces the joint distribution par Pθ,ϕ (x, β) := vϕ (β | x) · pθ (x | β),

(28)

with marginal distribution par Pθ,ϕ (x) =

vϕ (β | x) · pθ (x | β).

(29)

A.3. Heuristic Schedules Heuristic schedules can be formalized as a special case of our order policy framework where the policy vϕ is deterministic and non-trainable. These methods define a fixed policy that greedily selects the position i maximizing a predefined score computed from the model’s position-wise predictive distributions piθ (· | x̃(t−1) ). Recent works (Nie et al., 2025b; Kim et al., 2025) consider greedy schedules based on per-position uncertainty. Per-position uncertainty scores. At each sequential step t, for a masked position i ∈ Mt , let p(a) := piθ (a | x̃(t−1) ) be the local probability distribution. Let â := arg maxa p(a) denote the most likely token. We define the following scores (higher is better): Confidence: Margin: Entropy:

sconf (i) := p(â), smarg (i) := p(â) − max p(a), a̸=â X sent (i) := p(a) log p(a). a∈X

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

Algorithm 1 Unmasking-Order Policy Learning Require: Initial order policy vϕinit ; frozen diffusion model pθ ; dataset D of pairs (c, x⋆ ); hyperparameters ϵ, β, µ; group size G. 1: Initialize policy ϕ ← ϕinit 2: for iteration = 1, . . . , I do 3: reference policy vref ← vϕ 4: for step = 1, . . . , M do 5: Sample minibatch B ∼ D 6: old policy vold ← vϕ ⋆ ⋆ 7: Sample G orders {σ (i) }G i=1 ∼ vold (· | c, x ) for each (c, x ) ∈ B (teacher-forced rollout) 8: Compute self-aware returns {Ri } using Eq. (8) (frozen pθ ) bi } 9: Compute group-relative advantages {A 10: for GRPO update = 1, . . . , µ do 11: Update ϕ by maximizing the clipped GRPO objective 12: end for 13: end for 14: end for Ensure: Trained order policy vϕ

Sequential heuristics.

At sequential step t, given the current state x̃(t−1) , choose it ∈ arg max s(i; x̃(t−1) ), i∈Mt−1

where s ∈ {sconf , smarg , sent }. Parallel heuristics. At parallel step k with block size bk , given the current state x̃(k−1) , choose the block Bk ⊆ Mk−1 as the top-bk masked indices under the score: Bk := Top-bk { s(i; x̃(k−1) ) : i ∈ Mk−1 } . All positions in Bk are then sampled from piθ (· | x̃(k−1) ) using the same shared context. Discussion. These heuristics can be viewed as deterministic policies induced by the scores and the current decoding state. While they provide a strong baseline compared to random scheduling, they remain inherently myopic: they rely exclusively on local, instantaneous uncertainty in piθ (· | x̃(t−1) ) at the current state, and do not explicitly optimize a global objective over entire unmasking trajectories. As a result, they can prioritize “easy” positions early even when revealing a different position would better shape the future conditional landscape, leaving room for learned schedules that perform long-horizon credit assignment.

B. Theoretical Guarantees In this part, we provide theoretical analysis and detailed proofs of the proposed theorems and corollaries in this paper. In Appendix B.1, we first focus on the quantification of the parallel decoding errors mentioned in Remark 3.5 of the main article; in particular, we derive total variation bound in Corollary B.2. We then move to the proof of Theorem 3.2 and Corollary 3.3, in Appendices B.2 and B.3 respectively. B.1. Parallelization error and conditional total correlation In this section, we provide more details on the analysis of the error caused by parallel decoding. More specifically, we will seq par use KL(Pθ,ϕ , Pθ,ϕ ) to characterize the additional discrepancy introduced by parallel decoding. We first introduce the concept of conditional total correlation. 16

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

Conditional total correlation. For random variables X ∈ X N under a distribution P , and disjoint sets A, B ⊆ [N ], we define the conditional total correlation as X TCP (XB | XA ) := HP (Xi | XA ) − HP (XB | XA ) i∈B

(30)

! = KL P (XB | XA )

P (Xi | XA ) .

i∈B

This quantity is zero if and only if the variables in B are conditionally independent given XA under P . While parallel decoding updates a block Bk using a single shared context x̃(k−1) , which prevents within-block conditioning on newly sampled tokens. This induces an information loss whenever the true conditional dependence within Bk is nontrivial. Theorem B.1 (Parallelization error bounded by conditional TC). Consider a block schedule (B1 , . . . , BK ) (e.g., obtained by chunking a fine-grained order), and assume parallel updates sample each Xi for i ∈ Bk independently given the prefix XB<k . Then the additional mismatch introduced by parallelization is governed by the sum of within-block conditional total correlations under the sequential law: seq par KL(Pθ,ϕ , Pθ,ϕ ) =

K X

seq TC seq (X EX∼Pθ,ϕ Bk | XB<k ) . Pθ,ϕ

(31)

k=1 seq par Proof. For brevity, write P seq := Pθ,ϕ and P par := Pθ,ϕ .

Step 1: Factorizations under the block schedule. Fix ordered blocks (B1 , . . . , BK ) and define B<k := the chain rule grouped by blocks, P seq factorizes as P seq (x) =

K Y

P seq (xBk | xB<k ).

j<k Bj . By

(32)

k=1

Under the theorem assumption (within-block independent sampling given the prefix), Y P par (xBk | xB<k ) = P seq (xi | xB<k ),

(33)

i∈Bk

and therefore P par (x) =

K Y Y

P seq (xi | xB<k ).

(34)

k=1 i∈Bk

Step 2: Expand KL(P seq ∥P par ) and identify conditional total correlation. If P seq ̸≪ P par , then KL(P seq ∥P par ) = +∞ and the bound is trivial. Assume henceforth P seq ≪ P par . Then P seq (X) seq par seq KL P ∥P = EX∼P log par P (X) "K # K X X X seq seq = EX∼P seq log P (XBk | XB<k ) − log P (Xi | XB<k ) k=1 i∈Bk

k=1

(by (32) and (34)) =

K X k=1

# P seq (XBk | XB<k ) EX∼P seq log Q . seq (X | X i B<k ) i∈Bk P

By the KL form of conditional total correlation (with respect to P seq ), " # P seq (XBk | XB<k ) TCP seq (XBk | XB<k ) = EX∼P seq log Q , seq (X | X i B<k ) i∈Bk P 17

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

and hence we obtain the exact identity KL P seq ∥P par

K X

EX∼P seq TCP seq (XBk | XB<k ) .

(35)

k=1

The theorem explains why parallel decoding can incur an intrinsic error when blocks contain highly dependent variables (high conditional TC), motivating schedules that reveal strongly coupled variables in a dependency-aware manner. Quantification of error from data distribution Theorem B.1 quantifies the parallelization gap between the sequential and parallel output laws. We now relate this to the final deviation of parallel decoding from the data distribution π using total variation distance. Corollary B.2 (Parallel decoding error via sequential error + TC). For any (θ, ϕ) and any ordered block schedule (B1 , . . . , BK ) satisfying the within-block conditional independence assumption of Theorem B.1, v u K q h u X i par seq 1 t1 seq TC seq X TV π, Pθ,ϕ ≤ KL π∥P + E . (36) | X B B X∼Pθ,ϕ Pθ,ϕ k <k θ,ϕ 2 2 k=1

Proof. By the triangle inequality of total variation, par seq seq par TV π, Pθ,ϕ ≤ TV π, Pθ,ϕ + TV Pθ,ϕ , Pθ,ϕ . Applying Pinsker’s inequality to each term yields q seq seq , TV π, Pθ,ϕ ≤ 12 KL π∥Pθ,ϕ

q seq par seq par TV Pθ,ϕ , Pθ,ϕ ≤ 12 KL Pθ,ϕ ∥Pθ,ϕ .

Finally, Theorem B.1 gives the exact identity K h X i seq par seq TC seq X KL Pθ,ϕ ∥Pθ,ϕ = EX∼Pθ,ϕ Bk | XB<k , Pθ,ϕ k=1

which proves (36). Discussion. Corollary B.2 separates the parallel decoding error into (i) a sequential mismatch term that can be reduced by improving the denoiser and/or learning better schedules (as in Theorem 3.2), and (ii) an intrinsic parallelization term controlled by within-block conditional dependence. In particular, the parallelization term vanishes only when the tokens within each block are conditionally independent given the prefix (i.e., zero conditional total correlation), explaining why overly aggressive parallel updates can degrade accuracy. Remark B.3. The first term in (36) can be viewed as the statistical learning error, while the second term corresponds to the numerical error. The design of order schedules in discrete diffusion models can impact both sources of error. In contrast, for continuous diffusion models, scalar schedules affect only the numerical error and lead to statistically equivalent models under KL divergence in path space (Chen et al., 2025; Kingma et al., 2021). B.2. Proof of Theorem 3.2 We then proceed to prove that our proposed objective, LSAS , bounds the Kullback-Leibler divergence between the data distribution and the model-induced distribution, up to a constant term. Theorem 3.2 (Self-Awareness as Joint Distribution Matching). Let Qϕ (x, σ) := π(x)vϕ (σ|x) and Pθ,ϕ (x, σ) := pθ (x|σ)vϕ (σ|x) be the data-policy and model-policy joint distributions, respectively. The self-aware loss LSAS satisfies the following identity and bound: seq KL(π(x)∥Pθ,ϕ (x)) ≤ LSAS (θ, ϕ) − H(π)

= KL(Qϕ (x, σ)∥Pθ,ϕ (x, σ)), 18

(37)

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

where H(π) is the constant data entropy. Consequently, minimizing LSAS is equivalent to minimizing the joint distributional seq mismatch, which upper bounds the generation error KL(π(x)∥Pθ,ϕ (x)). Proof. Recall the induced sequential marginal seq Pθ,ϕ (x) =

vϕ (σ | x) pθ (x | σ),

σ∈Sn

and the joint distributions defined as Qϕ (x, σ) = π(x) vϕ (σ | x), Step 1: Identity KL(Qϕ ∥Pθ,ϕ ) = LSAS (θ, ϕ) − H(π). KL(Qϕ ∥Pθ,ϕ ) =

Pθ,ϕ (x, σ) = vϕ (σ | x) pθ (x | σ). By definition,

Qϕ (x, σ) log

x,σ

Qϕ (x, σ) Pθ,ϕ (x, σ)

= Ex∼π, σ∼vϕ (·|x) log

π(x)vϕ (σ | x) vϕ (σ | x)pθ (x | σ)

= Ex∼π [log π(x)] − Ex∼π, σ∼vϕ (·|x) [log pθ (x | σ)] = −H(π) + LSAS (θ, ϕ), where H(π) = −Ex∼π [log π(x)]. seq X Step 2: Bound KL(π∥Pθ,ϕ ) ≤ KL(Qϕ ∥Pθ,ϕ ). Let QX ϕ and Pθ,ϕ denote the marginals of Qϕ and Pθ,ϕ on x. By construction, X X seq X QX Qϕ (x, σ) = π(x), Pθ,ϕ (x) = Pθ,ϕ (x, σ) = Pθ,ϕ (x). ϕ (x) = σ

KL divergence contracts under marginalization (data processing), hence seq X KL(π∥Pθ,ϕ ) = KL(QX ϕ ∥Pθ,ϕ ) ≤ KL(Qϕ ∥Pθ,ϕ ).

Combining Step 1 and Step 2 yields Eq. (11), completing the proof. B.3. Proof of Corollary 3.3 We then move to prove the tightness of the bound derived in Theorem 3.2 can be achieved by matching a zero joint KL divergence. Corollary 3.3 (Joint realizability and tightness). If the model family is expressive enough to allow for joint realizability—i.e., there exist parameters (θ⋆ , ϕ⋆ ) such that KL(Qϕ⋆ (x, σ)∥Pθ⋆ ,ϕ⋆ (x, σ)) = 0—then the marginal mismatch seq KL(π∥Pθ,ϕ ) is zero, and the self-aware lower bound of Theorem 3.1 is tight, satisfying LSAS (θ⋆ , ϕ⋆ ) = H(π). Proof. Assume KL(Qϕ⋆ ∥Pθ⋆ ,ϕ⋆ ) = 0. Then Qϕ⋆ = Pθ⋆ ,ϕ⋆ almost everywhere, i.e., for all (x, σ) such that Qϕ⋆ (x, σ) > 0, π(x) vϕ⋆ (σ | x) = vϕ⋆ (σ | x) pθ⋆ (x | σ). Therefore, for any x with π(x) > 0 and any σ with vϕ⋆ (σ | x) > 0, pθ⋆ (x | σ) = π(x). (i) Zero marginal mismatch.

By Theorem 3.2, KL(π∥Pθseq ⋆ ,ϕ⋆ ) ≤ KL(Qϕ⋆ ∥Pθ ⋆ ,ϕ⋆ ) = 0,

seq hence KL(π∥Pθseq ⋆ ,ϕ⋆ ) = 0, i.e., Pθ ⋆ ,ϕ⋆ = π.

(38)

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

(ii) Tightness of Theorem 3.1. Fix any x with π(x) > 0. From Eq. (38), we know pθ⋆ (x | σ) = π(x) for all orders sampled from vϕ⋆ . Because this term is constant with respect to σ, the expectation over σ simply returns the value itself: Eσ∼vϕ⋆ (·|x) [log pθ⋆ (x | σ)] = Eσ∼vϕ⋆ (·|x) [log π(x)] = log π(x). Since Pθseq ⋆ ,ϕ⋆ (x) = π(x) (from step i), the LHS is also log π(x), proving the bound is tight. (iii) Value of LSAS .

By Theorem 3.2 and KL(Qϕ⋆ ∥Pθ⋆ ,ϕ⋆ ) = 0, LSAS (θ⋆ , ϕ⋆ ) − H(π) = 0

⇒

LSAS (θ⋆ , ϕ⋆ ) = H(π).

seq Corollary 3.3 states that under joint realizability this lower bound is achievable, yielding KL(π∥Pθ,ϕ ) = 0. This provides a principled justification for our training strategy: LSAS is not only a dense, model-consistent reward for learning the unmasking order, but fundamentally connects order optimization to distribution matching. Crucially, this implies that perfect distributional alignment does not require the diffusion model to match the data distribution across all possible unmasking orders; rather, it suffices to identify a single optimal ordering policy along which the model’s predictions align with the ground truth.

B.4. Proof of Remark 3.4 Remark 3.4 implies that a perfect model is order-invariant, i.e. if the model class contains a parameter θ⋆ whose conditionals match the true conditionals for every teacher-forced state (i.e., piθ⋆ (xi | x̃(t) ) = π(xi | x̃(t) ) for all valid masks). In this regime, the upper bound in Theorem 3.2 is tight with LSAS (θ⋆ , ϕ) = H(π) for any order policy ϕ. Proof. Fix any target sequence x ∈ X n and any order σ = (i1 , . . . , in ) ∈ Sn . Let x̃(t−1) (σ) be the teacher-forced state that reveals {i1 , . . . , it−1 } from x. By the chain rule of π along the order σ, π(x) =

n Y

π xit | x̃(t−1) (σ) .

(39)

t=1

By the assumed conditional correctness of θ⋆ , for every teacher-forced state and every currently masked position it , we have piθt⋆ xit | x̃(t−1) (σ) = π xit | x̃(t−1) (σ) . Multiplying over t and using the definition of the path likelihood, p (x | σ) = θ⋆

n Y t=1 n Y

piθt⋆ xit | x̃(t−1) (σ) π xit | x̃(t−1) (σ) = π(x),

t=1

where the last equality follows from (39). Hence for any order policy vϕ , X X Pθseq vϕ (σ | x) pθ⋆ (x | σ) = vϕ (σ | x) π(x) = π(x), ⋆ ,ϕ (x) = σ∈Sn

σ∈Sn

seq so Pθseq ⋆ ,ϕ = π and therefore KL(π∥Pθ ⋆ ,ϕ ) = 0.

Perfect denoiser renders order irrelevant. Since pθ⋆ (x | σ) = π(x) for all σ, we also have Pθ⋆ ,ϕ (x, σ) = vϕ (σ | x)π(x) = Qϕ (x, σ) for any ϕ. Thus KL(Qϕ ∥Pθ⋆ ,ϕ ) = 0 and by Theorem 3.2, LSAS (θ⋆ , ϕ) = H(π) for any order policy ϕ. However, since the pretrained model pθ is inevitably imperfect, it justifies our motivation for self-aware learning: improving either ϕ (finding easier generation paths) or θ (improving predictions) reduces the joint mismatch, as we observe empirically in Section 6.4. 20

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

B.5. Parallel Decoding for Unique Solution/Dirac Measure Tasks A key advantage of diffusion language models is their ability to decode multiple tokens in parallel. While our main theory focuses on sequential unmasking, the same self-aware objective also admits a simple interpretation for parallel decoding in unique-solution tasks, such as Sudoku. In these tasks, conditioned on the prompt or constraint c, the target distribution is a Dirac measure: π(· | c) = δx⋆ (c) . (40) Let β = (B1 , . . . , BK ) denote a block order, where each Bk is a set of positions decoded in parallel at step k. For a fixed block order β, define the teacher-forced parallel path likelihood as ⋆ ppar θ (x | β, c) =

K Y Y

piθ x⋆i | x̃(k−1) (β; c), c ,

(41)

k=1 i∈Bk

where x̃(k−1) (β; c) reveals the ground-truth tokens in the earlier blocks B<k and masks the remaining positions. For unique-solution tasks, each block conditional target is also a Dirac measure: Y π XBk | XB<k = x⋆B<k , c = δx⋆B = δx⋆i . k

(42)

i∈Bk

Since the output space is discrete and each model conditional distribution is parameterized by a softmax, every ground-truth token has positive probability: piθ x⋆i | x̃(k−1) (β; c), c > 0, ∀i, k. (43) Hence the teacher-forced parallel path likelihood is well-defined and strictly positive: ⋆ ppar θ (x | β, c) > 0.

(44)

Therefore, for any fixed block order β, the KL divergence from the Dirac target to the parallel teacher-forced model path reduces to the negative log-likelihood of the unique solution: par ⋆ KL(δx⋆ ∥ ppar θ (· | β, c)) = − log pθ (x | β, c) K X X =− log piθ x⋆i | x̃(k−1) (β; c), c .

(45)

k=1 i∈Bk

Equivalently, defining the parallel self-aware reward as Rθpar (x⋆ , β) :=

K X X

log piθ x⋆i | x̃(k−1) (β; c), c ,

(46)

k=1 i∈Bk

we obtain par ⋆ KL(δx⋆ ∥ ppar θ (· | β, c)) = − Rθ (x , β).

(47)

Thus, for unique-solution tasks, maximizing the parallel teacher-forced self-aware reward is exactly equivalent to minimizing the parallel path KL. This provides a direct justification for applying SAS to blockwise or parallel decoding in settings where the target solution is deterministic given the condition c.

C. Implementation Details This section details the concrete implementation of the decode order policy used in our experiments. We focus exclusively on architectural choices, data flow, masking logic, and training/inference mechanics that are not explicit from the mathematical formulation in the main paper. 21

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

C.1. Design Goals and Constraints The decode order policy is designed to satisfy the following practical constraints: • Backbone-agnostic: the policy operates on model outputs and auxiliary statistics without modifying or backpropagating through the diffusion language model. • Lightweight: the policy introduces minimal additional parameters and computational overhead relative to the frozen diffusion backbone. • Position-selective: at each diffusion step, the policy selects exactly one token position to reveal, subject to validity constraints. • Mask-aware: the policy must never select padding positions, prompt tokens, or already revealed tokens. These constraints motivate a per-position scoring formulation with explicit candidate masking. C.2. Policy Inputs At each diffusion step, the policy operates on a partially revealed sequence of length L. For each position j ∈ {1, . . . , L}, we construct a compact feature vector summarizing the diffusion model’s current belief about that position. Specifically, the policy input tensor has shape [B, L, 3], where B is the batch size, and each feature vector consists of: 1. Maximum token probability: the top-1 probability under the diffusion model’s predicted distribution at position j. 2. Token entropy: the entropy of the predicted token distribution, capturing uncertainty. 3. Normalized position index: j/L, providing weak positional context. 4. Current State Current state st with mask on future prediction positions. These features are computed on-the-fly during diffusion rollouts and are detached from the backbone model graph. C.3. Policy Architecture The decode order policy is implemented as a small Transformer encoder that maps per-position features to scalar reveal scores. Encoder. We use a Transformer encoder with (i) batch-first layout, (ii) nlayer self-attention layers, (iii) model dimension dmodel and (iv) nhead attention heads. Self-attention allows the policy to reason about relative confidence and uncertainty across positions rather than scoring positions independently. Output Head. The encoder output at each position is projected through a linear layer to produce a scalar λj , yielding a score tensor λ ∈ RB×L . No softmax is applied at this stage; normalization is deferred until candidate masking is applied. C.4. Candidate Masking and Valid Actions Not all positions are valid candidates for selection at every step. We construct a binary candidate mask M ∈ {0, 1}B×L indicating valid reveal positions. A position is considered invalid if it satisfies any of the following: • it corresponds to a padding token, • it lies in the prompt region, • it has already been revealed in a previous diffusion step. Invalid positions are assigned −∞ before normalization to ensure they receive zero probability mass. Only positions with Mj = 1 participate in action selection. 22

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

C.5. Action Selection At each diffusion step, the policy selects exactly one position per sequence. Given masked scores λ̃, we form a categorical distribution over valid positions and sample (or greedily select) an index: j ∗ = arg maxj λ̃j . The selected position is then revealed in the diffusion canvas, and the process repeats until all answer tokens are unmasked. C.6. Training-Time Supervision During training, we supervise the policy using ground-truth reveal orders extracted from reference diffusion trajectories. For each diffusion step: • the candidate mask is constructed, • logits are normalized only over valid positions, • GRPO loss is computed against the target reveal position. Loss is accumulated across steps and averaged over the batch. Padding positions and invalid candidates are fully excluded from loss computation. C.7. Inference-Time Integration At inference time, the policy replaces heuristic scheduling strategies. The diffusion model remains unchanged; only the reveal order differs. The policy is queried once per diffusion step, introducing negligible overhead compared to the backbone forward pass. All experiments use greedy selection for stability. Implementation Notes. All components are implemented in PyTorch. Attention masks and candidate masks are represented as boolean tensors and applied consistently in both training and inference. The policy network is trained independently and can be attached to any compatible diffusion language model without retraining the backbone.

D. Training and Evaluation Results In this section, we provide further details on the evaluation of our unmasking order policy across the Sudoku, GSM8K, and MBPP datasets. D.1. Sudoku

(a) Training Loss

(b) Avg Reward

Figure 6. Training and Validation Metrics. (a) Training loss decreases over time. (b) Average training reward stabilizes. (c) Validation total reward converges.

Training dynamics We show the training and evaluation dynamics in Figure 6. It illustrates the training and validation dynamics of the proposed order policy optimization. Panel (a) depicts the training loss, which exhibits a consistent downward trend over the course of training, indicating stable optimization of the objective function. Panel (b) tracks the average training reward, showing a rapid initial increase followed by stabilization, which reflects the policy successfully learning to prioritize high-likelihood reveal orders. Finally, Panel (c) demonstrates that the total reward on the validation set follows a similar trajectory, converging to a high value; this plateauing confirms that the learned policy generalizes robustly to unseen data without significant overfitting. 23

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

Equivalence-class Kendall’s τ We also conduct a systematic analysis of the order under different schedulings and compare them with human EXPERT order using equivalence class Kendall’s τ . In this paragraph, we show details on equivalence-class Kendall’s τ used for order analysis in Sudoku regime (Section 6.1). Standard rank correlation metrics penalize any deviation from the reference order. However, in Sudoku, many moves are strategically equivalent. For instance, a Naked Single (a cell with only one valid candidate) and a Hidden Single (a cell is the only position in a row/col/box for a specific number) are both deterministic logical steps. If a board state offers multiple Naked Singles simultaneously, they can be filled in any order without changing the underlying strategy. We propose a variant of Kendall’s τ that ignores permutations among such equivalent moves. Given a human expert order π and a model order σ, we iterate through the expert’s trajectory. Let st be the board state after t − 1 expert moves. For the current expert move πt and any future move πj (where t < j), we check if πj could have been solved at step t using the same deduction logic as πt (i.e., c(πj ; st ) = c(πt ; st )). If so, the choice between πt and πj is arbitrary; we treat this pair as unordered and exclude it from the evaluation. We calculate Kendall’s τ solely on the remaining strategy-relevant pairs—those where the expert prioritized a simpler or more fundamental deduction class over a complex one (e.g., prioritizing a Naked Single over a complex Intersection Removal). Let Prel be the set of relevant pairs. We compute the number of concordant pairs C ∗ (model agrees with expert hierarchy) and discordant pairs D∗ (model flips hierarchy) within Prel . The metric is defined as: τeq =

C ∗ − D∗ . C ∗ + D∗

A value of τeq = 1 implies the model perfectly matches the expert’s strategic hierarchy, even if the exact sequence differs on interchangeable steps. When does the model fail on Sudoku? We report an empirical observation on when sequential masked-diffusion decoding tends to make its first mistake on Sudoku. We first fine-tune the base diffusion model on the Sudoku corpus using the standard masked-diffusion pretraining objective (Section 2.1) so that it acquires basic Sudoku-solving capability. We then evaluate the resulting model by sequentially unmasking cells under a given decoding schedule (e.g., our learned order policy). To quantify failure timing, we define the first-step failure as the earliest decoding step t at which the generated value for the selected cell differs from the ground-truth value at that position (i.e., the first wrong fill along the rollout). Figure 7 plots the distribution of this first failure step on the evaluation set for random, confidence, margin, and learned order policy. Interestingly, for heuristic schedules, failures concentrate near the end of decoding (roughly the last few steps), rather than at the beginning. In contrast, we observe that a random schedule tends to fail much earlier, indicating that scheduling substantially affects not only the final accuracy but also the error onset during generation. Even with a trained order policy, however, the diffusion model still exhibits a nontrivial mass of late-stage failures, suggesting that the remaining unsolved cells are intrinsically harder and motivating further improvements to the diffusion head (e.g., our second-stage fine-tuning in Section 6.4). D.2. Math and code reasoning with LLaDA-8B Training Details & Computational Cost We trained the order policy using the Self-Aware Scheduling (SAS) framework with a group size of G = 6 and a batch size of B = 3 (total effective batch size of 18 trajectories per optimization step). The maximum sequence length was set to L = 512, with a learning rate of 1 × 10−4 on 4 NVIDIA H200 GPUs. Based on the training logs, the wall-clock time per step—including group sampling, reward computation, and backpropagation—averaged approximately 3.2 minutes. The policy typically converged within 200 steps (approx. 12 GPU-hours). The computational cost of SAS training is dominated by the rollout phase, where the policy samples G distinct ordering trajectories, and the reward phase, where the frozen diffusion model evaluates the pathwise log-likelihood of these trajectories via teacher-forced recomputation. The frozen denoiser’s likelihood evaluation (Reward) consumes the majority of FLOPs, significantly exceeding the policy’s gradient update cost. Evaluation We also report detailed evaluation result in Table 5. We use ARGMAX choice with order policy by setting the temperature 0 during evaluation. While heuristic baselines exhibit inconsistent performance across different settings, our learned order policy consistently demonstrates superior or competitive results across the entire spectrum of decoding regimes. This robustness confirms that our scheduling strategy captures fundamental reasoning structures, allowing it to transfer effectively well beyond its training configuration. 24

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

(a) Random Order

(b) Confidence-based order

(d) Policy-based Order

Figure 7. Distribution of failing steps across different ordering strategies. (a) Random ordering fails early. (b) Confidence strategies show late-failure patterns. (c) Margin and (d) Learned policy shifts failures later, but failure still concentrates on the last 5 steps.

Scheduling Thoughts: Learning the Order of Thought in Diffusion Language Models Table 5. Generalization across generation length and semi-autoregressive decoding. We report pass@1 accuracy with error bars on GSM8K (5-shot) and MBPP (3-shot). Notation: L--B denotes total generation length L and block size B; L--L is full diffusion decoding. We vary the number of decoding blocks K ∈ {1, 4, 8} across different generation lengths.

GSM8K Order

128–128

128–32

128–16

256–256

256–64

256–32

512–512

512–128

512–64

Confidence (Top-1) 0.620±0.018 0.728±0.017 0.744±0.017 0.628±0.018 0.760±0.016 0.758±0.017 0.765±0.018 0.810±0.017 0.810±0.016 0.608±0.019 0.742±0.017 0.740±0.018 0.640±0.018 0.794±0.016 0.796±0.017 0.628±0.020 0.800±0.017 0.815±0.016 Margin Entropy 0.600±0.020 0.740±0.018 0.700±0.019 0.620±0.019 0.790±0.017 0.790±0.017 0.792±0.016 0.798±0.017 0.802±0.016 Order Policy (Ours) 0.660±0.017 0.788±0.016 0.756±0.017 0.760±0.016 0.810±0.016 0.800±0.016 0.793±0.017 0.825±0.016 0.827±0.016 MBPP Order

128–128

128–32

128–16

256–256

256–64

256–32

512–512

512–128

512–64

Confidence (Top-1) 0.373±0.011 0.400±0.010 0.388±0.010 0.395±0.011 0.396±0.010 0.387±0.010 0.385±0.011 0.392±0.010 0.388±0.010 Margin 0.333±0.012 0.390±0.010 0.388±0.010 0.363±0.011 0.380±0.010 0.378±0.010 0.385±0.011 0.388±0.010 0.380±0.010 0.323±0.012 0.388±0.010 0.396±0.010 0.388±0.011 0.385±0.010 0.378±0.010 0.390±0.011 0.390±0.010 0.393±0.010 Entropy Order Policy (Ours) 0.375±0.010 0.405±0.009 0.416±0.009 0.410±0.010 0.407±0.009 0.400±0.009 0.407±0.010 0.395±0.009 0.395±0.009

Table 6. Inference-time overhead of SAS on GSM8K with generation length 128, averaged over 15 runs. Setting

Avg. decoding time

Backbone only Backbone + SAS Added time from SAS

338.753s 344.415s +5.662s (1.67%)

Inference cost. SAS introduces only a small inference-time overhead because the order policy is lightweight compared to the frozen backbone denoiser. On GSM8K with generation length 128, averaged over 15 runs, backbone-only decoding takes 338.753s, while backbone decoding with SAS takes 344.415s. Thus, SAS adds only 5.662s of decoding time, corresponding to a 1.67% total overhead. This suggests that the learned ordering mechanism improves decoding quality while preserving nearly the same inference efficiency as the original backbone. Direct generalization to parallel decoding. We also evaluate whether a sequentially trained SAS policy can be directly reused for parallel top-k decoding, where multiple positions are decoded at each step. Here the decoding setting is denoted by L-B-T , where L is generation length, B is block size, and T is the number of decoding steps. Table 7 reports the results on GSM8K and MBPP. These results suggest that a policy trained only for sequential decoding does not directly optimize the parallel top-k decoding regime. Compared with sequential decoding, direct parallel decoding leads to performance drops of different magnitudes across tasks and block configurations. This is expected because decoding multiple tokens simultaneously introduces additional dependencies that are not present in the sequential training objective. We therefore view direct parallel policy learning as an important extension of SAS: while sequentially trained policies can be applied to parallel decoding, fully exploiting the parallel advantage of diffusion language models likely requires training the order policy directly under blockwise or parallel decoding objectives. D.3. Open-ended Instruction-following Generation To evaluate whether learned scheduling also helps beyond structured reasoning tasks, we additionally train and test SAS on the open-ended instruction-following benchmark IFEval (Zhou et al., 2023). We compare the learned order policy against standard heuristic schedules under the same decoding setting, where L-B-T denotes generation length, block size, and number of decoding steps. Table 8 reports results for the 256-64-256 setting. The learned order policy outperforms all heuristic baselines, suggesting that the benefit of learned scheduling is not limited to Sudoku, math, or code generation, but also extends to open-ended instruction-following generation. D.4. Comparison with other order learning methods We compare SAS with prior order-learning methods (Hong et al., 2025; Jazbec et al., 2025) from two complementary perspectives. A fully controlled reimplementation of all prior methods is difficult because some codebases and training 26

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

Figure 8. Visualization of learned decoding order on a multi-step reasoning task. Prompts = [“A factory produces notebooks over several days. First, determine a key parameter: on Monday, the factory produces notebooks at a constant rate and makes 180 notebooks in 6 hours. 1. Using the rate from second question, how many notebooks will the factory produce in 14 hours? 2. What is the production rate in notebooks per hour? 3. if the factory must package notebooks in boxes of 28, how many full boxes can be packed from the notebooks produced in those 14 hours?”]. The text is colored by generation step, transitioning from blue (generated first) to red (generated last). Notably, the model prioritizes answering Question 2 (calculating the rate) first, identifying it as a prerequisite before solving Question 1, which depends on that rate. This non-sequential ordering demonstrates the model’s capacity for global semantic understanding, allowing it to plan and resolve logical dependencies effectively.

details are not publicly available, and because several prior methods rely on external verifiable terminal rewards rather than the self-aware likelihood reward used in SAS. Therefore, we avoid making claims of a definitive apples-to-apples benchmark. Instead, we provide: (i) a controlled reward ablation within our own framework, and (ii) matched-setting comparisons against reported numbers from prior work. Controlled terminal-reward comparison within our framework. To isolate the key methodological difference, we replace our dense self-aware reward with a sparse terminal reward similar to those used in prior verifier-based order-learning methods, while keeping the same backbone denoiser, order-policy network, and training pipeline fixed. This comparison is reported in Section 6.5 and Figure 3. Under this controlled setup, SAS performs better overall, suggesting that the dense self-aware reward provides a more effective learning signal than sparse terminal correctness feedback for learning the unmasking order. Matched-setting comparison to reported results. We also compare SAS against reported values from Hong et al. and Jazbec et al. under matched decoding settings when available. We denote each setting by L-B-T , where L is the generation length, B is the block size, and T is the number of decoding steps. The results are shown in Table 9. SAS is competitive overall and stronger in most matched reported settings, especially on GSM8K. On GSM8K, SAS improves over prior reported results in two settings and matches the best reported value in the third. On MATH500, SAS outperforms Jazbec et al. in the 256-256-256 setting, while falling below the best reported value in the other two settings. 27

Scheduling Thoughts: Learning the Order of Thought in Diffusion Language Models Table 7. Direct generalization of a sequentially trained SAS policy to parallel top-k decoding. Each setting is denoted by L-B-T , where L is generation length, B is block size, and T is the number of decoding steps. Task GSM8K MBPP

128-128-64

128-64-32

256-64-64

256-64-32

0.58 0.2567

0.60 0.14

0.61 0.14

0.58 0.1667

Table 8. Open-ended instruction-following evaluation on IFEval. The decoding setting is 256-64-256, where L-B-T denotes generation length, block size, and number of decoding steps. Method

IFEval Score

Confidence Margin Entropy Order Policy (Ours)

0.5067 ± 0.0289 0.5233 ± 0.0289 0.5267 ± 0.0288 0.5633 ± 0.0288

Sudoku comparison across problem scales. We further evaluate order learning on Sudoku, where prior work and SAS were originally tested on different problem scales. To make the comparison more informative, we train the available (Hong et al., 2025). codebase on 9 × 9 Sudoku and our codebase on 4 × 4 Sudoku, so that both methods are evaluated on both 4 × 4 and 9 × 9 settings. The results are shown in Table 10. SAS performs favorably on both scales. Interpretation. Overall, these comparisons suggest that SAS is competitive with existing order-learning methods and often improves over them in matched reported settings. However, we emphasize that the comparison is not perfectly controlled: methods may differ in backbone model, architecture, training pipeline, reward design, and task setup. We therefore view the reported comparisons as suggestive evidence rather than a definitive apples-to-apples benchmark. A fully controlled same-backbone comparison, for example on GSM8K with LLaDA, would be a valuable direction for future work.

E. Comparison with Hidden-State–Conditioned Decode Order Policy In addition to the feature-based order policy described in the main paper, we implement and evaluate an alternative decodeorder policy that directly conditions on the internal hidden states of the frozen LLaDA diffusion model. This variant is intended to test whether richer token-level representations can further improve reveal order selection in practice. The feature-based policy operates on lightweight summary statistics (e.g., entropy, maximum probability), which are inexpensive and model-agnostic. In contrast, hidden states encode substantially richer contextual information, including long-range dependencies and latent reasoning structure. However, they are also higher-dimensional, more entangled with the backbone model, and potentially misaligned with the decode-order decision boundary. E.1. Policy Inputs At each diffusion step, we extract the last-layer hidden states from the frozen LLaDA model for all token positions. For a batch of size B and sequence length L, this yields a tensor: H ∈ RB×L×dhid . Hidden states corresponding to padding tokens are masked out and never used as candidates. No gradients are propagated into the diffusion model. E.2. Architecture The hidden-state order policy is implemented as a lightweight projection network operating on H: • A linear projection from dhid to dproj • A non-linear activation • A final linear layer producing a scalar score per position This produces per-position scores λ ∈ RB×L , analogous to the feature-based policy. No additional self-attention layers are 28

Scheduling Thoughts: Learning the Order of Thought in Diffusion Language Models Table 9. Matched-setting comparison with reported results from prior order-learning methods. Each setting is denoted by L-B-T , where L is generation length, B is block size, and T is the number of decoding steps. Missing entries indicate that the corresponding setting was not reported in the prior work.

GSM8K Method

128-32-128

256-32-256

256-256-256

0.703 – 0.7875

– 0.800 0.800

– 0.330 0.760

Hong et al. Jazbec et al. SAS (ours)

MATH500 Method

128-32-128

256-32-256

256-256-256

0.284 – 0.260

– 0.355 0.300

– 0.200 0.2467

Hong et al. Jazbec et al. SAS (ours)

Table 10. Cross-scale Sudoku comparison of SAS and (Hong et al., 2025). Both methods are trained and then evaluated on 4 × 4 and 9 × 9 Sudoku.

Method Hong et al. SAS (ours)

4 × 4 Sudoku

9 × 9 Sudoku

80% 95%

76% 91%

used, as contextualization is already encoded in the backbone hidden states. E.3. Candidate Masking and Action Selection Candidate masking follows the same rules as in the feature-based policy: • prompt tokens are excluded, • already revealed tokens are excluded, • padding positions are excluded. At each step, the policy selects the highest-scoring valid position and reveals the corresponding token. All training uses the GRPO loss. E.4. Comparison with Feature-Based Policy Table 11 summarizes performance on GSM8K under varying decode budgets. Table 11. GSM8K accuracy under different decode budgets. Method

5-shot / 128–128

5-shot / 128–64

5-shot / 128–32

Top-1 Confidence Margin Entropy Order Policy (features) Left-to-Right Hidden-State Policy

0.620 0.608 0.600 0.660 0.754 ± 0.019 0.478 ± 0.022

0.711 0.752 0.756 0.748 0.756 ± 0.019 0.700 ± 0.021

0.728 0.742 0.740 0.788 0.756 ± 0.019 0.754 ± 0.019

Observations. Across settings, the hidden-state policy does not consistently outperform the simpler feature-based policy. While performance improves as the decode budget becomes more constrained, the hidden-state variant remains less robust than the feature-based order policy and often underperforms even heuristic schedules. 29

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

Discussion. These results suggest that directly exposing high-dimensional backbone representations to the order policy can be counterproductive. Although hidden states contain rich semantic information, they also entangle many factors—such as token identity, syntactic structure, and semantic content—that may not be directly relevant to the decode-order decision. This can make the hidden-state policy harder to optimize and less stable. In contrast, explicitly constructed confidence and uncertainty features, such as top-token probability, margin, and entropy, provide a more direct, interpretable, and stable signal for deciding which position to reveal next, while remaining lightweight and model-agnostic. We therefore view the hidden-state variant as an exploratory design study rather than a negative result about all representation-based policies. The feature-based policy is chosen not because the design space is exhausted, but because it is the most stable, lightweight, and empirically reliable option we found.

Record · ID 299942 · SHA-256 f16064a4842d269d
Retrieved via Conceptio — every document is proof-bundled with source, license, and retrieval metadata.