循环语言模型中循环计算的分配

# 循环语言模型中的循环计算分配
来源：https://arxiv.org/html/2608.18230
Yiyang Guo11footnotemark:1Rui\-jie ZhuHao YeJason Eshraghian致谢：通讯作者。单位：加州大学圣克鲁兹分校

###### 摘要

循环语言模型通过重复应用共享计算来提升推理和知识处理能力。现有系统通常重复整个层堆叠，尽管混合器和密集前馈网络执行不同的操作且成本不同。我们提出一个更具体的问题：*应该循环什么？*我们将循环视为状态更新的反复组合，并论证当一次应用能暴露新的跨位置影响方向，且该方向在任务输出端仍可观测时，该应用才具有价值。迭代传输秩描述累积影响轨迹；边际迭代传输秩描述连续应用贡献的非冗余影响。这一视角催生了MixerLoop，它对每个Gated DeltaNet混合器进行循环，而密集FFN仅执行一次。我们在相同数据、初始化和架构下，比较了15M和110M参数规模下MixerLoop与无循环和全循环方案。有限上下文关闭干预测试了后续混合器应用是否在最终语言模型输出端产生显著、非冗余且有益的变化。MixerLoop在15M参数下整体CORE评估超越FullLoop，在110M参数下保留其41.5%的CORE提升，同时将循环主干投影计算量减少45.9%。这些结果表明，循环深度的效益可以在不重复执行密集FFN的情况下得以保留。

## 1引言

权重共享循环在不添加唯一参数的情况下扩展语言模型的计算深度。从通用Transformer到近期如Ouro和LT2等模型，循环Transformer在最终词元预测前重复应用相同层，使固定参数执行多轮潜在计算。现有多数模型仍在块级定义循环，在每个循环步重复词元混合器和前馈网络。这种设计证明了潜在计算的作用，但未说明哪个操作符有帮助、如何在块内分配循环计算，或为何某操作符值得重复。SGT选择性循环高熵softmax注意力头，同时仅执行一次FFN，但其标准与单步softmax注意力分布绑定，无法解释LT2通过重复应用有状态线性混合器获得的增益。因此，操作符选择性循环的一般性解释仍然缺失：哪个组件应获得更多循环步，其边际贡献如何在循环深度上被衡量？

我们用MixerLoop回答这个问题，它将循环词元混合与重复前馈计算分离。MixerLoop反复应用Gated DeltaNet混合器，而密集FFN仅执行一次，并在相同唯一参数量、词元预算、数据顺序和优化设置下与匹配的NoLoop和FullLoop基线比较。为表征连续的混合器应用，我们引入迭代传输秩和边际迭代传输秩，衡量累积和新增的跨词元影响多样性。有限上下文关闭干预进一步测试后续混合器应用是否产生显著、非冗余、有益的最终输出变化。在15M和110M模型中，仅混合器循环持续优于NoLoop，在15M参数下整体CORE超越FullLoop，在110M参数下给出中间能力-计算权衡，并达到最高1.52× FullLoop预填充吞吐量。

（图注）图1：动机与架构概览。循环反复组合状态更新，有用的应用暴露的跨位置效应在最终输出端仍可观测。MixerLoop在每个物理层内重复共享GDN混合器，而密集FFN仅执行一次。
## 2相关工作

#### 循环Transformer与循环深度。

通用Transformer通过重复应用权重共享块（固定或自适应步数）引入深度维度Transformer循环。后续工作表明，这种归纳偏置可通过重复潜在计算改善系统性泛化并实现迭代算法。此观点现已扩展到大规模语言建模：Huginn添加可测试时增加计算量的循环核心；Ouro重复应用共享主干并在更大训练规模下展现良好参数扩展性；LT2用线性和稀疏词元混合器替代二次注意力；MELT减少循环解码内存消耗。其他模型用专家、注意力混合、多态层或部分参数特化修改循环块。额外方法在词元、层或计算预算上分配可变步数，包括SkipLayer、LoopFormer、Think-at-Hard和稳定性正则化循环。这些工作扩展了模型循环的时机、位置和时长，但多数仍以完整Transformer块作为循环单元，未解决块内分配问题。

