Tiny Recursive Reasoning with Mamba-2 Attention Hybrid

本文提出了一种将 Mamba-2 状态空间模型与递归推理框架相结合的混合架构，在保持参数规模相当的情况下，通过替换 Transformer 模块显著提升了模型在 ARC-AGI-1 任务上的候选解覆盖率和整体性能，证明了 SSM 算子适用于递归推理设计。

要点

这篇论文讲述了一个关于**“如何让小模型变得更聪明”的有趣故事。为了让你轻松理解，我们可以把人工智能（AI）想象成一个正在解谜题的“思考者”**。

1. 背景：小个子也能有大智慧？

以前，大家认为 AI 要变聪明，就得像盖摩天大楼一样，把模型做得巨大无比（参数越多越好）。但最近的研究发现，其实不需要那么大的“大脑”，只要给模型更多的“思考时间”，让它反复琢磨，小模型也能解决非常抽象的难题（比如 ARC-AGI 这种像智力测试一样的图形谜题）。

这就好比一个小侦探，虽然经验不如老侦探丰富，但如果给他足够的时间，让他把线索在脑海里反复推演、修正，他也能破案。

2. 核心问题：换一种“思考方式”行不行？

这篇论文的作者（来自 Intercom 的 Wenlong Wang 和 Fergal Reid）提出了一个大胆的想法：

• 原来的做法：小侦探（Tiny Recursive Model, TRM）在思考时，使用的是**“注意力机制”（Transformer）。这就像侦探在脑海里把线索“全盘扫描”**，同时看所有信息，互相联系。
• 新的尝试：他们想试试用**"Mamba-2"。Mamba-2 是一种新的技术，它的特点是“像流水一样”，信息是一个接一个流过的（状态空间模型）。这就像侦探“按顺序”**一步步检查线索，先看完第一个，再想第二个。

疑问是：如果把侦探的“全盘扫描”能力换成“按顺序检查”的能力，他还能保持那么聪明吗？甚至会不会更聪明？

3. 实验设计：给侦探换件“新衣服”

为了公平比较，作者做了一个非常精细的实验：

• 保持体重不变：他们确保新旧两个模型的大小（参数量）几乎一模一样（都是约 680 万参数），就像给两个侦探穿了一模一样重的衣服，只换里面的“思考肌肉”。
• 混合模式：他们没有完全抛弃“全盘扫描”，而是做了一个混合体（Hybrid）。这个新模型既有 Mamba-2 的“按顺序流”能力，又保留了部分“注意力”能力，确保它不会漏掉关键信息。

4. 实验结果：新侦探的表现如何？

他们在三个不同的“考场”进行了测试：

• 考场 A：ARC-AGI（抽象推理，最难）

  • 结果：新模型（Mamba-2 混合版）表现更好！
  • 具体表现：在官方评分标准（Pass@2，即给两次机会猜对）上，它比旧模型高了 2%。
  • 更惊人的是：如果给更多次机会（比如 100 次），新模型的优势会扩大到 4.75%。
  • 这意味着什么？ 旧模型虽然也能猜对，但新模型**“猜对的方案更多样”**。它就像是一个更博学的侦探，能提出更多种可能的解法，其中包含正确答案的概率更高。

• 考场 B：数独（Sudoku，规则严格）

  • 结果：旧模型（纯注意力）表现更好。
  • 原因：数独需要瞬间看清整个网格的关联，这时候“全盘扫描”比“按顺序流”更有效。

• 考场 C：迷宫（Maze，路径寻找）

  • 结果：新模型（Mamba-2 混合版）大获全胜，旧模型甚至完全失败。
  • 原因：迷宫需要一步步推演路径，Mamba-2 的“按顺序流”特性在这里简直是天作之合。

5. 核心发现：多样性 vs. 精准度

论文发现了一个有趣的**“权衡”**现象：

• 旧模型（TRM-attn）：像一个**“果断的决策者”**。它很自信，一旦觉得某个答案对，就把所有票数投给它。所以它第一次猜对的概率（Pass@1）很高，但备选方案比较单一。
• 新模型（TR-mamba2attn）：像一个**“博采众长的探索者”。它生成的答案更多样化**（Candidate Coverage 更高）。虽然它第一次猜对的概率和旧模型差不多，但它**“口袋里装的正确答案更多”**。当你给它更多次机会去挑时，它总能翻出正确答案。

比喻：
想象你要找一把钥匙开锁。

