Universal Reasoning Model
Universal Reasoning Model
- 作者/机构: Zitian Gao, Lynx Chen, Yihao Xiao, He Xing, Ran Tao, Haoming Luo, Joey Zhou, Bryan Dai
- 机构: Ubiquant
A1 主要贡献
近期,诸如通用 Transformer (Universal Transformers, UTs) 【5, Universal transformers, 2019】之类的循环模型【7, Hierarchical reasoning models: Perspectives and misconceptions, 2025】【11, Less is more: Recursive reasoning with tiny networks, 2025】【20, Hierarchical reasoning model, 2025】在解决 ARC-AGI 和数独【2, Arc prize 2024: Technical report, 2025】【3, Arc-agi-2: A new challenge for frontier ai reasoning systems, 2025】等复杂推理任务中展现了其有效性。尽管这些基于 UT 的小型模型是在这些任务上从头开始训练的,没有经过互联网规模的预训练,但它们的表现却持续大幅优于大多数标准的基于 Transformer 的大型语言模型 (LLMs)【20, Hierarchical reasoning model, 2025】。尽管这种对比凸显了 UTs 在深度密集型迭代推理方面的潜力,但门控机制的功能和影响除了其最初的直觉之外,尚未得到充分的探索。
以往的研究常常将性能提升归因于高层次的架构创新【7, Hierarchical reasoning models: Perspectives and misconceptions, 2025】【11, Less is more: Recursive reasoning with tiny networks, 2025】【20, Hierarchical reasoning model, 2025】,然而本文的分析揭示,核心的性能增益实际上源于通用 Transformer 固有的、常被忽视的循环归纳偏置。特别地,非线性的深度计算所起的作用远大于之前的认识,这表明增强循环处理的架构修改可以带来显著的下游性能改进。受此启发,本文进一步研究并加强了这种归纳偏置,通过对 UT 框架进行简化而有效的增强,从而在保持参数效率的同时,实现了更强的抽象能力。
本文的主要贡献如下:
* 通过广泛的消融研究,本文表明模型在 ARC-AGI 这类复杂推理任务上的性能主要源于其非线性。此外,本文揭示了超越标准 Transformer 的推理能力的真正来源是通用 Transformer 的循环机制,而非先前工作中过于精巧的设计。
* 通过在通用 Transformer 中引入短卷积和截断反向传播,本文在 ARC-AGI 1 上实现了 53.8% 的 pass@1 准确率,在 ARC-AGI 2 上实现了 16.0% 的 pass@1 准确率,达到了业界领先水平。
A3 背景知识
2.1 标准 Transformer
词汇表与嵌入函数定义。假设 $V$ 表示大小为 $V$ 的词汇表,$x = (x_1, . . . , x_N ) \in V^N$ 是一个长度为 $N$ 的输入序列。我们将词元嵌入函数定义为 $\phi : V^N \rightarrow R^{N \times d}$,它将离散的词元映射到一个 $d$ 维的连续表示。相反,反嵌入函数(或语言建模头)表示为 $\psi : R^{N \times d} \rightarrow R^{N \times V}$,它将隐藏状态投影回词汇表的 logit 空间。
单层 Transformer 定义。一个由 $\theta$ 参数化的单层 Transformer 被定义为一个函数 $T_{\theta} : R^{N \times d} \rightarrow R^{N \times d}$。这个函数通常由一个多头自注意力(MHSA)模块和一个逐位置的前馈网络(FFN)组成,每个模块都包裹着残差连接和层归一化。
$$\begin{aligned} \begin{aligned} \mathcal{T}_{\theta}(H) & = \text{FFN}(\text{LN}(H' + H)), \\ \text{where } H' & = \text{MHSA}(\text{LN}(H)) \end{aligned} \end{aligned}$$标准 Transformer 模型构建。一个深度为 $L$ 的标准、非递归 Transformer 模型 $M_{std}$ 是通过堆叠 $L$ 个具有不同参数 $\Theta = \{\theta_1, . . . , \theta_L\}$ 的层来构建的。其前向传播是这些层的复合:
$$\mathcal{M}_{\mathrm{std}}(\mathbf{x}) = \psi \circ \mathcal{T}_{\theta_L} \circ \cdots \circ \mathcal{T}_{\theta_1} \circ \phi(\mathbf{x})$$这里,操作符 $\circ$ 表示函数复合。计算成本和参数数量都与 $L$ 呈线性关系,这在模型容量和推理计算之间形成了一种刚性耦合。
2.2 通用 Transformer
UT 核心机制。通用 Transformer (UT)【5, Universal transformers, 2019】通过引入深度上的循环计算,扩展了标准 Transformer【18, Attention is all you need, 2017, NIPS】。UT 不是堆叠 $L$ 个不同的层,而是重复应用单个转换块来优化词元表示。对于一个输入序列 $x$,其嵌入矩阵为 $H_0 \in R^{n \times d}$,UT 的状态更新如下:
$$\mathbf{H}^{t+1}=\operatorname{LayerNorm}\left(\mathbf{H}^{t}+\operatorname{MHA}\left(\mathbf{H}^{t}\right)\right),$$共享转换函数。接着是一个共享的逐位置转换函数:
$$\mathbf{H}^{t+1} \leftarrow \text{LayerNorm}(\mathbf{H}^{t+1} + \text{Transition}(\mathbf{H}^{t+1})) , \quad t = 0, \ldots, T-1,$$其中 Transition 可以是一个前馈网络或可分离卷积。为了同时编码位置和优化深度,UT 在每一步都添加了二维正弦嵌入。
2.2.1 参数共享
跨深度权重绑定。UT 的一个关键设计是跨深度的权重绑定。注意力和转换参数
$$\Theta_{\mathrm{UT}}=\left\{\mathbf{W}_{h}^{Q}, \mathbf{W}_{h}^{K}, \mathbf{W}_{h}^{V}, \mathbf{W}^{O}, \Theta_{\text {Transition }}\right\}$$在所有时间步 $t$ 中都被重用。因此,该模型使用灵活的步数 $T$ 进行迭代式表示优化,这使得 (i) 推理时可以进行深度自适应,以及 (ii) 理论上比固定深度的 Transformer 具有更高的表达能力。
2.2.2 自适应计算时间 (ACT)
动态计算分配。通过自适应计算时间 (ACT)【9, Adaptive computation time for recurrent neural networks, 2017】,不同的词元可以在不同的循环步骤停止。在步骤 $t$,每个位置都会预测一个停止概率:
$$p_{t,i}=\sigma(\mathbf{w}^{\top}\mathbf{h}_{t,i}+b),$$该概率会一直累加,直到达到阈值 $1 - \epsilon$。
最终表示生成。最终的词元表示是一个加权混合:
$$\mathbf{h}_{i}^{\text{final}} = \sum_{t} \Delta_{t,i} \mathbf{h}_{t,i},$$其中 $\Delta_{t,i}$ 是截断后的分配量。ACT 允许 UT 为复杂的词元分配更多的计算资源,而为简单的词元分配较少的计算资源。
A2 方法细节
模型基础架构。我们的通用推理模型 (URM) 的基础架构与通用 Transformer【5, Universal transformers, 2019】紧密相连,不同之处在于它采用了仅解码器(decoder-only)的设计。这一点与之前的研究如 HRM【20, Hierarchical reasoning model, 2025】和 TRM【11, Less is more: Recursive reasoning with tiny networks, 2025】是一致的。我们的工作与以往模型【11, Less is more: Recursive reasoning with tiny networks, 2025】【20, Hierarchical reasoning model, 2025】的不同之处在于引入了以下的 ConvSwiGLU 模块和截断循环反向传播机制。
3.1 ConvSwiGLU
设计动机与机制。为了增强通用 Transformer 的非线性,我们引入了 ConvSwiGLU(动机见 4.6 节),它通过一个深度可分离的短卷积来增强标准的 SwiGLU 前馈模块。与传统的逐点 SwiGLU【16, Glu variants improve transformer, 2020】(它独立处理每个词元)不同,我们的设计明确地将局部上下文交互注入到门控机制中,从而在不增加序列级复杂性的情况下,引入了词元空间中轻量级的通道混合【1, Physics of language models: Part 4.1, architecture design and the magic of canon layers, 2025, NeurIPS】【22, Metaformer is actually what you need for vision, 2022, CVPR】。
输入投影。给定一个输入序列 $X \in R^{T \times d}$,我们首先将其投影到一个扩展的中间表示:
$$[\mathbf{G}, \mathbf{U}] = X W_{\text{up}} \in \mathbb{R}^{T \times 2m}$$门控表示。SwiGLU 激活函数产生一个门控表示:
$$\mathbf{H}_{\mathrm{ffn}}=\mathrm{SiLU}(\mathbf{G}) \odot \mathbf{U} .$$短程交互集成。为了整合短程词元交互,我们在门控特征上应用了一个深度可分离的一维卷积:
$$\mathbf{H}_{\mathrm{conv}} = \sigma(\mathbf{W}_{\mathrm{dwconv}} * \mathbf{H}_{\mathrm{ffn}}),$$其中 $W_{dwconv} \in R^{m \times 1 \times k}$ 是一个大小为 $k=2$ 的深度可分离卷积核。
输出投影。最后,输出被投影回隐藏维度:
$$\mathbf{Y} = [\sigma(\mathbf{W}_{\text{dwconv}} * (\text{SiLU}(\mathbf{G}) \odot \mathbf{U}))] W_{\text{down}}.$$3.2 截断循环反向传播
设计动机。当循环推理的次数变得很大时,从早期循环传播的梯度可能会因为噪声累积和不稳定性而阻碍优化(经验证据见 4.5 节)。为了缓解这个问题,我们采用了截断循环反向传播(TBPTL),只计算后期循环的梯度。
模型定义。考虑一个在训练期间展开 M 次迭代循环的 D 层通用推理模型。设 $h^{(d)}_t$ 表示在第 $t \in \{1, . . . , M \}$ 次迭代时第 $d \in \{1, . . . , D\}$ 层的隐藏表示。循环转换定义为:
$$\mathbf{h}_t^{(d)}=F_\theta^{(d)}(\mathbf{h}_t^{(d-1)}, \mathbf{h}_{t-1}^{(d)}),$$其中 $F^{(d)}_{\theta}$ 表示第 $d$ 层的参数化变换,其可训练参数为 $\theta$。
循环划分。我们不通过所有 M 个循环进行反向传播,而是将展开过程划分为仅前向传播和可训练两部分。具体来说,对于一个截断索引 $N < M$:
梯度计算。在训练期间,我们只对后 $(M - N)$ 个循环中累积的损失计算梯度:
$$\mathcal{L}_{\mathrm{TBPTL}}(\theta)=\sum_{t=N+1}^{M} \mathcal{L}\left(\mathbf{h}_{t}^{(D)}, y\right),$$其中 $L(·)$ 是交叉熵损失函数。
梯度公式。因此,关于 $\theta$ 的梯度为:
$$\nabla_{\theta} \mathcal{L}_{\mathrm{TBPTL}}=\sum_{t=N+1}^{M} \frac{\partial \mathcal{L}}{\partial \mathbf{h}_{t}^{(D)}} \frac{\partial \mathbf{h}_{t}^{(D)}}{\partial \theta}.$$示例。例如,在一个配置为 $D = 4$ 层和 $M = 8$ 个内部循环的设置中,我们选择 $N = 2$ 个仅前向传播的循环。因此,只有最后 6 个循环(即 $t = 3$ 到 $t = 8$)对梯度计算有贡献。
A4 实验环境
我们的实验设置主要遵循 HRM 和 TRM【11, Less is more: Recursive reasoning with tiny networks, 2025】【20, Hierarchical reasoning model, 2025】。
* 数据集: 使用与先前工作相同的 ARC-AGI 1, ARC-AGI 2, Sudoku 数据集和增强数据。
* 软件配置: 遵循【11, Less is more: Recursive reasoning with tiny networks, 2025】,对模型参数应用指数移动平均 (EMA) 以提高训练稳定性。所有模型均使用 AdamAtan2 优化器【6, Scaling exponents across parameterizations and optimizers, 2024, ICML】进行训练。
* 学习率:
* ARC-AGI 1 和 ARC-AGI 2: 主模型学习率分别为 $1 \times 10^{-4}$ 和 $3 \times 10^{-4}$,谜题嵌入学习率为 $1 \times 10^{-2}$。
* Sudoku: 谜题嵌入学习率为 $1 \times 10^{-4}$。
* **权重衰减**: * ARC-AGI 1 和 ARC-AGI 2: 主模型和谜题嵌入均为 0.1。 * Sudoku: 1.0,与先前工作一致。
- 模型架构:
- 模型包含 4 个层,隐藏层大小为 512,有 8 个注意力头。
- 内部循环运行 8 步,其中前两步仅进行前向传播。
- 外部循环采用自适应计算时间 (ACT)【9, Adaptive computation time for recurrent neural networks, 2017】,最大步数为 16。
A4 实验结果
4.2 主要结果
如表1所示,通用推理模型 (URM) 在所有基准测试中均取得了相较于先前基于UT的方法的显著改进。
- ARC-AGI 1: URM 达到了 53.8% 的 pass@1,大幅超越了 TRM (40.0%) 和 HRM (34.4%)。
- ARC-AGI 2: URM 获得了 16.0% 的 pass@1,几乎是 HRM 的三倍,是 TRM 的两倍多。
- Sudoku: URM 实现了 77.6% 的准确率,同样超过了 TRM 和 HRM。
- 分析: 值得注意的是,在更大的采样预算下(例如 pass@1000),URM 的优势进一步扩大,这表明迭代优化能够生成更丰富的候选答案,而非脆弱的单步预测。
4.3 为什么是通用 Transformer?
实验内容: 表2比较了标准 Transformer 和通用 Transformer (UTs) 在不同模型深度、隐藏层大小和循环次数下的性能,报告了在 ARC-AGI 1 上的 pass@n 结果。
实验结果:
* 参数效率: 表2显示,UTs 在 ARC-AGI 1 上的性能提升源于其极高的参数效率,而非模型规模或计算预算的增加。一个仅有4倍参数的UT模型达到了40.0的pass@1得分,远超使用多达32倍参数但性能依然较弱的标准 Transformer。简单地增加标准 Transformer 的深度或宽度会带来递减的回报,甚至导致性能下降。
* 计算效率: 即使在计算量(FLOPs)相同的情况下,UTs 的优势依然存在。在32倍FLOPs下,将计算从深的、非共享的层重新分配到循环优化中,pass@1 从标准 Transformer 的23.75提升到UTs的40.0。
分析:
* 这种卓越的效率是由跨深度参数共享引入的循环归纳偏置驱动的。通过重复应用共享的变换,UTs 实现了与算法推理结构更好对齐的迭代优化,同时避免了参数数量的增加。
* 此行为与先前研究【15, Reasoning with latent thoughts: On the power of looped transformers, 2025, ICLR】的分析一致,该研究认为许多推理任务更多地受益于迭代计算而非增加独立层数。在标准 Transformer 中,额外的 FLOPs 通常用于高层中冗余的优化,而循环计算则将相同的预算转化为增加的有效深度【15, Reasoning with latent thoughts: On the power of looped transformers, 2025, ICLR】【23, Investigating layer importance in large language models, 2024, BlackboxNLP】。
* 因此,在固定参数和固定 FLOPs 预算下,UTs 在推理任务上始终优于标准 Transformer,使其特别适合于像 ARC-AGI 这样多步抽象比纯粹规模更关键的推理密集型环境。
4.4 短卷积
实验内容: 为了增强通用 Transformer 的非线性归纳偏置,引入了一个由 $W_{dwconv}$ 参数化的深度可分离短卷积模块。实验评估了在UT转换模块内不同位置插入该模块对性能的影响。
实验结果与分析:
* 插入位置: 如图3左图所示,在六个插入点中,将 $W_{dwconv}$ 模块插入注意力路径内部,即位置 (a)至(d),并未带来性能提升,反而常常导致性能下降,这表明局部扰动干扰了注意力线性投影的几何结构。在位置 (e) 有轻微增益,此时扰动仅作用于聚合后的多头特征。
* 最佳位置: 最显著的效果出现在位置 (f),即 MLP 扩展之后,表明短程混合在已经是非线性的子空间内应用时最为有益。这支持了一种功能性解释,即 MLP——而非注意力——构成了模型表达性非线性的主要来源;用 $W_{dwconv}$ 增强它,显著提升了模型的非线性表示能力。
* 注意力可视化: 如图4所示,将短卷积整合到 MLP 中显著增强了通道混合。标准通用 Transformer 表现出相对稀疏和同质的注意力模式,而带有 ConvSwiGLU 的模型则产生了更多样化和结构化的注意力矩阵。这表明短卷积促进了更有效的通道间信息流动,从而提高了注意力机制的表达能力。
4.5 截断循环反向传播
实验内容: 在总内部循环次数固定为8的情况下,通过实验评估了截断循环反向传播(TBPTL)的效果。实验改变了不计算梯度的前向传播循环次数,并记录了在 ARC-AGI 1 上的性能。
实验结果: 如表3所示,当截断前两个循环的梯度时——即前两个内部循环迭代以仅前向模式运行——模型取得了最佳性能。pass@1 和 pass@1000 都在此截断设置下达到峰值,而更短或更长的截断范围则导致较差的结果。
分析:
* 这种趋势与循环神经网络中的截断时间反向传播(TBPTT)非常相似,其基本动机大体相同。在完全的时间反向传播中,梯度会通过整个序列传播,这会带来高昂的计算和内存成本,并且由于梯度消失或爆炸行为,常常导致长程梯度无效。因此,实际应用通常将梯度传播限制在一个固定的近期窗口内【14, On the difficulty of training recurrent neural networks, 2013, ICML】【17, Unbiasing truncated backpropagation through time, 2017】。
* 同样,在通用 Transformer 中,跨所有内部循环迭代传播梯度可能导致优化不稳定,而过于激进的截断则限制了模型协调多步优化的能力。因此,适度截断梯度传播在优化稳定性和有效的长时程学习之间提供了一个有利的平衡。
* 需要注意的是,该实验中的所有结果都是使用一个没有短卷积模块的两层 URM 获得的,这与之前报告的完整 URM 模型不同。
4.6 Transformer 的非线性
实验内容: 对通用推理模型的非线性架构组件进行了消融研究,并在 ARC-AGI 1 上报告了 pass@n 结果。
实验结果: 如表4所示,随着模型中非线性组件的逐步移除,ARC-AGI 1 的性能单调下降。在这些组件中,MLP 中的激活函数扮演了尤其关键的角色:
* 将 SwiGLU 替换为 SiLU 或 ReLU 等更简单的非线性函数,会导致性能大幅下降。
* 完全移除注意力 softmax 会导致性能急剧崩溃。
分析:
* 这种清晰的单调趋势凸显了强非线性变换对于解决复杂抽象推理任务的重要性。
* 这些结果表明,ARC-AGI 所需的表达能力主要来自丰富的非线性映射。削弱非线性可能会系统性地限制模型表示复杂推理技能的能力。
* 需要注意的是,模型中仍保留了一些本研究未进行消融的非线性形式,例如每层后应用的 RMSNorm 以及注意力中查询和键之间的点积交互。然而,这些组件要么难以在不引起训练不稳定的情况下移除,要么与显式激活函数相比,其非线性效应相对较弱。由于消融它们通常会导致训练失败,因此它们不在本次分析的范围之内。
4.7 Muon 优化器
实验内容: 为了评估通用推理模型 (URM) 的训练效率,将 Muon (Momentum Updated Orthogonal Newton) 优化器【12, Muon: An optimizer for hidden layers in neural networks, 2024】与标准自适应基线 AdamAtan2【6, Scaling exponents across parameterizations and optimizers, 2024, ICML】进行了比较。
实验结果:
* 收敛速度: 在 ARC-AGI 1 和 ARC-AGI 2 基准测试中,Muon 表现出显著更快的收敛速度。在 ARC-AGI 2 上,使用 Muon 优化的模型在大约 600,000 个训练步骤内达到了 11.5% 的 pass@1 准确率,而 AdamAtan2 基线则需要超过 1,300,000 步才能达到相同的性能,这对应于优化速度近两倍的提升(如图5所示)。
* 最终性能: 尽管在早期训练中存在优势,但两种方法最终收敛到相似的准确率(ARC-AGI 1 上约 53.8%,ARC-AGI 2 上约 16.0%),表明它们的渐近性能相当。
分析:
* 这些结果表明 URM 中优化效率与架构能力是分离的。Muon 优化器通过近似二阶曲率来应用正交更新,以更好地处理深度循环结构引起的复杂损失景观【8, What makes looped transformers perform better than non-recursive ones (provably), 2025】,从而预处理了循环权重矩阵具有挑战性的谱特性【13, Muon is scalable for llm training, 2025】并降低了训练成本,但它并未带来最终泛化能力的提升。
A7 补充细节
5.1 ARC-AGI
先前工作的分类。先前在 ARC-AGI 基准【2, Arc prize 2024: Technical report, 2025】【3, Arc-agi-2: A new challenge for frontier ai reasoning systems, 2025】上的工作涵盖了基于视觉的构想、大型语言模型(LLM)的适配以及循环推理架构。以视觉为中心的方法,如 Vision ARC【10, Arc is a vision problem!, 2025】,将 ARC 重新表述为图像到图像的转换问题,并表明标准的视觉归纳偏置可以取得有竞争力的性能,特别是在集成和测试时扩展的情况下。基于 LLM 的方法探索了微调和测试时训练,证明了瞬时参数更新在类似 ARC 的任务上优于静态的上下文学习。除了语言和视觉模型,循环架构强调将迭代计算作为抽象的核心机制。分层推理模型(HRM)【7, Hierarchical reasoning models: Perspectives and misconceptions, 2025】【20, Hierarchical reasoning model, 2025】引入了多时间尺度的循环,并取得了优异的 ARC-AGI 结果,而随后的分析表明其增益可能主要源于循环本身而非显式的层次结构。微型递归模型(TRM)【11, Less is more: Recursive reasoning with tiny networks, 2025】进一步简化了这一范式,表明单个轻量级网络递归应用可以匹配或超过更复杂的层次设计。
5.2 通用 Transformer (循环 Transformer)
UT 的发展与特性。通用 Transformer (UT),也被称为循环 Transformer,由 Dehghani 等人【5, Universal transformers, 2019】提出,作为标准 Transformer 的扩展,增加了循环计算和自适应计算时间。后续工作表明,UTs 展现出比标准 Transformer 显著更强的多步推理能力,因为循环优化机制有助于克服多跳推理任务中的架构限制【4, The devil is in the detail: Simple tricks improve systematic generalization of transformers, 2021, EMNLP】【19, Grokked transformers are implicit reasoners: A mechanistic journey to the edge of generalization, 2024, NeurIPS】。此外,UTs 展示了改进的算法学习能力,能够更有效地建模迭代和基于规则的计算【21, Looped transformers are better at learning learning algorithms, 2024, ICLR】。通过在优化步骤中重用参数,UTs 还实现了更高的参数效率,允许在不增加模型大小的情况下进行更具表达力的计算【15, Reasoning with latent thoughts: On the power of looped transformers, 2025, ICLR】。
A5 结论
本文系统地研究了通用 Transformer 模型在复杂推理任务上性能增益的来源。广泛的消融研究表明,这些增益主要源于 Transformer 的循环归纳偏置和强大的非线性组件,而非过于复杂的架构设计。受此启发,本文提出了通用推理模型 (URM),该模型通过短卷积门控增强了非线性深度计算,并通过截断循环反向传播提高了优化稳定性。URM 在 ARC-AGI 1 和 2 上均取得了当前最佳性能。
💬 评论讨论
欢迎在这里分享您的想法和见解!