FlashLoop循环Transformer：用循环深度重构推理效率的帕累托前沿

循环结构在神经网络里不是新东西，但把循环深度和推理效率放在同一个帕累托前沿上讨论，是最近才被认真对待的方向。FlashLoop 这个工作之所以值得单独拿出来聊，是因为它没有走"堆参数、堆层数"的老路，而是把目光投向了另一个维度——让同一组权重被反复调用，用计算步数换模型容量。这个思路听起来简单，但真正落地时会撞上一堆工程问题：KV cache 怎么复用、循环次数怎么定、训练时梯度怎么传、推理时延迟和吞吐怎么权衡。我花了几天时间把它的设计逻辑和实现细节捋了一遍，下面把这些东西拆开讲清楚，适合对 Transformer 推理优化感兴趣、想理解"循环深度"这条路线的读者。

1. 为什么循环深度突然成了推理效率的焦点

1.1 从"堆层数"到"复用层"的范式转移

过去几年 Transformer 的 scaling 基本沿着一条路走：层数越多、参数越大、效果越好。但这条路在推理侧遇到了硬墙。一个 70B 级别的模型，光是权重加载就要占满显存带宽，每生成一个 token 都要把全部层跑一遍，延迟随层数线性增长。更麻烦的是，KV cache 的大小和层数成正比，长上下文场景下显存直接被吃光。

循环深度的思路是：与其训练 100 层不同的权重，不如训练 12 层权重然后循环跑 8 次。参数量降下来了，但有效计算深度还是 96 层。这里的关键洞察是——模型的表达能力不完全来自参数数量，也来自计算的迭代次数。就像数值分析里迭代求解一个方程，迭代次数够多就能逼近精确解，不一定需要更复杂的公式。

这个思路在理论上早有支撑。Deep Equilibrium Model 那类工作证明了，一个无限深度的网络可以等价于某个不动点方程的解。循环 Transformer 本质上是把这个"无限深度"截断成有限次循环，用可控的计算量去逼近那个不动点。

1.2 帕累托前沿到底在权衡什么

"帕累托前沿"这个词用在这里很准确。推理效率不是一个单一指标，它至少包含四个维度：

传统 dense 模型在这四个维度上是一个固定点，你没法在不重训的情况下调整。循环模型则提供了一个旋钮：循环次数。推理时你可以根据硬件条件和延迟要求，动态选择跑 4 次还是 8 次。这就形成了一个可调节的帕累托前沿——同样的权重，不同的效率/质量权衡点。

FlashLoop 的核心贡献，就是把这个旋钮做得足够精细，让循环带来的额外开销尽可能小，同时保证质量下降在可接受范围内。

1.3 和 KV cache 的纠葛是绕不开的

循环结构最反直觉的地方在于 KV cache 的处理。标准 Transformer 里，每一层有自己的 KV cache，层与层之间不共享。循环模型里，同一组权重被调用多次，那这多次调用是共享一份 KV cache，还是每次循环都写一份新的？

如果共享，那循环就退化成了对同一份上下文的重复注意力，表达能力受限。如果每次循环都写新的 KV，那 cache 大小就随循环次数线性膨胀，显存优势荡然无存。FlashLoop 在这两者之间找了一个折中方案，这也是它区别于早期循环 Transformer 的关键。具体怎么做的，后面第 3 节会展开。

2. FlashLoop 的核心机制拆解

2.1 权重共享的具体形式：哪些层参与循环

FlashLoop 不是把整个模型都做成循环的。它的结构可以理解为：前若干层是独立的"预热层"，中间一段是共享权重的"循环块"，最后若干层是独立的"收尾层"。这种设计有明确的动机。

预热层负责把原始输入映射到一个合适的表示空间。如果一上来就循环，第一遍循环的输入分布和后续循环的输入分布差异很大，权重很难同时适配。预热层把这个分布"预热"到一个稳定区域，后续循环的输入就落在相近的分布上，共享权重才有意义。

