LT2: Linear-Time Looped Transformers——线性时间循环变换器-北京尧图网络

一、研究背景与问题循环变换器Looped Transformer, LT通过多次重用同一组网络层权重共享来模拟更深的网络在参数固定的情况下提升模型推理能力。但其核心瓶颈在于每一轮循环都需重新执行全注意力Full Attention机制其复杂度为O(L²)导致训练时的注意力FLOPs随循环次数和序列长度急剧增长推理时的KV缓存内存占用同样呈二次增长长上下文场景下效率极低难以实用扩展。二、核心贡献LT2 架构作者提出LT2Linear-Time Looped Transformer将循环变换器中原有的二次Softmax全注意力替换为次二次线性或稀疏注意力机制实现训练和推理复杂度接近线性O(L)保持甚至超越原有循环变换器的性能显著提升长上下文下的解码吞吐量和内存效率。主要变体变体注意力类型特点LT2-linear线性注意力如GDN、KDA循环记忆可迭代精炼支持DPLR状态更新LT2-sparse稀疏注意力如DSA、NSA循环扩展有效感受野从窗口 w 扩展到 TwLT2-hybrid混合策略线性稀疏/全注意在深度或循环维度上混合不同注意力类型三、理论分析作者从理论上证明循环 × DPLR线性注意力单次循环仅能实现秩-1的状态更新多次循环后状态转移矩阵可达到秩-T显著提升表达能力甚至可模拟任意正交变换借助Cartan-Dieudonné定理。循环 × 稀疏注意力单次窗口注意力覆盖 w 个令牌T 次循环后有效感受野扩展为 O(Tw)即计算换上下文。残差连接影响残差会稀释远距离信息实际有效感受野增长弱于理论但仍显著优于单次稀疏注意。四、实验验证1. 语言建模FineWeb-Edu, 0.6B/1.3B, 100B tokens线性/稀疏循环模型GDN、KDA、DSA性能接近全注意力循环模型差距在1个百分点以内混合模型 FullGDN在1.3B规模下平均零样本准确率提升至62.89%基准59.27%困惑度PPL降至9.12基准9.87混合模型 GDNDSA无全注意力完全匹配基准性能且完全线性时间。2. 长上下文效率1k~32k tokens全注意力循环模型在8k以上解码吞吐量急剧下降LT2线性/混合模型在32k上下文下仍保持平坦的高吞吐量线性模型能支持更大批量bs8而全注意力模型在8k即OOM。3. 消融实验混合比例Full:GDN 1:4 最优混合模式交错分布优于集中放置混合级别深度级混合优于循环级混合SDPA输出门控有效缓解注意力汇聚Attention Sink随循环累积的问题提升性能。4. 训练稳定性缺少门控或delta规则的模型如RetNet训练不稳定GDN同时具备数据依赖门控和delta规则训练最稳定稀疏注意力变体训练平稳但性能略逊混合模型继承了两者的优点性能优且稳定。5. 合成任务状态跟踪长距离召回循环显著提升次二次模型的表现纯全注意力循环在第4阶段n64卡住混合模型GDNNSA/Window达到第5阶段n128完全线性时间模型首次解决该任务。6. 真实长上下文检索知识召回大海捞针循环提升所有模型的知识召回能力2~4个百分点混合模型 GDNDSA 在无二次注意力的前提下逼近全注意力性能混合模型 FullGDN 为最强配置。五、模型蒸馏与迁移作者展示了将预训练的全注意力循环模型Ouro蒸馏为LT2混合模型的有效方法仅需约10亿令牌的持续训练转换后的Ouro-Hybrid-1.4B优于同等规模的业界1B模型媲美4B级别模型保持线性时间推理效率提出了多阶段蒸馏策略基础蒸馏 → 长上下文 → 推理强化并强调每循环监督的重要性。六、主要结论循环与次二次注意力具有独特协同效应可提升记忆精炼与感受野扩展混合架构如 FullGDN 和 GDNDSA是实用的帕累托最优方案兼顾性能与效率权重共享高效注意力使循环模型成为参数高效、可扩展的语言建模新方向现有全注意力循环模型可通过蒸馏低成本转换为线性时间模型具备实际部署价值。七、局限性未探索完全的循环级混合不同循环使用不同注意力家族未设计显式的跨循环状态传递机制未来可进一步优化长上下文建模与计算效率。这里是自己的论文阅读记录感兴趣的话可以参考一下如果需要阅读原文的话可以看这里如下所示项目地址在这里如下所示预训练模型发布地址在这里如下所示循环变换器LT通过在解码最终令牌前多次迭代其层已成为一种强大的架构。然而它们与全注意力的结合保留了二次复杂度使得计算成本高昂且速度缓慢。我们引入了LT2线性时间循环变换器这是一系列循环架构用具有线性时间复杂度的次二次注意力替代了二次softmax注意力。我们研究了两种变体具有线性注意力的LT2-linear和具有稀疏注意力的LT2-sparse。我们发现循环与这些变体具有独特的协同作用它能在线性注意力中实现迭代记忆精炼并在稀疏注意力中逐步扩展有效感受野。我们从理论上阐述了这些优势并在受控的召回、状态跟踪和语言建模任务中展示了一致的实证提升。接着我们探索了LT2-hybrid这是一种在循环设置中结合不同注意力变体的混合架构。我们发现了两种有前景的架构变体(1) LT2-hybrid (GDNDSA)它交错使用线性和稀疏注意力以最大化效率在完全线性时间成本下达到了标准循环变换器的质量(2) LT2-hybrid (FullGDN)它将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]。图21.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系列模型的成本壁垒。表1LT2支持的令牌混合器。蓝色突出显示门控/留存赭色突出显示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变体在我们的研究中稳定性较差这最终限制了其性能。线性-稀疏混合循环以极低成本匹配全注意力循环。循环混合模型 (GDNDSA) 不包含全注意力在两个规模上都匹配了全注意力参考模型1.3B时PPL为9.72 vs. 9.87。它还提供了最大的效率加速在32k上下文下解码吞吐量提升了2.9倍。我们认为这是一个有趣的混合设置其中线性注意力有助于全局压缩而稀疏注意力有助于精确的KV位置选择。表2FineWeb-Edu上两个规模的零样本下游性能K 4次循环。D-Gate 数据依赖门控Δ DPLR线性变体。奶油色高亮表示没有全注意力的最佳LT2模型。每个规模每列最佳值用粗体表示次佳值用下划线表示。令牌预算相对于Chinchilla计算最优缩放[29]。循环混合模型 (FullGDN) 将帕累托前沿推向新高度。这是总体最强的配置在两个规模上都提高了通用语言建模性能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 (GDNFull) 是唯一在所有批大小下都能达到32k并在整个范围内保持解码吞吐量平坦的变体混合LT2 (GDNDSA) 由于其top-k KV读取而紧随其后。线性时间混合器消除了长上下文解码悬崖。循环Transformer在4k到32k之间损失了一半以上的解码吞吐量因为其KV缓存每循环迭代都在增长。循环GDN、混合LT2 (GDNFull) 和混合LT2 (GDNDSA) 在整个范围内和所有批大小下都保持了平坦的解码速率在bs132k时它们的解码速度比LT快约3倍在bs8时它们达到了32k而LT在8k时就已内存溢出。这种优势随批大小而累积因为每个额外的批次元素都会增加一个固定大小的GDN状态但会增加一个长度成比例的KV缓存。表3混合LT2消融实验。1.3B / T 4 / 100B FineWeb-Edu令牌。奶油色高亮显示每组中的最佳行。平均值是表2中八个任务的平均值。线性时间混合器也扩展了内存溢出边界。随着批次增长循环Transformer的内存溢出越来越早bs4上限接近16kbs8上限接近8k而循环GDN在所有批大小下都能达到32k混合LT2 (GDNFull) 在bs8时也能达到32k。混合LT2 (GDNDSA) 介于两者之间因为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比例下bookendFull在顶部和底部GDN在中间略微优于均匀交错暗示在输入编码和最终读出时都使用注意力有小幅益处。将注意力层集中在一端——前端加载或后端加载——平均准确率损失超过0.7个百分点。或许结论是任何沿深度的合理分散都远优于任何集中。级别跨迭代异质性没有帮助。我们尝试了三种循环级调度来替代深度级混合粗到细 (Full→SWA-512→SWA-256→SWA-128)细到粗相反以及一个随机基线每步重新采样1:4深度级混合在评估时对K5个样本进行多数投票。粗到细在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中学到的汇聚点被重新注入循环t1而不是被重置因此每个后续迭代开始时已经偏向汇聚点。最大残差激活遵循相同的复合模式与已有文献[58]一致。残差流RMS范数沿层内深度和跨循环迭代增长到第4次循环结束时达到约20倍。表4SDPA输出门控对三个包含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报告了稀疏注意力循环的相同诊断。所有三种稀疏变体WindowNSADSA训练平稳它们的梯度范数在整个运行过程中处于或低于循环Transformer并且没有出现全注意力循环在训练中期偶尔出现的尖锐尖峰。代价是在语言建模损失上存在一个小但持续的差距将每次迭代限制在稀疏感受野限制了每循环计算并相对于密集注意力减慢了收敛速度。在稀疏选择中循环DSA是最强的这就是我们在本文其余部分将其采纳为LT2的稀疏组件的原因。图7稀疏循环变体训练时没有全注意力循环中出现的尖峰但最终损失略高于循环Transformer。混合混合器结合了稳定性和能力。图8显示了两种混合配置继承了两者的优点。循环混合模型 (GDNDSA) 和循环混合模型 (FullGDN) 从训练一开始就在损失上跟踪循环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次迭代中共享权重T1是标准非循环模型。我们使用AdamW训练峰值学习率3×10^-4批次32。循环的效果。图9报告了每个架构T解决的最高阶段。显著的模式是循环对次二次混合器的帮助大于对全注意力的帮助。循环Transformer和循环FullWindow在第4阶段n_max 64达到平台期并且在任何T下都未达到第5阶段。相比之下三种次二次变体——循环NSA、循环GDNWindow和循环GDNNSA——都达到了第5阶段n_max 128循环FullGDN也做到了但仅因为其一半的块已经是线性时间的GDN。循环GDNWindow最为显著在T ≤ 4时为第3阶段在T 8时达到第5阶段。跨架构比较。参考点是循环Transformer。相对于它几种次二次混合器在这个联合任务上既获得了表达能力又提高了召回能力。循环Transformer在所有T下都止步于第4阶段但循环NSA、循环GDNWindow和循环GDNNSA都达到了第5阶段——在相同参数预算下n_max比全局注意力基线翻倍。其中循环GDNNSA和GDNWindow是解决该问题的完全线性时间模型这与上面讨论的出色语言建模性能一致。图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表51.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个令牌上重新归一化softmaxw_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一个线性时间循环变换器家族用线性、稀疏和混合注意力机制取代了循环架构中的二次令牌混合瓶颈。特别是我们的混合变体恢复或超过了全注意力循环变换器的质量同时显著提高了推理效率。这些结果表明高效的令牌混合器可以使递归深度成为未来语言模型的一个实用扩展轴。局限性。有两个方向尚未探索。首先我们研究了深度级混合和简单的循环级调度但没有研究完全的循环级混合其中不同的迭代可能使用不同的注意力家族而不仅仅是改变掩码。其次我们没有设计显式的跨循环循环状态传递机制原则性的跨循环状态共享可能进一步改善长上下文建模、内存重用和计算效率。

