一、研究背景与问题
循环变换器(Looped Transformer, LT)通过多次重用同一组网络层(权重共享)来模拟更深的网络,在参数固定的情况下提升模型推理能力。但其核心瓶颈在于:每一轮循环都需重新执行全注意力(Full Attention)机制,其复杂度为O(L²),导致:
训练时的注意力FLOPs随循环次数和序列长度急剧增长;
推理时的KV缓存内存占用同样呈二次增长;
长上下文场景下效率极低,难以实用扩展。
二、核心贡献:LT2 架构
作者提出LT2(Linear-Time Looped Transformer),将循环变换器中原有的二次Softmax全注意力替换为次二次(线性或稀疏)注意力机制,实现:
训练和推理复杂度接近线性(O(L));
保持甚至超越原有循环变换器的性能;
显著提升长上下文下的解码吞吐量和内存效率。
主要变体:
| 变体 | 注意力类型 | 特点 |
|---|---|---|
| LT2-linear | 线性注意力(如GDN、KDA) | 循环记忆可迭代精炼,支持DPLR状态更新 |
| LT2-sparse | 稀疏注意力(如DSA、NSA) | 循环扩展有效感受野(从窗口 w 扩展到 Tw) |
| LT2-hybrid | 混合策略(线性+稀疏/全注意) | 在深度或循环维度上混合不同注意力类型 |
三、理论分析
作者从理论上证明:
循环 × DPLR线性注意力:
单次循环仅能实现秩-1的状态更新;
多次循环后,状态转移矩阵可达到秩-T,显著提升表达能力,甚至可模拟任意正交变换(借助Cartan-Dieudonné定理)。
循环 × 稀疏注意力:
单次窗口注意力覆盖 w 个令牌;
T 次循环后,有效感受野扩展为 O(Tw),即计算换上下文。
残差连接影响:
残差会稀释远距离信息,实际有效感受野增长弱于理论,但仍显著优于单次稀疏注意。
四、实验验证
1. 语言建模(FineWeb-Edu, 0.6B/1.3B, 100B tokens)
线性/稀疏循环模型(GDN、KDA、DSA)性能接近全注意力循环模型,差距在1个百分点以内;
混合模型 Full+GDN在1.3B规模下:
平均零样本准确率提升至62.89%(基准59.27%);
困惑度(PPL)降至9.12(基准9.87);
混合模型 GDN+DSA(无全注意力)完全匹配基准性能,且完全线性时间。
2. 长上下文效率(1k~32k tokens)
全注意力循环模型在8k以上解码吞吐量急剧下降;
LT2线性/混合模型在32k上下文下仍保持平坦的高吞吐量;
线性模型能支持更大批量(bs=8),而全注意力模型在8k即OOM。
3. 消融实验
混合比例:Full:GDN = 1:4 最优;
混合模式:交错分布优于集中放置;
混合级别:深度级混合优于循环级混合;
SDPA输出门控:有效缓解注意力汇聚(Attention Sink)随循环累积的问题,提升性能。
4. 训练稳定性
缺少门控或delta规则的模型(如RetNet)训练不稳定;
GDN同时具备数据依赖门控和delta规则,训练最稳定;
稀疏注意力变体训练平稳,但性能略逊;
混合模型继承了两者的优点:性能优且稳定。
5. 合成任务(状态跟踪 + 长距离召回)
循环显著提升次二次模型的表现;
纯全注意力循环在第4阶段(n=64)卡住;
混合模型(GDN+NSA/Window)达到第5阶段(n=128),完全线性时间模型首次解决该任务。
6. 真实长上下文检索(知识召回 + 大海捞针)
循环提升所有模型的知识召回能力(+2~4个百分点);
混合模型 GDN+DSA 在无二次注意力的前提下逼近全注意力性能;
混合模型 Full+GDN 为最强配置。
五、模型蒸馏与迁移
作者展示了将预训练的全注意力循环模型(Ouro)蒸馏为LT2混合模型的有效方法:
仅需约10亿令牌的持续训练;
转换后的Ouro-Hybrid-1.4B:
优于同等规模的业界1B模型;
媲美4B级别模型;
保持线性时间推理效率;
提出了多阶段蒸馏策略(基础蒸馏 → 长上下文 → 推理强化),并强调每循环监督的重要性。
六、主要结论
循环与次二次注意力具有独特协同效应,可提升记忆精炼与感受野扩展;
混合架构(如 Full+GDN 和 GDN+DSA)是实用的帕累托最优方案,兼顾性能与效率;
权重共享 + 高效注意力使循环模型成为参数高效、可扩展的语言建模新方向;
现有全注意力循环模型可通过蒸馏低成本转换为线性时间模型,具备实际部署价值。
七、局限性
未探索完全的循环级混合(不同循环使用不同注意力家族);
未设计显式的跨循环状态传递机制,未来可进一步优化长上下文建模与计算效率。
这里是自己的论文阅读记录,感兴趣的话可以参考一下,如果需要阅读原文的话可以看这里,如下所示:
项目地址在这里,如下所示:
预训练模型发布地址在这里,如下所示:
循环变换器(LT)通过在解码最终令牌前多次迭代其层,已成为一种强大的架构。然而,它们与全注意力的结合保留了二次复杂度,使得计算成本高昂且速度缓慢。我们引入了LT2(线性时间循环变换器),这是一系列循环架构,用具有线性时间复杂度的次二次注意力替代了二次softmax注意力。我们研究了两种变体:具有线性注意力的LT2-linear和具有稀疏注意力的LT2-sparse。我们发现循环与这些变体具有独特的协同作用:它能在线性注意力中实现迭代记忆精炼,并在稀疏注意力中逐步扩展有效感受野。我们从理论上阐述了这些优势,并在受控的召回、状态跟踪和语言建模任务中展示了一致的实证提升。接着,我们探索了LT2-hybrid,这是一种在循环设置中结合不同注意力变体的混合架构。我们发现了两种有前景的架构变体:(1) LT2-hybrid (GDN+DSA),它交错使用线性和稀疏注意力以最大化效率,在完全线性时间成本下达到了标准循环变换器的质量;(2) LT2-hybrid (Full+GDN),它将GDN与一小部分全注意力层交错以最大化质量,在性能和效率上都超越了标准循环变换器。此外,我们还展示了如何将预训练的LT转换为LT2-hybrid模型。仅需约10亿令牌的训练,我们的转换模型(Ouro-hybrid-1.4B)便超越了业界1B级模型,并与业界4B级模型具有竞争力,同时保持了线性时间注意力的速度优势。这两方面共同展示了使循环变换器成为语言建模更可扩展架构,并推动高效、高性能小语言模型发展的清晰路径。
图1:(左)LT2引入的新参数效率前沿。(右)转换后的LT2-Hybrid优于同等规模的业界1B模型,同时媲美4B模型。
1. 引言
沿着参数轴扩展神经语言模型推动了现代NLP的大部分进展[7, 34, 29]。一个互补的轴——通过权重共享递归扩展深度——最近已成为一种有前景的替代方案。这些架构,通常被称为循环变换器(LT,最初为通用变换器[16]),在解码最终预测令牌之前,跨多个步骤重用相同的权重[23, 69, 76]。实际上,重复计算变成了有效深度:模型在保持独特参数数量固定的同时执行多轮潜在计算,使得循环变换器成为一种实现参数高效推理的有吸引力的方法。
然而,当前的循环变换器扩展性较差,因为每个循环都必须重复地将二次全注意力应用于整个序列。因此,其成本和推理时的存储随序列长度增长,并随每次循环迭代而叠加。结果是,尽管参数被重用,训练时的注意力FLOPs和推理时的KV缓存使用量随循环次数的增加而扩展性不佳,使得注意力成为扩展循环变换器的主要瓶颈[62, 76]。如图2所示,对每个令牌进行T次注意力处理导致训练时的注意力FLOPs和推理时的KV缓存内存都大幅增长。在长上下文中,二次注意力项占主导地位,增加循环步骤很快变得不切实际[76]。
图2:1.3B模型的注意力FLOPs和推理缓存内存与序列长度的关系。
我们引入了LT2(线性时间循环变换器),这是一系列循环架构,用次二次令牌混合原语替代了二次softmax注意力。我们主要研究两种不同的变体,LT2-linear和LT2-sparse,它们分别用线性注意力[35, 72, 63]和稀疏注意力[68, 15]替代了二次注意力。我们展示了循环操作可以将计算转化为上下文:它能够对线性注意力中的循环记忆进行更精细的控制,并扩大稀疏注意力中的感受野;我们在§2.2提供直观理解,在附录B.1提供详细的理论分析。此外,我们探索了LT2-hybrid,这是一种通过混合循环设置中不同注意力变体,将性能-效率前沿推向新水平的混合架构。我们证明,LT2-hybrid (GDN [70] + DSA [15])——在循环设置中结合了线性和稀疏注意力——达到了标准循环变换器的质量(平均零样本准确率59.3%),同时在8k上下文、批大小为8时提供约5.7倍更高的解码吞吐量(125 vs. 22 tokens/s),完全没有二次注意力。LT2-hybrid (Full + GDN),它将GDN与一小部分全注意力层交错,更进一步:在相同设置下,其平均零样本性能比标准循环变换器提高了+2.1个百分点(61.4% vs. 59.3%),同时仍实现约5倍更高的解码吞吐量,并在语言建模、召回、状态跟踪和效率基准上持续优于标准循环变换器(§3)。
最后,我们探索将预训练的循环变换器(具体为Ouro [76])蒸馏为LT2模型。如图1(右)所示,仅需约10亿令牌的持续训练,我们转换后的Ouro-Hybrid-1.4B保留了其全注意力教师模型的质量,同时继承了LT2的线性时间效率。由此产生的模型在标准零样本基准上,与业界开源1B-4B参数范围的模型相比具有竞争力,达到或超过了1B级基线,并在多项任务上接近3B-4B模型。这表明实践者无需从头开始重新训练:现有的循环变换器可以高效地转换为线性时间变体,从而降低了采用LT2系列模型的成本壁垒。
表1:LT2支持的令牌混合器。蓝色突出显示门控/留存,赭色突出显示DPLR风格操作。训练FLOPs按每层、序列长度L报告;缓存/状态内存是推理时每层的。w表示稀疏注意力的窗口/预算大小,且w << L。
2. LT2:线性时间循环变换器
2.1. 架构
循环变换器 (LT)。设L表示序列长度,d表示隐藏维度;我们将隐藏状态序列记为h ∈ R^(L×d),位置t处的状态记为h_t ∈ R^d。深度为N的标准Transformer堆叠了N个独立参数化的块{F_ℓ}_{ℓ=1}^N,每个块由一个令牌混合器和一个带残差连接的位置前馈网络(FFN)组成:
F_ℓ(h) = h' + FFN_ℓ(h'), h' = h + MHA_ℓ(h), (1)
其中MHA_ℓ是多头自注意力(为简洁起见,我们省略了前置归一化)。循环变换器(LT)将这N个共享块重用T次迭代:
h^(0) = Emb(x), h^(τ) = (F_N ∘ … ∘ F_1)(h^(τ-1)), τ = 1,…,T, ŷ = Dec(h^(T)), (2)
得到有效深度T·N,但仅有N个独特参数集——与同等深度的Transformer相比,参数减少了T倍。遵循Ouro [76],我们在整个预训练过程中使用固定的T,并在附录A中讨论自适应计算时间。
LT2。LT2简单地将公式(1)中的MHA子层替换为次二次令牌混合器,因此每个共享块变为:
F_ℓ(h) = h' + FFN_ℓ(h'), h' = h + LinearMixer_ℓ(h), (3)
其中LinearMixer_ℓ是表1中的任何线性或稀疏注意力原语。在整个过程中,q_ℓ, k_ℓ ∈ R^(d_k) 和 v_ℓ ∈ R^(d_v) 表示h_ℓ的查询/键/值投影;S_ℓ ∈ R^(d_k × d_v) 是线性注意力混合器的循环状态。我们另外插入了一个零初始化、每通道学习的门控ρ_τ ∈ R^d,作为跨循环迭代的残差,h^(τ) = \tilde{h}^(τ) + ρ_τ ⊙ h^(τ-1),其中\tilde{h}^(τ) 是迭代τ时循环块堆栈的输出(即 \tilde{h}^(τ) = (F_N ∘ … ∘ F_1)(h^(τ-1)))。因此,我们的设置包括两个层次的残差连接:传统的每块恒等残差连接和学习的每循环残差。
2.2. 超越效率:循环的优势
次二次注意力提供了明显的效率提升。一个更有趣的问题是循环为这些注意力变体增加了什么。我们提出两个主张:通过T次循环迭代,一个对角加低秩(DPLR)线性注意力块将其秩-1状态更新转变为秩-T更新,而一个滑动窗口块将其大小为w的窗口转变为大小为Tw的有效感受野。
循环 × DPLR线性注意力:循环记忆上的秩-T更新。前沿的线性注意力架构现在使用DPLR混合器,例如GDN [70],KDA [63]和RWKV7 [47]。我们以KDA作为运行示例,它通过以下方式在序列位置t维护一个循环状态S_t ∈ R^(d_k × d_v):
S_t = A_t S_{t-1} + β_t k_t v_t^T, A_t = Diag(α_t)(I - β_t k_t k_t^T), (4)
其中α_t ∈ [0, 1]^(d_k) 是一个对角门控,因此A_t是单位矩阵(I)加上一个秩-1 (k_t k_t^T) 扰动。先前的工作表明,单个这样的块每个令牌只能模拟两个元素的排列,并且无法在有限精度下解决S_n单词问题(n ≥ 3)[27]。当相同的共享块被循环T次时,每个循环迭代τ ∈ {1, …, T} 都会在位置t处对循环状态贡献一个新的DPLR因子A_t^(τ),因此每令牌累积状态转移算子变为:
A_t^eff = ∏{τ=1}^T A_t^(τ) = ∏{τ=1}^T Diag(α_t^(τ)) (I - β_t^(τ) k_t^(τ) k_t^(τ)T)。 (5)
DeltaProduct [57] 表明,循环DPLR的表达能力增益取决于循环特定键 {k_t^(τ)}_{τ=1}^T 之间的关系。在一种极端情况下,如果所有键都相同,那么每次循环迭代都会擦除循环记忆中相同的方向,与无循环情况相比没有表达能力增益。在另一种极端情况下,如果来自不同循环迭代的键是正交的,那么循环会沿着T个不同方向擦除历史信息。在这种情况下,包含单个秩-1扰动的原始转移被替换为具有秩-T记忆擦除子空间的有效转移。我们在附录B.1中详细讨论了证明。
循环 × 稀疏注意力:感受野扩展。一个窗口大小为w的单个滑动窗口块允许位置t的每个查询只关注最后w个令牌,
I_t^(1) = {t - w + 1, …, t},
因此每循环的感受野是O(w),超出w个令牌的任何信息都是不可见的。循环该块会重新在序列上应用相同的窗口,信息随着每次循环迭代而传播得更远:在循环迭代τ时,位置t关注循环-(τ-1)状态的窗口,而这些状态已经在循环迭代τ-2时吸收了来自它们自己窗口的信息,依此类推。归纳地链接这个论证(附录B.2)得出T次循环迭代后的感受野:
I_t^(T) ⊇ {max(1, t - Tw + 1), …, t}, |I_t^(T)| = O(Tw)。 (6)
换句话说,窗口为w的块的T次循环达到了与T层堆叠的窗口-w注意力[9]一样远的范围,但参数却少了T倍。因此,循环将计算转化为上下文:一旦T适中,一个小的固定窗口就足以覆盖长序列,这使得稀疏混合器成为长上下文设置中循环的自然伙伴。
2.3. 混合LT2:跨深度和循环混合混合器
我们进一步探索循环设置中的混合架构。混合模型中的常见做法是将线性块与全注意力块交错,以实现强大的语言建模性能,同时恢复召回能力[41, 43, 64, 45]。我们表明,循环变换器为这种混合开辟了第二个轴:除了沿深度变化混合器外,我们还可以跨循环迭代变化混合器。我们在以下第3.3节中探索这两种选项(图3)。
图3:混合LT2的两种方式。(a) 深度级混合在共享块内的线性层之间交错全注意力层。(b) 循环级混合在循环迭代间变化混合器,例如,先是全注意力循环,然后是窗口缩小的滑动窗口循环 (256 → 128)。
3. 实验
我们围绕四个问题组织主要实验。首先,我们测试LT2在标准语言建模规模(§3.1)和现实长上下文检索(§3.7)下是否具有竞争力。然后,我们对混合设计选择进行消融研究:混合应用于何处、混合器如何沿深度排列,以及使用何种混合器比例(§3.3)。接着,我们研究SDPA输出门控,它减轻了循环处理下的注意力汇聚积累(§3.4)。此外,我们还进行了涵盖合成召回/状态跟踪、长上下文效率和训练稳定性的实验(§3.5, 3.2, 和 3.6)。
3.1. 语言建模
我们在FineWeb-Edu [46]上以0.6B和1.3B参数规模预训练所有模型,令牌预算为100B,每个循环变体使用T = 4次循环。混合比例对于(Full:Linear)和(GDN:DSA)变体均为1:4。完整设置见附录C。表2总结了结果。我们在第3.2节提供详细的效率比较。
次二次混合器几乎匹配全注意力循环。循环GDN、KDA和DSA在两个规模上的平均性能均在循环Transformer参考模型的一个百分点左右,同时避免了二次复杂度。在较小的0.6B规模下,循环GDN仍略落后于循环Transformer;然而,在较大的1.3B规模下,它在保持线性时间复杂度的同时超越了循环Transformer。在循环设置中,我们发现门控和DPLR线性注意力都很重要,门控似乎比DPLR风格更新公式的作用更大。相比之下,循环纯DeltaNet变体在我们的研究中稳定性较差,这最终限制了其性能。
线性-稀疏混合循环以极低成本匹配全注意力循环。循环混合模型 (GDN+DSA) 不包含全注意力,在两个规模上都匹配了全注意力参考模型(1.3B时PPL为9.72 vs. 9.87)。它还提供了最大的效率加速:在32k上下文下,解码吞吐量提升了2.9倍。我们认为这是一个有趣的混合设置,其中线性注意力有助于全局压缩,而稀疏注意力有助于精确的KV位置选择。
表2:FineWeb-Edu上两个规模的零样本下游性能,K = 4次循环。D-Gate = 数据依赖门控;Δ = DPLR线性变体。奶油色高亮表示没有全注意力的最佳LT2模型。每个规模每列最佳值用粗体表示,次佳值用下划线表示。令牌预算相对于Chinchilla计算最优缩放[29]。
循环混合模型 (Full+GDN) 将帕累托前沿推向新高度。这是总体最强的配置,在两个规模上都提高了通用语言建模性能(1.3B时为61.39 vs. 59.27),在更难的任务上增益最大。由于只有一小部分层是二次的,它仍然提供了×2.7的解码加速。这两种混合模型共同构成了新的帕累托前沿:一种以接近线性的成本匹配全注意力,另一种则超越它,同时明显快于全注意力循环模型。
3.2. 长上下文下的效率
我们测量了四种循环LT2候选模型在1k到32k令牌、批大小{1, 2, 4, 8}下的预填充和解码吞吐量,使用单个H100 (80 GB) GPU,对softmax注意力使用FlashAttention-2 [13],对GDN使用融合分块内核。所有变体在匹配的参数计数下使用T = 4。空心方块表示每种配置在内存溢出前的最后一个长度。
图4:长上下文下跨批大小的效率。行是批大小 (1, 2, 4, 8);列是 (a) 预填充和 (b) 解码吞吐量与序列长度的关系。空心方块标记每种配置在耗尽80 GB HBM之前达到的最后一个序列长度。循环GDN和混合LT2 (GDN+Full) 是唯一在所有批大小下都能达到32k,并在整个范围内保持解码吞吐量平坦的变体;混合LT2 (GDN+DSA) 由于其top-k KV读取而紧随其后。
线性时间混合器消除了长上下文解码悬崖。循环Transformer在4k到32k之间损失了一半以上的解码吞吐量,因为其KV缓存每循环迭代都在增长。循环GDN、混合LT2 (GDN+Full) 和混合LT2 (GDN+DSA) 在整个范围内和所有批大小下都保持了平坦的解码速率:在bs=1,32k时,它们的解码速度比LT快约3倍;在bs=8时,它们达到了32k,而LT在8k时就已内存溢出。这种优势随批大小而累积,因为每个额外的批次元素都会增加一个固定大小的GDN状态,但会增加一个长度成比例的KV缓存。
表3:混合LT2消融实验。1.3B / T = 4 / 100B FineWeb-Edu令牌。奶油色高亮显示每组中的最佳行。平均值是表2中八个任务的平均值。
线性时间混合器也扩展了内存溢出边界。随着批次增长,循环Transformer的内存溢出越来越早(bs=4:上限接近16k,bs=8:上限接近8k),而循环GDN在所有批大小下都能达到32k,混合LT2 (GDN+Full) 在bs=8时也能达到32k。混合LT2 (GDN+DSA) 介于两者之间,因为DSA仍然为top-k选择维护KV缓存,但每个查询只读取其中的一小部分。在实践中,这决定了是在有用的批大小下服务长上下文,还是无法做到。
3.3. 消融实验:混合比例、模式和混合级别
第3.1节中的混合LT2同时固定了三个设计选择:循环中有多少注意力,它沿深度位于何处,以及在何种级别混合混合器。在本节中,我们对此进行了仔细的消融研究(表3)。
比例:清晰的倒U型,最优在1:4。在循环Transformer (1:0) 和循环GDN (0:1) 之间扫描Full:GDN比例,内部呈现清晰的倒U型,1:4最优。过多的注意力挤占了第3.5节中记录的正则化效果;过少则使循环缺乏精确检索能力。1:4是仍然能够恢复完整检索质量的最小注意力量,与标准混合Transformer基线[38]匹配。
模式:分散优于集中。在固定的1:4比例下,bookend(Full在顶部和底部,GDN在中间)略微优于均匀交错,暗示在输入编码和最终读出时都使用注意力有小幅益处。将注意力层集中在一端——前端加载或后端加载——平均准确率损失超过0.7个百分点。或许结论是,任何沿深度的合理分散都远优于任何集中。
级别:跨迭代异质性没有帮助。我们尝试了三种循环级调度来替代深度级混合:粗到细 (Full→SWA-512→SWA-256→SWA-128),细到粗(相反),以及一个随机基线(每步重新采样1:4深度级混合,在评估时对K=5个样本进行多数投票)。粗到细在PPL上胜出,但在下游任务上失败,它在最终迭代中过度拟合局部统计。细到粗则相反。随机投票总体最好,但推理计算量增加5倍,与简单的固定交错相比难以证明其合理性。
图5:循环Transformer (T = 4, 24层) 的展开诊断。x轴贯穿展开的计算,虚线标记循环边界。(a) 首个令牌注意力质量形成锯齿状并随每个循环加剧——汇聚点被重新注入而非重置。(b) 最大FFN残差激活遵循相同的复合模式(对数刻度)。(c) 残差流RMS范数随层内深度和跨循环迭代增长。SDPA输出门控使(a)/(b)平坦化,并显著缓解——但未消除——(c)中的跨循环增长。
3.4. 消融实验:注意力汇聚和SDPA输出门控
权重共享循环的一个自然担忧是,底层注意力块(特别是注意力汇聚[68])中的病态,即一小部分令牌吸收不成比例的softmax质量[58],可能会在循环迭代中复合:相同的softmax块被应用于已经携带前一轮汇聚点的残差流。门控注意力[51]表明,在缩放点积注意力(SDPA)之后添加一个头特定的sigmoid门控可以消除标准Transformer中的汇聚点。我们对循环模型提出同样的问题并采用相同的修复方法,将其应用于循环块内部,使得W_θ在每次迭代中被重用。我们将此门控添加到三个LT2变体中,在参数计数匹配的情况下保持FFN宽度,并在1.3B / T = 4设置下在100B FineWeb-Edu令牌上重新训练。
汇聚点是真实存在的,并且会跨循环复合。图5沿轨迹(循环1,层1)→(循环1,层24)→(循环2,层1)→… 展开了循环Transformer。首个令牌的注意力质量呈现锯齿状,并随着循环而加剧:在循环t中学到的汇聚点被重新注入循环t+1而不是被重置,因此每个后续迭代开始时已经偏向汇聚点。最大残差激活遵循相同的复合模式,与已有文献[58]一致。残差流RMS范数沿层内深度和跨循环迭代增长,到第4次循环结束时达到约20倍。
表4:SDPA输出门控对三个包含softmax的LT2变体的影响。1.3B / T = 4 / 100B令牌。我们报告表2中八个零样本基准的平均值。
3.5. 训练稳定性
循环模型的一个实际问题是,重复应用相同的块可能会放大激活并破坏优化稳定性。我们跟踪了整个预训练过程中的语言建模损失和全局梯度范数,发现循环内部混合器的选择对稳定性有显著影响。
门控和delta规则保持线性循环有界。图6比较了循环Transformer与四种次二次混合器。循环RetNet在训练过程中表现出持续较大的梯度范数和频繁的尖峰,与其在表2中的发散一致。循环DeltaNet和循环Mamba2明显更好,但仍然显示出偶尔的尖峰,这些尖峰传播到损失中。相比之下,循环GDN在整个运行过程中梯度范数低于循环Transformer,并产生了所有变体中最平滑的损失曲线。有两个因素似乎很重要:
图6:循环GDN在所有线性和全注意力变体中训练损失最平滑,梯度范数最小;同时缺乏数据依赖门控和delta规则的循环RetNet则发散。
数据依赖门控,它让循环忘记陈旧状态而不是让它在迭代中积累;以及delta规则,它限制了循环记忆的更新。只具备两者之一的混合器(Mamba2有门控但没有delta规则;DeltaNet有delta规则但门控较弱)是稳定的,但噪音明显比GDN大,而两者都没有的RetNet则不稳定。
稀疏注意力稳定但能力略逊。图7报告了稀疏注意力循环的相同诊断。所有三种稀疏变体(Window,NSA,DSA)训练平稳:它们的梯度范数在整个运行过程中处于或低于循环Transformer,并且没有出现全注意力循环在训练中期偶尔出现的尖锐尖峰。代价是在语言建模损失上存在一个小但持续的差距,将每次迭代限制在稀疏感受野限制了每循环计算,并相对于密集注意力减慢了收敛速度。在稀疏选择中,循环DSA是最强的,这就是我们在本文其余部分将其采纳为LT2的稀疏组件的原因。
图7:稀疏循环变体训练时没有全注意力循环中出现的尖峰,但最终损失略高于循环Transformer。
混合混合器结合了稳定性和能力。图8显示了两种混合配置继承了两者的优点。循环混合模型 (GDN+DSA) 和循环混合模型 (Full+GDN) 从训练一开始就在损失上跟踪循环Transformer,并在结束时略微领先,同时它们的梯度范数始终保持较小,并且没有全注意力循环偶尔产生的尖峰。因此,将循环混合器与稀疏或密集注意力配对似乎可以正则化循环:线性分支在迭代中保持梯度范数有界,而注意力分支提供了纯线性模型所缺乏的精确检索。
图8:两种混合变体在损失上匹配或超越循环Transformer,同时在训练过程中产生更小、更平滑的梯度范数。
在所有三个比较中都出现了相同的图景。具有数据依赖门控和delta规则的混合器(GDN,以及包含它的混合模型)在循环下训练得比普通全注意力更稳定,而稀疏注意力虽然表达能力较弱,但从未导致不稳定。这激发了本文其余部分使用的两种LT2实例化:当稳定性是优先事项时使用LT2-sparse(带有DSA的循环混合模型),当能力是优先事项时使用LT2-linear(带有GDN的循环混合模型)。
3.6. 合成任务:状态跟踪 + 召回
状态跟踪和长距离检索通常被视为对立的压力测试:状态跟踪有利于循环深度,而检索有利于精确关注长历史。我们用来探测LT2的第一个合成实验将这两种压力置于同一任务上。我们遵循Olmo-hybrid [43]的基于状态的召回构建。
该任务结合了长距离召回(从约Θ(m)个令牌前获取bits[.],对压缩RNN来说很难)和状态跟踪(将运行的置换序列应用于指针,对TC^0中的固定深度Transformer来说很难)。我们将n = m,并沿着课程 {8, 16, 32, 64, 128, 256} 一起增长它们;一旦在100k步预算内评估准确率达到0.90,模型就晋级。标题指标是n_max,即解决的最大n = m。所有模型共享一个4层、256宽、4头骨干网络(RoPE);循环变体在此骨干网络的T次迭代中共享权重,T=1是标准非循环模型。我们使用AdamW训练(峰值学习率3×10^-4,批次32)。
循环的效果。图9报告了每个(架构,T)解决的最高阶段。显著的模式是循环对次二次混合器的帮助大于对全注意力的帮助。循环Transformer和循环Full+Window在第4阶段(n_max = 64)达到平台期,并且在任何T下都未达到第5阶段。相比之下,三种次二次变体——循环NSA、循环GDN+Window和循环GDN+NSA——都达到了第5阶段(n_max = 128);循环Full+GDN也做到了,但仅因为其一半的块已经是线性时间的GDN。循环GDN+Window最为显著:在T ≤ 4时为第3阶段,在T = 8时达到第5阶段。
跨架构比较。参考点是循环Transformer。相对于它,几种次二次混合器在这个联合任务上既获得了表达能力又提高了召回能力。循环Transformer在所有T下都止步于第4阶段,但循环NSA、循环GDN+Window和循环GDN+NSA都达到了第5阶段——在相同参数预算下,n_max比全局注意力基线翻倍。其中,循环GDN+NSA和GDN+Window是解决该问题的完全线性时间模型,这与上面讨论的出色语言建模性能一致。
图9:关于循环次数T解决的最高课程阶段s。纯混合器在白线之上,混合模型在下方。
3.7. 现实召回和长上下文检索
我们现在转向现实的长上下文召回,其中模型必须从远长于循环状态所能舒适容纳的自然文本中检索特定事实。我们遵循Mamba-3 [38]的评估协议。所有模型都在1.3B规模。纯模型堆叠单个混合器家族,而混合模型则以固定的4:1比例将该混合器与全注意力交错,循环变体在匹配的参数计数下,在相应的非循环骨干网络的T = 4次迭代中共享权重。我们评估两个互补套件:在2048个令牌上的知识型召回(SWDE [1], SQuAD [52], FDA [1], TriviaQA [33], Natural Questions [37] 和 DROP [17]),以及在1024、2048和4096个令牌上的精确大海捞针检索(NIAHSingle-1/2/3)[30],其中4096列信息量最大,因为所有模型都在2048上进行了预训练,必须外推。
三个发现突出(表5)。循环在匹配的参数计数下持续改善底层混合器:循环Transformer、循环GDN和循环Mamba-2在知识型套件上的平均得分分别比其非循环对应物提高了大约2-4个百分点,任务级别的波动在任一方向(某些列如TQA已经饱和,而FDA和NQ受益更多于额外迭代),并且它们保留了其基础版本的定性NIAH行为——循环骨干优雅地外推到4096,而密集注意力骨干则不然。
循环混合模型 (GDN + DSA) 在知识套件上紧密跟踪循环Transformer,尽管不包含二次组件,并且在NIAH-4096上外推效果显著更好。循环混合模型 (Full + GDN) 是总体最强的配置,优于循环Transformer
表5:1.3B参数的长上下文评估。知识型基准 (SWDE-DROP) 在2048个令牌下评估;NIAH-Single-1/2/3 在1024、2048和4096个令牌下评估(模型在2048上预训练,必须外推到4096)。纯模型使用单个混合器;混合模型以1:1的比例与全注意力交错。循环变体在相应的非循环骨干网络的T = 4次迭代中共享权重。粗体标记每列最佳结果;下划线标记次佳。
图10:蒸馏后Ouro-Hybrid-1.4B变体的能力保留。Ouro-Hybrid (均匀) 以固定的交错模式交错线性和全注意力层,而Ouro-Hybrid (先前SoTA) 使用先前最佳方法[40]选择保留全注意力的层。
L_KD = Σ_{τ=1}^T w_t^(τ) KL( σ_top-k( o_t^(τ) / T_kd ) || σ_top-k( s_t^(τ) / T_kd ) ), (7)
其中σ_top-k 在教师的前k个令牌上重新归一化softmax,w_t ∈ Δ^(T-1) 控制步骤t时每个循环接收多少监督。我们逐步预热循环级监督,在训练步骤的一半时间内设置均匀权重以平等监督每个循环,然后切换到仅最终输出监督(6亿令牌,长度4096)。
阶段3(长上下文延续)。然后,我们将阶段2扩展到在35K OpenThoughts-v3 [28]上的延续阶段,序列长度为32k,使用与公式(7)相同的KD损失和恒定的学习率(6亿令牌,长度32768)。
4.2. 结果
将预训练的全注意力模型蒸馏为线性时间变体仍然不简单。我们发现,通用和数学基准性能的主要驱动力是蒸馏数据的组成。阶段1和2利用通用数据集如DCLM [39],在常识下游任务上取得了高分。在阶段3整合推理特定数据如OpenThoughts显著缩小了数学推理上的差距。从算法的角度来看,有两个因素很重要。如图11所示,渐进式长度扩展对于保持长上下文性能至关重要。此外,每循环监督提供了比仅监督最终循环更稳定的梯度信号。我们向希望蒸馏循环架构的研究人员推荐这种多阶段方案。
图11:不同蒸馏模型的Ruler子任务性能。关键任务差异在于多键检索,这从每循环监督中获益更多。
最后,我们将蒸馏后的模型与业界小语言模型进行比较(图1),展示了极具竞争力的性能。我们将向公众发布完整的模型检查点,以促进对高效、高能力小模型的进一步研究。
5. 相关工作
循环变换器及通往可扩展递归的路径。通用变换器[16]通过捆绑每一层的权重[5]并将同一块迭代固定或自适应步数,将深度维度的递归重新引入Transformer。早期的后续工作表明,这种简单的归纳偏置改善了组合基准上的系统泛化[11],并且更广泛地说,跨层的参数共享是一种可行的设计选择,而非特例[60]。在理论方面,循环架构在温和假设下是图灵完备的[49],并且可以被编程以模拟迭代算法,如多步梯度下降[23, 69, 21, 20, 18],这已被形式化为循环的“潜在思维”观点,即每次传递精炼一个内部计算[53, 25, 6, 8]。这种观点推动了最近一波以递归为中心推理模型的发展,包括HRM [66]、TRM [32]、通用推理模型[19]以及大规模训练的循环语言模型[76]。然而,核心障碍是可扩展性:将同一块循环T次会将计算量乘以T而不增加参数,因此天真的通用变换器在匹配FLOPs下性能不如标准Transformer[62, 48]。有三项工作直接针对这一效率差距。第一种是将稀疏性和专家混合注入共享块,使得容量增长无需按比例增加计算量,如稀疏通用变换器[61]、MoEUT [12]、通用专家混合[10]和参数高效的FFN重用[44]。第二种是用低秩增量放松严格的权重绑定,使得每次迭代可以廉价地特化[3]。第三种,最符合自适应计算时间的精神,是为每个令牌分配可变数量的递归步骤:递归混合学习动态的每令牌深度,在一个令牌级路由框架中[4],而弹性和深度递归变体将此思想扩展到视觉和注意力感知的潜在推理[26, 36, 74, 56]。最近的工作进一步通过并行采样加速了递归深度模型的推理[22]。我们的工作延续了这一轨迹,专注于如何使循环计算在固定的计算预算下可预测地扩展。
次二次注意力。并行研究路线用成本为线性或接近线性的序列混合器替代softmax注意力。线性注意力[35]将注意力表达为核特征图,并将推理重新表述为循环状态更新,这被解释为快速权重编程,并可追溯到早期关于联想快速权重的工作[55, 31, 2]。从这个基础出发,涌现了一系列高效递归,包括RetNet [59]、门控线性注意力[71]、HGRN2 [50]、DeltaNet及其并行和门控变体[72, 70]、Mamba-2 [14]及其后继Mamba-3 [38]的SSM-注意力对偶性,以及RWKV-7 [47]。最近的工作还通过负特征值[27]和Householder乘积[57]解决了这些递归有限的状态跟踪能力。由于纯线性注意力模型在需要精确召回的任务上仍落后于softmax注意力[1],一个互补的路线是在混合堆栈中交错线性和softmax层,如Jamba [41]、Kimi Linear [63]和Olmo Hybrid [43],或将预训练的Transformer蒸馏为混合或线性后继者[24, 40]。第三条路线保留softmax注意力但强制稀疏性以减少其二次成本,范围从用于流式推理的注意力汇聚[68]到原生可训练的块稀疏模式[75, 9, 15]。这些方法在很大程度上与深度维度的递归正交:它们降低了单次前向传递的成本,而循环则在多次传递中重用参数。将两者结合是一个自然的方向,也是我们在这项工作中探索的方向。
6. 结论
我们提出了LT2,一个线性时间循环变换器家族,用线性、稀疏和混合注意力机制取代了循环架构中的二次令牌混合瓶颈。特别是,我们的混合变体恢复或超过了全注意力循环变换器的质量,同时显著提高了推理效率。这些结果表明,高效的令牌混合器可以使递归深度成为未来语言模型的一个实用扩展轴。
局限性。有两个方向尚未探索。首先,我们研究了深度级混合和简单的循环级调度,但没有研究完全的循环级混合,其中不同的迭代可能使用不同的注意力家族,而不仅仅是改变掩码。其次,我们没有设计显式的跨循环循环状态传递机制;原则性的跨循环状态共享可能进一步改善长上下文建模、内存重用和计算效率。