🤖 machine learning

Latent Recurrent Transformer: Architecture Exploration, Training Strategies, and Scaling Behavior

本文介绍了潜在循环 Transformer（LRT），这是一种轻量级架构，通过在不增加推理深度的情况下将高层隐藏状态复用为循环记忆来增强自回归模型，并提出了一种交错并行训练策略以高效扩展该方法，从而在最小参数开销下实现了语言建模和上下文学习性能的提升。

原作者： Zeyi Huang, Xuehai He, LiLiang Ren, Yiping Wang, Baolin Peng, Hao Cheng, Shuohang Wang, Pengcheng He, Jianfeng Gao, Yong Jae Lee, Yelong Shen

发布于 2026-05-27

📖 1 分钟阅读☕ 轻松阅读

CC BY 4.0

原作者： Zeyi Huang, Xuehai He, LiLiang Ren, Yiping Wang, Baolin Peng, Hao Cheng, Shuohang Wang, Pengcheng He, Jianfeng Gao, Yong Jae Lee, Yelong Shen

原始论文采用 CC BY 4.0 许可（http://creativecommons.org/licenses/by/4.0/）。 ✨ 这是对下方论文的AI生成解释。它不是由作者撰写或认可的。如需技术准确性，请参阅原始论文。阅读完整免责声明

想象一下，你正在教一个机器人逐字编写故事。在标准做法（称为“自回归 Transformer"）中，机器人会查看它已经写下的单词，思考这些单词，然后选择下一个单词。一旦选定该单词，它便立即继续前进。这就像工厂的装配线：每个物品只经过机器一次快速处理，随后便完成。

本文提出了一种名为**潜在循环 Transformer（LRT）**的新思路。以下是其工作原理，辅以简单的类比：

1. 问题：装配线的“失忆症”

在标准机器人中，一旦完成对某个单词的处理，它就会忘记关于该单词的深层、高层级思考。它仅保留一个基本的“笔记”（即隐藏状态），以辅助处理下一个单词。如果机器人需要记住几个单词之前的某个复杂概念来理解当前句子，它就必须在一堆基本笔记中翻找。这就像试图仅通过查看电影票根来回忆剧情反转，而非回顾实际场景。

2. 解决方案：“智能笔记本”（LRT）

LRT 为机器人配备了一本智能笔记本。

工作原理：当机器人完成对第 1 个单词的处理时，它不会丢弃其深层思考，而是将这些思考的“高层级摘要”写入笔记本。
神奇之处：当它开始处理第 2 个单词时，会立即打开笔记本，读取来自第 1 个单词的摘要。它将此摘要作为“记忆增强”，以帮助处理第 2 个单词。
结果：机器人无需停止装配线或进行额外工作，就能从过去的自己那里获得“第二意见”。这就像一位厨师在切下一块蔬菜时，瞥一眼刚写下的关于上一块蔬菜的笔记，以确保风味特征保持一致。

3. 技巧：如何在学习时避免停滞

你可能会问：“如果机器人需要读取过去的笔记来学习，这是否意味着它必须等待过去完成才能开始未来？那会让训练变得超级缓慢！”
通常确实如此。如果你试图让机器人按步骤学习（先单词 1，再单词 2，然后单词 3），你就会失去并行处理所有任务的能力，导致训练耗时极长。

作者发明了一种巧妙的训练技巧，称为交错并行训练。这就像一场带有转折的接力赛：

第 1 步（热身）：机器人以正常速度跑完全程（整个句子），仅为了生成一份粗略的笔记草稿。
第 2 步（接力）：机器人不再重新跑完全程，而是将比赛分为两组跑者：“偶数”跑者和“奇数”跑者。
- 首先，“偶数”跑者有机会查看“奇数”跑者（已完成热身）的笔记，并据此改进自己的表现。
- 接着，“奇数”跑者查看“偶数”跑者改进后的笔记，并据此改进自己的笔记。
优势：这样，每个单词都有机会从邻居那里学习，同时计算机仍能并行执行大量工作。这就像一群学生一起学习：他们不会等待某人读完整本书；而是通过小组交换笔记来加速学习。

4. 结果：以更低的成本实现更智能

本文将这种新机器人与旧的标准机器人进行了测试。

效率：即使给予相同的“脑力”（计算时间），新机器人（LRT）也比旧机器人学会了写得更好（错误率更低），并且对上下文的理解更佳（回答问题能力更强）。
微小升级：新机器人仅需增加极少量的额外内存（约 0.3% 的体积）即可获得这些显著改进。这就像给汽车加装一个小型而强大的 GPS，而非购买全新的引擎。
最佳平衡点：他们发现，最佳的“笔记本”并非机器人最后产生的那个想法（那过于专注于下一个单词），而是来自其大脑中部的某个想法。它的高层级程度足以发挥作用，但又不过于专业化。

总结

潜在循环 Transformer是一种让 AI 模型变得更聪明的方法，它允许模型复用前一个单词的“高层级思考”来帮助处理下一个单词。他们通过一种巧妙的“接力赛”方法训练模型，从而在不减慢训练速度的情况下实现了这一目标。最终得到的模型学习更快、表现更优，且无需显著增大体积。

技术摘要：潜在循环 Transformer（LRT）

问题陈述

自回归 Transformer 是语言建模的标准，但它们对每个生成的令牌都依赖固定深度的前馈计算。虽然循环机制理论上可以通过迭代细化增加计算能力，但现有方法会引入显著的推理成本。在深度上增加循环（对同一令牌重复应用模块）或在时间上增加循环（插入暂停或思考令牌）的方法，都会通过额外的解码步骤或模块应用增加推理延迟。

