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

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

Related papers

• Constrained Decoding for Diffusion Language Models via Efficient Inference over Finite Automata [57.27430779838529]
  We present an exact and tractable algorithm for sampling from the constrained mean-field posterior under any constraint expressible as a finite automaton.<n>The approach guarantees constraint satisfaction by construction, supports both greedy and sampling-based decoding, and is compatible with parallel and block-wise decoding.<n> Empirical evaluations on Dream-7B and LLaDA-8 show substantial accuracy gains across various tasks.
  arXiv  Detail & Related papers  (2026-07-08T05:48:57Z)
• JUMP: Single-Pass Membership Inference on Fine-Tuned Diffusion Language Models [5.43345665278304]
  We study membership inference attacks (MIAs) for fine-tuned discrete diffusion language models (dLLMs)<n>We propose JUMP (Joint Uncertainty-Guided Mask Probing), a single-pass scoring attack that exploits both distinctive properties of dLLMs.<n>On fine-tuned LLaDA-8B-Base across six MIMIR domains, JUMP improves mean ROC-AUC from 0.82 to 0.90 over SAMA and substantially improves low-FPR detection.
  arXiv  Detail & Related papers  (2026-05-09T05:40:00Z)
• DPRM: A Plug-in Doob h transform-induced Token-Ordering Module for Diffusion Language Models [76.12556589212666]
  We introduce DPRM (Doob h-transform Process Reward Model), a plug-in token-ordering module for diffusion language models.<n>DPRM starts from confidence-driven progressive ordering and gradually shifts to Doob h transform Process Reward guided ordering through online estimates.<n>Under tractable optimization assumptions, DPRM also yields a sample-complexity advantage over random and confidence-only ordering.
  arXiv  Detail & Related papers  (2026-04-27T11:50:26Z)
• LogicDiff: Logic-Guided Denoising Improves Reasoning in Masked Diffusion Language Models [0.0]
  We introduce LogicDiff, an inference-time method that replaces confidence-based unmasking with logic-role-guided unmasking.<n>A dependency-ordered scheduler then unmasks tokens in logical dependency order.<n> LogicDiff improves LLaDA-8B-Instruct accuracy from 22.0% to 60.7% on GSM8K and from 23.6% to 23.6% on MATH-500, with less than 6% speed overhead.
  arXiv  Detail & Related papers  (2026-03-24T13:08:10Z)
• Adaptation to Intrinsic Dependence in Diffusion Language Models [5.185131234265025]
  Diffusion language models (DLMs) have emerged as a promising alternative to autoregressive (AR) approaches.<n>We introduce a distribution-agnostic unmasking schedule for DLMs that adapts to the (unknown) dependence structure of the target data distribution.<n>Our results significantly improve upon prior convergence theories and yield substantial sampling acceleration for low-complexity distributions.
  arXiv  Detail & Related papers  (2026-02-23T18:41:34Z)
• Search or Accelerate: Confidence-Switched Position Beam Search for Diffusion Language Models [24.78455014605002]
  Diffusion Language Models generate text by iteratively denoising a masked sequence.<n>Standard decoding follows a greedy rule: unmask the most confident positions.<n>We present SOAR, a training-free decoding algorithm that adapts its behavior to the model's uncertainty.
  arXiv  Detail & Related papers  (2026-02-11T15:41:09Z)
• Lookahead Unmasking Elicits Accurate Decoding in Diffusion Language Models [51.12873073612084]
  Masked Diffusion Models (MDMs) as language models generate by iteratively unmasking tokens, yet their performance depends on the inference time order of unmasking.<n>We propose Lookahead Unmasking (LookUM), which addresses these concerns by reformulating sampling as path selection over all possible unmasking orders.<n>LookUM requires only two to three paths to achieve peak performance, demonstrating remarkably efficient path selection.
  arXiv  Detail & Related papers  (2025-11-04T02:37:37Z)
• Improving Discrete Diffusion Unmasking Policies Beyond Explicit Reference Policies [47.6755955972232]
  We cast denoising as a KL-regularized Markov decision process (MDP) with an explicit reference policy and optimize a regularized objective.<n>We prove that the optimized policy under this framework generates samples that more closely match the data distribution than schedules.
  arXiv  Detail & Related papers  (2025-10-07T09:44:24Z)
• Blockwise SFT for Diffusion Language Models: Reconciling Bidirectional Attention and Autoregressive Decoding [60.06816407728172]
  Discrete diffusion language models have shown strong potential for text generation.<n>Standard supervised fine-tuning misaligns with semi-autoregressive inference.<n>We propose Blockwise SFT, which partitions responses into fixed-size blocks.
  arXiv  Detail & Related papers  (2025-08-27T02:49:33Z)
• Large Language Diffusion Models [93.26422905620008]
  Large language models (LLMs) are widely regarded as relying on autoregressive models (ARMs)<n>We introduce LLaDA, a diffusion model trained from scratch under the pre-training and supervised fine-tuning paradigm.<n>Across extensive benchmarks on general tasks, math, code, and so on, LLaDA demonstrates strong scalability and performs comparably to our self-constructed ARM baselines.
  arXiv  Detail & Related papers  (2025-02-14T08:23:51Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.