HRM-Text: Efficient Pretraining Beyond Scaling
论文链接:https://arxiv.org/pdf/2605.20613
代码链接:https://github.com/sapientinc/HRM-Text
摘要
当前大语言模型的预训练范式依赖于海量计算资源和互联网规模的原始文本,这给基础研究带来了巨大的障碍。相比之下,生物系统通过多时间尺度处理展现出高效的样本学习能力,例如额顶叶环路的功能组织。受此启发,我们提出了 HRM-Text,它用分层循环模型(HRM)取代了标准的 Transformer 模型,将计算解耦为缓慢演化的策略层和快速演化的执行层。为了稳定这种用于语言建模的深度循环,我们引入了 MagicNorm 和预热深度信用分配。此外,我们没有采用标准的原始文本预训练,而是使用任务完成目标和 PrefixLM MASK,专门针对指令-响应对进行训练。作为高效预训练的实证,一个仅使用 40B 个唯一 token 和 1500 美元预算从头开始训练的 1B 参数 HRM-Text 模型,在 MMLU 测试中取得了 60.7% 的准确率,在 ARC-C 测试中取得了 81.9% 的准确率,在 DROP 测试中取得了 82.2% 的准确率,在 GSM8K 测试中取得了 84.5% 的准确率,在 MATH 测试中取得了 56.2% 的准确率。尽管训练 token 的数量比标准基线模型少了约 100-900 倍,预计计算量也少了 96-432 倍,HRM-Text 的性能仍然能够与 2B 至 7B 参数的开源模型相媲美。这些结果表明,协同设计架构和目标可以显著降低计算性能比,使从头开始进行预训练的方法能够被更广泛的研究群体所接受。
1.Introduction

大语言模型(LLM)的显著成功目前主要源于一种单一的模式:大规模、多阶段的流水线训练,首先在互联网规模的原始文本上进行广泛的无监督预训练。这种蛮力式的扩展范式虽然无疑有效,但在数据量有限的情况下效率极低。大量的计算资源被用于预测提示性文本或与任务无关的文本,仅仅是为了构建通用表示。因此,这种极高的计算门槛在很大程度上阻碍了更广泛的研究群体进行基础预训练的探索。普遍的观点是,如果没有庞大的计算集群和数万亿个 token,探索新的架构或从头开始训练都是徒劳的。
这种对数据的蛮力渴求与人类智能形成鲜明对比,人类智能仅凭少量示例就能掌握运行规律并进行启发式搜索。在我们之前的研究中,我们引入了 Hierarchical Recurrent Model (HRM),这是一种受生物额顶叶环路功能组织启发的双时间尺度架构。HRM 将决策过程解耦为缓慢演化的策略层和快速演化的执行层,提供了一种结构性的归纳偏置,有助于避免局部停滞,并成功指导组合任务上的符号搜索。
然而,将循环架构扩展到语言建模这种开放式复杂领域会引入严重的梯度不稳定性风险。仅靠结构先验是不够的;要实现具有竞争力的开放领域性能,需要进行整体的协同设计。在本文中,我们再次证明架构和训练方法至关重要。我们探索了实现这种高效采样引擎的两个主要协同方向:
- Architectural Exploration:为了在不导致参数数量呈比例爆炸的情况下实现深度计算,我们基于 HRM 的模块化、多时间尺度递归机制。快速的 L-module 执行局部迭代细化,而慢速的 H-module 则在循环中保持稳定的语义上下文。为了使这种深度递归机制在数学上适用于语言,我们引入了诸如 MagicNorm 和预热深度信用分配等稳定化技术,这些技术在保持反向优化稳定性的同时,限制了前向激活方差。
- Objective Exploration:我们挑战了在原始文本上进行自回归预训练的传统做法。由于模型主要用于推理时的条件生成,我们直接从头开始在指令-响应对上预训练 HRM-Text。我们优化了一个任务完成目标,仅针对响应计算负对数似然损失:。我们将其与 PrefixLM 注意力掩码相结合,从而在保留响应的标准因果生成的同时,允许在指令 token 上实现完全的双向(类似编码器的)注意力。
当这两个方向结合起来时,其结果便是一个实证存在性证明,挑战了当前的扩展性教条。HRM-Text 从零开始训练,仅使用 40B 个唯一 token,在大多数基准测试中都取得了优于 Llama、Qwen、Gemma、OLMo、Ouro 和 Huginn 等当代开源模型的优异性能。值得注意的是,如图 1 和表 4 所示,它仅使用比这些基线模型少约 100-900 倍的训练 token 和 96-432 倍的估计训练计算量,就达到了类似的性能水平。
我们并不将 HRM-Text 视为最终或最优的语言模型,而是将其作为证明,证明特定的结构先验和针对性的训练目标能够显著改变计算性能比。由于入门门槛大幅降低,这种方法使基础人工智能研究更加普及。从零开始的预训练再次成为可能——我们诚邀社区成员加入我们,共同探索智能架构和聚焦目标能够带来怎样的改变。
2.Methods

