← 返回 2026-05-21

HRM-Text:超越规模化的高效预训练 HRM-Text: Efficient Pretraining Beyond Scaling

Guan Wang, Changling Liu, Chenyu Wang, Cai Zhou, Yuhao Sun, Yifei Wu, Shuai Zhen, Luca Scimeca, Yasin Abbasi Yadkori 📅 2026-05-20 👍 323 2026-07-13 08:36
PrefixLM 任务完成目标 小算力训练 层级推理 循环神经网络 高效预训练

1B 循环模型仅用 40B token 预训练,算力省百倍即达 7B 水平

前置知识

PreNorm 与 PostNorm

残差 Transformer 中归一化层(Norm)的放置位置决定训练稳定性。PreNorm 把 Norm 放在残差分支内部,保留恒等路径、梯度顺畅但深层激活方差会累积;PostNorm 放在残差分支外,能严格约束方差但容易梯度消失。这是网络加深时必须权衡的核心设计点。

MagicNorm 正是同时利用这两种范式的优点,理解该权衡是读懂方法的关键。

循环深度 Transformer 与 BPTT

循环(looped/recurrent)Transformer 在多个步骤上复用同一组参数重复计算,不增加参数量却能提升有效深度。训练时梯度需沿时间反向传播(BPTT),但深循环里雅可比矩阵连乘会让梯度方差爆炸,因此常用截断反向传播(TBPTT),只在最后若干步回传梯度。

HRM-Text 的 MagicNorm 与 warmup 深度信用分配都建立在 TBPTT 截断这一前提之上。

PrefixLM 掩码

PrefixLM 是介于 encoder-decoder 与 decoder-only 之间的注意力掩码:对提示词(prefix)段用双向注意力,让模型像编码器那样充分理解指令;对响应段保持因果掩码以自回归生成。因为不需要预测指令,去掉指令段因果约束能让模型更高效地编码上下文。

它是 HRM-Text 提升每条训练样本产出的关键工具,与任务完成目标天然配套。

任务完成目标 (Task-completion Objective)

传统自回归预训练对所有 token 算 NLL 损失,含提示词本身。任务完成目标只在响应 token 上算条件 NLL,把全部梯度花在生成正确答案而非重建已知输入。HRM-Text 把它当作从头预训练的唯一目标,直接跳过大规模原始文本预训练阶段。

理解它才明白为何 40B token 就够:模型不浪费算力去预测无关文本。

HRM 层级循环模型

HRM 受大脑额顶环启发,用两种时间尺度的循环模块:慢速 H 模块维护稳定的语义策略,快速 L 模块做局部迭代细化。典型配置先执行多次 L 再做一次 H(如 H2L3 表示 2 个 H 外循环、每循环 3 个 L 步)。这种结构先验能在符号搜索等组合任务上避免局部停滞。

HRM-Text 正是把 HRM 从符号任务扩展到开放语言建模的核心架构基础。

研究动机

当前大语言模型的预训练范式高度依赖规模法则(scaling laws),需要从互联网抓取数万亿 token 的原始文本进行无监督自回归训练,再叠加高质量 mid-training 与后训练阶段。这种暴力堆数据的做法在数据受限场景下极不高效:模型把大量算力花在预测类似提示词的文本或与下游任务无关的文本上,仅仅是为了构建通用表征。由此带来的算力门槛(动辄上千张 H100、数百万美元)实际上把广大学术团队、小实验室和个人研究者排除在「从头预训练」的探索之外。主流观点甚至默认没有海量算力和万亿 token 就根本无法验证新架构或新训练方法。例如 Olmo3 7B 用了约 6T token 和 $252\times10^{21}$ FLOPs,Qwen3.5 2B 用了 36T token 和 $432\times10^{21}$ FLOPs,这些数字对小型研究团队几乎不可企及。

本文的目标是本文要证明:在一个极小的预算(约 1500 美元、16 张 H100、1.9 天)和仅 40B 唯一 token 的条件下,仍可从随机初始化训练出一个 1B 参数的模型,使其在主流基准(MMLU、ARC-C、DROP、GSM8K、MATH 等)上达到与 2–7B 开源模型相当的水平。更具体地,作者希望把训练 token 数相对标准基线减少 100–900 倍、估计算力减少 96–432 倍,从而把「从头预训练」重新变成普通研究者可参与的实验。HRM-Text 并不追求成为最优模型,而是作为一个存在性证明(existence proof):特定的结构先验与定向训练目标可以从根本上改变算力-性能比。