收尾层负责把循环后的表示映射回输出空间。循环块本身不需要关心最终输出格式，它只负责迭代精炼表示。这种分工让循环块的权重可以专注于"如何让表示更好"这一件事。

循环块内部通常是标准的 Transformer block：自注意力 + FFN + 残差 + LayerNorm。循环时，输入是上一轮循环的输出，权重不变。可以写成：

h_0 = Preprocess(x) for i in range(n_loops): h_{i+1} = LoopBlock(h_i) y = Postprocess(h_n)

这里有个细节：循环块是否带自己的位置编码？如果带，那每轮循环位置编码都会重新加一遍，可能造成位置信息被过度强调。FlashLoop 的做法是在预热层注入位置信息，循环块内部不再重复注入，靠残差连接保持位置信息的传递。

2.2 循环次数怎么定：训练时随机、推理时固定

循环次数 n 是个超参数。训练时如果固定 n，模型会过拟合到这个特定深度，推理时改成别的 n 效果会崩。FlashLoop 采用了训练时随机采样 n 的策略——每次前向传播随机选一个 n，比如在 4 到 12 之间均匀采样。

这样做的好处是模型见到过各种循环深度，对 n 的变化鲁棒。推理时你就可以根据延迟预算自由选择 n，不用重训。这有点像 dropout 的思路：训练时随机丢，推理时用全部，模型学会了在多种配置下工作。

但随机采样 n 会带来训练不稳定的问题。因为不同 n 对应的计算图深度不同，梯度回传的路径长度差异很大。n 小的时候梯度容易爆炸，n 大的时候梯度容易消失。FlashLoop 用了两个手段来缓解：一是对循环块的输出做尺度归一化，二是对梯度做裁剪。这两个都是常规操作，但在这里是必需的，不是可选项。

2.3 循环块内部的残差缩放

循环结构有个天然的风险：残差连接在多次循环后会累积成很大的值。假设每轮循环残差分支的输出量级是 O(1)，循环 12 次后主分支的量级可能到 O(12)，LayerNorm 虽然能拉回来，但数值精度会受影响。

FlashLoop 在循环块的残差连接上引入了一个缩放因子，通常是 1/sqrt(n_loops) 或者一个可学习的标量。这个缩放让每轮循环的增量被压制，多轮累积后总量级保持在合理范围。这个技巧在深度很深的 Transformer 里也常见，但循环结构里它从"优化技巧"变成了"必需组件"。

实测中这个缩放因子的初始化很关键。如果初始化得太小，循环块几乎不起作用，模型退化成只有预热层和收尾层；如果太大，训练初期就会发散。比较稳的做法是初始化为 1/sqrt(期望循环次数)，然后让它跟着训练慢慢调整。

3. KV cache 在循环结构里的三种处理方案

3.1 方案一：每轮循环独立 cache

这是最直观的做法。循环块每跑一轮，就把这一轮的 K 和 V 追加到 cache 里。第 i 轮循环的注意力可以看到前面所有轮次的 KV，包括预热层的 KV 和之前循环轮次的 KV。

优点是表达能力强，每一轮循环都能"看到"历史迭代的信息，相当于一个隐式的迭代记忆。缺点是 cache 大小随循环次数线性增长。如果循环 8 次，cache 就是标准模型的 8 倍（假设循环块只有一层）。长上下文场景下这是灾难性的。

这个方案适合短序列、对质量要求极高的场景。比如一些推理任务，输入不长但需要深度思考，cache 膨胀可以接受。

3.2 方案二：所有循环轮次共享一份 cache

另一个极端是循环块只维护一份 KV cache，每轮循环覆盖写。第 i 轮循环的注意力只能看到当前轮的 KV，看不到历史轮次。

优点是 cache 大小和标准模型一样，显存优势最大化。缺点是循环退化成了"对同一份表示反复做同样的注意力"，表达能力提升有限。实测中这种方案的收益递减很快，循环 4 次以上基本没有额外提升。

