Hierarchical Reasoning Model

关键思想：

HRM在单次前向传递中执行顺序推理任务，无需对中间过程进行显式监督
- 两个相互依赖的循环模块 (two interdependent recurrent modules)
  - 高层模块 (high-level module) 负责缓慢、抽象的规划
  - 低层模块 (low-level module) 处理快速、详细的计算

背景知识

Transformer模型本身由一个有限（常数）数量的层组成，每层包含多个注意力头（Attention Heads）和一个前馈网络（Feed-Forward Network）。这个层的数量（例如12层）是固定的，它不随输入句子的长度$n$变化。论文《Average-Hard Attention Transformers are Constant-Depth Uniform Threshold Circuits》的核心在于证明：Transformer模型中的每一个操作，即使是复杂的注意力机制，都可以被$TC^0$电路家族模拟。

$TC^0$电路无法解决需要 多项式深度（Polynomial Depth） 的复杂顺序推理问题，例如许多需要复杂逻辑、计数和无限长推理链的问题。
这就是为什么有人认为Transformer——尽管它在实践中非常强大，但在理论上，它的计算深度被限制在了一个非常低且不变的复杂度类$TC^0$中，这给它最被期望的能力——复杂推理——设置了一个基础的理论上限。

更进一步 $TC^0$ 局限性在于：多项式时间 (Polynomial Time, $P$) 是指解决一个问题所需要的计算步骤随着输入规模 $n$ 以多项式 (如 $n^2$, $n^3$)增长。很多我们认为需要 “算法” 来解决的问题都在这个类别中，例如排序、最短路径等。而 $TC^0$ 的能力严格弱于$P$ 。这意味着，任何需要真正多项式时间算法才能解决的问题，Transformer从理论上都无法完美解决

另外一点：LLM不具有图灵完备性（Turing-completeness）。图灵完备性是指一个系统（如编程语言或者计算机）能够模拟任何图灵机，也就是执行任何算法。图灵完备的系统必须具有循环、条件判断或者无限内存等能力。Transformer时一个固定的前馈网络，接受输入，经过固定的层数计算，然后输出结果。它无法根据输入动态地决定“再循环计算次数”或者“创建一个新的变量来存储中间结果”。因此LLM无法执行通用的算法。这就意味着对于需要严谨算法步骤的任务，比如深思熟虑的规划（如制定一个多步骤，有依赖关系的复杂计划）或者符号操作（如解代数方程），LLM无法像传统计算机程序那样可靠地执行。它智能通过模式匹配来“模仿”这些过程，不是真正的“执行”

数独（Sudoku）作为例证: 数独是一个绝佳的测试案例。解决数独需要严格的逻辑推理和约束满足，这是一个典型的算法问题。它不是一个模糊的语言模式匹配任务。
- 实验结果: 研究发现，即使把Transformer的层数（物理深度）增加很多，它在解决数独问题上的表现虽然有所提升，但远未达到完美，并且会很快遇到性能瓶颈。
对“缩放范式”的挑战: 这个结果直接挑战了“只要模型越大、数据越多，能力就越强”（即缩放范式，Scaling Paradigm）的信念。它表明，对于某些类型的推理问题，简单地“缩放”模型规模可能无法克服其底层的架构限制。模型的性能提升可能只是因为它记住了更多解题模式，而不是学会了通用的解题算法。

为了提高LLM的这种解决复杂问题的能力，大多数都依靠 CoT（Chain-of-Thought）。

CoT主要是通过外化推理（externalizes reasoning）– 也就是说土里并不是在模型内部一个抽象的逻辑空间发生的。而是模型把“思考”的过程写出来，变成文字。
具体而言：CoT通过将复杂任务分解为更简单的中间步骤，并使用一个浅层模型（也就是Transformer-based LLM）顺序地生成文本，从而将推理过程外化为词元级别的语言。然而，用于推理的CoT是一个‘拐杖’，而非一个令人满意的解决方案。它依赖于脆弱的、由人类定义的分解方式，其中任何一个步骤的错误或顺序错乱都可能让整个推理过程完全脱轨。这种对明确语言步骤的依赖，将推理束缚在了词元级别的模式上。其结果是，CoT推理通常需要大量的训练数据，并在处理复杂推理任务时生成海量词元，导致响应时间缓慢。我们需要一种更高效的方法来最小化这些数据需求。

