分层循环模型HRM-Text:用极少量计算颠覆Transformer预训练范式
大语言模型的预训练正越来越依赖少数拥有充足算力和数据资源的机构。训练一个有竞争力的基础模型,往往需要数万亿token、数千张GPU,甚至上千万美元的算力投入。然而,当前的训练模式并不高效,大量计算都消耗在了提示词、格式填充等无关内容上。
清华团队提出的分层循环模型HRM-Text重新设计了架构和训练目标。该模型采用双时间尺度的分层结构,把计算拆分为慢速的H模块和快速的L模块。标准Transformer对每个token只做一次前向传播,HRM则会在同一token上进行多轮递归更新,在不增加参数量的前提下提高了计算深度。训练目标也不再沿用标准的全文自回归预训练,而是直接在指令-回答对上训练,只对回答部分计算损失,并配合PrefixLM掩码,让指令部分双向注意、回答部分按因果掩码生成。
为了提升递归训练的稳定性,研究团队引入了MagicNorm和Warmup Deep Credit Assignment。MagicNorm是一种混合归一化策略,利用截断反向传播下前向与反向计算深度的不对称性,在模块内部采用PreNorm并在模块出口额外加入归一化,从而提升深层递归训练的稳定性。Warmup Deep Credit Assignment则在训练初期仅对最后2个递归步骤回传梯度,随后线性扩展至最后5步,让模型在较短的信用路径上稳定收敛,再逐步引入更长的依赖关系。
实验结果表明,在仅使用比标准模型少约100至900倍的训练token、96至432倍的估计计算量的情况下,HRM-Text依然实现了可媲美2B至7B参数开源模型的性能表现。使用1B参数、40B非重复token,并以约1500美元的训练成本,该模型在MMLU、ARC-C、GSM8K等基准测试上分别取得了60.7%、81.9%、84.5%的成绩。