这个方案适合显存极度受限、循环次数很少（2-3 次）的场景。

3.3 方案三：分组共享 + 滑动窗口

FlashLoop 实际采用的是折中方案。它把循环轮次分成若干组，组内共享 cache，组间独立。比如循环 8 次，分成 2 组，每组 4 次循环共享一份 cache。这样 cache 大小是标准模型的 2 倍，但表达能力比完全共享强很多。

更进一步，组内的注意力可以用滑动窗口，只关注最近若干轮的 KV，进一步压缩 cache。滑动窗口的大小是个可调参数，窗口越大质量越好但 cache 越大。

这个表格里的 g 是组数，通常取 2 到 4。选择 g 的时候要考虑：g 太小退化成完全共享，g 太大退化成独立 cache。经验值是让每组覆盖 3-4 轮循环，这样组内既有足够迭代，组间又有足够差异。

3.4 cache 的预分配和复用

工程实现上还有个细节：循环结构的 cache 预分配策略和标准模型不同。标准模型每层 cache 大小固定，可以一次性预分配。循环模型如果采用分组共享，每组的 cache 大小取决于该组实际用到的轮次，需要动态管理。

FlashLoop 的做法是预分配最大可能的 cache，然后用一个 mask 控制实际使用范围。这样避免了动态分配的开销，代价是显存占用按最坏情况算。如果你的场景循环次数固定，可以按实际次数预分配，省一些显存。

提示：分组共享方案下，组与组之间的 cache 切换点要小心处理。如果切换点正好在一个注意力窗口的中间，可能导致某些 token 的 KV 被重复计算或遗漏。建议把切换点对齐到序列的边界或者固定的 chunk 边界。

4. 训练循环 Transformer 时最容易踩的坑

4.1 梯度检查点不能随便开

循环结构天然适合梯度检查点，因为同一组权重被反复调用，中间激活值可以只存一份，反向传播时重算。但这里有个陷阱：如果循环次数是随机采样的，检查点的重算必须用相同的循环次数，否则梯度对不上。

我见过有人为了省显存，在检查点重算时用了不同的随机种子，结果训练 loss 震荡得厉害，查了半天才发现是梯度不一致。正确做法是把循环次数作为检查点的一部分存下来，重算时严格复用。

另外，循环块的检查点粒度要选好。如果整个循环块作为一个检查点单元，重算开销是 O(n_loops)；如果把循环块内部再细分，重算开销降低但管理复杂度上升。一般建议循环块整体作为一个检查点单元，因为循环块本身不大。

4.2 学习率要按有效深度调

循环模型的有效深度是 n_loops 乘以循环块层数。如果循环块有 2 层，循环 8 次，有效深度就是 16 层。这个有效深度决定了合适的学习率范围。

直接用标准 Transformer 的学习率往往偏大，因为循环结构的梯度路径更长，更容易不稳定。经验做法是按有效深度反比缩放学习率，或者用 warmup 把学习率慢慢拉上去。warmup 步数也要相应增加，给模型时间适应循环带来的梯度变化。

还有个细节：循环块的权重和预热层、收尾层的权重，学习率可以不同。循环块因为被反复调用，梯度累积效应更强，学习率可以设小一点。预热层和收尾层只调用一次，可以用标准学习率。这种分层学习率在实测中能明显提升训练稳定性。

4.3 循环次数的采样分布很讲究

前面说训练时随机采样循环次数，但采样分布不是随便定的。如果均匀采样 4 到 12，模型在 4 和 12 附近的表现会比较好，中间反而一般。这是因为均匀分布下，端点附近的样本相对密度更高（边界效应）。

更好的做法是用一个偏向中间的正态分布，或者用分段采样：大部分时候采中间值，偶尔采极端值。这样模型在常见循环次数上表现最好，同时保持对极端值的鲁棒性。