LT2: Linear-Time Looped Transformers——线性时间循环变换器

相关推荐

零基础搞定抖音直播间弹幕抓取：三步跑通DouyinLiveWebFetcher

技术写作的知识整理方法

免费的QQ空间说说备份工具GetQzonehistory：扫码免密，六张表把十年历史完整搬回本地

更多资讯

完全二叉树、非完全二叉树与哈希表的实现与应用

Python中文词云7案例：从能出图到能控图

Conda环境清理全攻略：释放磁盘空间与提升性能

阶梯碳交易与电制氢综合能源系统24小时调度建模

冒泡排序与稀疏数组：算法优化与高效存储技术

拆开AI黑箱：从梯度下降到SHAP的模型可解释性实践

今日更新

PLFM_RADAR：从被动告警到主动感知的微服务健康监控系统

SPI隔离设计实战：MAX14431BAEE+与STM32F407工业级应用

数字隔离器如何根治SPI通信地环路干扰：MAX14431与TM4C129实战

本周精选

面向复杂光照与角度的宠物面部识别鲁棒性增强策略

基于DRV8818与STM32的双极步进电机驱动方案实战解析

DRV8818PWPR+PIC18F45K80双极步进电机控制方案实战

本月热门

面向复杂光照与角度的宠物面部识别鲁棒性增强策略

基于DRV8818与STM32的双极步进电机驱动方案实战解析

DRV8818PWPR+PIC18F45K80双极步进电机控制方案实战

读完文章，想聊聊您的网站？