Hierarchical Reasoning Model

那么Hierarchical Reasoning Model是如何做的？HRM模型本质在于探索其内在的抽象逻辑空间(latent reasoning)。换句话说：HRM的本质在于在内部的隐藏状态空间内完成推理或者逻辑的计算。这种模式的优势在于其真正地模拟了人在用语言去沟通交流的模式。

HRM

HRM由四个可学习的模块组成:

输入网络(input network): $f_{I}(\cdot, \theta_{I})$
浅层循环网络(low-level recurrent module): $f_{L}(\cdot, \theta_{L})$
深层循环网络(high-level recurrent module): $f_{H}(\cdot, \theta_{H})$
输出网络(output network): $f_{O}(\cdot, \theta_{O})$

此时，我们假设一个前馈过程的总步骤数为 $i=1, \cdots, N \times T$

这里的 $N\times T$ 其含义是：模型会动态的折叠为 $N$ 个深层循环，而每个深层循环对应一个$T$ 步骤长度的浅层循环
在这个过程中:浅层循环网络 $f_{L}(\cdot, \theta_{L})$ 会保持一个隐层状态 ($z_{L}^{i}$)。与之对应的，深层循环网络 $f_{H}(\cdot, \theta_{H})$ 也会维护一个隐层状态($z_{H}^{i}$)
- 基于此，那么初始的隐层状态就分别为 $z_{L}^{0}$ 和 $z_{H}^{0}$

HRM的通俗理解：CEO与工程师的“隐式推理”

HRM模型提出模型不应该“大声思考”（就是CoT），而应该在内部完成推理，即在它的隐藏状态中进行思考。这被称为 “隐式推理” 。受人脑分层处理信息的启发，HRM采用了两个相互协作的模块：

高层模块（“CEO”或“战略家”，也就是high-level recurrent module）：这个模块负责制定战略。它运行缓慢，进行抽象思考，并为解决问题设定总体规划。
底层模块（“工程师”或“执行者”，也就是high-level recurrent module）：这个模块负责执行。它运行速度极快，接收CEO的当前计划，并进行快速、细致的计算来完成任务。这种结构让模型可以在一次前向传播中完成“思考”，而无需生成任何中间文本。下面是它解决一个任务的流程：

接收任务：输入（例如一个数独棋盘）被送入模型。
CEO制定战略：CEO模块创建一个初始的宏观计划。
工程师开始工作：在固定的步数（T步）内，工程师模块快速更新其状态，试图执行CEO的计划，解决谜题的局部。这是一个高强度、专注的计算过程。
提交进度报告：在这轮“爆发式”工作后，工程师将其进展汇报给CEO。
战略复盘：CEO根据新的信息评估进展，优化其宏观战略，并下达新的指令。
循环往复：这种“慢速战略规划”与“快速细节执行”的循环会重复固定的次数（ $N$ 次）。
输出最终答案：所有循环结束后，CEO的最终战略状态被用来一次性生成完整的解决方案。

HRM的内部网络结构

Step 1: 输入$x$会先经过一个映射的过程，论文中为 working representation：

$$ \tilde{x}=f_I\left(x ; \theta_I\right) $$

Step 2: 在每一个内部时间步$i$ (从1 到$N \times T$)，两个模块会更新他们的隐层状态($z_L$可以看作是工程师, $z_H$可以看作是CEO)

工程师($z_L$)在每一步都会更新。它会利用自己的前一步状态$z_L^{i-1}$，CEO当前的战略$z_H^{i-1}$（注意CEO的当前战略是本循环中是保持不变的），以及输入信息