与已有工作不同的是,绝大多数效率改进工作要么只优化数据配比,要么只压缩已有 Transformer,默认接受「原始文本自回归预训练 + 后续指令微调」的两阶段范式。本文的切入点在于同时改造两个维度:架构上不再用单次前向的标准 Transformer,而是用受生物多时间尺度启发的层级循环结构(HRM)增加有效深度而不爆炸式增长参数;目标函数上直接抛弃原始文本预训练,从第一天起就在指令-响应对上用任务完成损失训练。两个方向协同——循环带来的深度需要 MagicNorm 和 warmup 深度信用分配才能在语言规模上稳定,而 PrefixLM 加上响应-only 损失让每条样本的产出最大化。这种「架构加目标的联合设计」正是被规模化主流忽视的角度。

核心方法

可以把 HRM-Text 想象成一个会「反复思考」的小型大脑:它不是一次读完全文就吐答案,而是把输入编码成一个慢变的策略状态(H 模块),再用一个快变的执行状态(L 模块)反复迭代细化,每轮迭代都在做内部推理。技术上,HRM-Text 把标准 Transformer 替换为层级循环模型:前向传播由输入 token embedding 初始化高层状态 $z_0^H$,低层状态 $z_0^L$ 固定;核心处理是 2 个高层循环,每个循环先执行 3 次快速 L 模块更新,再做 1 次慢速 H 模块更新(记为 H2L3,共 8 步 H/L),最后用线性头从最终 H 状态输出 token logits。为让这种深循环在语言建模里稳定,作者引入 MagicNorm(PreNorm 加末端 Norm 的混合)和 warmup 深度信用分配(梯度回传步数从 $K=2$ 线性升到 $K=5$)。训练上完全抛弃原始文本,直接在指令-响应对上最小化响应的 NLL,并用 PrefixLM 掩码让指令段双向注意、响应段因果生成。

全文最关键的创新是 MagicNorm 对「前向-反向计算视界不对称」的利用。在截断反向传播(TBPTT)下,前向要走过全部 $N$ 步循环,而反向梯度只回传 $K$ 步($K \ll N$)。MagicNorm 在每个循环模块内部用 $L$ 个 PreNorm 子层(保留恒等路径利于梯度),但在模块出口再加一个 Norm(约束每步激活方差)。于是前向看,每个循环步都被模块级 Norm 约束,方差不会无限增长,获得 PostNorm 式的前向稳定性;反向看,由于截断,误差信号只穿过模块级 Norm 共 $K$ 次,同时仍能利用模块内部的 PreNorm 恒等连接,所以优化时更像稳定的 PreNorm。这种「前向像 PostNorm、反向像 PreNorm」的设计是让深循环在语言规模上首次稳定训练的钥匙。与之配套的「直接在指令-响应对上从头预训练加 PrefixLM」从目标侧把每条样本的信息产出最大化,二者协同才达成 40B token 的极致效率。

方法步骤详情

训练流程分五步。(1) 数据构造:从 FLAN、Tasksource、SYNTH、AceReason、OpenThoughts2、DMMath、Sudoku 等 18 个开源数据集组成约 176.5B token 初始语料,用分层采样控制每数据集上限、小数据集上采样,最终采样出 40B 唯一 token、总训练量 60B token 的混合,并剥离所有 ... 思维链标记。(2) 前向:token embedding 初始化 $z_0^H$,$z_0^L$ 固定,执行 H2L3 循环(2 轮,每轮 3 次 L 加 1 次 H),每个 H/L 模块是 16 层 MagicNorm Transformer(hidden 1536、head 128、RoPE、SwiGLU、sigmoid 门控自注意力)。(3) 损失:仅对响应 token 算 $-\log P(x_a\mid x_q)$。(4) 反向:warmup 深度信用分配,$K$ 从 2 线性升到 5。(5) 优化:Adam-atan2、LR $2.2\times10^{-4}$ 恒定、batch 196608 token、EMA 0.9999、bfloat16、FSDP2 单连续训练、无梯度裁剪。

技术新颖性

