🤖 machine learning

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

本文介绍了自我感知调度（Self-Aware Scheduling, SAS），这是一个通过推导解码失配的可解上界来训练轻量级策略，从而优化掩码扩散语言模型中标记去掩码顺序的原则性框架，与启发式调度相比，该框架显著提升了在数独和数学推理等任务上的生成质量。

原作者： Jiawei Xu, Minghui Liu, Aakriti Agrawal, Yifan Chen, Furong Huang

发布于 2026-06-23

📖 1 分钟阅读☕ 轻松阅读

CC BY 4.0

原作者： Jiawei Xu, Minghui Liu, Aakriti Agrawal, Yifan Chen, Furong Huang

原始论文采用 CC BY 4.0 许可（http://creativecommons.org/licenses/by/4.0/）。 ✨ 这是对下方论文的AI生成解释。它不是由作者撰写或认可的。如需技术准确性，请参阅原始论文。阅读完整免责声明

核心思想：机器如何“思考”至关重要

想象一下，你正在尝试拼凑一个巨大的拼图，但你看不见包装盒上的图案。你必须猜测每一块碎片应该放在哪里。

大多数 AI 模型（比如那些写文章或写代码的模型）的工作方式就像一个人在读书：它们一个词、一个词、一个词地按顺序写出来，严格遵循从左到右的直线顺序。这被称为**自回归（Autoregressive）**生成。

然而，一种新型的 AI——扩散模型（Diffusion Model），其工作方式则完全不同。想象一下，你有一个拼图，每一块碎片都被一个黑方块（“掩码/Mask”）盖住了。AI 的任务就是逐一揭开这些碎片，直到整幅图像显现出来。

问题在于： AI 有多种选择。它可以按任何顺序揭开碎片。

它可以先从角落的碎片开始（容易）。
它可以先从中间的碎片开始（难）。
它也可以随机挑选碎片。

过去，科学家们告诉 AI 根据简单的规则来挑选，比如“总是先选出现在看起来最确定的那个碎片”。论文将这些称为启发式调度（Heuristic Schedules）。作者认为这些规则是“近视”的（短视的）。它们先挑容易的碎片，但这可能会让剩下的拼图变得更难解决。

解决方案：“自我感知调度”（SAS）

作者创造了一种名为**自我感知调度（Self-Aware Scheduling, SAS）**的新方法。AI 不再遵循僵化的规则，而是学习它自己的“思考顺序”。

这就像一个学生在参加考试：

旧方法（启发式）： 学生看着试卷，立即回答所有他百分之百确定的问题，把难的问题留到最后。如果这些简单的问题无法提供足够的线索来解决难题，学生就会陷入困境。
SAS 方法： 学生看着整张试卷并问道：“如果我先做第 5 题，它能否为我解决第 10 题提供必要的线索？”学生学会了挑选那个能让剩余部分更容易完成的问题，即使这个题目本身现在并不容易回答。

它是如何工作的（“秘诀”）

论文使用了一些复杂的数学公式，但概念很简单：

“自我感知”奖励： AI 有一个“大脑”（模型），它已经训练好了。新的部分是一个“管理者”（策略），负责决定顺序。
测试： 管理者尝试不同的揭开拼图碎片的顺序。
评分： 与其等到最后看拼图是否完成（这很慢且很难成功），管理者会检查：“给定我刚刚选定的顺序，我的‘大脑’现在能多好地解释这个答案？”
- 如果这个顺序让“大脑”感到自信且逻辑通顺，管理者的得分就高。
- 如果这个顺序让“大脑”感到困惑，得分就低。
学习： 管理者利用这个分数来学习哪些顺序效果最好。这就像教练告诉球员：“不要只是跑得快；要朝着能帮助整个团队获胜的方向跑。”

研究发现（结果）

研究人员在三种不同类型的“拼图”上测试了该方法：