还有个进阶技巧：训练后期逐渐收窄采样范围，让模型在目标循环次数上精调。比如前 80% 的训练用宽范围采样，后 20% 收窄到目标值附近。这样既保证了鲁棒性，又保证了最终质量。

4.4 位置编码的注入位置要实验

前面提到位置编码在预热层注入，循环块不重复注入。但这个不是唯一选择，具体哪种好要做实验。

另一种做法是每轮循环都注入位置编码，但用一个衰减因子控制注入强度。第一轮注入强，后续轮次逐渐减弱。这样模型在早期循环有明确的位置信息，后期循环更关注内容而非位置。

还有一种做法是完全不在循环块注入位置编码，靠预热层的位置信息通过残差传递。这种最简洁，但实测中对长序列效果会差一些，因为位置信息在多次循环后会被稀释。

我的建议是：短序列（<2K）用第一种，长序列用第二种，超长序列（>16K）考虑在循环块里加轻量的相对位置编码。

5. 推理部署时的效率实测与调优

5.1 循环次数和延迟的非线性关系

理论上循环次数翻倍，延迟也翻倍。但实测中不是这样。因为循环块的计算可以和其他操作重叠，而且 GPU 的利用率在循环结构下往往更高（同一组权重反复用，cache 命中率高）。

实测数据显示，循环次数从 4 增加到 8，延迟增加大约 70% 而不是 100%。从 8 增加到 12，延迟增加约 50%。这个非线性来自几个因素：权重加载只发生一次，后续循环复用；kernel 启动开销被摊薄；批处理效率提升。

这意味着循环结构的帕累托前沿比理论预测的更靠前。你可以在延迟增加不多的情况下，显著提升有效深度。

5.2 批处理策略要重新设计

标准模型的批处理是静态的：凑够一个 batch 就跑。循环模型的批处理可以更灵活，因为不同样本可以用不同的循环次数。

比如一个 batch 里有简单样本和困难样本，简单样本跑 4 次循环就够了，困难样本跑 12 次。如果统一跑 12 次，简单样本浪费了计算。FlashLoop 支持 per-sample 的循环次数，但实现上要注意：不同循环次数的样本在同一个 batch 里，需要 padding 到相同长度，padding 部分不参与计算。

这个特性在 serving 场景下很有价值。你可以根据请求的难度动态分配计算资源，简单请求快速返回，困难请求多花点时间。整体吞吐量比统一循环次数高不少。

5.3 和投机解码的配合

投机解码用一个小的 draft 模型生成候选 token，大的 target 模型验证。循环结构天然适合做 draft 模型：用少量循环次数快速生成候选，然后用完整循环次数验证。

具体做法是：draft 阶段用 2-4 次循环，生成一批候选 token；验证阶段用 8-12 次循环，一次性验证所有候选。因为 draft 和 target 共享权重，不需要额外训练 draft 模型，这是循环结构相对标准投机解码的一个优势。

实测中这种自投机解码能带来 1.5 到 2 倍的加速，具体取决于任务。代码生成这类确定性强的任务加速比高，开放域对话加速比低。

5.4 显存和吞吐的权衡表

部署时最常问的问题是：给定显存预算，怎么配置循环次数和 batch size 最优。下面是一个经验参考：

这个表是粗粒度的，实际要结合序列长度调整。序列越长，KV cache 占比越大，能用的 batch size 越小。长序列场景下建议优先保证循环次数，牺牲 batch size，因为循环次数对质量的影响比 batch size 大。

注意：循环次数不是越多越好。实测中超过某个阈值后，质量提升趋于平缓，但延迟继续线性增长。这个阈值通常在 8 到 16 之间，具体取决于任务和模型规模。建议在自己的任务上做一次循环次数扫描，找到质量-延迟的拐点。

6. 循环 Transformer 的适用边界与不适用场景

6.1 什么任务最适合循环结构