$$ z_L^i=f_L\left(z_L^{i-1}, z_H^{i-1}, \tilde{x} ; \theta_L\right) $$

CEO（$z_H$）只在工程师完成一整个$T$步的计算周期后，才更新一次。在周期内战略会保持不变，以指导工程师的工作。简而言之：只有在步骤数量达到$T$之后才会更新，否则就是保持和前一个状态一致

$$ z_H^i= \begin{cases}f_H\left(z_H^{i-1}, z_L^{i-1} ; \theta_H\right) & \text { if } i \equiv 0(\bmod T), \\ z_H^{i-1} & \text { otherwise } .\end{cases} $$

Step 3: 在经过 $N$个完整的循环后（也就是$N\times T$个步骤后），预测的$\hat y$最终由CEO的隐藏状态$z_H^{N T}$决定

$$ \hat{y}=f_O\left(z_H^{N T} ; \theta_O\right) $$

HRM的优化策略

理解HRM的优化策略，将其分为了四步：

第一步

首选训练循环神经网络RNN，通用的方法叫做BPTT（backpropagation through time）。简单理解就是假设RNN为了解决一个问题，内部思考了1000步，BPTT为了更新RNN，则需要记录1000步的隐藏状态。在进行backpropagation的时候，需要从最后的结果去一步一步到这追溯回去。

弊端：需要巨大的显存

第二步

研究 Deep Equilibrium Model (DEQ)提供的思路

核心思想：如果一个循环系统最终会达到一个稳定状态（均衡点或定点），那么我们其实不需要关心它到达那里的过程。我们只需要分析那个最终的稳定状态本身。
好比：一滴墨水滴入清水中。要预测最终均匀混合的状态，你不需要模拟每一个墨水分子在水中的随机碰撞路径。你可以直接用描述最终均衡状态的物理定律来分析。
数学工具：隐函数定理（Implicit Function Theorem, IFT) 就是这个“物理定律”。它允许我们直接计算出在那个“均衡点”上，模型参数应该如何更新，而完全不需要进行时间上的反向追溯。

DEQ的数学背景

我们已经知道通过BPTT有 $O(T)$ 的内存成本，其中 $T$ 是时间步长。这阻碍了我们训练能够进行深度推理( $T$ 非常大)的模型。我们的目标是找到一种内存成本为 $O(1)$ 的训练方法。

DEQ的核心思想是不动点(Fixed Point): 一个良好性能的循环神经网络，在接受输入 $x$ 后，其隐藏状态 $z$ 经过多次迭代，最终会收敛到一个均衡状态。也就是fixed point，记作 $z^{*}$
用数学公式表示不动点为 $z^{*} =f(z^{*}, x;\theta)$，其中$f$是RNN的一个网络层，$x$是输入，$\theta$是模型参数
所以DEQ的核心观点: 既然网络最终会收敛到 $z^{*}$ 那么我们只需要关心这个最终的均衡状态 $z^{*}$，而不需要关心它是如何一步一步到达这里。如果我们能直接找到 $z^{*}$，并计算损失函数对$z^*$ 的梯度，然后将这个梯度直接传递给 $\theta$，那么就可以完全绕过BPTT

那么如何计算最终的 $z^*$ 或者如何计算 $\frac{\partial z^*}{\partial \theta}$ ?

推理过程就需要隐函数定理：首先，从不动点的定义开始

$$ z^* - f(z^*, x; \theta) = 0 $$

这是一个关于 $z^*$ 和 $\theta$ 的隐式方程。我们想要 $z^*$ 对 $\theta$ 的导数，可获得

$$ \frac{\partial z^*}{\partial \theta} - \left( \frac{\partial f}{\partial z^*} \frac{\partial z^*}{\partial \theta} + \frac{\partial f}{\partial \theta} \right) = 0 $$