数独（逻辑谜题）：
- 他们使用了一个 10 亿参数的 AI。
- 旧方法： 最好的“基于规则”的方法解决了约 82% 的谜题。
- SAS 方法： AI 学会了自己的顺序，解决了 91.8% 的谜题。
- 惊喜： 甚至人类专家的“逻辑”顺序（先解最简单的数字）也比不上 AI 学到的顺序！AI 找到了一条更适合其特定“大脑”的路径。
数学问题 (GSM8K)：
- 他们使用了一个更大的 80 亿参数 AI。
- 旧方法： 正确解决了 64% 的问题。
- SAS 方法： 正确解决了 76% 的问题。
编程 (MBPP)：
- 旧方法： 解决了 39.5% 的问题。
- SAS 方法： 解决了 41% 的问题。

“第二阶段”加成

论文还发现，一旦 AI 学会了最佳思考顺序，你可以针对这条路径给 AI 进行额外的专门训练。

这就像音乐家学习了一首曲子的最佳练习方式，然后只通过这种方式练习直到精通。
在数独任务中，这一额外步骤将准确率从 91.8% 提升到了 97.5%。

为什么这很重要

论文声称，AI 如何思考（它揭示信息的顺序）与它知道什么同样重要。通过教 AI 规划自己的思考路径，而不是强迫它遵循僵化的规则，我们可以让它在解决数学、逻辑和编程等复杂问题时变得更加聪明。

简而言之：这篇论文教 AI 不仅仅是“猜下一个词”，而是开始“规划通往答案的最佳路径”。

技术摘要：调度思维 (SAS)

1. 问题陈述

掩码扩散模型 (MDMs) 通过迭代地取消掩码（unmasking）标记来生成离散序列，这一过程为自回归模型固定的从左至右分解提供了一种灵活的替代方案。然而，这种灵活性引入了一个关键的自由度：思维顺序（即取消掩码的调度）。

目前的方法依赖于启发式调度（例如，根据局部置信度、边际概率或熵来选择标记）。作者认为这些启发式方法是“近视”的：它们优化的是紧接着的下一步，而非全局解码轨迹。因此，它们经常选择那些局部确定性高但无法为后续步骤有效传播约束的“容易”标记。此外，即使是专家设计的逻辑顺序（例如，人类解决数独的策略），对于特定的预训练扩散模型来说也可能是次优的，因为模型的条件预测可能在与人类逻辑不同的路径上最为可靠。

核心问题在于如何学习一种最优的取消掩码顺序，以在不重新训练基础扩散模型的情况下最大化生成质量，同时解决搜索 $n!$ 种可能顺序的难解性问题。

2. 方法论：自我感知调度 (SAS)

本文提出了自我感知调度 (Self-Aware Scheduling, SAS)，这是一个即插即用的框架，它在保持基础扩散模型 ( $p_\theta$ ) 冻结的同时，学习一个轻量级的顺序策略 ( $v_\phi$ )。

理论基础

作者将扩散解码建模为一个潜在顺序生成过程。其目标是最小化真实数据分布 $\pi(x)$ 与由解码过程 $P^{seq}_{\theta,\phi}(x)$ 诱导的分布之间的 Kullback–Leibler (KL) 散度。

难解性： 由于需要对所有 $n!$ 种解码顺序求平均，直接最小化该 KL 散度是难以实现的。
可解上界： 作者推导出了该失配的一个可解上界。他们证明，最小化 KL 散度等价于（在常数范围内）最小化数据-策略联合分布与模型-策略联合分布之间的联合 KL 失配。
自我感知奖励： 这一推导产生了一个稠密的、模型自我感知的奖励函数，定义为目标序列在冻结去噪器下的路径对数似然：
$R_\theta(x, \sigma) = \log p_\theta(x | \sigma) = \sum_{t=1}^n \log p_{i_t}^\theta(x_{i_t} | \tilde{x}^{(t-1)}(\sigma))$
该奖励根据模型在被迫按特定顺序揭示标记时解释数据的能力，为每个揭示决策分配信用。