HRM-Text 基于改进的 HRM 架构,采用双时间尺度循环。前向传播初始化时,使用由输入 token 嵌入生成的高级状态 和一个固定的低级状态 。核心处理序列包含两个高级循环。每个高级循环执行三次快速的 L-module 更新,然后执行一次慢速的 H-module 更新。token logits 通过对最终 H-module 状态的输出应用线性头生成。我们采用了一种预热深度信用分配策略:梯度最初仅在最后两个循环步骤中进行反向传播,随着训练的进行扩展到最后五个步骤。
在内部, 和 循环模块均采用 MagicNorm 进行结构化。此外,我们还利用了无参数 RMSNorm(省略了可学习的 参数)、SwiGLU 激活函数、旋转位置嵌入 (RoPE) 和 sigmoid 门控自注意力机制。
与基于原始文本的标准自回归预训练不同,我们优化的是一个任务完成目标。该模型直接从头开始,在指令-响应对 上进行预训练,使用仅基于响应计算的负对数似然 (NLL) 损失函数 。该目标函数自然地与 PrefixLM 注意力掩码相结合,从而实现了指令 token 之间的完全双向注意力。
以下各节将详细介绍 HRM-Text 实现极高效率的具体机制。第 2.1 节深入探讨了我们提出的新型稳定化技术,而第 2.2 节则探讨了任务完成预训练目标和 PrefixLM 掩码策略。
2.1 Scaling to language with recurrence
2.1.1 Stabilization via MagicNorm
尽管最初的 HRM 在符号任务上表现出色,但将循环架构扩展到语言建模会引入严重的梯度不稳定性风险。Transformer 的设计本身就需要在归一化层的放置上做出妥协;循环会放大这种妥协,因为同一个变换会在多个步骤中重复应用。
PostNorm。PostNorm 将归一化操作放在残差分支之外:
这有效地限制了激活方差,可以提高表达能力,但会破坏清晰的恒等路径,并可能导致更深网络中的梯度消失。
PreNorm。PreNorm 将归一化操作放在残差分支内部:
这保持了直接的恒等路径,即 ,使得梯度能够更直接地流向早期层。然而,未归一化的残差累积会导致隐藏状态方差随深度增加,这可能会导致表示崩溃或相对于 PostNorm 的性能下降。
MagicNorm。为了解决循环模型中的这种权衡,我们引入了 MagicNorm,它利用了由时间截断反向传播 (TBPTT) 引起的正向和反向计算范围之间的不对称性。
令 表示循环前向步骤的总数, 表示截断的后向范围,其中 。在 MagicNorm 中,每个循环模块由 个内部 PreNorm 块组成,但在其出口处以最终归一化层进行封顶:
在前向传播过程中,循环状态 会经历 次模块级归一化操作。由于这些归一化操作直接作用于主循环路径,因此它们限制了每个循环步骤结束时的激活方差。这避免了纯粹的 PreNorm 导致的方差无界增长,并赋予循环核心类似于 PostNorm 的前向稳定性。
相反,在反向传播过程中,截断的梯度范围意味着误差信号仅经过模块级归一化 次。在同一范围内,梯度还流经 个内部 PreNorm 恒等连接。由于 相对于完整的递归深度 很小,因此 MagicNorm 在优化过程中表现得更像一个稳定的 PreNorm 架构。
2.1.2 Warmup deep credit assignment