#### 操作符选择性循环。

SGT是块边界下循环的最接近先例。它识别高熵softmax注意力头，进行重复，并在循环注意力后执行一次FFN。这表明循环增益不需要每步所有子层，但熵标准不能直接扩展到循环状态混合器。循环稀疏FFN是互补方案：Sparse Layers和LoopMoE表明当连续通道激活不同专家时，重复FFN计算仍可有用。这些结果共同表明，循环价值取决于连续应用暴露的计算，而非固定操作符身份。这促使我们比较无循环、仅混合器和全块循环，并选择状态转移可在循环深度上被追踪的混合器。

#### 线性注意力作为循环记忆。

线性循环模型天然适合此研究，因其用固定大小状态（跨序列更新）替代增长的注意力图。S4为长序列建模建立结构化状态空间层；RetNet将循环状态更新与类注意力并行表示关联。Mamba引入选择性状态空间层S6，其输入依赖动态决定什么进入、持续于或离开循环状态。相关研究将线性注意力解释为快速权重记忆：每个词元将键值关联写入紧凑矩阵状态，后续词元查询该状态而无需存储完整注意力历史。DeltaNet使用增量规则更新在写入新信息前修正现有关联；Gated DeltaNet添加自适应遗忘以控制长程状态演化。Gated DeltaNet-2分离擦除和写入操作，独立控制移除哪些先前记忆组件和提交哪些新值组件。从结构化循环到输入选择性和覆写感知记忆的演进，使线性注意力特别适合操作符级循环：连续应用可细化写入、遗忘和检索，同时保持线性序列复杂度。我们使用GDN作为循环混合器，并测试状态细化是否提供有效的循环边界，而无需重复执行密集FFN。

## 3方法

### 3.1 MixerLoop架构

我们在Gated DeltaNet主干上构建MixerLoop，改变循环计算应用的边界。令\(\mathcal{A}_{i}\)和\(\mathcal{F}_{i}\)表示物理层\(i\)中的残差词元混合器和密集前馈网络：
\[\mathcal{A}_{i}(h)=h+\operatorname{GDN}_{i}\!\left(\operatorname{Norm}_{A,i}(h)\right),\]
\[\mathcal{F}_{i}(h)=h+\operatorname{FFN}_{i}\!\left(\operatorname{Norm}_{F,i}(h)\right).\]
非循环模型对每个物理层的混合器和FFN各应用一次：
\[h^{\mathrm{NoLoop}}=\left(\mathcal{F}_{L}\circ\mathcal{A}_{L}\circ\cdots\circ\mathcal{F}_{1}\circ\mathcal{A}_{1}\right)(h^{0}).\]
现有全循环语言模型则重复使用完整主干\(T\)个循环步：
\[h^{\mathrm{FullLoop}}=\left(\mathcal{F}_{L}\circ\mathcal{A}_{L}\circ\cdots\circ\mathcal{F}_{1}\circ\mathcal{A}_{1}\right)^{T}(h^{0}).\]
一次主干传递的输出成为下一次的输入，因此每个循环步都执行词元混合器和密集FFN。

MixerLoop将循环移入每个物理层，仅应用于GDN混合器：
\[h^{\mathrm{MixerLoop}}=\left(\mathcal{F}_{L}\circ\mathcal{A}_{L}^{T}\right)\circ\cdots\circ\left(\mathcal{F}_{1}\circ\mathcal{A}_{1}^{T}\right)(h^{0}).\]
\(\mathcal{A}_{i}\)的\(T\)次应用共享参数，而不同物理层保留独立的混合器和FFN权重。因此MixerLoop增加了词元混合的有效深度，而不增加唯一参数数量或重复执行密集FFN。尽管相同的GDN参数被重用，连续应用并非重复相同计算。每次应用接收前一次通道产生的隐藏序列，并重新计算其输入依赖的键、值、衰减门和增量规则更新。因此它在因果循环状态上引发新的记忆写入、擦除和读取序列。MixerLoop允许这种有状态的词元混合过程在表示通过密集FFN之前被细化数次。

