像 GPT-4 和 Gemini 这样的大型语言模型 (LLMs) 是计算领域的巨擘,能够编写代码、创作诗歌并回答各种各样的问题。但尽管它们功能强大,却有一个致命弱点: 复杂的多步推理谜题。像解决棘手的数独或破译 ARC-AGI 基准测试中的抽象模式这样的任务,即使是最先进的 LLM 也可能出错。它们自回归、逐个 token 的生成过程意味着一个错误就可能导致整个解答脱轨,并且没有简单的方法可以回溯和修正。
研究人员已经开发出像 思维链 (CoT) 提示 这样的技术,来引导 LLMs 进行循序渐进的“思考”,这有所帮助但并未解决核心问题。那么,如果我们不去构建越来越大的模型,而是设计出更小、更高效、并且 擅长 这种迭代推理的系统,会怎么样呢?
最近的一篇论文《少即是多: 使用微型网络进行递归推理》正是探讨了这个问题。作者们引入了 微型递归模型 (TRM) ,这是一个异常小而简单的模型,在那些困扰大型 LLM 的谜题上取得了惊人的性能。TRM 的参数量少至 700 万——不到 GPT-3 等模型规模的 0.01%——却在数独、迷宫和 ARC-AGI 等基准测试中创造了新的最先进记录。
本文将深入探讨 TRM 的工作原理。我们首先会了解它的前身——分层推理模型 (HRM) ,以理解其构建的基础。然后,我们将剖析使 TRM 如此高效的精妙简化设计。最后,我们将审视那些令人惊叹的结果,它们证明了有时候,少即是多。
背景: 分层推理 (HRM) 的前景与复杂性
TRM 并非凭空出现——它直接演进自一个名为 分层推理模型 (HRM) 的模型。HRM 是一种新颖的方法,它通过使用两个相互递归调用的小型神经网络来优化解答,展现出了巨大潜力。其设计灵感来源于关于大脑如何处理不同频率信息的复杂生物学论证。
让我们来分解一下 HRM 的关键组成部分:
递归分层推理
HRM 使用两个网络: 一个 低层 网络 \(f_L\) 和一个 高层 网络 \(f_H\)。\(f_L\) 以高频率递归处理细粒度细节,而 \(f_H\) 则以较低频率递归来整合信息。这两个网络在两个潜在特征向量上操作: \(z_L\) 和 \(z_H\)。深度监督
HRM 并非训练模型一次性得到正确答案,而是采用了一个迭代过程。在多达 16 个 监督步骤 中,模型将其先前的输出和潜在特征作为输入,并尝试改进它们。这模拟了一个非常深的网络,却没有单次前向传播所带来的巨大内存开销。单步梯度近似
\[ z_L^* \approx f_L\left(z_L^* + z_H + x\right) \]
在 HRM 中,一个完整的前向传播涉及多次递归调用。对所有这些调用进行反向传播的计算成本会非常高昂。为了解决这个问题,HRM 的作者使用了一个聪明 (但可能存在缺陷) 的捷径:** 隐函数定理 (IFT)** 。在特定条件下,该定理允许你仅通过对最后一步进行反向传播来近似梯度。这依赖于一个假设,即递归会收敛到一个不动点:
\[ z_H^* \approx f_H\left(z_L + z_H^*\right) \]自适应计算时间 (ACT)
为了使训练更高效,HRM 使用一个 Q-learning 机制来决定一个解答何时“足够好”,从而模型可以停止对特定训练样本的迭代,避免在每个样本上都花费全部 16 个步骤。
HRM 是一项突破,它在其他模型难以解决的谜题上取得了高准确率。但它也很复杂——依赖于不确定的生物类比和不保证成立的不动点假设——这使得它难以理解和改进。而这正是 TRM 登场的契机。
核心方法: 剖析微型递归模型 (TRM)
TRM 的创造者们仔细审视了 HRM 并提出了一个问题: 我们能否通过剥离其复杂性来达到相同或更好的结果? 答案是响亮的 是。TRM 是简化的典范——它所做的改变在减小模型尺寸的同时,显著提升了性能。
TRM 的整体架构如下图所示。
图 1. TRM 在输入问题
x
的引导下,通过迭代其潜在推理状态z
来递归地改进其预测答案y
。
1. 抛弃不动点定理,采用完全反向传播
TRM 放弃了 HRM 的单步梯度捷径。不动点假设在仅仅几次递归后不太可能被满足,因此 TRM 定义了一个 完整的递归过程,并对整个过程进行反向传播:
一个完整的递归过程包括对推理向量的 n
次更新和对答案向量的一次更新:
…
为了保持训练效率,TRM 执行 T-1
次这样的递归过程而不计算梯度,用它们来优化潜在状态,然后再进行最后一次 带 梯度的递归用于学习更新。
这一改变产生了巨大影响: 在消融研究中,从单步梯度切换到完全反向传播,将 Sudoku-Extreme 的准确率从 56.5% 提升到 87.4%。
表 1. 设计选择对 TRM 最终性能的贡献。
2. 对潜在特征更简单、更直观的理解
HRM 的两个潜在特征 \(z_L\) 和 \(z_H\) 是通过生物学的“层次结构”来解释的。TRM 提供了一个更简单的解释:
- \(z_H\) →
y
— 当前嵌入的答案。 - \(z_L\) →
z
— 潜在的推理或“草稿纸”。
为了优化一个解答,模型需要三样东西: 原始问题 (x
) 、先前的答案 (y
) ,以及导致该答案的推理链 (z
) 。忽略其中任何一个都会削弱其改进解答的能力。
一个数独可视化 (图 6) 清楚地说明了这一点: 解码 y
会得到一个几乎正确的网格,而解码 z
则得到一个难以辨认的数字图——这证明了 z
确实是潜在的推理。
图 6a. 一个 Sudoku-Extreme 谜题的输入
x
。
图 6b. Token 化后的
z_H
直接对应于预测的解答。
从经验上看,双特征 (y
和 z
) 设计的性能优于单特征和多特征的变体。
表 2. 两个独立的特征带来了最高的准确率。
3. 一个网络统一全局
HRM 使用了两个网络,使参数数量加倍。TRM 将 \(f_L\) 和 \(f_H\) 统一到一个单一网络中,该网络学习两种任务,通过输入 (是否存在 x
) 来区分。这不仅将参数减半,还将数独的准确率从 82.4% 提高到了 87.4%。
4. “少即是多”: 微型更优
将网络从 2 层扩展到 4 层反而降低了准确率——更大的模型在数据有限的情况下会过拟合。TRM 小巧的 2 层网络通过递归展开,在推理过程中达到了必要的深度,同时保持了泛化能力。
5. 其他巧妙的改进
- 更简单的 ACT: 用单个二元交叉熵损失取代了 Q-learning 的停止机制,从而无需进行第二次前向传播。
- 对小输入使用无注意力机制: 对于 9x9 的数独,用 MLP-Mixer 替换注意力机制提升了性能。更大的网格仍然能从注意力机制中受益。
- 权重的指数移动平均 (EMA) : 平滑训练过程,防止在小数据集上崩溃。
精简的 TRM 伪代码反映了其高效性:
图 3. TRM 的循环结构: 在每个监督步骤的梯度更新之前进行递归更新。
实验与结果: 微型网络,巨大性能
这些模型在以下任务上进行了测试:
- Sudoku-Extreme: 9x9 数独,仅有 1000 个训练谜题。
- Maze-Hard: 在复杂的 30x30 迷宫中寻找路径。
- ARC-AGI-1/2: 抽象几何推理谜题。
数独与迷宫谜题
TRM 不仅仅是略胜 HRM 一筹——它简直是 完胜。
表 4. TRM 以微小的参数量在谜题基准测试中占据主导地位。
多达数十亿参数的 LLM 得分为 0.0%。HRM 在数独上达到 55.0%,在迷宫上达到 74.5%。而 TRM 的 MLP 变体在数独上达到了 87.4%,基于注意力的变体在迷宫上得分 85.3%——参数量仅为 500-700 万。
ARC-AGI 挑战
ARC-AGI 被认为是一项重大的推理挑战。TRM 再次大放异彩。
表 5. TRM-Att 超越了 HRM 和几个强大的 LLM。
700 万参数的 TRM-Att 在 ARC-1 上得分 44.6%,在 ARC-2 上得分 7.8%,击败了 HRM 的 40.3%/5.0%,超越了 Gemini 2.5 Pro,并挑战了规模远大于它的定制模型。
即使在“有效深度” (层数 × 递归次数) 相同的情况下,TRM 也比 HRM 提供了更高的准确率:
表 3. 在相似的计算深度下,TRM 的设计被证明更优越。
结论: AI 推理的新路径?
微型递归模型 (TRM) 有力地证明了更大并非总是更好。通过将一个复杂的模型 (HRM) 进行毫不留情的简化,研究人员创造出了一个远比之前更优雅、更高效、也更有效的模型。
核心要点:
- 直接优于捷径: 完全反向传播的性能优于不动点梯度近似。
- 直觉优于教条: 简单的“答案 + 推理”设计同时提升了清晰度和结果。
- 少即是多: 微型递归网络避免了过拟合,并通过递归利用了计算深度。
TRM 的成功揭示了一条构建先进 AI 推理系统的激动人心的替代路径。我们不必扩展通用巨型模型,而是可以构建小型的专用模型,通过迭代优化它们的解答。这种方法不仅参数高效,而且对于解决那些构成真正智能基础的逻辑谜题可能至关重要。
征程远未结束,但 TRM 已经表明,有时候,最深刻的解答来自最微小的封装。