Scheduling Thoughts: Learning the Order of Thought in Diffusion Language Models
作者: Jiawei Xu, Minghui Liu, Aakriti Agrawal, Yifan Chen, Furong Huang
分类: cs.LG, cs.AI
发布日期: 2026-06-22
🔗 代码/项目: PROJECT_PAGE
💡 一句话要点
提出自我感知调度以优化扩散语言模型的思维顺序
🎯 匹配领域: 支柱四：生成式动作 (Generative Motion)
关键词: 扩散语言模型 自我感知调度 策略优化 数独 数学推理 生成质量 Kullback-Leibler散度 轻量级策略网络
📋 核心要点
- 现有的掩码扩散语言模型在解码过程中，解掩码顺序通常依赖启发式选择，导致生成质量不稳定。
- 本文提出自我感知调度（SAS），通过优化解掩码顺序来提升生成质量，将其视为策略优化问题。
- 实验结果表明，SAS在数独任务中将准确率从82.0%提升至91.8%，在数学推理任务中也显著提高了通过率。
📝 摘要（中文）
掩码扩散语言模型通过逐步解掩码令牌进行解码，其中解掩码顺序定义了“思维顺序”，对生成质量有显著影响，但通常是启发式选择。本文推导了顺序解码不匹配的可处理上界，利用Kullback-Leibler散度进行测量，并在模型表达能力足够的情况下保持紧密性。该界限引入了基于有序轨迹的密集自我感知奖励，将顺序选择视为一个原则性的策略优化问题。我们将这一思想实例化为自我感知调度（SAS），使用组相对策略优化学习轻量级顺序策略，能够无缝应用于任意顺序和半自回归解码。在数独和数学推理任务中，SAS显著提高了生成准确性。
🔬 方法详解
问题定义：本文旨在解决掩码扩散语言模型中解掩码顺序选择的不足，现有方法依赖启发式策略，导致生成质量不一致。
核心思路：提出自我感知调度（SAS），通过引入基于有序轨迹的自我感知奖励，将顺序选择转化为策略优化问题，从而系统性地学习最佳解掩码顺序。
技术框架：SAS的整体架构包括：首先，通过模型的路径对数似然推导出解码不匹配的上界；其次，利用组相对策略优化算法学习解掩码顺序；最后，将学习到的顺序应用于任意顺序和半自回归解码。
关键创新：SAS的主要创新在于将解掩码顺序选择视为一个优化问题，利用自我感知奖励机制来指导学习过程，这与传统的启发式方法有本质区别。
关键设计：在设计上，SAS采用了轻量级的策略网络，并通过密集的自我感知奖励来优化顺序选择，确保在不同任务和生成长度下均能有效提升性能。
🖼️ 关键图片
📊 实验亮点
在数独任务中，SAS将准确率从82.0%提升至91.8%，并通过二次微调达到97.5%。在数学推理任务中，SAS在GSM8K数据集上的通过率从64%提升至76%，在MBPP数据集上从39.5%提升至41%，表现优于传统启发式调度。
🎯 应用场景
该研究在自然语言生成、对话系统及其他需要序列生成的任务中具有广泛的应用潜力。通过优化解掩码顺序，能够显著提升生成内容的质量和准确性，未来可能推动更高效的语言模型开发。
📄 摘要（原文）
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