新颖性集中在三点。第一,MagicNorm 是首个显式利用 TBPTT 前向/反向视界不对称的归一化方案:纯 PreNorm 反向稳但前向方差爆炸,纯 PostNorm 前向稳但反向梯度消失,MagicNorm 在截断设定下同时拿到两者优点,这是循环深度模型走向语言规模的关键缺失环节,附录 C 还给出了梯度对数散布由 H 循环深度主导的机制证据。第二,warmup 深度信用分配把「信用分配路径长度」当作课程来调度,既受生物时序学习启发,也实测能减少早期长梯度路径的优化病态(图6/图7 显示完整 BPTT 会出现罕见但巨大的梯度尖峰)。第三,把任务完成损失加 PrefixLM 当作从头预训练的唯一阶段(而非微调阶段),与原始文本自回归预训练的工业范式正面对抗。相较 Huginn、Ouro 等仍依赖大 token 预算或多阶段训练的循环模型,HRM-Text 是首个在小算力下端到端联合设计架构与目标的循环语言模型。

HRM-Text architecture: dual-timescale recurrence, MagicNorm, gated attention, PrefixLM mask
Figure 2: HRM-Text architecture: dual-timescale recurrence, MagicNorm, gated attention, PrefixLM mask
Task-completion and PrefixLM improve response modeling
Figure 3: Task-completion and PrefixLM improve response modeling

实验结果

HRM-Text 1B 在 40B 唯一 token、约 $1.0\times10^{21}$ FLOPs(2×8 H100、46 小时、约 1472 美元)下取得 MMLU 60.7%、ARC-C 81.9%、DROP 82.2%、GSM8K 84.5%、MATH 56.2%。同等算力对照(表1):在 $1.0\times10^{21}$ FLOPs 匹配下,HRM 1B 几乎全面击败 Looped Transformer 1B(MMLU 56.5、MATH 48.3)、RINS 1B(MMLU 56.1)及更大的 Transformer 3B Deep(MMLU 56.7)/ Wide(54.5)。消融(表3)显示从「Transformer+$P(x)$+Causal」到「HRM+$P(x_a|x_q)$+PrefixLM」,MMLU 从 40.6 升到 60.7、GSM8K 从 48.4 升到 84.5,三个组件各自显著。对比当代开源模型(表4):HRM-Text 在 DROP、GSM8K、MATH 上甚至超过算力大 96–432 倍的 Gemma3 4B、Llama3.2 3B、Qwen3.5 2B;MMLU 60.7 接近 Olmo3 7B 的 65.8。有效深度分析显示 HRM 相邻层 $\|\Delta h\|_2$ 均值 34.12,明显高于 Looped Transformer(29.65)与 RINS(25.42),深层仍有贡献。

Performance comparison across architectures (Transformer vs HRM), objectives and attention
Table 3: Performance comparison across architectures (Transformer vs HRM), objectives and attention
Evaluation results of HRM-Text 1B vs contemporary fully open / open-weight models
Table 4: Evaluation results of HRM-Text 1B vs contemporary fully open / open-weight models
Pretraining efficiency: FLOPs and tokens vs benchmark average
Figure 1: Pretraining efficiency: FLOPs and tokens vs benchmark average
Effective depth analysis: block-wise Δh norm and cosine similarity
Figure 4: Effective depth analysis: block-wise Δh norm and cosine similarity
查看结构化数据
任务指标本文基线提升
MMLU(综合知识) accuracy (%) 60.7 Olmo3 7B 65.8(252×FLOPs)/ Qwen3.5 2B 64.5(432×FLOPs) 用约 96–252 倍更少算力达到 7B 模型约 92% 的水平
ARC-C(科学推理) accuracy (%) 81.9 Olmo3 7B 81.6 / Qwen3.5 2B 81.0 以 1B 参数、40B token 持平甚至略超 7B/2B 模型
GSM8K(小学数学) accuracy (%) 84.5 Llama3.2 3B 77.7 / Gemma3 4B 38.4 +6.8 / +46.1 个百分点,循环推理优势明显
MATH(竞赛数学) accuracy (%) 56.2 Llama3.2 3B 48.0 / Olmo3 7B 40.0 +8.2 / +16.2 个百分点
DROP(离散推理) F1/EM (%) 82.2 Qwen3.5 2B 30.8 / Llama3.2 3B 45.2 大幅领先,差距达 37–51 个百分点