• 旧模型手里只有一把钥匙，它非常确定这把是对的，直接插进去（Pass@1 高）。
• 新模型手里有一串钥匙，它不确定哪把是对的，但它把所有可能的钥匙都试了一遍。虽然它第一次插进去可能没开，但只要给它多试几次（Pass@2 或 Pass@100），它肯定能打开那把锁。

6. 总结与意义

这篇论文告诉我们：

1. Mamba-2 可以进入“递归推理”的领域：以前大家觉得只有 Transformer（注意力机制）才能做这种反复思考的任务，现在证明 Mamba-2 这种新技术也能胜任，甚至在某些方面更强。
2. 思考方式决定了解题风格：不同的“思考肌肉”（混合策略）适合不同类型的难题。有的擅长瞬间全局判断，有的擅长一步步推演。
3. 未来的方向：既然 Mamba-2 本身就是一种“迭代更新”的过程，也许未来的 AI 可以把“思考”和“记忆”结合得更紧密，让模型在内部状态里直接完成推理，而不需要像现在这样一层层地“想”。

一句话总结：
作者给一个小侦探换了一种新的“思考肌肉”（Mamba-2 混合技术），发现它虽然思考方式变了，但不仅没变笨，反而在解决复杂谜题时，想出的“好点子”更多了，证明了这种新技术是未来构建更聪明 AI 的有力候选者。

技术摘要

这是一篇发表于 ICLR 2026 隐式与潜在思维研讨会 (Latent & Implicit Thinking Workshop) 的论文，标题为 《Tiny Recursive Reasoning with Mamba-2 Attention Hybrid》（基于 Mamba-2 注意力混合架构的微型递归推理）。

以下是对该论文的详细技术总结：

1. 研究背景与问题 (Problem)

• 递归推理的兴起： 现有的大型语言模型（LLM）通常通过思维链（Chain-of-Thought, CoT）在 Token 空间进行显式推理，但这会消耗大量计算资源且可能产生与推理无关的冗余 Token。相比之下，潜在递归（Latent Recursion）（如 Tiny Recursive Model, TRM）通过在隐藏表示空间进行迭代 refinement（细化），不生成中间 Token，仅用极小的参数量（如 7M）就能在抽象推理任务（如 ARC-AGI）上取得优异表现。
• 核心问题： 现有的 TRM 架构完全依赖 Transformer 块（注意力机制） 作为每一步的更新算子。然而，Mamba-2 作为一种状态空间模型（SSM），其核心机制本身就是一种基于状态更新的迭代过程（$h_t = a_t h_{t-1} + B_t x_t$）。
• 研究动机： 将 Mamba-2 引入递归推理框架是否可行？Mamba-2 的线性复杂度和内在的递归特性是否能替代或增强 Transformer 块，从而在保持推理能力的同时提升效率或性能？目前尚不清楚 SSM 算子是否适合嵌入递归推理的“外层循环”中。

2. 方法论 (Methodology)

作者提出了一种新的 TRM 变体，用 Mamba-2 + Attention 混合算子 替换了原始 TRM 中的纯 Transformer 块，同时严格保持参数量一致。

• 架构设计 (TR-mamba2attn)：
  • 保留了 TRM 的双层潜在状态结构（$z_H$ 高层状态和 $z_L$ 低层状态）以及递归调度（3 个外层循环，4-6 个内层循环）。
  • 核心替换： 将更新函数 $f$ 中的 Transformer 块替换为混合块堆栈。具体顺序为：Mamba-2 $\rightarrow$ Mamba-2 $\rightarrow$ Attention $\rightarrow$ MLP。
  • 混合动机： Mamba-2 擅长单向序列依赖，但 ARC-AGI 等任务需要双向空间依赖。因此，引入 Attention 块（或 MLP-t）来提供必要的跨位置信息流，弥补纯 Mamba 的因果性限制。
  • 备选方案： 实验了另一种变体 TR-mamba2mlpt，用 MLP-t（转置维度上的全连接）替代 Attention，但在大网格任务上表现不佳。
• 参数匹配 (Parameter Matching)：
  • 为了公平比较，混合模型与原始 TRM 的参数量几乎完全一致（6.86M vs 6.83M）。
  • 隐藏层维度设为 512，Mamba-2 的具体参数为 $d_{state}=128, head_{dim}=64, expand=2$。
• 归一化策略 (Post-Norm)：
  • 论文强调在递归展开中必须使用 Post-Norm（残差连接后归一化），而非现代 LLM 常用的 Pre-Norm。
  • 原因： 在重复应用同一残差模块时，Pre-Norm 可能导致残差流幅度随递归深度 $t$ 增长（$\approx \sqrt{t}$），引发数值不稳定（NaN）。Post-Norm 能限制隐藏状态尺度，确保递归稳定性。

