自主思维学习：抽象令牌课程

# 学习你自己的思维：抽象令牌课程  
来源：https://arxiv.org/html/2609.19717  

Khashayar Gatmiry 加州大学伯克利分校 [email protected]&Avrajit Ghosh* 加州大学伯克利分校 [email protected]&Parsa Mirtaheri* 加州大学圣地亚哥分校 [email protected] 以及 Jason D. Lee 加州大学伯克利分校 [email protected]&Nika Haghtalab 加州大学伯克利分校 [email protected]&Emmanuel Abbe EPFL [email protected] 以及 Peter Bartlett Google DeepMind 和加州大学伯克利分校 [email protected]††致谢：同等贡献。  

#### 摘要  
大型语言模型（LLMs）通过利用思维链（CoT）作为思考中间阶段的草稿本，已经取得了卓越的推理能力。然而，CoT 技术需要对思维令牌进行显式监督，这需要丰富且特定于任务的数据。在这项工作中，我们提出了抽象令牌课程（ATC），这是一种新颖的课程学习框架，能够在没有直接监督或人工草稿本设计的情况下，引出有效的连续中间表示。ATC 通过一系列分布逐步增加问题复杂性，训练模型在连续表示空间中发展内部抽象的“思维”。本文提供了理论与实验证据，证明了 ATC 的益处及其在训练连续思维方面相对于先前方法的优势。理论上，我们表明，对于使用单层 softmax 注意力通过 ATC 学习奇偶校验函数的情况，注意力自然会集中在上下文中为预测下一个令牌提供“最简路径”的 CoT 令牌上。实验上，我们展示了 ATC 在图可达性和算术学习任务上的有效性。  

