LooM：循环不变内存，让循环Transformer推理预算可调

循环Transformer通过重复执行共享层来增加有效深度，但内存随循环次数增长、计算只能整块调整、训练后难以改变循环次数——这三大限制让它的落地一直不顺。一篇题为《LooM: Loop-Invariant Memory for Causal Encoder-Decoder Language Models》的研究提出了针对性的解法：只在因果编码器-解码器（CED）的解码器层做循环，并且让内存写入与循环次数解耦。

循环Transformer的三个卡点

循环Transformer的基本思路是让同一组层被反复执行，从而以较少的参数获得更深的网络效果。但研究指出，这条路有三个绕不开的限制：

• 内存随循环次数增长：每一轮循环都要保留相应的缓存，循环越多内存占用越大；
• 计算粒度太粗：计算量只能按整个块来增减，无法精细调节；
• 循环次数训练后难改：模型训练完成后，想调整循环次数并不容易。

这三条限制叠在一起，使得循环Transformer在部署阶段缺乏灵活性——而部署恰恰是最需要按预算调节计算和内存的地方。

KV 缓存（Key-Value Cache） 是自回归语言模型推理时的核心内存结构。生成每个 token 时，模型需要用到此前所有 token 的键（Key）和值（Value）向量来计算注意力，为避免重复计算，这些向量会被缓存下来。标准 Transformer 的 KV 缓存大小与序列长度线性相关；而在循环 Transformer 中，如果每一轮循环都各自维护一份缓存，缓存就会再乘以循环次数，这正是「内存随循环次数增长」的根源。LooM 的解法则利用了循环层共享同一套注意力参数这一事实：同一个 token 的 KV 条目在多次循环中是相同的，因此只需写入一次，后续循环直接复用。理解这一点，才能明白为何「循环不变内存」能够在增加有效深度的同时把缓存控制在单次前向的水平。

LooM 的核心思路：内存只写一次

LooM（Loop-invariant Memory，循环不变内存）的做法是分而治之。在因果编码器-解码器架构中，它只对第一个解码器层之后的解码器层做循环。这些层读取一个全局内存，而这个全局内存对每个 token 只写入一次。

关键在于 KV 缓存的处理：每个循环层只在第一次访问某个 token 时发布该 token 的键值（KV）条目，因此缓存不会随循环次数增长。这正是“循环不变内存”名字的由来——循环次数变了，内存占用不变。

这一设计直接化解了第一个限制。相比让整个模型反复执行，只循环解码器层、且内存写入与循环解耦，是更轻量的路径。

预算可调：从 3 到 21 次层执行

LooM 的第二个设计是让这些循环层独立循环。这样一来，每个 token 可以获得从 3 到 21 次层执行的任意预算，计算量不再是整块跳变，而是可以连续调节。

更实用的是，一个在随机采样预算上训练出来的模型，就能覆盖上述所有预算，不需要额外的路由器（router）来动态选择。这意味着部署时可以按需选择计算量，而不必为每个预算单独训练模型。对于需要在不同硬件或延迟要求下切换的场景，这种灵活性有实际价值。

效果与训练预算的影响

研究给出的结论是：在循环次数足够时，LooM 在语言建模、零样本准确率和长上下文困惑度上都优于在相同数据上训练的单次前向基线，同时使用的缓存大小相同。

论文的分析还揭示了一个值得关注的细节：训练时的预算分布决定了额外深度能带来什么收益。换句话说，循环次数不是越多越好，训练阶段如何采样预算，会影响模型从额外深度中获得的增益类型。这为后续研究提供了一个可调的维度。

为什么说 CED 是循环的天然场所

论文的落脚点是一个架构判断：因果编码器-解码器是循环的自然场所——内存只写一次，深度则成为部署时选择的预算。

这个观点的价值在于把“深度”从训练时的固定结构，转变为推理时的可调资源。传统模型一旦训练完成，深度就固定了；而 LooM 让深度变成一个可以在部署阶段按预算选择的变量，同时不牺牲缓存效率。对于追求推理成本弹性的应用来说，这是一个方向性的提示。

当然，这一方案的前提是采用因果编码器-解码器架构，且循环仅作用于解码器层。它在其他架构上的适用性、以及训练预算分布的具体设计策略，仍有待进一步验证。

因果编码器-解码器（Causal Encoder-Decoder，CED） 是一种把编码器与解码器都设计为因果（即每个位置只能看到左侧上下文）的架构，区别于 T5 这类编码器双向可见的经典编码器-解码器。在 CED 中，编码器负责处理输入并产出表示，解码器则以自回归方式逐 token 生成输出。这种结构对循环特别友好：编码器只运行一次、无需循环，循环仅发生在解码器内部，且解码器天然需要 KV 缓存，因此「每 token 只写一次内存」的约束可以严格成立。相比之下，在纯解码器模型中循环全部层会迫使每一轮都重新处理缓存，在编码器-解码器混合模型中又难以界定循环边界。CED 因此成为 LooM 设计的架构前提，也解释了论文为何强调该方案的适用范围受架构约束。

要点回顾

• LooM 只在 CED 的解码器层做循环，全局内存每 token 只写一次，KV 缓存不随循环次数增长；
• 循环层独立运行，支持每 token 3 到 21 次层执行的连续预算，单模型覆盖全部预算且无需路由器；
• 足够循环次数下，语言建模、零样本准确率和长上下文困惑度均优于同数据的单次前向基线，缓存大小相同；
• 训练预算分布决定额外深度的收益类型，深度由此成为部署时可选的预算。