原始的 HRM 采用固定的单步梯度策略,仅在最后两个循环步骤(最后一个 和最后一个 )中进行反向传播。我们通过预热深度信用分配扩展了这种方法。该方案的灵感来源于时间课程原则:早期优化仅限于较短的信用分配路径,只有在模型达到更稳定的状态后才会引入更长的路径。这种设计也与生物学中关于时间学习的解释相一致,例如局部痕迹可以支持延迟积分分配,奖赏预测信号可以从奖赏附近的事件转移到更早的线索,以及发展性课程可以通过在学习者接触长程依赖关系之前让他们接触短程结构来改善序列学习。
在操作层面上,我们会动态调整反向梯度范围 。在预训练初期,我们仅计算最后两个循环步骤的梯度(),然后线性地将范围加深至最后五个步骤()。这种逐步加深的方式使得模型能够利用更长的循环计算,同时减少初始化阶段因梯度路径过长而导致的优化问题。由于预热阶段反向传播的循环步骤少于最终设置,因此也减少了平均反向传播计算量,从而加速了早期训练。
2.2 Task-completion objective and PrefixLM
目前训练基础模型的主流范式依赖于资源密集型的多阶段流程。从 T5 模型到现代大语言模型,训练通常以广泛的无监督预训练开始,随后进行更高质量的中期训练。
在预训练阶段,模型使用互联网规模的原始语料库进行训练,以学习通用的语言表示。在训练中期(或退火)阶段,模型在高质量文本(通常是指令类数据)上进行优化。在这两个阶段,模型都针对所有 token 优化自然语言学习(NLL)目标函数。
虽然这种方法有效,但在数据和资源有限的情况下效率低下。广泛的原始文本预训练消耗了大部分计算和数据,而且大部分 token 级损失都用于预测提示性文本或与任务无关的文本。然而,在推理阶段,模型主要应用于条件生成:给定一个 query 或指令,它们必须产生适当的响应。
为了提高样本效率,HRM-Text 省略了广泛的原始文本预训练,而是完全从头开始训练指令-响应对。给定一个包含指令和响应 的示例,我们优化响应在指令条件下的 NLL:
由于模型不预测指令 token,因此其参数更新的重点在于生成准确的响应。图 3-(a) 展示了这一效果。虽然有无任务完成目标时总损失相近,但与响应部分相关的误差显著降低。
此外,这种单阶段条件目标自然地与 PrefixLM 注意力掩码相契合。由于模型无需自回归预测指令 ,我们移除了指令段的因果掩码:所有指令 token 之间双向关联,而响应序列则保持标准的因果掩码。这使得 HRM-Text 在解码器式实现中实现了类似编码器-解码器的分离。指令段首先被整合为一个完全可见的上下文,类似于编码器端的表示;而响应段则以自回归的方式生成,类似于解码器。
图 3(b) 显示,PrefixLM 能够带来更高的注意力 softmax 熵,表明其注意力覆盖了更多样化的 token 集合。图 3(c) 显示,因果注意力更偏向局部化,而 PrefixLM 注意力则更偏向全局化和多样化。综上所述,仅响应条件损失和 PrefixLM 注意力共同提高了在数据和计算资源受限情况下的样本效率。
3.Results
本文的核心问题是:在预训练预算有限的情况下,从随机初始化训练的模型能否达到有意义的开放模型性能水平?为此,我们采用小预算设计探索的方法:
- 首先,架构选择能否提高固定训练计算资源的利用率;
- 其次,目标函数和输入结构能否提高每个训练样本的产出。
- 最后,我们将 HRM-Text 与当前的完全开放模型和开放权重模型进行比较,以量化其相对于当前预训练实践的效率,并分析循环架构是否能增加有效深度。
所有模型的训练细节均在第 4 节中提供。
在这些实验中,HRM-Text 模型从头开始训练,训练数据集为第 4.1 节所述的任务格式混合数据集,仅使用 40B 个唯一 token。我们报告了单个 HRM-Text 检查点的所有性能。
3.1 Architecture efficiency under matched training compute