表1：本研究中研究的循环边界。\(\mathcal{A}\)表示GDN词元混合器，\(\mathcal{F}\)表示密集FFN。所有模型使用相同物理主干，仅在循环计算位置上不同。
### 3.2训练设置

我们比较具有匹配唯一参数数、数据、处理词元预算、优化设置、GDN主干、分词器、上下文长度和语料的架构。标称15M/110M配置分别具有15.7M/116.9M唯一参数、6/12层、宽度288/768、4/12个GDN头，以及FFN宽度768/2,048；两种循环架构均使用\(T=4\)，与LT2循环数匹配。

训练使用32K SentencePiece分词器，上下文长度1024，相同的170分片ClimbMix子集（10.69B语料词元），100,000流式更新（每个模型52.43B处理词元），全局批大小512。我们使用AdamW优化器，\((\beta_{1},\beta_{2})=(0.9,0.95)\)，峰值学习率\(5\times 10^{-4}\)，1000次预热更新后余弦衰减，权重衰减0.1，梯度裁剪1.0，相同种子和数据顺序，无架构特定超参数搜索。

我们评估DataComp-LM的22任务CORE套件，将每个任务针对其随机基线居中，并报告CORE平均值和固定操作子集，因为循环模型在信息操作上倾向于比知识存储获得更大增益。架构比较仅改变循环边界，而模型类别、唯一参数、数据、处理词元和评估保持固定。

### 3.3估计计算减少

令\(C_{A}\)和\(C_{F}\)表示物理层中一个GDN混合器和一个密集FFN的投影计算量，\(C_{H}\)表示最终语言模型头的成本。忽略共享嵌入查找，一次前向传递的投影成本为：
\[C_{\mathrm{NoLoop}}=L\left(C_{A}+C_{F}\right)+C_{H},\]
\[C_{\mathrm{MixerLoop}}=L\left(TC_{A}+C_{F}\right)+C_{H},\]
\[C_{\mathrm{FullLoop}}=TL\left(C_{A}+C_{F}\right)+C_{H}.\]
相对于NoLoop，MixerLoop增加\(L(T-1)C_{A}\)投影计算量，而FullLoop增加\(L(T-1)(C_{A}+C_{F})\)。

在循环主干内，相对投影成本为：
\[\frac{C_{\mathrm{MixerLoop}}-C_{H}}{C_{\mathrm{FullLoop}}-C_{H}}=\frac{TC_{A}+C_{F}}{T(C_{A}+C_{F})}.\]
由于密集FFN占每层投影计算量的61.2–61.3%，式9在\(T=4\)时得0.541：
\[\frac{TC_{A}+C_{F}}{T(C_{A}+C_{F})}=0.541.\]
因此MixerLoop保留FullLoop主干投影计算量的54.1%，将其减少45.9%。

由于每个架构仅执行一次LM头，端到端节省对15M为33.9%，对110M为43.0%；32K词表头在15M中占比更大，而循环主干计算在110M中占主导。预填充延迟和吞吐量在第4节测量。

### 3.4有限迭代传输秩

我们接下来问：为何在移除重复FFN计算后，后续混合器应用仍然有用？对此直接比较隐藏状态是不够的。即使中间表示变化很大，若这些变化被后续层消除或对模型预测影响甚微，也可能无效。因此我们通过有限干预测量迭代传输，并观察其在最终语言模型输出端的影响。

对于每个物理层\(\ell\)，我们为其原生GDN混合器构建一个*上下文关闭*对应物。原生混合器通过其因果卷积和循环矩阵状态联合处理序列。上下文关闭混合器保留相同的输入投影、门控、归一化、残差路径和词元局部非线性计算，但独立处理每个词元。其因果卷积