Note -「Hierarchical Reasoning Model」思考，像脑子一样！

\[\mathscr{Lorain~wy~Lora~blea.} \newcommand{\DS}[0]{\displaystyle} % operators alias \newcommand{\opn}[1]{\operatorname{#1}} \newcommand{\card}[0]{\opn{card}} \newcommand{\lcm}[0]{\opn{lcm}} \newcommand{\char}[0]{\opn{char}} \newcommand{\Char}[0]{\opn{Char}} \newcommand{\Min}[0]{\opn{Min}} \newcommand{\rank}[0]{\opn{rank}} \newcommand{\Hom}[0]{\opn{Hom}} \newcommand{\End}[0]{\opn{End}} \newcommand{\im}[0]{\opn{im}} \newcommand{\tr}[0]{\opn{tr}} \newcommand{\diag}[0]{\opn{diag}} \newcommand{\coker}[0]{\opn{coker}} \newcommand{\id}[0]{\opn{id}} \newcommand{\sgn}[0]{\opn{sgn}} \newcommand{\Res}[0]{\opn{Res}} \newcommand{\Ad}[0]{\opn{Ad}} \newcommand{\ord}[0]{\opn{ord}} \newcommand{\Stab}[0]{\opn{Stab}} \newcommand{\conjeq}[0]{\sim_{\u{conj}}} \newcommand{\cent}[0]{\u{\degree C}} \newcommand{\Sym}[0]{\opn{Sym}} \newcommand{\Var}[0]{\opn{Var}} \newcommand{\wg}[0]{\wedge} \newcommand{\Wg}[0]{\bigwedge} \newcommand{\sq}[0]{\opn{\square}} % symbols alias \newcommand{\E}[0]{\exist} \newcommand{\A}[0]{\forall} \newcommand{\l}[0]{\left} \newcommand{\r}[0]{\right} \newcommand{\ox}[0]{\otimes} \newcommand{\lra}[0]{\leftrightarrow} \newcommand{\llra}[0]{\longleftrightarrow} \newcommand{\iso}[1]{\overset{\sim}{#1}} \newcommand{\eps}[0]{\varepsilon} \newcommand{\Ra}[0]{\Rightarrow} \newcommand{\Eq}[0]{\Leftrightarrow} \newcommand{\d}[0]{\mathrm{d}} \newcommand{\e}[0]{\mathrm{e}} \newcommand{\i}[0]{\mathrm{i}} \newcommand{\j}[0]{\mathrm{j}} \newcommand{\k}[0]{\mathrm{k}} \newcommand{\Ex}[0]{\mathbb{E}} \newcommand{\D}[0]{\mathbb{D}} \newcommand{\oo}[0]{\infty} \newcommand{\tto}[0]{\rightrightarrows} \newcommand{\mmap}[0]{\hookrightarrow} \newcommand{\emap}[0]{\twoheadrightarrow} \newcommand{\actl}[0]{\curvearrowright} \newcommand{\actr}[0]{\curvearrowleft} \newcommand{\nsubg}[0]{\triangleleft} \newcommand{\nsupg}[0]{\triangleright} \newcommand{\lin}[0]{\lim_{n\to\oo}} \newcommand{\linf}[0]{\liminf_{n\to\oo}} \newcommand{\lsup}[0]{\limsup_{n\to\oo}} \newcommand{\ser}[0]{\sum_{n=1}^\oo} \newcommand{\serz}[0]{\sum_{n=0}^\oo} \newcommand{\isoto}[0]{\overset\sim\to} \newcommand{\F}[0]{\mathbb F} \newcommand{\x}[0]{\times} \newcommand{\M}[0]{\mathbf{M}} \newcommand{\T}[0]{\intercal} \newcommand{\Co}[0]{\complement} \newcommand{\alp}[0]{\alpha} \newcommand{\lmd}[0]{\lambda} \newcommand{\mmid}[0]{\parallel} \newcommand{\loop}[0]{{\circlearrowleft}} \newcommand{\go}[0]{\triangleright} % symbols with parameters \newcommand{\der}[1]{\frac{\d}{\d #1}} \newcommand{\ul}[1]{\underline{#1}} \newcommand{\ol}[1]{\overline{#1}} \newcommand{\wt}[1]{\widetilde{#1}} \newcommand{\br}[1]{\l(#1\r)} \newcommand{\bk}[1]{\l[#1\r]} \newcommand{\ev}[1]{\l.#1\r|} \newcommand{\wh}[1]{\widehat{#1}} \newcommand{\eval}[1]{\l[\!\l[#1\r]\!\r]} \newcommand{\abs}[1]{\l|#1\r|} \newcommand{\bs}[1]{\boldsymbol{#1}} \newcommand{\dat}[1]{\bs{\mathrm{#1}}} \newcommand{\env}[2]{\begin{#1}#2\end{#1}} \newcommand{\ALI}[1]{\env{aligned}{#1}} \newcommand{\CAS}[1]{\env{cases}{#1}} \newcommand{\pmat}[1]{\env{pmatrix}{#1}} \newcommand{\algo}[1]{\begin{array}{r|l}#1\end{array}} \newcommand{\dary}[2]{\l|\begin{array}{#1}#2\end{array}\r|} \newcommand{\pary}[2]{\l(\begin{array}{#1}#2\end{array}\r)} \newcommand{\pblk}[4]{\l(\begin{array}{c|c}{#1}&{#2}\\\hline{#3}&{#4}\end{array}\r)} \newcommand{\u}[1]{\mathrm{#1}} \newcommand{\t}[1]{\text{#1}} \newcommand{\ts}[1]{\textsf{#1}} \newcommand{\tb}[1]{\textbf{#1}} \newcommand{\os}[2]{\overset{#1}{#2}} \newcommand{\lix}[1]{\lim_{x\to #1}} \newcommand{\ops}[1]{#1\cdots #1} \newcommand{\seq}[3]{{#1}_{#2}\ops,{#1}_{#3}} \newcommand{\dedu}[2]{\u{(#1)}\Ra\u{(#2)}} \newcommand{\prv}[3]{\DS{{\DS #1} \over {\DS #2}}~(#3)} \]

  一个老生常谈的话题是, 以 transformer 架构为代表的典型 LLMs 都是固定深度的网络, 直觉上, 它们的计算能力天生限制在 \(\tb{AC}^0\) 或者 \(\tb{TC}^0\) 中. 为获得强大的推理能力, LLMs 选择用 CoT 来将复杂问题拆解到 token 级别, 并以高延迟, 高脆弱性为代价链式地在 token sequence 上用浅网络 "推理". 就像是一个只能通过打草稿-读草稿来进行思考的解题者, 听上去也不赖, 但的确不算高明.

  理想的思考会是怎样的? 在脑子里想, 也就是所谓的 latent reasoning! 你会想到 RNN, 但 RNN 的过早收敛 (early convergence) 问题和昂贵的 BPTT 代价也常被鞭尸. RNN 喊着什么类脑喊着什么下一代网络就冲过来了, 但我们的大脑怎么可能做 BPTT?

网络结构

  慢速迭代的高层 "指导" 快速迭代的低层.

  HRM 有四个可学习网络, 分别是:

• 输入网络 \(f_I(\cdot;\theta_I)\);
• 低层循环网络 \(f_L(\cdot;\theta_L)\);
• 高层循环网络 \(f_H(\cdot;\theta_H)\);
• 输出网络 \(f_O(\cdot;\theta_O)\).

  每一步高层循环对应 \(T\) 步低层循环, \(N\) 步高层循环共对应 \(N\x T\) 个时间步. 设低层和高层的隐状态分别为 \(z_L^i\) 和 \(z_H^i\), 输入为 \(x\), 则首先由

\[\tilde x=f_I(x;\theta_I) \]

得到模型输入, 接着在循环网络中迭代

\[\ALI{ z_L^i &= f_L(z_L^{i-1},z_H^{i-1},\tilde x;\theta_L);\\ z_H^i &= \CAS{ z_H^{i-1},&i\not\equiv0\pmod T;\\ f_H(z_L^{\color{red}{i}},z_H^{i-1};\theta_H),&i\equiv 0\pmod T. } } \]

(原文红色 \(i\) 为 \(i-1\), 疑似笔误: 高层应该接受最新的低层状态作为反馈.) 最后通过高层读出

\[\hat y=f_O(z_H^{NT};\theta_O). \]

  细节上, \(f_L\) 和 \(f_H\) 是维度与架构都相同的 encoder-only transformer blocks (RoPE + GLU + RMSNorm, 去除 linear bias); \(f_I\) 做 tokens 到 vector representations 的 embedding; \(f_O\) 是简单的 linear read-out + 概率归一化.

一步梯度近似 (one-step gradient)

  只在最后一次 forward 开启梯度, 这居然 work? 我们来编个理由.

  先回顾一下隐函数定理, 设有方程 \(G(z,\theta)=0\), \((z^*,\theta^*)\) 在其轨迹上, 且 \(G\) 关于 \(z\) 的 Jacobi 矩阵 \(\frac{\part G}{\part z}\) 在这点可逆, 那么在 \((z^*,\theta^*)\) 附近, 可以把 \(G=0\) 轨迹上的 \(z\) 视为关于 \(\theta\) 的可微函数 \(z=z^*(\theta)\), 那么对 \(G(z^*(\theta),\theta)=0\) 两侧取 \(\frac{\part}{\part\theta}\) 得到

\[\frac{\part G}{\part z}\frac{\part z^*}{\part\theta}+\frac{\part G}{\part \theta}=0. \]

因而

\[\frac{\part z^*}{\part\theta}=-\br{\frac{\part G}{\part z}}^{-1}\frac{\part G}{\part\theta}. \]

我们需要求 \(\frac{\part \mathcal L}{\part\theta}\), 而 \(\frac{\part\mathcal L}{\part z}\) 来自上游梯度, 还需要 \(\frac{\part z}{\part\theta}\), 所以就可以使用隐函数定理.

  接着, 我们研究类似 RNN 的迭代形式 \(z^i=f(z^{i-1};\theta)\). 假设 \(z\) 已经取到不动点 \(z^*\), 即 \(z^*=f(z^*;\theta)\), 那么令 \(G(z,\theta)=z-f(z;\theta)\), 就有 \(G(z^*,\theta^*)=0\). 此时 \(\frac{\part G}{\part z}=\bs 1-\frac{\part f}{\part z}\). 记 \(J_f=\ev{\frac{\part f}{\part z}}_{z^*}\), 那么

\[\frac{\part z^*}{\part \theta}=-\br{\bs 1-J_f}^{-1}\cdot -\frac{\part f}{\part \theta}=(\bs 1-J_f)^{-1}\frac{\part f}{\part\theta}. \]

这就 Deep Equilibrium Model (DEQ) 中常见的不动点梯度.

  现在我们的瓶颈是对 \((\bs 1-J_f)^{-1}\) 的处理. 我们假设 \(\rho(J_f)<1\), 则由 Neumann 级数,

\[(1-J_f)^{-1}=\bs 1+J_f+J_f^2+\cdots. \]

额... 我们直接取零阶近似 \((\bs1-J_f)^{-1}\approx\bs 1\)... 论文声称是这样, 这样有 \(\mathcal O(\|J_f\|)\) 的误差, 但论文似乎没有细致分析其安全性; 而且诸如一阶近似 \(\bs 1+J_f\) 的计算开销其实也可接受? 我们暂且承认这是一个工程上合理的近似吧, 取 \((\bs 1-J_f)^{-1}\approx\bs 1\) 得到

\[\frac{\part z_H^*}{\part\theta_H}\approx\frac{\part f_H}{\part \theta_H},\quad \frac{\part z_H^*}{\part\theta_L}\approx\frac{\part f_H}{\part z_L^*}\cdot\frac{\part z_L^*}{\part\theta_L},\quad \frac{\part z_H^*}{\part\theta_I}\approx\frac{\part f_H}{\part z_L^*}\cdot\frac{\part z_L^*}{\part\theta_I}. \]

其中再用同样的近似

\[\frac{\part z_L^*}{\part\theta_L}\approx\frac{\part f_L}{\part\theta_L},\quad \frac{\part z_L^*}{\part \theta_I}=\frac{\part z^*_L}{\part\tilde x}\frac{\part\tilde x}{\part\theta_I}\approx\frac{\part f_L}{\part\theta_I}. \]

我们就得到了快速反向传播的策略. 这种近似称为一步梯度近似.

  你会发现用 PyTorch 实现起来特别简单: 前若干步推理直接关掉梯度, 最后一步推理再打开梯度, 然后用优化器优化即可.

深度监督 (Deep Supervision)

  把一个 \(N\x T\) 的 forward 视为一个 segment, 每个 segment 视为独立的 input-output 做训练; 后续推理时也以 segment 为单位, 决定 "要不要再推理一个 segment".

  如上所述, 记 \((z^m,\hat y^m)\) 为第 \(m\) 个 segment 结束的隐状态和输出, \(\theta^m\) 为第 \(m\) 个 segment 更新后的参数, 则

\[\ALI{ (z^m,\hat y^m)&\gets\t{HRM}(z^{m-1},x;\theta^{m-1}),\\ \mathcal L_m&\gets\t{Loss}(\hat y^m,y),\\ \theta^m&\gets\t{OptStep}(\theta^{m-1},\nabla_{\theta^{m-1}}\mathcal L_m). } \]

  注意 \(z^m\) 在传给下一个 segment 前会 detach. 也就是说，第 \(m+1\) 个 segment 的梯度不会传回第 \(m\) 个 segment. 这相当于对 "递归深监督过程" 再做一次 one-step 近似. 工程上, 这样能让高层模块更频繁收到反馈, 训练更稳定; 类脑解释上, 周期性学习窗口类似神经振荡调控何时发生学习.

自适应思考时长 (Adaptive Computational Time, ACT)

  用 Q-learning 决策 continue or halt 行为.

  Q-value \(\hat Q^m=\pmat{\hat Q^m_{\t{halt}}&\hat Q^m_{\t{cont}}}\) 通过

\[\hat Q^m\gets\t{sigmoid}(\theta_Q^\T\x z_{H}^{m}) \]

给出; 最大推理步数 \(M_{\max}\) 作为超参数取定, \(M_{\min}\) 以 \(\eps\) 的概率取 \(\t{Unif}([2:M_{\max}])\) 以鼓励长思考, 以 \(1-\eps\) 的概率取 \(1\). Q-target 自然地给出

\[\ALI{ \hat G_{\t{halt}}^m &= \mathbb 1\{\hat y^m=y\},\\ \hat G_{\t{cont}}^m &= \CAS{ \hat Q_{\t{halt}}^{m+1},&{\color{red}{m+1\ge M_{\max}}};\\ \max\{\hat Q_{\t{halt}}^{m+1},\hat Q_{\t{cont}}^{m+1}\},&\t{otherwise}. } } \]

(原文红色部分条件为 \(m\ge N_{\max}\), 是 \(m\ge M_{\max}\) 的笔误; 但 \(m=M_{\max}-1\) 时就应该强制下一步的 reward 取自 halt.)

  最终, 综合 sequence-to-sequence loss 和 Q-head loss, 我们给出 ACT 的损失:

\[\mathcal L_{\t{ACT}}^m = \t{Loss}(\hat y^m,y)+\t{BinaryCE}(\hat Q^m,\hat G^m). \]

  Clarify: 实现时, 每当 batch 中有 sample halt, 我们就将其替换为一个新的 sample. 此外, 带 ACT 的 HRM 在深度监督的同时做 ACT 训练, 跨 segment 的隐状态被 detach; 同一 segment 内, prediction loss 与 Q-head loss 联合优化共享表示.

  关于 ACT, 还有一个值得关心的问题是它作为一个 deep Q-learning 的稳定性问题, 我们知道这类 learning 通常是不够稳定的, 需要配合 replay buffer (缓存经验, 采样 mini-batch 更新) 和 target network (保存一个旧网络计算 target 维持其稳定性). 但 Gallici et al. 指出, 若网络参数有界, 训练有 weight decay, 网络中使用 post-norm 结构, 则 Q-learning 可以收敛. HRM/ACT 中采用 post-norm 的 RMSNorm 和 AdamW (实际训练使用的 Adam-atan2) 优化器, 其稳定性被上述论断支持.

实验结果

  谁想读论文的 result? 本节是 codex 写的.

  实验部分的核心结论是: HRM 在不使用预训练, 不使用 CoT 监督, 仅依赖 input-output pairs 的情况下, 用约 \(27\mathrm M\) 参数和约 \(1000\) 个训练样本, 在多类需要深度推理的任务上显著超过直接预测式 transformer 和若干 CoT baselines.

  论文主要报告了四个主结果:

benchmark	training examples	HRM accuracy	任务要点
ARC-AGI-1	960	\(40.3\%\)	从少量输入输出网格示例中归纳抽象规则.
ARC-AGI-2	1120	\(5.0\%\)	更强调组合推理, 多步逻辑, 上下文规则和符号抽象.
Sudoku-Extreme	1000	\(55.0\%\)	\(9\times 9\) 数独, 要求完整预测唯一解.
Maze-Hard	1000	\(74.5\%\)	\(30\times 30\) 迷宫最短路, 预测必须 valid 且 optimal.

  对 ARC-AGI, 论文使用数据增强来提高泛化: 对训练和测试输入应用平移, 旋转, 翻转和颜色置换; 测试时对每个 test input 生成并求解 \(1000\) 个增强版本, 再反变换回原坐标, 选出现频率最高的两个预测作为最终答案. 在 ARC-AGI-1 上, HRM 的 \(40.3\%\) 高于文中列出的 o3-mini-high 的 \(34.5\%\) 和 Claude 3.7 8K context 的 \(21.2\%\). 但 ARC-AGI-2 的绝对准确率仍只有 \(5.0\%\), 因而这个结果更应理解为 "在低样本, 无预训练设定下相对更强", 而不是已经解决 ARC-AGI-2.

  对 Sudoku-Extreme, 论文特意强调常用 Sudoku 数据集并不够难: Kaggle Sudoku 可以用 elementary single-digit techniques 解决, 17-clue puzzles 虽然 clue 少, 但 clue 往往较正交, 反而会产生许多直接路径. 他们构造的 Sudoku-Extreme 来自 easy puzzles 与社区公认困难 puzzles 的混合, 并用 tdoku solver 所需 backtracks 衡量难度. Sudoku-Extreme 平均需要 \(22\) 次 backtracks, 明显高于 Sudoku-Bench 的 \(0.45\) 次. 在主实验中, HRM 只用 \(1000\) 个训练样本达到 \(55.0\%\); 在用于分析的 Sudoku-Extreme-Full 上, Figure 2 显示单纯增加 transformer 宽度收益很小, 而增加计算深度才关键, HRM 能更有效利用深度并接近完美表现.

  对 Maze-Hard, 任务是寻找 \(30\times 30\) 迷宫中的最短路径. 数据生成遵循 Lehnert et al. 的流程, 但额外保留 shortest path length 超过 \(110\) 的困难实例; 训练集和测试集各 \(1000\) 个样本. 论文指出, 这些需要长推理轨迹的 Sudoku/Maze 任务对 CoT baselines 和 direct-pred transformer 都很困难: 与 HRM 同尺寸, 同训练设置的 \(8\) 层 transformer direct prediction baseline 在 \(1000\) 样本设定下几乎无法解决这些任务; 即使在 Sudoku-Extreme-Full 上也只有 \(16.9\%\).

  ACT 的实验主要在 Sudoku-Extreme-Full 上验证. Figure 5 显示, 与固定 segment 数 \(M\) 的模型相比, ACT 在 \(M_{\max}\) 增大时平均实际计算步数仍保持较低且稳定, 但准确率接近固定计算模型. 也就是说, ACT 学到的不是盲目多想, 而是根据样本难度分配计算量. 进一步地, 训练时给定某个 \(M_{\max}\) 的模型, 推理时增大 \(M_{\max}\) 仍可能提升 Sudoku accuracy, 说明 HRM 具有一定 inference-time scaling; 但论文也指出, 对 ARC-AGI 这类通常只需少数变换的任务, 额外推理计算带来的收益较小.

  论文还可视化了中间 timestep 的预测轨迹. 在 Maze-Hard 中, HRM 看起来会先并行探索多条可能路径, 再去掉 blocked 或低效路径并逐步形成最终路线; 在 Sudoku 中, 轨迹类似 depth-first search, 会尝试候选解并在遇到矛盾时 backtrack; 在 ARC 中, 轨迹更像 hill-climbing, 即逐步调整网格直到接近解. 这些可视化支持 "HRM 在不同任务上学到不同推理策略" 的解释, 但论文也明确说这还不是对内部算法的完整刻画.

---

Reference.

• G. Wang et al., “Hierarchical Reasoning Model,” Aug. 04, 2025, arXiv: arXiv:2506.21734. doi: 10.48550/arXiv.2506.21734.