3. 实验设置 (Experiments)

• 基准测试：
  • ARC-AGI-1： 抽象推理核心基准，使用官方指标 Pass@K（K=1, 2, 5, 10, 100, 1000）。输入经过二面体变换和颜色置换增强（每个输入约 880 个增强样本）。
  • Sudoku-Extreme： 9x9 数独，测试约束满足能力。
  • Maze-30x30-Hard： 30x30 迷宫，测试大规模空间推理。
• 对比模型：
  • TRM-attn：原始基于 Attention 的 TRM。
  • TR-mamba2attn：本文提出的 Mamba-2 + Attention 混合模型。
  • TRM-mlp-t / TR-mamba2mlpt：基于 MLP 的变体（作为对照）。

4. 关键结果 (Key Results)

A. ARC-AGI-1 性能提升

• Pass@2 (官方指标)： 混合模型达到 45.88%，比原始 TRM (43.88%) 提升 +2.0%。
• Pass@K 趋势： 随着 K 值增加，混合模型的优势扩大。在 Pass@100 时提升 +4.75%，在 Pass@1000 时提升 +4.25%。
• Pass@1 (Top-1 选择)： 两者基本持平（混合模型仅低 0.25%），说明混合模型在保持最佳选择质量的同时，显著扩大了候选解的覆盖范围。

B. 其他任务表现

• Sudoku (数独)： 纯 MLP 变体表现最好（87.4%），混合模型略低于纯 MLP 但优于 Attention 模型。这表明小网格约束满足任务受益于密集的全连接通信。
• Maze (迷宫)： 纯 MLP 变体在 30x30 大网格上完全失败（0%），而混合模型达到 80.6%（原始 Attention 模型为 60.8%）。这证明了在大规模空间推理中，Mamba-2 的序列处理能力结合 Attention 的跨位置混合至关重要。

C. 深入分析：覆盖度与选择性的权衡 (Coverage vs. Selection Trade-off)

• 候选多样性： 混合模型生成的唯一候选解数量更多（每道题 339.5 vs 266.6），且投票熵（Vote Entropy）更高。这意味着 Mamba-2 的序列处理引入了不同的解路径，增加了候选池的多样性。
• 决策集中度： 原始 TRM 的投票更集中在 Top-1 候选上（41.1% vs 32.9%），且 Top-1 的置信度差距（Margin）更大。
• 结论： 混合模型通过更广泛的探索（Coverage）提高了 Pass@K，而原始模型通过更果断的选择（Selection）维持了 Pass@1。两者在难度分层上表现出互补性：混合模型在“困难”样本上优势明显，而原始模型在“简单”样本上略优。

5. 主要贡献 (Contributions)

1. 首个 Mamba 混合递归模型： 首次将 Mamba-2 引入潜在递归推理框架，验证了 SSM 算子作为递归更新算子的可行性。
2. 性能验证： 在 ARC-AGI-1 上实现了官方指标 Pass@2 的显著提升（+2.0%），并在高 K 值下表现出更强的候选覆盖能力。
3. 机制洞察： 揭示了递归推理中“候选覆盖度”与“最佳选择”之间的权衡。Mamba-2 的引入并未牺牲 Top-1 质量，而是通过生成更多样化的解路径提升了整体成功率。
4. 架构建议： 强调了在递归架构中使用 Post-Norm 对稳定性的关键作用，并指出了不同任务（如数独 vs 迷宫）对混合算子（Attention vs MLP）的不同需求。

6. 意义与展望 (Significance)

• 理论意义： 证明了递归推理的核心在于“思考时间”（迭代过程）而非单纯的模型规模或特定的注意力机制。SSM 的内在递归性与外部递归循环的结合是有效的。
• 实践价值： 为设计更高效、更小的推理模型提供了新方向。Mamba-2 的线性复杂度可能为未来的测试时计算扩展（Test-time Compute Scaling）提供比纯 Transformer 更优的基底。
• 未来方向： 论文提出下一步可以探索将递归循环内化到 SSM 的状态更新中（即利用 Mamba 的内在递归 $h_t$ 来模拟外层循环），而不是仅仅依赖外层的迭代调用，这可能进一步释放 SSM 在递归推理中的潜力。

总结： 该论文成功地将 Mamba-2 集成到微型递归推理模型中，不仅保持了原有的推理能力，还通过增加候选解的多样性显著提升了在抽象推理基准上的表现，确立了 SSM 混合算子在递归推理设计空间中的重要地位。