## 1 引言  
大型语言模型（LLMs），通常实例化为基于 Transformer 的架构（Vaswani 等人，2017）并在互联网规模的语料库上进行训练，在推理和指令遵循任务中展现了强大的泛化能力。一种特别有效的提示策略是引导出自然语言中的中间推理步骤，即所谓的思维链（CoT），这已被证明可以提高一系列推理基准测试的性能（Wei 等人，2022；Kojima 等人，2022）。直观上，CoT 为中间论证和辅助计算提供了一个草稿空间，有助于分解复杂任务。最近的一些工作进一步通过强加结构来引导这种中间推理，例如，通过符号表示（Xu 等人，2023）或将部分推理委托给可执行代码（Chen 等人，2022）。另一种强大的方法是提供特定任务的草稿本（Nye 等人，2021；Zhou 等人，2022；Besta 等人，2023），这些草稿本指定了中间步骤的明确格式。虽然在特定领域对专门数据集进行微调已被证明是有效的，但设计可靠可学习的特定任务草稿本或结构化思维链提示可能需要大量人力和专业知识（Nye 等人，2021；Besta 等人，2023）。一个自然的问题以及潜在的困难来源是，哪些形式的中间推理是本身可高效学习的。Lotfi 等人（2024）通过“全局性程度”的概念将此形式化，并推测可学习的思维链具有恒定的全局性程度。大致而言，他们推测，对于那些下一个令牌可以仅从上下文中的少数几个思维令牌预测的 CoT，模型应学会在训练期间有效地模仿思维链，将其注意力集中在少数令牌上。我们称之为“注意力简单性偏差”。另一方面，直接在固定、全难度目标分布上训练可能会带来严重的优化障碍（Bengio 等人，2009）。一种互补的方法是课程学习：在一系列逐渐变难的分布上进行训练，这可以提高在挑战性推理问题上的优化和样本效率（Bengio 等人，2009；Hase 等人，2024）。然而，虽然课程学习相对于固定分布训练提高了优化和样本效率，但它本身并不能克服底层模型计算能力的固有局限性。最近的理论工作表明，标准 Transformer 架构的计算能力在被限制为不产生中间推理步骤的输出时，其计算能力是固有的。特别是，没有思维链（CoT）或草稿本令牌的有界深度 Transformer 可以用相对浅层的电路类别来表征，无法高效执行某些固有的顺序或组合计算（Hahn，2020；Merrill 和 Sabharwal，2023）。相反，用 CoT 或中间令牌增强 Transformer 可证明能增强其计算能力，通过有效增加其顺序计算深度，并在足够中间步骤下能够表示更复杂的函数类别，直至多项式时间计算（Merrill 和 Sabharwal，2023；Feng 等人，2023）。除了表达能力，CoT 还被证明通过结构化中间依赖关系和促进优化，可以提高某些任务的可学习性和样本复杂度（Kim 等人，2024）。综合来看，思维链提示的优势和特定任务草稿本的局限性提出了一个根本问题：是否存在一种方法，无需为每个新问题显式设计特定任务的草稿本，就能从模型中引出有效的思维链？似乎提示技术，如使用思维链示例的小样本学习，在引出这些推理痕迹方面能力有限；它们主要仅在模型已经对特定任务拥有大量基础知识时才有效（Min 等人，2022；Madaan 等人，2023）。在这项工作中，我们提出了抽象令牌课程（ATC），它利用输入分布上的课程技术，在无需显式监督思维链令牌的情况下，为给定任务引出抽象的思维链（CoT）。对于我们的理论贡献，我们解决了 Lotfi 等人（2024）提出的关于使用单层注意力模型通过 ATC 学习 n 个比特的奇偶校验的猜想，从而在此设置下建立了注意力简单性偏差。已知在单次设置中，使用梯度下降学习超立方体上的高阶函数（如奇偶校验）是不可能的（Kearns，1998）。在此设置中，我们表明 ATC 很可能将注意力分数集中在具有预测输出的最简信号的思维令牌上。我们注意到，虽然 ATC 适用于比奇偶校验预测更通用的设置，但由于我们证明的复杂性以及奇偶校验作为研究挑战性学习问题的测试平台的重要性，我们的理论分析（定理 1）集中在奇偶校验的 ATC 上。在此设置中，我们为 ATC 的有效性和注意力简单性偏差的有效性提供了理论证据。我们进一步进行了广泛的实验，证明了 ATC 在奇偶校验之外的任务（包括图推理和多位数算术）上的有效性。定理 1 表明，对注意力键执行一个大的梯度下降步骤足以使注意力简单性偏差显现：以恒定概率，注意力分数集中在最近的思维链令牌上。然而，定理 1 的一个关键假设是，在 ATC 早期阶段学习的表示在之前的思维令牌中保持不变，以便它们可以在当前阶段可靠使用。在我们的实验中，我们发现这一特性并非自动成立。具体来说，逐步训练更长的连续思维链带来了两个挑战。1. 随着课程推进，在更难问题上的训练可能会改变早期潜在思维的表示。2. 全反向传播需要存储每个潜在步骤的激活，导致内存和计算成本随着推理链长度的增长而增加。为解决第一个挑战，我们引入课程回溯作为一般保留机制：模型在所有先前完成的阶段上重新评估，并返回到准确率低于课程阈值的早期阶段。因此，回溯使先前已解决的问题保持在训练标准内，并防止在早期推理被遗忘时课程继续推进。为解决第二个挑战，我们使用截断反向传播，它仅通过最近思维的固定窗口传播梯度。在我们的实验中，全反向传播已经可以保留早期阶段，而在截断反向传播下，这一保护变得至关重要。当输入提示中的答案候选被模型复制作为快捷方式而不执行底层计算时，会出现进一步的挑战。为此，我们引入候选丢弃，它在部分训练示例中隐藏候选，并鼓励模型根据问题输入计算答案。我们的实验表明，课程回溯、截断反向传播和候选丢弃共同实现了在诸如图可达性和算术运算等任务上的推理。它们的消融实验揭示了各自的局限性：没有候选丢弃，训练会停滞在复制快捷方式上；没有回溯，在图可达性上的截断训练无法保留先前计算并崩溃。截断反向传播在两项任务中都降低了内存和计算成本，而回溯则在性能退化时恢复早期阶段的性能。这些发现证明了防止快捷方式、降低训练成本和保留推理所依赖的计算之间的互补作用。  