循环结构的本质是"迭代精炼表示"，所以最适合那些需要多步推理的任务。数学题、逻辑推理、代码生成、多跳问答，这些任务的标准 Transformer 往往需要很深的层数才能做好，循环结构用较少的参数就能达到类似效果。

另一个适合的场景是资源受限的部署。边缘设备上显存有限，循环结构可以用较小的权重占用达到可接受的质量。推理时根据设备负载动态调整循环次数，负载低时多循环几次提升质量，负载高时少循环几次保证响应。

还有一个不太被提及但很有价值的场景：需要可控计算量的场景。有些应用要求推理延迟严格可控，标准模型没法在不换模型的情况下调整延迟。循环模型可以通过调整循环次数精确控制延迟，这在实时系统里很有用。

6.2 什么任务不适合

循环结构不适合那些"一遍过"的任务。比如简单的分类、检索、模式匹配，这些任务不需要迭代精炼，循环带来的额外计算是浪费。实测中这类任务上循环结构的收益接近零，甚至因为训练难度增加而略有下降。

也不适合对首 token 延迟极度敏感的场景。循环结构的首 token 延迟比同参数量的标准模型高，因为要跑完所有循环才能出第一个 token。如果你的应用要求首 token 在几十毫秒内返回，循环结构可能不合适。

还有一个边界是超长序列。虽然循环结构在长序列上有理论优势（参数少，cache 管理灵活），但实测中超长序列（>32K）下循环结构的注意力计算占比很高，循环带来的收益被注意力开销抵消。这种场景下稀疏注意力或者线性注意力可能更合适。

6.3 和 MoE、稀疏化的关系

循环结构和 MoE 是可以叠加的。MoE 在 FFN 层做稀疏激活，循环结构在深度方向做权重共享。两者结合可以同时降低参数和计算。但叠加后训练难度也叠加，需要更仔细的超参调优。

一个务实的做法是先用循环结构把参数量降下来，如果还不够再考虑 MoE。不要一上来就两个都用，否则出问题很难定位是哪个组件导致的。

稀疏化方面，循环结构本身可以看作一种时间维度的稀疏——不是所有层都独立计算。如果再叠加空间维度的稀疏（比如稀疏注意力），要注意两者的交互。循环结构下，稀疏注意力的 mask 每轮循环可能不同，需要仔细设计。

7. 自己动手实现一个最小循环 Transformer

7.1 核心代码结构

下面是一个最小实现的骨架，用 PyTorch 风格伪代码展示。重点看循环块的定义和循环逻辑。

class LoopBlock(nn.Module): def __init__(self, d_model, n_heads, dropout=0.1): super().__init__() self.attn = MultiHeadAttention(d_model, n_heads) self.ffn = FeedForward(d_model) self.norm1 = LayerNorm(d_model) self.norm2 = LayerNorm(d_model) self.res_scale = nn.Parameter(torch.tensor(1.0)) def forward(self, h, kv_cache=None): # 自注意力，带残差缩放 attn_out = self.attn(self.norm1(h), kv_cache) h = h + self.res_scale * attn_out # FFN，带残差缩放 ffn_out = self.ffn(self.norm2(h)) h = h + self.res_scale * ffn_out return h class FlashLoopModel(nn.Module): def __init__(self, d_model, n_heads, n_pre, n_post, max_loops): super().__init__() self.pre_layers = nn.ModuleList([ TransformerBlock(d_model, n_heads) for _ in range(n_pre) ]) self.loop_block = LoopBlock(d_model, n_heads) self.post_layers = nn.ModuleList([ TransformerBlock(d_model, n_heads) for _ in range(n_post) ]) self.max_loops = max_loops def forward(self, x, n_loops=None): if n_loops is None: n_loops = random.randint(4, self.max_loops) # 预热层 h = x for layer in self.pre_layers: h = layer(h) # 循环块 for i in range(n_loops): h = self.loop_block(h) # 收尾层 for layer in self.post_layers: h = layer(h) return h