算法

策略参数化： 顺序策略 $v_\phi$ 是一个轻量级的神经网络（例如 MLP 或 Transformer），它输出当前被掩码位置的分类分布。
优化： 使用群体相对策略优化 (GRPO) 来训练策略。该过程被公式化为一个有限时界的马尔可夫决策过程 (MDP)：
- 状态： 部分揭示的序列。
- 动作： 选择下一个要取消掩码的索引。
- 奖励： 当前上下文（教师强制）下该位置处真实标记的对数似然。
两阶段框架：
1. 顺序学习： 在 $\theta$ 冻结的情况下优化 $\phi$ ，以寻找现有模型的最佳轨迹。
2. 第二阶段微调： （可选）冻结学习到的顺序策略，并沿着学习到的轨迹微调扩散模型 $\theta$ ，以进一步降低误差。

3. 核心贡献

理论： 本文推导了一个将解码失配（KL 散度）与去噪器的累积路径对数似然联系起来的可解上界。这为使用模型自身的似然作为顺序选择的稠密奖励提供了理论依据。
算法 (SAS)： 引入了一个通过 GRPO 学习取消掩码顺序策略的即插即用框架。它适用于任意顺序序列解码和半自回归（块）解码。
实验验证： 证明了优化思维顺序可以在多种推理任务（数独、数学、代码）和模型规模（1B 到 8B 参数）上，相比启发式和专家调度实现显著且一致的改进。

4. 实验结果

数独 (1B 模型)

性能： SAS 将谜题准确率从 82.0%（最佳启发式：置信度）提升至 91.8%。
对比： 学习到的策略显著优于人类专家顺序（76.6%），这表明对于不完美的扩散模型，最优路径未必是人类的逻辑路径。
分析： 可视化显示，该策略采用了一种“由易到难”的课程学习方式，首先解决高确定性的标记（如唯余数/naked singles）以传播约束，但在对整体有利之处会偏离人类逻辑。
第二阶段： 沿着学习到的轨迹对模型进行微调，进一步将准确率提升至 97.5%。

数学推理 (LLaDA-8B)

GSM8K： Pass@1 从 64% 提高到 76%。
MBPP (代码)： Pass@1 从 39.5% 提高到 41%。
泛化性： 学习到的策略在不同生成长度（128 到 512）和块大小（半自回归解码）的变化下保持稳健，始终优于包括从左至右启发式在内的基准模型。

跨领域迁移

在数学任务上训练的策略可以有效地迁移到代码任务（反之亦然），在所有测试设置中均优于最强的启发式基准。这表明该策略学习的是关于信息密度的可迁移调度原则，而非记忆特定任务的模式。

奖励对比

与稀疏的 0/1 终端奖励（正确/错误答案）相比，稠密的自我感知奖励提供了更稳定且通常更优的性能，突显了稠密、模型一致的信号相对于稀疏终端反馈的优势。

5. 意义与主张

论文将解码顺序定位为一种可控且可优化的“推理原语”。作者声称，除了学习“生成什么”之外，还可以学习“何时”对每项信息进行提交。

原则性优化： 不同于优化局部不确定性的启发式调度，SAS 优化的是全局轨迹似然，使解码路径与模型的预测优势相一致。
高效性： 该方法计算效率高，仅需训练轻量级策略，而无需对完整的扩散骨干网络进行昂贵的强化学习更新。
自我改进： 该框架为通过迭代调度和模型细化来实现扩散大语言模型（Diffusion LLMs）的自我改进提供了一条切实可行的路径，即模型学习其自身最可靠的推理路径。

作者承认存在局限性，包括教师强制训练与自由运行推理之间可能存在的差距，以及目前对单任务训练的依赖，并指出虽然跨领域迁移前景广阔，但尚未得到充分表征。