## 2 设置  
我们考虑从序列输入 x=(x₁,…,xₙ) 学习目标 y 的问题，其中每个 xᵢ 来自字母表 Σ，y 属于 Y（可以是离散或连续）。我们的方法，抽象令牌课程（ATC）如图 1 所示，由以下组件定义：  
- • 一系列分布 D=(𝒟₁,…,𝒟ₖ) 在 Σⁿ 上，其中每个 𝒟ᵢ 代表用于第 i 阶段课程的输入序列的分布。  
- • 一个嵌入模型 fₜ𝒻: Σⁿ ⊗ (ℝᵈ)ᵏ → ℝᵈ，它接受输入 x=(x₁,…,xₙ) 的令牌序列 (x₁,…,xₙ,z₁,…,zₖ) 和连续思维链向量 z=(z₁,…,zₖ)，并将其映射到一个连续向量，我们将其作为 Transformer 的最后一层隐藏嵌入，对应于上下文中的最右侧令牌。具体来说，Transformer fₜ𝒻 接受 Σ 中的离散令牌（初始输入）和 ℝᵈ 中的连续令牌（学习到的思维链表示）。在本文的其余部分，我们假设 fₜ𝒻 选择为 Transformer。我们关注的 fₜ𝒻 的一个常规选择是 Transformer 架构，其中初始嵌入层仅应用于离散输入令牌 xᵢ∈Σ，而对于连续令牌，我们跳过初始嵌入层，直接将它们馈送到第一层注意力（在层归一化之后）。输出向量 fₜ𝒻(x₁,…,xₙ,z₁,…,zₖ) 是应用于上下文中最后一个令牌的 Transformer 的最后一层隐藏嵌入。  
- • 一系列确定性变换 (𝒯ᵢ)ᵢ₌₁,…,ₖ，其中 𝒯ᵢ 应用于 fₜ𝒻(x,z₁,…,zᵢ) 以获得下一个连续思维链，即 z_{i+1} := 𝒯ᵢ(fₜ𝒻(x,z₁,…,zᵢ))。例如，𝒯ᵢ 可以是恒等函数，直接传递隐藏嵌入 fₜ𝒻(x,z₁,…,zᵢ) 作为下一个连续思维链 z_{i+1}。  
- • 一系列头 (f_headⁱ)ᵢ₌₁,…,ₖ，其中每个 f_headⁱ: ℝᵈ → Y 是一个线性层，将 d 维隐藏嵌入 fₜ𝒻(x,z₁,…,zᵢ) 映射到输出空间 Y。  

ATC 的核心思想是将训练组织在分布 𝒟ᵢ 上，其支持实例 x 的固有难度 i（随 i 增加），并运行 k 个阶段，每个阶段在逐渐增加实例复杂性的同时向上下文追加一个新的连续表示。具体来说，我们首先在最简单的实例 𝒟₁ 上训练，并使用应用于最后一个隐藏嵌入（在 xₙ 处）的头 f_head¹ 预测目标 y，即 fₜ𝒻(x₁,…,xₙ)；当 y 是离散的，f_head¹ 可以是语言头，而当 y 是连续的，则可以是线性头。在 𝒟₁ 上训练以预测 y 迫使 fₜ𝒻(x) 包含在 𝒟₁ 上对 y 有预测性的表示，这可能有助于对更复杂的 𝒟₂ 中的实例进行预测。为了利用这一点，一旦在 𝒟₁ 上验证损失饱和，我们就将 fₜ𝒻(x) 追加到上下文的末尾（在最后一个输入令牌 xₙ 之后）作为第一个连续思维链令牌 z₁。我们继续这个过程；在阶段 r，对于采样实例 x∼𝒟ᵢ，我们首先将输入 x 放入上下文，然后是连续思维链令牌 z₁,…,z_{r-1}，并训练模型使用 f_headⁱ 从最后一个隐藏嵌入 fₜ𝒻(x,z₁,…,z_{r-1}) 预测 y。通过这种方式，我们将预测输出的复杂性分解为更简单的中间训练步骤，从形成预测 𝒟ᵢ 的表示到预测 𝒟_{r+1}。