作者观察到，自回归解码已经“免费”计算了高层循环信号：即前一个令牌的隐藏状态。具体而言，针对下一个令牌预测训练的上层状态包含了有用的潜在上下文。核心研究问题是：这些从前一个令牌计算出的高层表示，能否在不增加额外解码步骤的情况下，被重用作下一个令牌的循环记忆？

方法论

架构：潜在循环 Transformer（LRT）

LRT 是对标准自回归 Transformer 的一种轻量级增强，引入了跨层、跨令牌的潜在路径。

机制：在令牌位置 $t$ ，LRT 重用来自前一位置的源层隐藏状态 $m_{t-1} = h^{\ell_{src}}_{t-1}$ 作为循环记忆。其中， $\ell_{src}$ 是选定的源层（通常是中间偏上的层）。
注入：记忆 $m_{t-1}$ $m_{t - 1}$ 通过两种互补机制注入到当前令牌 $t$ $t$ 的 Transformer 层中：
1. KV 投影：将记忆投影到注意力键值空间（ $\tilde{k}_{rec}, \tilde{v}_{rec}$ ），并使用输入依赖的门控机制将其与局部键/值结合。这为前一个令牌的表示进入注意力机制提供了一条直接路径。
2. 残差注入：将记忆直接添加到残差流中（ $\bar{x}^{(\ell)}_t = \alpha_\ell x^{(\ell)}_t + \gamma_\ell m_{t-1}$ ），使整个模块暴露于循环信号之下。
推理效率：关键在于，LRT 保留了标准的仅解码器骨干网络、注意力机制和 KV 缓存接口。它不插入记忆令牌、暂停令牌或深度循环。每个生成的令牌仍然仅利用一次标准的 Transformer 前向传播，循环状态仅被更新并传递给下一步。默认变体仅增加 0.3% 的参数。

训练策略：交错并行训练

在大规模上预训练这种循环机制是一个主要挑战。精确训练需要沿整个序列顺序展开 Transformer（时间反向传播），这将破坏 Transformer 高效预训练所必需的并行性。分块训练（顺序处理连续块）保留了一定的并行性，但仅能在块边界传播记忆，仅提供粗略的近似。

作者引入了交错并行训练，以在保留并行计算的同时近似令牌级循环：

初始化传递：进行一次完整的序列前向传播，构建一个共享缓冲区，包含所有位置的键、值以及源层循环状态。
细化传递：将序列位置划分为 $S$ 个不相交的子集（例如，步长划分的子集）。这些子集并行进行细化。
写回：来自某个子集的更新状态被写回共享缓冲区。后续子集可以消耗来自先前子集的细化记忆。
目标：损失函数是初始化损失与每个子集细化损失的平均值。

该策略确保每个令牌都接收到一次感知循环记忆的监督步骤。在计算成本方面，它大约花费标准 Transformer 更新 2 倍 的成本（一次初始化 + 一次有效细化），但比块状训练更紧密地模拟了令牌级循环。

主要贡献

LRT 架构：一种轻量级循环扩展，重用来自前一个令牌的高层源层隐藏状态作为记忆，实现了在不增加额外解码步骤或改变 KV 缓存接口架构的情况下进行跨层循环。
交错并行训练：一种并行训练近似方法，在固定的训练预算下细化不相交的令牌子集以模拟令牌级循环，避免了标准 RNN 式训练所需的顺序展开。
实证验证：证明了 LRT 在各种模型规模和每参数令牌预算下改善了计算与质量的权衡，默认共享投影变体仅增加了极小的参数开销（0.3%）。

结果

实验在 nanochat 风格的 GPT 骨干网络（20 层和 24 层）上进行，这些模型在 FineWeb-Edu 数据集上进行了预训练。

扩展行为：在基线等效的训练计算量下（考虑交错训练的 2 倍成本），LRT 将扩展曲线推向更低的每字节比特数（BPB）和更高的CORE（衡量 22 项任务中上下文学习能力的指标）。
性能提升：
- 在 24 层骨干网络、基线等效计算量 80 的条件下，LRT-shared 实现了 0.695 BPB，而基线为 0.699。
- 上下文学习（CORE）也有所提升；在 20 层骨干网络、计算量 80 的条件下，LRT-shared 达到 0.274，而基线为 0.271。
消融研究结果：
- 源层：中间偏上的层（例如 20 层模型中的第 12 层）表现优于最终层，表明有用的记忆应该是高层的，但不应过度专门化于 logits。
- 时间源：紧邻的前一状态（ $m_{t-1}$ ）是保持单次前向传播约束的最强单一记忆源。
- 注入：KV 投影和残差注入是互补的；同时使用两者可获得最佳结果。
- 训练：交错并行训练显著优于块状训练，后者仅提供更粗略的循环近似。

意义与主张

该论文声称，LRT 成功改变了自回归 Transformer 的计算与质量权衡。通过重用前一个令牌的“免费”高层信号，并通过并行近似对其进行训练，LRT 实现了迭代细化和更好的上下文利用，同时避免了深度循环或暂停令牌方法带来的推理延迟惩罚。

作者强调，LRT 是一种轻量级增强，保留了标准的 Transformer 接口。其收益归功于循环记忆机制本身，而非仅仅增加额外的参数或门控，因为在同一强基线上添加通用架构组件（如门控注意力或层缩放）仅产生了微不足道的改进。该工作表明，有用的循环记忆应当是高层的，但不应是最终专门化于 logits 的层，并且令牌级循环可以在预训练期间通过并行方式有效近似。