注意

$\frac{\partial f}{\partial \theta}$ 是函数 $f$ 对 $\theta$ 的梯度
$\frac{\partial f}{\partial z^*}$ 是函数 $f$ 对$z^*$ 的雅可比矩阵（Jacobian Matrix），这里记作$J_f$

最后可的

$$ \frac{\partial z^*}{\partial \theta} - J_f \frac{\partial z^*}{\partial \theta} = \frac{\partial f}{\partial \theta} $$

$$ (I - J_f) \frac{\partial z^*}{\partial \theta} = \frac{\partial f}{\partial \theta} $$

最后，我们得到了IFT的最终结果，也就是计算理想梯度的公式：

$$ \frac{\partial z^*}{\partial \theta} = (I - J_f)^{-1} \frac{\partial f}{\partial \theta} $$

这个公式有什么意义呢？直观理解，在不动点$z^*$，我们可以通过计算一个雅可比矩阵的逆，然后乘一个局部梯度，来直接得到$z^*$对$\theta$的精确梯度。从而完全避免BPTT

这个公式很完美，但是如何计算$(I - J_f)^{-1}$是一个巨大的问题：

需要计算雅可比矩阵$J_f$
对$(I-J_f)$求逆这两步的计算成本很高，尤其是对高纬度的隐藏状态

有没有什么近似的办法呢？数学上的诺伊曼级数，告诉我们如果矩阵$A$的谱半径小于1，那么

$$ (I - A)^{-1} = I + A + A^2 + A^3 + \dots $$

矩阵的谱半径

矩阵都有对应的特征值对吧。那么谱半径就是一个矩阵所有特征值的绝对值中，最大的那一个。

此时将$A$替换为雅可比矩阵$J_f$，可得

$$ (I - J_f)^{-1} = I + J_f + J_f^2 + J_f^3 + \dots $$

这个公式提供了一个近似的方法：我们可以只取这个无穷级数的前几项作为近似。例如$(I - J_f)^{-1} \approx I + J_f$

第三步

HRM的设计受到了DEQ的启发，但它采取了一种极其大胆且务实的简化。

核心简化：HRM说：“那个基于IFT的复杂矩阵求逆过程太慢了，并且大胆的假设近似。干脆假设那个复杂的矩阵就是单位矩阵 I。” 这在数学上意味着 $(I - J_F)⁻¹ ≈ I$。
理想梯度的大胆近似：一步梯度近似（1-step gradient） $$ \frac{\partial z^*}{\partial \theta} = (I - J_f)^{-1} \frac{\partial f}{\partial \theta} \approx I \cdot \frac{\partial f}{\partial \theta} = \frac{\partial f}{\partial \theta} $$ 这个结果 $\frac{\partial z^*}{\partial \theta}\approx \frac{\partial f}{\partial \theta}$ 的直观解释是：在计算梯度时，我们假装不动点 $z^*$ 与参数 $\theta$ 之间没有复杂的历史依赖关系。我们把它当作一个普通的、固定的输入来对待。因此，$z^*$ 对 $\theta$ 的梯度，就近似等于更新函数 $f$ 本身对 $\theta$ 的梯度，就好像 $z^*$ 是一个常数一样。
这在实践中意味着什么？
这意味着在计算梯度时，我们完全忽略了所有跨越时间的长期依赖关系。梯度的“传播路径”被极大地缩短了。

让我们用一个“责任分配”的比喻来理解：

BPTT的责任分配：最终答案错了，BPTT会说：“这是因为第1000步的状态错了，而第1000步错了是因为第999步错了，而第999步错了又是因为……” 一直追溯到第一步。责任链条非常长。
HRM近似梯度的责任分配：最终答案错了，HRM会说：
1. 这个错，只怪罪于$H$模块的最后一个状态 $z_H$。
2. H模块最后一个状态 $z_H$ 的“锅”，只甩给L模块的最后一个状态 $z_L$。
3. L模块最后一个状态 $z_L$ 的“锅”，只甩给最初的输入嵌入 $\tilde{x}$。梯度的传播路径变成了一条极短的链条：输出 → $H$模块最终状态 → $L$模块最终状态 → 输入。