局限与改进

作者承认三点局限。第一,推理时因循环展开,单次前向计算量约为单次 Transformer 的 4 倍(H2L3),文中明确把这视为代价并寄希望于自适应计算时间(ACT)来缓解。第二,在 MMLU 这类广覆盖知识基准上仍落后于 7B 级模型(Olmo3 7B 65.8 vs 60.7),因为事实知识对模型规模和数据广度更敏感,40B token 不足以覆盖。第三,DROP 数据集在 n=13 的污染检测上达到统计显著(clean 子集 81.1 vs dirty 85.5,$|Z_k|>2$),虽然 n=20 不显著且 clean 子集仍很强,但作者仍坦诚报告。我自己的观察:论文未提供中等到大规模(>3B)的扩展证据,1B 到更大规模是否保持效率优势未验证;PrefixLM 在多轮对话的 KV-cache 处理需要定制工程,作者也承认;基线对比中部分模型(如 Qwen3.5 2B 在 DROP 仅 30.8)分数异常偏低,让人怀疑评测设置是否完全公平,缺乏统一的 lm-eval-harness 复跑说明。

独立分析的弱点

第一个弱点是推理算力开销:H2L3 配置使前向计算约为单次 Transformer 的 4 倍,部署成本不容忽视,改进方向是引入 ACT,让简单问题提前停止循环(auto-guidance 中 $w<0$ 表现更好正是可早停信号)。第二个弱点是知识覆盖天花板:仅 40B token 且知识源(SYNTH 改写维基)只占 21.7B,MMLU 难追万亿 token 模型,改进方向是把循环推理核与外部检索/可学习记忆(如 Engram 条件记忆)解耦,让小模型专注推理、知识外挂。第三个弱点是扩展性证据不足:论文最大只到 HRM 1B / Transformer 3B,scaling law 在更大尺度是否保持未知,建议补做 3B–7B 扩展实验。第四个弱点是评测公平性:部分基线(Qwen3.5 DROP 仅 30.8)偏低,缺统一复跑协议,建议用同一 lm-eval-harness 加同一 few-shot 设置复跑所有基线。第五个弱点是污染披露不完整:仅 DROP 触发显著,其它基准的 $Z$ 统计未披露,建议逐基准补充污染分析。

未来方向

作者明确点名的方向有三:(1) 自适应计算时间(ACT),让循环深度按问题难度动态调整,回收推理开销并支持测试时扩展;(2) 知识与推理解耦,把 HRM-Text 作为紧凑推理核,搭配检索增强或条件记忆(如 Engram)供给事实;(3) PrefixLM 推理工程化,尤其是多轮对话中用户/助手段的 KV-cache 逻辑。基于成果可延伸的方向包括:把 MagicNorm 加 warmup 深度信用分配迁移到 Huginn、Ouro 等其它循环深度模型验证通用性;在更大规模(3B–7B)补全扩展曲线;探索循环深度与思维链/RLVR 的关系——论文已剥离 标记,是否让循环隐式承担了 CoT 角色值得专门研究;把 auto-guidance 推广为可微的自适应深度门控,实现训练-推理一致的动态深度。

复现评估

复现门槛在同类论文里属于最低一档。代码已在 github.com/sapientinc/HRM-Text 开源;训练数据全部来自公开数据集(FLAN、Tasksource、SYNTH、OpenMathInstruct2、NuminaMath、AceReason、OpenThoughts2 等),表5/表6 给出了完整数据配比与分层采样上限;算力仅需 2 个 8×H100 节点、46 小时、约 1472 美元(按 2 美元/H100·小时计),单次连续训练、无中间检查点、无崩溃恢复;关键超参(LR $2.2\times10^{-4}$、batch 196608、EMA 0.9999、Adam-atan2、bfloat16、LeCun 初始化、无梯度裁剪)全部公开,模型配置(每模块 16 层、hidden 1536、head 128、vocab 65536、context 4096)也完整。主要复现难点在于 MagicNorm 和 warmup 深度信用分配的实现细节,以及 HRM 的 FSDP2 并行;论文未提供中间检查点或训练日志,但从零复现一个 1B 模型对中型团队完全可行。