本研究的第一部分探讨了架构设计在多大程度上能够提升固定训练预算的利用率。我们通过比较标准 Transformer 模型、更大匹配 FLOPs 的 Transformer 模型、循环 Transformer 模型、RINS 和 HRM 模型在匹配训练计算条件下的性能来验证这一假设。
表 1 将训练浮点运算次数 (FLOPs) 匹配的循环架构(包括 HRM、循环 Transformer 和 RINS)与标准 Transformer 进行了比较。对于递归模型,递归次数列中的值表示每次前向传播的总计算量,以不使用递归时所需计算量的倍数表示。例如,H2L3 表示 2 个外部 循环,每个外部循环内有 3 个 模块循环,总共有 个 模块步数。由于每个 H 或 L 模块包含完整 HRM 循环核心一半的非嵌入参数,因此表中对应的递归次数为 。对于标准 Transformer 模型,该值为 1。
循环 Transformer 和 RINS 通常优于相同规模的 Transformer 模型,这表明循环或循环计算是一种有效的架构方向。然而,在训练浮点运算预算相同的情况下,与规模更大的 Transformer 相比,它们的优势并不那么明显。HRM 是这种架构设计空间的一个典型例子,并且在与列出的基线模型(包括规模更大的深度 Transformer)的比较中表现良好。
在循环设计中,我们进一步将 HRM 与 TRM 进行比较,以区分分层双时间尺度循环和共享参数双时间尺度循环变体。

TRM 是 HRM 的一个变体,它共享 和 模块的参数,旨在小规模符号推理问题上取得优异的性能。表 2 对 HRM 和 TRM 进行了比较。由于 TRM 在 H-L 模块之间共享参数,因此有两种方法可以近似匹配 FLOPs:保持参数总数不变并减少递归次数,或者保持递归结构不变并减少参数数量。在第一种情况下,TRM 的训练稳定性较差,这可能是由于递归次数的减少削弱了预期的迭代计算。在第二种情况下,增加递归次数可以稳定训练并提高性能,但模型的性能仍然落后于 FLOPs 匹配的 HRM。在此比较中,HRM 在 FLOPs 显著低于 TRM 的情况下,通常可以达到相当或更强的性能。
这些结果支持了小预算设计探索的第一部分:循环和循环架构可以在固定的训练计算量下提高基准测试的收益,而 HRM 是这一更广泛的架构设计空间中的一个有效点。
3.2 Task-completion objective and PrefixLM yield

本研究的第二部分探讨了训练目标和输入结构是否能够提高每个训练样本的产出。我们通过增量消融实验来验证这一点。实验首先使用一个基于完整问答对并采用因果注意力机制训练的标准 Transformer 模型,然后依次添加任务完成目标、PrefixLM 注意力机制,最后添加 HRM 架构。所有实验均以 FLOPs 为单位进行性能匹配。
如表 3 所示,任务完成目标、PrefixLM 训练和 HRM 架构均对整体性能有显著贡献。引入任务完成目标在所有基准测试中均取得了初步提升,而 PrefixLM 训练相比标准因果掩码进一步增强了这些结果。最终,从标准 Transformer 架构过渡到 HRM 架构在所有测试中均实现了性能的全面且稳定的提升。
3.3 Comparison with contemporary open models

在探索了小预算设置下的架构、目标和输入结构之后,我们将得到的 HRM-Text 检查点与使用更大预算训练的当代完全开放和开放权重模型进行了比较。
图 1 和表 4 将 HRM-Text 1B 与当代完全开放模型和开放权重模型进行了比较,包括 Llama、Qwen、Gemma、OLMo 以及循环模型 Huginn 和 Ouro。HRM-Text 在大多数基准测试中表现优异,即使在参数数量较少且预训练预算仅为 40B 唯一 token 的情况下,在 MMLU 测试中也保持了竞争力。这种模式与 HRM-Text 的定位相符:循环深度和任务完成预训练能够提升推理和任务执行能力,而广泛的事实知识覆盖则对模型规模和数据广度更为敏感。HRM-Text 达到此性能水平所需的估计训练计算量比对比的开放基线模型少 96-432 倍,训练 token 数量也少约 100-900 倍。这一比较支持了本文的核心问题,即小型、面向任务完成的预训练运行可以达到使用更大 token 和计算预算训练的开放模型的性能水平。
我们报告的扩展实验将 Transformer 模型的参数扩展到 3B,将 HRM-Text 模型的参数扩展到1B。在此范围内,结果表明,使用有限数据训练的模型可以与使用更大规模数据集(高达 36T 个 token)的当前工业级预训练方法相媲美。在更大的模型规模下证明类似的效率提升仍是未来工作的研究方向。
3.4 Effective depth analysis