这个骨架省略了 KV cache 管理和分组共享的逻辑，但核心的循环结构已经体现出来了。实际实现中，KV cache 的管理是最复杂的部分，需要根据第 3 节的方案做具体设计。

7.2 训练脚本的关键配置

训练循环模型时，下面几个配置项和标准模型不同，需要特别注意：

# 优化器配置 optimizer = AdamW([ {'params': model.pre_layers.parameters(), 'lr': 1e-4}, {'params': model.loop_block.parameters(), 'lr': 5e-5}, # 循环块学习率减半 {'params': model.post_layers.parameters(), 'lr': 1e-4}, ], weight_decay=0.01) # 梯度裁剪 torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0) # 循环次数采样 def sample_loops(step, total_steps): if step < 0.8 * total_steps: # 前期宽范围采样 return random.randint(4, 12) else: # 后期收窄 return random.randint(6, 10)

学习率分层是循环模型训练的一个关键技巧。循环块的梯度累积效应更强，用相同的学习率容易发散。梯度裁剪的阈值也要比标准模型小一些，因为循环结构的梯度范数天然更大。

7.3 验证循环结构是否真的在工作

训练完之后，怎么确认循环结构真的在起作用，而不是退化成普通 Transformer？有几个诊断方法。

第一个方法是看循环次数对输出的影响。固定输入，改变循环次数，看输出变化。如果循环 4 次和循环 12 次的输出几乎一样，说明循环块没起作用，可能残差缩放太小或者循环块权重被训练成了恒等映射。

第二个方法是看循环块权重的变化。如果循环块的权重和初始化时差不多，说明它没学到东西。正常情况下循环块的权重应该有明显变化，尤其是注意力和 FFN 的权重。

第三个方法是消融实验。把循环次数固定为 1，看质量下降多少。如果下降很少，说明循环结构对这个任务是多余的。如果下降明显，说明循环确实在贡献。

提示：循环块退化成恒等映射是常见问题，尤其在训练数据量小的时候。如果发现这个问题，可以尝试增大残差缩放的初始化值，或者给循环块加一个辅助损失，强制它的输出和输入有差异。

8. 循环路线的局限和可能的演进方向

循环 Transformer 不是万能药，它有自己的天花板。最大的局限是循环块内部的权重共享限制了表达多样性。标准 Transformer 每一层可以学到不同的特征，循环块只有一组权重，所有循环轮次共享，这限制了模型能表达的变换种类。

一个可能的演进方向是"部分共享"：循环块内部不是所有子层都共享，比如注意力共享但 FFN 不共享，或者反过来。这样在参数增加不多的情况下提升表达多样性。具体共享哪些、不共享哪些，需要实验确定。

另一个方向是条件循环：不是所有 token 都循环相同次数，而是根据 token 的难度动态决定。简单 token 少循环，困难 token 多循环。这需要设计一个难度评估机制，可以在循环块里加一个轻量的分类头，预测当前 token 是否还需要继续循环。

还有个方向是把循环结构和状态空间模型结合。状态空间模型天然是循环的，但它的循环是沿着序列维度的，和循环 Transformer 沿着深度维度的循环不同。两者结合可能产生新的架构，但这还是开放问题。

从工程角度看，循环结构最大的价值是给了推理时一个可调节的旋钮。在标准模型里，你要么用大模型（高质量高延迟），要么用小模型（低质量低延迟），没有中间选项。循环模型让你可以用同一套权重覆盖整个延迟-质量谱系，这在产品部署时非常实用。你可以根据用户等级、设备性能、当前负载动态调整，而不需要维护多个模型。

我在实际测试中的体会是，循环结构的收益在中等规模模型（1B 到 10B）上最明显。太小的模型循环收益有限，太大的模型循环带来的训练难度增加超过了收益。如果你的场景正好在这个区间，值得认真考虑循环结构。如果不在，可以先观望，等工具链更成熟再上手。