这就是论文中图4（Figure 4）伪代码的精髓：

在 with torch.no_grad(): 的代码块里，模型进行了大量的循环计算。no_grad() 明确告诉PyTorch：“别记录这些步骤的历史，我不需要BPTT！”
在循环结束之后，模型才在没有 no_grad() 的情况下重新计算了最后一步。这样，PyTorch只会为这最后一步建立一个极短的计算图，梯度也只能沿着这个短路径传播。

第四步

你可能会问：“这么粗糙的近似，能学到东西吗？” 答案是肯定的，并且好处巨大。

O(1) 恒定内存：因为只需要存储每个模块的最后一个状态，所以无论模型“思考”多少步（1000步还是10000步），内存占用都是一个固定的常数。这彻底解决了BPTT的内存瓶颈，让训练进行超深度推理的模型成为可能。
极高的计算效率：没有了时间展开和复杂的矩阵运算，训练过程变得飞快。
生物学上的合理性：这种“只根据近期活动更新”的机制，更接近我们对大脑局部学习规则的理解。
为什么它依然有效：虽然这个梯度是“近似”的，但它仍然为模型的学习提供了正确的“方向”。更重要的是，HRM将这个技术与 “深度监督”（Deep Supervision） 相结合。深度监督会在多个时间点为模型提供学习信号，这种频繁的、直接的反馈在很大程度上弥补了近似梯度损失掉的“精确性”，确保了模型能够稳定、有效地学习。

Appendix

诺伊曼级数 (Neumann Series) 是什么?

诺伊曼级数本质上是高中数学里等比数列求和公式的矩阵版本。

回忆一下等比数列，我们都学过，对于一个数 $r$，如果它的绝对值 $|r| < 1$，那么下面这个无穷级数的和是收敛的：

$$ 1 + r + r^2 + r^3 + \dots = \frac{1}{1-r} $$

例如，如果 $r = 0.5$，那么 $1 + 0.5 + 0.25 + … = 1 / (1 - 0.5) = 2$。如果 $|r| ≥ 1$，这个级数就会发散到无穷大，公式就不成立了。

推广到矩阵，诺伊曼级数将这个思想从单个数字推广到了方块矩阵 $A$。

数字$1$变成了单位矩阵$I$。
除法 $1/x$ 变成了矩阵求逆 $X⁻¹$。

于是，公式就变成了：

$$ (I - A)^{-1} = I + A + A^2 + A^3 + \dots $$

这个公式的直观意义是： 一个矩阵的逆 $(I - A)⁻¹$ 可以被一个无穷的矩阵幂次之和来近似。

就像等比数列有收敛条件 $|r| < 1$ 一样，诺伊曼级数也有一个收敛条件。这个条件就是由“谱半径”来定义的。

谱半径 (Spectral Radius) 是什么？

谱半径是衡量一个矩阵“威力”或“放大能力”的核心指标。

什么是特征值 (Eigenvalue)？对于一个矩阵 $A$，它的一些“特殊向量”（称为特征向量 $v$）在经过 $A$ 的线性变换后，方向保持不变，只会被拉伸或压缩。这个拉伸/压缩的比例就是特征值 $\lambda$。
$$ Av = \lambda v $$
一个矩阵通常有多个特征值。
谱半径的定义 谱半径，记为 $\rho(A)$，就是一个矩阵所有特征值的绝对值中，最大的那一个。
$$ \rho(A) = \max_i |\lambda_i| $$
谱半径的直观意义：系统的稳定性指标想象一个简单的动态系统，它的状态在每一步都通过乘以矩阵 $A$ 来更新：$z_t = A * z_{t-1}$。
- 如果 $\rho(A) < 1$：这意味着矩阵 $A$ 在其“最强”的方向上也是一个“收缩”变换。无论初始状态是什么，经过多次迭代后，状态向量 $z$ 都会衰减至零。系统是稳定的。
- 如果 $\rho(A) > 1$：这意味着矩阵$A$ 在其“最强”的方向上是一个“扩张”变换。一个微小的初始信号会经过多次迭代后被放大到无穷。系统是不稳定的，会发生“爆炸”。
- 如果 $\rho(A) = 1$：系统处于临界状态，可能保持稳定也可能发散。