我们假设 HRM 的有效性源于其循环特性,从而增加了内部有效计算量(增加了参数的利用率)。我们通过检验 HRM 是否比标准 Transformer 基线和循环 Transformer 基线展现出更大的有效深度来验证这一假设。
图 4 从两个角度展示了有效深度:(a) 相邻循环块之间差异的范数,以及 (b) 块级表征的余弦相似度。这两个指标都表明,与标准 Transformer 和其他循环模型相比,HRM 在深度方向上保持了更活跃的表征变化。
参照 Hu et al. 的方法,我们也使用 logit lens 分析来评估模型输出分布何时开始趋于稳定。我们使用模型的输出投影头解码不同层的隐藏状态,然后计算每个探测预测与最终模型分布之间的 KL 散度。如图 5 所示,标准 Transformer 和循环 Transformer 都在相对较早的层中收敛到稳定的输出分布,表明它们更深层的增量贡献较小。相比之下,HRM 在更深的层中保持了更大的 KL 值,表明其有效深度更大。

4.Training details
4.1 Dataset

我们仅使用开源数据集训练 HRM-Text 模型,这些数据集包含通用指令、重写知识、数学和符号任务、教科书练习以及从网络提取的问题。初始语料库包含约 176.5B 个 token,分布在 593.7M 份文档中。我们从中抽取 40B 个唯一 token,使总训练量达到 60B 个 token,重复次数遵循下文所述的分层抽样方案。表 5 总结了数据集的组成。
为了在推理过程中控制响应属性,我们根据目标响应风格在指令前添加特定的条件标签。我们使用四种主要条件:直接(仅答案)、cot(思维链)、synth(合成答案风格)和 noisy(格式不统一的网络爬虫文本)。如表 5 “Condition” 列所示,该方法利用条件训练,从而能够在生成时显式选择模型的输出格式。
为了将训练信号集中于最终任务完成情况,我们在训练前移除所有包含在 <think>...</think> 边界内的文本。这消除了主要由强化学习与可验证奖励 (RLVR) 训练产生的显式长 CoT 轨迹,符合我们希望 HRM-Text 依赖于其内部分层计算而非显式推理步骤的目标。
我们采用 SeqIO 式分层抽样,将每个数据集或任务视为一个独立的层,而不是从合并的语料库中均匀抽样。为了确保训练样本的均衡性并防止大规模数据集过度代表,我们对每个任务或每个数据集的文档数量设置了上限,同时对较小的数据集进行上采样。具体的抽样限制和乘数详见表 6。

