Transformers Provably Learn Chain-of-Thought Reasoning with Length Generalization¶
会议: NeurIPS 2025
arXiv: 2511.07378
代码: 无
领域: LLM推理
关键词: CoT理论, 长度泛化, 注意力集中, 状态追踪, 复杂度类
一句话总结¶
从优化理论角度证明了一层 Transformer 通过梯度下降在合成状态追踪任务上能学会 CoT 推理并实现长度泛化，首次为常数深度 Transformer 学习 \(\mathsf{NC}^1\)-complete 问题（超越之前局限于 \(\mathsf{TC}^0\) 的理论）提供了收敛保证。
研究背景与动机¶
- 领域现状：CoT 推理是当前 LLM 突破复杂推理任务的关键技术——通过生成中间推理步骤来解决超出模型单步处理能力的问题。理论上已知带 CoT 的 Transformer 在表达力上可以解决 \(\mathsf{NC}^1\) 级别的问题（超越不带 CoT 时只能做 \(\mathsf{TC}^0\) 的限制）。
- 现有痛点：现有理论工作要么只研究表达力（"能不能表示"），不涉及优化（"能不能学到"）；要么只在极简单的 \(\mathsf{TC}^0\) 任务（如 parity）上证明了优化收敛保证。关键缺口：没有人证明 Transformer 能通过训练学会产生可长度泛化的 CoT。
- 核心矛盾：表达力结果不等于可学性结果——即使存在一组权重让 Transformer + CoT 能解决某问题，梯度下降是否能找到这组权重？找到后能否泛化到更长的输入？
- 本文要解决什么：为 Transformer 通过梯度下降学习 CoT 推理提供严格的收敛和长度泛化保证，且任务复杂度达到 \(\mathsf{NC}^1\) 级别。
- 切入角度：选择 LEGO（基于群操作的状态追踪）作为分析对象——它是分离 \(\mathsf{TC}^0\) 和 \(\mathsf{NC}^1\) 的标准理论基准，且具有丰富的代数结构可以利用。
- 核心 idea 一句话：通过"注意力集中"机制将 attention 层的检索鲁棒性与状态追踪任务的代数结构联系起来，证明群结构决定了长度泛化的程度。
方法详解¶
整体框架¶
分析对象是一层 Transformer（softmax attention + smoothed ReLU MLP），不使用位置编码（NoPE），在 LEGO 状态追踪任务上用 CoT loss 训练。LEGO 任务要求模型执行一系列群操作 \(g_1, g_2, ..., g_L\) 并逐步输出中间状态 \(y_1, y_2, ..., y_L\)，其中 \(y_{i+1} = g_{i+1}(y_i)\)。
关键设计¶
- 注意力集中（Attention Concentration）机制:
- 做什么：描述 attention 层如何实现对相关上下文的精确检索
- 核心思路：Transformer 通过注意力矩阵 \(Q\) 学会两个关键检索：(1) \(Q_{4,3}\) 抓取上下文中对应的群操作 \(g_{i+1}\)；(2) \(Q_{4,4}\) 抓取前一步的答案 \(y_i\)。当这两个注意力足够"集中"时，MLP 可以正确计算 \(g_{i+1}(y_i)\)
-
设计动机：长度泛化的关键在于——当输入变长时，attention 能否仍然聚焦在正确的位置上？这取决于任务的代数结构
-
两种群结构下的泛化分析:
- 简单传递群操作（如循环群）：学到的特征具有高"纯度"，attention 集中度极高（对正确子句的 attention 为 \(1 - 1/d^c\)），因此可以泛化到 \(\text{poly}(d)\) 的长度（训练长度的多项式倍）
-
对称群操作（\(S_n\)）：学到的特征不够"纯净"，直接训练只能泛化到训练长度的常数倍。但通过递归自训练（progressive length doubling）可以不断扩展可解长度
-
递归自训练方案（Recursive Self-Training）:
- 做什么：通过课程学习逐步扩展模型可处理的输入长度
- 核心思路：先在长度 \(L\) 上训练到收敛，用模型的贪心输出生成长度 \(2L\) 的训练数据，再继续训练。每轮 doubling 后模型精度达到 \(1 - O(2^{-k/\text{poly}(d)})\)
- 设计动机：对于不能直接长度泛化的复杂群结构，自训练提供了一条渐进式扩展能力的路径，且模型能获得超出训练数据覆盖范围的能力
训练策略¶
- 两阶段课程学习：先训练 MLP 学习单步群操作（one-step reasoning），再冻结 MLP 训练 attention 实现上下文检索
- 不使用位置编码——理论分析表明标准 PE（如 RoPE）因位置依赖性固有地阻碍长度泛化
实验关键数据¶
主实验（简单传递/循环群 \(C_6\)）¶
| 训练长度 | 测试长度5 | 测试长度10 | 测试长度20 | 测试长度30 | 测试长度40 | 测试长度50 |
|---|---|---|---|---|---|---|
| CoT (L=5) | 100% | 100% | 100% | 100% | 99.7% | 98.1% |
| Non-CoT (L=5) | 16.4% | 16.9% | 15.1% | 18.1% | 19.2% | 19.0% |
| Non-CoT (L=10) | 16.8% | 18.4% | 16.6% | 17.9% | 17.4% | 17.9% |
消融/对比¶
| 配置 | 关键发现 |
|---|---|
| CoT vs non-CoT | Non-CoT 模型甚至无法在训练分布内学会任务（接近随机猜测） |
| 循环群 vs 对称群 | 循环群可泛化到 10× 训练长度；对称群只能泛化到 ~2-3× |
| 无 PE vs 有 PE | 不使用 PE 是长度泛化的必要条件（标准 PE 偏向局部位置） |
关键发现¶
- CoT 是必要的：没有 CoT 的 Transformer 即使在训练分布内也接近随机猜测（~17%），说明状态追踪确实超出了 \(\mathsf{TC}^0\) 的能力
- 代数结构决定泛化：简单传递群（循环群）的泛化远优于对称群，与理论预测一致
- 递归自训练有效：对称群 \(S_5\) 设置下，每轮 doubling 训练长度后，模型确实能在新长度上达到高精度
- 注意力集中可验证：实验中观察到训练后 attention 确实集中在正确的子句上
亮点与洞察¶
- 首个 \(\mathsf{NC}^1\) 级别 CoT 优化保证：之前理论只能处理 parity 等 \(\mathsf{TC}^0\) 任务，本文一举扩展到 \(\mathsf{NC}^1\) 完备问题，这是理论复杂度上的本质性突破
- 注意力集中 = 长度泛化：给出了一个简洁优美的图景——长度泛化取决于 attention 的检索鲁棒性，而检索鲁棒性取决于学到的特征的"纯度"，后者又取决于群结构
- 无 PE 的理论支持：从理论角度解释了为什么去掉位置编码有助于长度泛化——标准 PE 的位置偏置会破坏 attention 的长距离检索能力
- 自训练超越训练数据：证明了自训练不只是"重复学会的东西"，而是能真正扩展模型的能力边界
局限性 / 可改进方向¶
- 任务特殊性：只在 LEGO 状态追踪上证明，虽然是标准理论基准但距离自然语言推理仍有距离
- 一层 Transformer：分析限于单层，尽管作者论证这已是当前理论前沿，但多层设置的结论可能不同
- 无位置编码：实际 LLM 都使用 PE（RoPE 等），去掉 PE 的理论结论能否指导实践存疑
- 课程学习假设：需要先学 MLP 再学 attention 的两阶段课程，与实际端到端训练不同
- 对称群限制：Theorem 4.2 要求 \(|Y| < \log\log d\)，对更一般的群结构未覆盖
相关工作与启发¶
- vs Merrill et al. (ICLR 2024): 他们证明了 CoT 的表达力可达 \(\mathsf{NC}^1\)（甚至 P），但未涉及优化；本文首次从可学性角度给出保证
- vs Kim & Suzuki (ICLR 2025): 他们证明 Transformer + CoT 能高效解决 parity（\(\mathsf{TC}^0\)），本文将复杂度扩展到 \(\mathsf{NC}^1\)
- vs Wen et al.: 他们从样本效率角度分析 CoT 对稀疏依赖问题的帮助；本文聚焦长度泛化，互为补充
评分¶
- 新颖性: ⭐⭐⭐⭐⭐ 首个 \(\mathsf{NC}^1\) 级别 CoT 可学性+长度泛化保证，注意力集中机制和递归自训练的理论化都是全新贡献
- 实验充分度: ⭐⭐⭐ 作为理论论文实验已足够验证理论预测，但只有合成任务，且非 CoT baseline 略显粗糙
- 写作质量: ⭐⭐⭐⭐ 论文清晰严谨，符号规范，但有 reviewer 指出图表可读性不佳
- 价值: ⭐⭐⭐⭐ 深化了对 CoT 推理本质的理论理解，"注意力集中↔长度泛化"的洞见对实践有启发意义