结论： 诺伊曼级数 $(I - A)⁻¹ = I + A + A² + …$的收敛条件，就是矩阵 $A$ 的谱半径必须小于1 ($\rho(A) < 1$)。这和等比数列的条件 $|r| < 1$ 是完美对应的。

为什么可以假设 $J_f$ 的谱半径小于1?

这是整个近似方法能够成立的关键前提。在HRM的上下文中，这个假设 $\rho(J_f) < 1$ 并非凭空而来，而是基于以下几个强有力的理由：

理由一：这是一个设计目标，也是一个成功训练的结果

DEQ和HRM的前提：整个模型框架的基本假设就是网络在处理输入后，其隐藏状态会收敛到一个稳定的不动点 $z^*$。
稳定性的必然要求：根据我们对谱半径的理解，一个动态系统要能够收敛，其状态转移的雅可比矩阵 $J_f$ 在不动点附近的谱半径必须小于或等于1。如果 $\rho(J_f) > 1$，系统就是不稳定的，任何微小的扰动都会被放大，隐藏状态会发散而不是收敛。
因此，一个能够成功解决问题并收敛的HRM模型，其内部动态必然满足这个条件。我们不是在凭空假设，而是在描述一个成功模型应该具备的内在属性。

理由二：训练过程本身会隐式地推动谱半径减小

梯度爆炸与损失函数：如果 $\rho(J_f) > 1$，网络就是不稳定的。在训练中，这通常表现为梯度爆炸。一个微小的输入变化会导致输出的巨大变化，这会产生一个极大的损失值。
优化器的作用：像AdamW这样的优化器，其目标是最小化损失。当它看到一个因梯度爆炸而产生的巨大损失时，它会调整网络权重（参数 $\theta$），以“平息”这种不稳定的动态。这个调整过程，实际上就是在隐式地迫使 $J_f$ 的谱半径减小，使网络变得更稳定。

理由三：现代网络架构中的内置稳定器

我们并非完全依赖训练过程的隐式推动。现代深度学习架构中有很多组件本身就有助于维持系统的稳定性：

归一化层 (Normalization Layers)：像 RMSNorm 或 LayerNorm 这样的层会强制性地将每一层神经元的激活值重新缩放到一个固定的范围内。这就像给系统装上了一个“稳压器”，极大地限制了状态值发生爆炸的可能性。
激活函数 (Activation Functions)：像 Tanh 或 Sigmoid 这样的激活函数会将输出值压缩到 [-1, 1] 或 [0, 1] 的区间内，天然地具有抑制“爆炸”的作用。
权重衰减 (Weight Decay)：这是优化器（如AdamW）中的一个正则化项，它会惩罚过大的权重值，鼓励模型使用更小、更平滑的参数，这同样有助于降低 J_f 的谱半径。

总结

我们之所以可以假设 $\rho(J_f) < 1$，是因为：

这是模型能够稳定收敛的数学前提。 一个不满足此条件的模型，根本无法正常工作。
这是模型训练的自然结果。 优化器会通过最小化损失来惩罚不稳定的动态，从而间接压低谱半径。
这是现代网络架构设计的内在优势。 归一化等技术为系统的稳定性提供了强有力的保障。

因此，这个假设在实践中是一个非常合理且通常会成立的假设，它为HRM使用高效的“一步梯度近似”提供了坚实的理论基础。