4.2 Dataset Contamination
虽然我们的预训练数据来源于广泛使用的公共数据源,并且许多数据源都采取了去污措施,但考虑到预训练的规模,残留污染仍然可能存在。为了严格评估我们模型的基准性能是否人为地受益于测试样本,我们采用了一种改编自 Llama 家族的统计检验方法。
我们对所有评估基准测试(不包括少样本示例)中的问题进行分词,并与完全分词的预训练语料库进行 n-gram 匹配。样本的污染率定义为该样本中参与这些匹配 n-gram 的 token 比例。
为了确定污染是否会提升性能,我们根据污染百分比将评估数据划分为四个重叠的子集:“Clean”(< 20%)、“Not Clean”(≥ 20%)、“Not Dirty”(< 80%)和 “Dirty”(≥ 80%)。要使污染被认为具有显著影响,“干净”样本的性能必须明显低于平均水平,而“脏”样本的性能必须明显高于平均水平。对于每个大小为 的子集,我们计算经验平均性能 和检验统计量 ,其中 和 分别是大小为 的抽样分布的均值和标准差。我们得出结论,只有当所有四个子集中的 时,数据集污染才能带来统计学意义上的显著性能提升。
我们使用 和 的样本量,对 HRM-Text 0.6B 和 1B 进行了测试,测试涵盖了表 4 中所示的所有基准测试。HRM-Text 0.6B 在两种设置下均未表现出明显的污染。HRM-Text 1B 在 n = 13 的 DROP 基准测试中显示出统计学意义(如表 7 所示),但在 n = 20 时则未显示出统计学意义。尽管如此,HRM-Text 1B 在 DROP 的严格纯净子集(平均污染率为 0%,包含 5904 个样本)中仍然获得了 81.1 分,表明尽管污染带来的潜在益处有限,但其基线泛化能力仍然很强。
4.3 Architecture and optimization details
HRM-Text 1B 模型在两台 8×H100 节点上预训练耗时 46 小时,成本约为 1472 美元(假设每小时 H100 费用为 2 美元)。下文将概述模型、优化和基础设施设置。
Tokenizer。我们采用字节对编码 (BPE) 分词器,词表为 65,536,使用分词器库进行训练。
Model configuration。每个模块都是一个包含 16 层的 Transformer 模型,隐藏层大小为 1536,头部大小为 128。我们使用 4096 的上下文大小和 RoPE 位置编码,其中 。RMSNorm 的 参数设置为 。所有模型均以 bfloat16 精度进行训练,所有模型权重均使用 LeCun 正态分布进行初始化。
Optimization。我们使用 Adam-atan2 优化器,,,权重衰减率为 。学习率在 2000 步内线性预热,然后保持在 。未应用梯度裁剪。batch size 为 196,608 个 token。我们没有采用标准的学习率衰减,而是使用衰减率为 0.9999 的指数移动平均 (EMA) 来维护模型权重。我们的最终评估和公开发布的模型权重均使用此 EMA 检查点。
Infrastructure。并行化框架基于 PyTorch FSDP2。所有模型均在一次连续运行中完成训练。我们不使用中间检查点、崩溃恢复或跳过损失峰值。
5.Discussion
5.1 Toward decoupling knowledge and reasoning
我们的研究结果表明,可以尝试将事实覆盖与推理计算部分解耦。HRM-Text 仅使用 40B 个唯一 token 进行训练,并且显式面向知识的数据源仅占任务格式混合数据源的一小部分。尽管如此,该模型在 MATH 和 GSM8K 等推理密集型基准测试中表现出色,同时在 MMLU 等更广泛的知识基准测试中也保持了不俗的性能。这种模式表明,紧凑的循环模型无需像通常需要万亿 token 预训练那样记忆大量事实,即可学习有效的任务执行和推理行为。
这一观察启发了未来系统将紧凑的推理核心与事实存储分离。在这样的系统中,像 HRM-Text 这样的循环模型可以专注于计算、规划和任务执行,而事实广度则由精心整理的语料库、检索增强型存储或学习到的记忆模块提供。最近的条件记忆方法,例如 Engram,也指向了类似的方向:它们没有强制 Transformer 主干网络通过密集计算来模拟静态模式查找,而是引入了可扩展的记忆查找作为补充的稀疏性轴,从而将神经计算解放出来,用于全局上下文整合和推理。HRM-Text 目前尚未实现检索或条件记忆,但其结果表明,将小型循环推理模型与外部或学习到的知识存储相结合,是未来研究的一个有前景的方向。
5.2 Adaptive computation time (ACT)
Wang et al. 为 HRM 配备了一个自适应计算模块,该模块允许较简单的问题提前终止,从而在保持接近最优性能的同时减少计算量。为了简化设计和训练过程,我们在 HRM-Text 中没有使用该组件,但它仍然是提高性能和计算效率的一个很有前景的方向。当前的循环调度提供了额外的有效串行深度,但相对于单遍 Transformer,它也增加了推理时间。ACT 允许简单的提示或 token 在更少的循环周期后停止,同时将全部循环预算保留给更复杂的情况,从而有可能收回相当一部分推理开销。因此,我们认为 ACT 是对 HRM-Text 循环深度设计的自然补充:HRM 在需要推理时提供深度,而 ACT 可以使该深度成为条件性的,而不是固定的。
5.3 PrefixLM with inference frameworks
PrefixLM 可以在诸如 vLLM 之类的标准文本生成推理框架内运行,而无需采用截然不同的服务栈。其主要要求是在预填充期间进行自定义注意力掩码处理,以便指令 token 可以双向参与,而响应 token 保持自回归特性。
在多轮聊天中使用 PrefixLM 式的注意力模式也需要精心设计的 key-value 缓存逻辑:用户 token 需要在每个用户片段内保持完全的注意力,而助手 token 必须保留因果关系。这是一种工程上的限制,而非概念上的局限,但在生产环境中的推理系统中应该明确地加以解决。