让扩散语言模型自动学习最优的解码顺序，提升生成质量。
Scheduling Thoughts: Learning the Order of Thought in Diffusion Language Models
- 基于路径似然推导出可优化的自注意奖励，指导解码顺序选择。
- 在数独任务中准确率从82.0%提升至91.8%，数学推理任务提升至76%。
- 适用于多种解码模式，适合关注生成质量与可控性的研究者。
掩码扩散语言模型通过迭代解码来恢复被遮蔽的词元，解码顺序定义了‘思维顺序’，对生成质量有显著影响，但通常由启发式方法决定。本文推导出一种可计算的上界，用于衡量序列解码不匹配度（以KL散度表示），该上界由模型的路径似然表达，在模型表达能力强时紧致。这一上界催生了一个密集的自注意奖励信号，将顺序选择建模为一个带冻结去噪器的策略优化问题。我们提出Self-Aware Scheduling（SAS）方法，使用分组相对策略优化学习轻量级顺序策略，并可无缝应用于任意顺序和半自回归解码。在10亿参数的掩码扩散模型（1B MDM）上的数独任务中，准确率从82.0%（最佳启发式）提升至91.8%，经第二阶段微调后达97.5%；在LLaDA-8B的数学推理任务中，GSM8K的pass@1从64%升至76%，MBPP从39.5%升至41%，在不同生成长度与块大小下均优于或等同于启发式调度。
原文摘要 · Abstract (English)
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
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。