Parallel Loop Transformer for Efficient Test-Time Computation Scaling

作者/机构: ByteDance Seed

A1 主要贡献

核心问题: 大型语言模型（LLMs）功能强大，但在实际部署中因推理成本（如延迟和内存带宽）过高而受到限制。一种参数高效的架构是循环 Transformer（Looped Transformer），它通过在多个计算步骤（或称“循环”）中重复使用同一组参数，以较小的存储占用实现更大的有效深度和更强的性能。然而，传统循环 Transformer 存在一个致命缺陷：其循环必须严格按顺序执行（如图 1a 所示），导致每个词元的计算量、时钟延迟和 KV 缓存大小都随着循环次数（L）线性增长（O(L)），这使得它们在对延迟敏感的应用中不切实际。

研究目标: 本文旨在解决循环 Transformer 的有效深度与推理速度之间的矛盾，提出一种名为并行循环 Transformer（Parallel Loop Transformer, PLT）的新型架构。该架构的目标是既能获得深度循环计算带来的性能优势，又能保持与标准非循环 Transformer 相当的低推理延迟。

创新点:
为了实现上述目标，PLT 架构引入了两项关键技术：

  1. 跨循环并行（Cross-Loop Parallelism, CLP）: 该技术通过重新设计训练和推理流程，打破了循环之间的顺序依赖。在解码过程中，CLP 可以在单次前向传播中，将当前词元 ti 的第 l 个循环与前一个词元 ti-1 的第 l+1 个循环等并行计算。这种重叠计算方式将 L 个顺序步骤有效地折叠为一个，从而将模型的有效深度与时钟延迟解耦（如图 1b 所示）。

  2. 高效表示增强（Efficient Representation Enhancement）: 该策略旨在解决 KV 缓存随循环次数线性增长的问题。

    • 首循环 KV 缓存共享: 所有后续循环共享第一个循环的 KV 缓存，将 KV 缓存的内存占用恢复到单循环的水平。
    • 门控滑动窗口注意力（G-SWA）: 为了在共享全局信息的同时保留局部上下文信息以维持高精度，在非首循环中引入了门控滑动窗口注意力机制。

主要贡献总结:

  • 提出了并行循环 Transformer（PLT）架构，据我们所知，这是首个成功并行化循环 Transformer 计算的架构，以可忽略的延迟开销实现了可扩展的测试时计算。
  • 引入了跨循环并行（CLP）和高效表示增强技术（KV 缓存共享与门控 SWA），克服了传统循环模型的关键延迟和内存瓶颈。
  • 通过在内部和开源模型上的大量实验证明，PLT 在准确性上显著优于普通 Transformer 基线，同时只增加了极小的延迟。
  • 证明了 PLT 远比普通循环 Transformer 高效，甚至能让一个更浅、更高效的 PLT 模型（例如，1.7B 激活参数）在性能和延迟上优于一个更大的普通模型（例如，2.5B 激活参数）。


图 1 计算流程图。(a) 普通循环 Transformer，其中每个词元的每个循环都必须串行计算。(b) 并行循环 Transformer (PLT)，其中同一蓝色虚线框内的 Transformer 循环可以并行计算。


(a) 普通循环 Transformer 的推理流程。(b) 所提出的 PLT 的推理流程。

A2 方法细节


图 2 循环次数 L=3 的 PLT 训练和推理流程。训练（左）：相同颜色的框追踪输入词元如何遍历循环以预测其目标（例如，词元 T1 经过三个循环来预测 T4，与图 1b 一致）。训练在词元维度上并行，在循环维度上串行。推理（右）：在使用 L=3 的循环 Transformer 解码 T4 和 T5 时，PLT 的并行化前向传播。因为在训练期间没有水平（同一步骤、跨循环）的激活依赖，所以在解码期间，同一步骤内的计算（每行；见蓝色虚线框）可以并行运行。

2.1 预备知识：普通循环 Transformer

普通循环 Transformer 的定义。我们考虑一个具有 L 个循环的普通循环 Transformer【11，Universal transformers，2019，International Conference on Learning Representations】。设 T = ($t_1$, $t_2$, . . . , $t_n$) 是一个词元序列，E = ($e_1$, $e_2$, . . . , $e_n$) 是经过嵌入后的词元序列。对于词元索引 i ∈ {1, . . . , n} 和循环索引 l ∈ {1, . . . , L}，令 $h_i^{(l)}$ 表示位置 i 在经过 l 次前向循环后的隐藏状态，其中 $h_i^{(0)}$ 是初始状态。

计算流程。第 i 个词元在普通循环 Transformer 中的计算流程如下：

$$h_{i}^{(0)}=t_{i}, \quad h_{i}^{(l)}=f^{(l)}\left(h_{i}^{(l-1)}\right), \quad l=1, \ldots, L,$$

其中 $f^{(l)}$ 表示第 l 次前向循环。最终，$h_i^{(L)}$ 会被送入分类器头以预测第 (i + 1) 个词元。

普通循环 Transformer 的挑战。在普通的循环 Transformer 中，循环是严格按顺序运行的，因此每个词元的计算量、时钟延迟和 KV 缓存大小都随着循环次数 L 呈 O(L) 扩展，如表 1 所示。虽然权重共享使其参数高效（用更少的存储权重实现了更大的有效深度），但这并不能减少延迟。因此，它主要在同等参数量比较下有帮助；而在同等延迟预算下——这在实际推理中很典型——普通循环 Transformer 没有内在优势，甚至可能因为更长的解码路径、更高的内存带宽压力和更大的 KV 缓存而表现更差。

2.2 并行循环 Transformer

动机。普通循环 Transformer 参数效率高，但依赖于每个词元的严格串行计算，这增加了前向传播次数和 KV 缓存占用。我们的目标是在保持参数节省的同时，通过在循环和词元之间实现并行化来缩短前向传播时间，且不破坏因果关系。

核心设计。如图 1b 所示，我们将普通的串行循环 Transformer 转换为一个并行循环 Transformer，该 Transformer 可以并行执行 L 个循环。该设计包含两个关键组成部分：(i) 用于并行循环推理的跨循环并行（第 2.2.1 节）；(ii) 用于以轻量级时间和 KV 缓存开销提高准确性的高效表示增强（第 2.2.2 节）。

2.2.1 跨循环并行

算法 1 循环次数为 3 的 PLT 解码
需要:
循环次数 L = 3
Transformer 模块函数 f；
输入序列 T = ($t_1$, . . . , $t_n$)；
最大新生成词元数 M。
1: $K_{share}, V_{share}, h_{1_n}, h_{2_{n-1}}, h_{3_{n-2}} \leftarrow f(T)$ ▷ 预填充
2: $logits \leftarrow ClassifierHead(h_{3_{k-2}})$
3: $t_{k+1} \leftarrow argmax(logits)$ ▷ 预测下一个词元
4: for i = n + 1 to n + M - 1 do
5: $e_i = Embedding(t_i)$
6: B ← $B_0, B_1, B_2 \leftarrow e_i, e_i + h_{1_{i-1}}, e_i + h_{2_{i-2}}$
7: $K_{share}, V_{share}, h_1, h_2, h_3 \leftarrow f(B, K_{share}, V_{share})$
8: $logits \leftarrow ClassifierHead(h_{3_{i-2}})$
9: $t_{i+1} \leftarrow argmax(logits)$

算法 2 PLT 的训练
需要:
循环次数 L
Transformer 模块函数 f
分类头 ClassifierHead(·)
输入词元序列 T = ($t_1$, . . . , $t_n$)
1: E ← Embedding(T)
2: $K_{share}, V_{share}, H^{(1)} \leftarrow f(E)$
3: for i = 2 to L do
4: $H^{(i-1)} \leftarrow concat(0, H^{(i-1)}[:-1])$
5: $B \leftarrow E + H^{(i-1)}$
6: $H^{(i)} \leftarrow f(B; K_{share}, V_{share})$
7: end for
8: $logits \leftarrow ClassifierHead(H^{(L)})$
9: $loss = CrossEntropy(logits, T)$
10: return loss

算法 3 PLT 中非首循环的门控滑动窗口注意力
需要:
输入隐藏状态 H
QKV 线性层 $f_{qkv}$, 输出线性层 $f_o$
门控线性层 $f_{gate}$
来自首循环的共享键值缓存 ($K_{share}$, $V_{share}$)
滑动窗口注意力 (SWA) 的窗口大小 w
1: Q, K, V = $f_{qkv}(H)$
2: $y_{global} = Attn(Q, K_{share}, V_{share})$
3: $y_{local} = SWA(Q, K, V, w)$
4: $g = Sigmoid(f_{gate}(Q))$
5: $\tilde{y} = g \odot y_{local} + (1-g) \odot y_{global}$
6: $o = f_o(\tilde{y})$
7: return o

跨循环并行（CLP）的定义。我们引入了跨循环并行（CLP）：将对较早词元的后续循环计算与对较晚词元的早期循环计算进行重叠。如图 1b 所示，以解码具有 3 个循环的词元 t4 为例，对词元 t3 的第一个循环、对词元 t2 的第二个循环和对词元 t1 的第三个循环是同时执行的，后续词元也执行类似的跨循环并行。该方法的详细推理和训练流程如下：

训练流程。PLT 的训练流程如图 2 和算法 2 所示。在第一个循环中，PLT 的行为与普通的 Transformer 一致：它将词元嵌入 H = E = ($e_1$, $e_2$, . . . , $e_n$) 输入模型，并获得最后一层的隐藏状态 $H^{(1)} = (h_1^{(1)}, h_2^{(1)}, . . . , h_n^{(1)})$。在第二个循环之前，PLT 将 $H^{(1)}$向右移动一个位置，从 ($h_1^{(1)}, . . . , h_n^{(1)}$) 变为 ($0, h_1^{(1)}, . . . , h_{n-1}^{(1)}$)，然后加回原始的嵌入：B = E + shift($H^{(1)}$)。这种移位操作消除了跨连续循环的相同索引状态之间的直接依赖，从而能够在解码期间进行并行处理；我们称之为跨循环并行。PLT 对 i = 2, . . . , L 重复此过程：在每个循环中，它将前一个循环的状态右移一个位置，加上嵌入，然后应用 Transformer，同时重用来自第一个循环的共享 $K_{share}$ 和 $V_{share}$。

表 1 普通 Transformer、循环 Transformer 和所提出的 PLT 的复杂性比较。普通 Transformer 作为基线，其参数量为 P，每个词元的计算成本为 C，在内存受限设置下单词元解码延迟为 t。注意力 KV 缓存的规模为 O(nd)，其中 n 是序列长度，d 是嵌入维度。这里，L 表示循环次数，w 表示 SWA（滑动窗口注意力）使用的滑动窗口大小。

推理流程。如图 2（右）所示，PLT 中的跨循环并行通过单次前向传播处理 L 个词元。这种并行设计利用了 LLM 解码的内存受限特性【44，LLM inference unveiled: Survey and roofline model insights，2024，CoRR】：增加并行的测试时计算 FLOPs 可以提高准确性，而额外的解码延迟则可以忽略不计。算法 1 展示了当 L=3 时 PLT 的解码流程。在解码步骤 i，我们构建一个错位的微批次 B={$B_0, B_1, B_2$}，其中 $B_0$ 是词元 i 的第一个循环，$B_1$ 是词元 i−1 的第二个循环，$B_2$ 是词元 i−2 的第三个循环。然后我们根据最终循环状态 $h_{3_{i-2}}$ 来预测词元 (i+1)。与普通循环 Transformer（每个词元应用 L 次顺序传递）相比，PLT 用每个词元一次的并行传递执行相同的逻辑计算。与标准的非循环解码器相比，它增加了能够提高准确性的 FLOPs，但延迟增加可以忽略不计，因为解码是内存受限的【44，LLM inference unveiled: Survey and roofline model insights，2024，CoRR】。

2.2.2 高效表示增强

问题与解决方案。跨循环并行解决了普通循环 Transformer 中推理时间随循环次数 L 扩展的问题，但它仍然会产生 L 倍的 KV 缓存内存成本，这限制了长上下文的使用。我们引入了高效表示增强技术，该技术包含两个部分：(i) 共享首循环的 KV 缓存以提供单一的全局表示，以及 (ii) 使用门控滑动窗口注意力来加强局部上下文。具体细节如下。

共享首循环的 KV 缓存。在标准的循环设计中，每个循环都维护自己的 KV 缓存，因此内存会随着 L 线性增长。为了减少 KV 内存，我们让所有后续循环共享第一个循环的 KV 缓存。如算法 3 所示，非首循环保留其私有的查询（queries），但在第一个循环的 $K_{share}$ 和 $V_{share}$ 上执行全局注意力。因此，只需要存储一个全局 KV 缓存。这种设计为非首循环保留了全局信息，并消除了依赖于 L 的 KV 缓存增长。

非首循环中的门控滑动窗口注意力（G-SWA）。为了在共享的全局表示之上进一步增强局部信息，非首循环在其私有的 Q、K 和 V 上应用滑动窗口注意力，窗口大小为 w。在我们的实验中，我们设置 w=64，并且它不随整体序列长度的增加而增加。如算法 3 的第 3-5 行所示，我们随后使用一个 sigmoid 门来融合 SWA（局部）的输出 $y_{local}$ 和对共享 KV（全局）的全注意力输出 $y_{global}$：

$$g \leftarrow \operatorname{Sigmoid}(f_{\text{gate}}(Q)), \quad \tilde{y} \leftarrow g \odot y_{\text{local}} + (1-g) \odot y_{\text{global}}.$$

G-SWA的效率。门控线性层 $f_{gate}$ 是逐头的，每个头输出一个标量，因此增加的参数和计算量可以忽略不计。此外，由于使用了 SWA，非首循环在解码期间只缓存最近的 w 个 KV 条目。总的来说，门控滑动窗口注意力使得所提出的 PLT 比普通循环 Transformer 在 KV 缓存方面更高效，因为非首循环只存储 w 个最近的 KV 条目，同时通过结合全局和局部信息来保持准确性。

2.3 推理效率分析

复杂度对比。如表 1 所示，我们在内存受限的设置下分析了单次词元推理的效率。普通 Transformer（基线）的参数量为 P，单次词元计算量为 C，KV 缓存为 O(nd)，解码延迟为 t。引入 L 个循环后，计算量增加到 LC，如果每个循环都保留自己的缓存，内存会增加到 O(Lnd)；由于循环是顺序运行的，延迟会增长到 Lt。采用我们提出的跨循环并行后，循环迭代可以并发运行，因此延迟约为 t，计算量仍为 LC，KV 缓存仍为 O(Lnd)。一种 KV 缓存共享的变体则重用第一个循环的缓存，并且不为非首循环存储单独的缓存，这将 KV 缓存减少到 O(nd)，但代价是准确性略有下降。最后，我们没有在非首循环中使用全上下文注意力或禁用注意力，而是在这些循环中应用大小为 w 的滑动窗口注意力，这使得总的 KV 缓存为 O(nd + (L − 1)wd)。由于通常 w ≪ L，额外的 KV 缓存和计算开销很小。总的来说，PLT 保持了接近基线的解码延迟，并避免了 KV 缓存的 L 倍增长。

表 2 激活参数数量相同的内部 Seed-MoE 和 PLT 的性能评估。Lat. 和 KV cache 分别表示在批量大小为 4 下的解码延迟（毫秒）和 KV 缓存内存开销。

A4 实验环境

模型架构：
* 内部模型:
* Seed-MoE 680M/13B（6.8亿激活参数，130亿总参数）
* Seed-MoE 2.5B/60B（25亿激活参数，600亿总参数）
* PLT 模型: Seed-MoE 1.7B/40B（17亿激活参数，400亿总参数）

  • 开源模型 (基于 OLMo 和 OLMoE):
    • Dense 模型 (1.2B): 16层，隐藏维度2048，MLP隐藏维度16384，使用GQA（32个查询头，8个键/值头）。
    • MoE 模型 (1B/7B): 16层，隐藏维度2048，使用SwiGLU专家和“8 in 64”配置，使用MHA（16个注意力头）。
    • Ablation Study 模型: 较小规模的Dense和MoE模型。

数据集:

  • 训练:

    • 内部模型: 680M模型使用150B高质量词元进行训练；2.5B模型使用1T词元进行训练。
    • 开源模型: 主要实验使用400B词元进行训练；消融研究使用100B词元。

  • 评估:

    • 内部模型评估: MMLU, CEval, AGIEval, MMLU-Pro, BBH, DROP, MBPP, HumanEval, MATH, GSM8k。
    • 开源模型评估: MMLU, Hellaswag, ARC-Challenge, ARC-Easy, PIQA, Winogrande, Commonsense QA。

硬件配置:

  • 评估在单张GPU上进行。附录中提到消融实验使用了H800 GPU。

软件配置:
* 实现: 开源模型实现基于OLMo和OLMoE。
* 推理框架: 效率评估基于VLLM 【23，Efficient memory management for large language model serving with pagedattention，2023，Proceedings of the ACM SIGOPS 29th Symposium on Operating Systems Principles】。
* 量化: 为了模拟真实服务场景，推理时使用了FP8自注意力【33，Flashattention-3: Fast and accurate attention with asynchrony and low-precision，2024，Advances in Neural Information Processing Systems】和W4A8线性层量化【19，Liquidgemm: Hardware-efficient w4a8 gemm kernel for high-performance llm serving，2025，arXiv preprint arXiv:2509.01229】。

A4 实验结果

3.1 同等参数下的准确性比较

实验内容:
本节在内部 Seed-MoE 680M/13B 模型上，比较了普通 Transformer、普通循环 Transformer 以及 PLT 的不同变体。通过逐一添加跨循环并行（CLP）、KV 缓存共享和门控滑动窗口注意力（G-SWA）来评估每个组件的效果，并研究了循环次数 L 从2扩展到3的可行性。

实验结果与分析:
根据表 2 的数据，主要发现如下：

  • 观察 1：跨循环并行（CLP）在保持准确性的同时降低了延迟。

    • 与普通模型相比，一个包含两个循环的 Transformer（第(2)行）使平均准确率提升了5.0个点（34.7→39.7），但延迟和内存也大幅增加（延迟4.8→9.4 ms，+96%；内存280M→560M，+100%）。
    • 在相同循环次数下加入CLP后（第(3)行），准确率几乎不变（39.7→39.6），但通过并行化循环计算，延迟从9.4 ms大幅削减至5.9 ms（-37%）。CLP 在保留了循环带来的大部分准确率增益的同时，仅引入了适度的延迟开销（约1.23倍）。

  • 观察 2：高效表示增强以极小的准确率损失减少了 KV 缓存。

    • KV 缓存共享（第(4)行）消除了循环带来的额外 KV 缓存占用（560M→280M，-50%），并因 KV 缓存加载时间减少而降低了延迟（5.9→4.8 ms，-19%），但代价是平均准确率下降了3.4个点（39.6→36.2）。
    • 进一步加入 G-SWA（第(5)行），通过局部窗口恢复了每个循环的特异性，将准确率从36.2提升回39.7（+3.5），而 KV 缓存开销仅增加1.4%（280M→284M），延迟几乎无惩罚（4.8→4.9 ms）。

  • 观察 3：PLT 实现了循环级别的准确率和接近普通模型的效率，且扩展性良好。

    • 使用两个循环的 PLT-2（第(5)行）在准确率上与朴素循环 Transformer 持平（39.7），但效率接近普通模型（延迟+2%，KV缓存+1.4%）。
    • 扩展到 PLT-3（第(6)行）后，平均准确率进一步提升至40.8（比 PLT-2 高1.1），延迟和 KV 缓存仅有微小增加。这表明 PLT 成功地将延迟和内存与循环次数解耦。

推理效率分析:
如表 3 所示，在不同批量大小（batch size, bs）下，PLT 的延迟优势依然显著。

  • 低吞吐量场景 (bs=4, 8): KV 共享是降低延迟的关键，将延迟减少了19%。最终的 PLT-2 模型延迟仅比普通模型高约2%。
  • 高吞吐量场景 (bs=32, 64): CLP 的并发性是主要驱动力。与普通循环 Transformer 相比，PLT-2 将单次词元延迟削减了47%，与普通模型的延迟差距也缩小到4-6%。

表 3 PLT 变体在不同批量大小下的推理效率评估。

3.2 同等准确率下的延迟比较

实验内容:
本节旨在通过调整 PLT 模型的大小，使其准确率与一个更大的普通 Seed-MoE 基线模型相匹配，同时提升推理效率。实验对比了一个 2.5B/60B 的普通 Seed-MoE 模型和一个层数减少（为基线的三分之二）、采用 PLT-2 配置的 1.7B/40B MoE 模型。

实验结果与分析:
* 准确率: 如表 4 所示，1.7B/40B 的 PLT-2 MoE 模型平均准确率达到 62.6，比 2.5B/60B 的普通 MoE 模型高出 0.5 个点。
* 推理效率: 如图 3 所示，在从 4 到 64 的所有批量大小下，1.7B/40B 的 PLT-2 MoE 模型比 2.5B/60B 的普通模型延迟低约 30%。同时，由于层数减少，其 KV 缓存也约为基线的三分之二。
* 结论: 实验证明，PLT 提升了 Transformer 的可扩展性，能用更少的参数、更低的 KV 缓存和更低的推理延迟，实现相近甚至更高的准确率。

表 4 内部 Seed-MoE 2.5B 和 PLT-2 1.7B 的性能比较。


图 3 Seed-MoE (2.5B/60B) 和 PLT-2 (1.7B/40B) 的批量大小与延迟关系图。

A7 补充细节：相关工作

4.1 潜在推理

垂直潜在推理：循环 Transformer。Dehghani 等人【11，Universal transformers，2019，International Conference on Learning Representations】和 Lan 等人【25，Albert: A lite bert for self-supervised learning of language representations，2020，International Conference on Learning Representations】的研究引入了跨层参数共享和循环机制，为循环 Transformer 奠定了基础。Saunshi 等人【32，Reasoning with latent thoughts: On the power of looped transformers，2025，arXiv preprint arXiv:2502.17416】强调了循环 Transformer 在复杂推理任务中的强大能力。同时，Fan 等人【13，Looped transformers for length generalization，2024，arXiv preprint arXiv:2409.15647】和 Yang 等人【43，Looped transformers are better at learning learning algorithms，2023，arXiv preprint arXiv:2311.12424】展示了其在数据拟合方面的参数效率以及在 RASP-L 任务上优越的长度泛化能力。近期的工作如 Chen 等人【7，Inner thinking transformer: Leveraging dynamic depth scaling to foster adaptive internal thinking，2025，arXiv preprint arXiv:2502.13842】和 Geiping 等人【14，Scaling up test-time compute with latent reasoning: A recurrent depth approach，2025，arXiv preprint arXiv:2502.05171】通过循环为复杂词元分配更大的计算深度，使模型能够在不增加参数数量的情况下扩展推理时的推理能力。

水平潜在推理：潜在思维链（Latent CoT）。先前的工作【15，Think before you speak: Training language models with pause tokens，The Twelfth International Conference on Learning Representations】【46，Quiet-star: Language models can teach themselves to think before speaking，2024，First Conference on Language Modeling】通过在序列中插入特殊的离散、非语义词元来增强模型的推理能力，从而分配额外的计算步骤来优化其中间表示。相比之下，最近的研究【16，Training large language models to reason in a continuous latent space，2024，arXiv preprint arXiv:2412.06769】【35，Llm pretraining with continuous concepts，2025，arXiv preprint arXiv:2502.08524】将推理步骤压缩到连续的潜在思维中，实现了更高效和更具表达力的内部推理。这些潜在 CoT 方法在 GSM8K【9，Training verifiers to solve math word problems，2021，arXiv preprint arXiv:2110.14168】、NaturalQuestions【22，Natural questions: a benchmark for question answering research，2019，Transactions of the Association for Computational Linguistics】和 CommonsenseQA【36，Commonsenseqa: A question answering challenge targeting commonsense knowledge，2018，arXiv preprint arXiv:1811.00937】等推理基准上展示了显著的改进。

与 PLT 的区别。无论是循环 Transformer 还是潜在 CoT 推理，都因其固有的串行计算（无论是跨循环还是按词元）而导致推理效率低下，而我们提出的 PLT 创新地将不同词元中不同循环的计算重叠起来，从而实现了极高的推理效率。

4.2 并行化计算

PHD。PHD【42，Efficient pretraining length scaling，2025，arXiv preprint arXiv:2504.14992】提出利用并行计算可以通过并行化重复词元的前向传播来实现可扩展的性能提升。为了在保持最小内存访问开销的同时获得更好的性能，PHD 引入了 KV 缓存共享和分块滑动窗口注意力。PHD 的缺点在于词元重复的方法，这是一种低效的并行计算利用方式，因为前几层 Transformer 中的隐藏表示非常相似。在高吞吐量服务场景下，PHD 提升的性能无法弥补因解码计算增加而导致的吞吐量损失。

ParScale。随后，ParScale【6，Parallel scaling law for language models，2025，arXiv preprint arXiv:2505.10475】提出使用序列重复和前缀微调【27，Prefix-tuning: Optimizing continuous prompts for generation，2021，Proceedings of the 59th Annual Meeting of the Association for Computational Linguistics and the 11th International Joint Conference on Natural Language Processing (Volume 1: Long Papers)】也可以实现可扩展的性能提升。ParScale 的缺点在于其效率低下，当激活 P 个推理流时会引入 P 倍的 KV 缓存，导致 KV 缓存占用和推理延迟方面的开销，尤其是在高吞吐量服务场景下。

StagFormer。StagFormer【10，Stagformer: Time staggering transformer decoding for runninglayers in parallel，2025，arXiv preprint arXiv:2501.15665】提出了一种时间交错解码机制，通过将层拆分为多个堆栈来沿深度轴并行化 Transformer 层的执行，其中上层堆栈通过交叉注意力关注前一时间步的下层堆栈的隐藏状态。这实现了部分层级并行，但在注意力计算和内存访问方面仍然存在不完全并行的问题，导致效率提升有限。具体来说，其权重分离的变体虽然硬件使用量翻倍，但吞-吐量提升不到50%；而权重共享的变体虽然参数较轻，但与我们的 loop2+CLP 设计相比，会产生额外的 KV 缓存成本和交叉注意力开销，导致推理效率更差。

PLT 的优势。本文提出的 PLT 通过完整的并行化和在循环 Transformer 上的共享策略，进一步展现了更令人印象深刻的推理效率，在不降低性能的情况下将 KV 缓存占用减少了50%以上。PLT 将并行计算的利用率提升到了前所未有的水平。通过发现循环 Transformer 天然适合利用并行计算，我们提出跨循环并行来改进循环 Transformer，在保持循环 Transformer 性能的同时，实现了尤其出色的推理效率。

A5 结论

本文提出了并行化循环 Transformer（PLT），该模型提出了跨循环并行（CLP）技术，将先前词元的后续循环与后续词元的先前循环的计算和内存访问进行重叠。此外，还提出了门控滑动窗口注意力（gated-SWA），以在不降低性能的情况下实现 KV 缓存的并行访问。与普通 Transformer 相比，PLT 在性能上取得了显著提升，而延迟开销可以忽略不计。与具有相同推理计算预算的普通循环 Transformer 相比，PLT 在性能相似甚至更好的情况下，展现出明显更高的推理效率，显示了并行计算利用的巨大潜力。

A6 附录

A 开源实验

实验安排。我们的开源实验分为两部分：
* 我们在 A.2 节中展示了在激活参数超过10亿的稠密模型系列和 MoE 模型系列上的开源实验。
* 我们在 A.3 节中对稠密模型进行了类似的消融研究，作为我们内部 Seed-MoE 实验的补充。

A.1 评估数据集

数据集列表。我们在评估中使用了以下开源数据集，包括 MMLU 【17，Measuring massive multitask language understanding，2020，arXiv preprint arXiv:2009.03300】，HellaSwag 【47，Hellaswag: Can a machine really finish your sentence?，2019，arXiv preprint arXiv:1905.07830】，ARC 【8，Think you have solved question answering? try arc, the ai2 reasoning challenge，2018，arXiv preprint arXiv:1803.05457】，PIQA 【2，Piqa: Reasoning about physical commonsense in natural language，2020，Proceedings of the AAAI conference on artificial intelligence】，Winogrande 【31，Winogrande: An adversarial winograd schema challenge at scale，2021，Communications of the ACM】，以及 CommonsenseQA 【36，Commonsenseqa: A question answering challenge targeting commonsense knowledge，2018，arXiv preprint arXiv:1811.00937】。

A.2 主要实验

A.2.1 训练和评估设置

实现基础。我们的开源实现基于 OLMo 和 OLMoE。我们将我们的方法与具有相同激活参数的普通 Transformer 和普通循环 Transformer 进行了比较。

稠密模型设置。对于稠密模型，我们将 Transformer 层数设置为 16。隐藏维度设置为 2048。对于这些模型中的 MLP 模块，我们将 MLP 隐藏维度设置为 16384。我们使用 FFN 输出归一化。对于这些模型中的注意力模块，我们使用 GQA，有 32 个查询头和 8 个键/值头。我们同时使用查询/键/值层归一化，但不使用注意力输出归一化。对于 LM 头模块，我们使用权重绑定，与嵌入层共享参数。

训练动态。我们总共在 400B 词元上训练模型，全局训练批量大小设置为 1024。我们使用余弦学习率调度，峰值学习率设为 3e-4，预热步数为 2000 步。

设置说明。需要注意的是，与公开可用的开源模型 OLMo-1B 的原始设置相比，出于训练稳定性的考虑【50，Hybridnorm: Towards stable and efficient transformer training via hybrid normalization，2025，CoRR】，我们对归一化设置做了一些微调。

MoE 模型设置。对于 MoE 模型，我们将 Transformer 层数设置为 16。隐藏维度设置为 2048。对于这些模型中的 MLP 模块，我们使用 SwiGLU 专家和 8 in 64 的配方。对于这些模型中的注意力模块，我们总共使用 16 个注意力头的 MHA。对于 LM 头模块，我们也使用与嵌入层共享参数的权重绑定。

训练动态。我们总共在 400B 词元上训练模型，全局训练批量大小也设置为 1024。我们使用余弦学习率调度，峰值学习率设为 4e-4，预热步数为 2500 步。

效率评估设置。基于 VLLM【23，Efficient memory management for large language model serving with pagedattention，2023，Proceedings of the 29th Symposium on Operating Systems Principles, SOSP 2023】，我们分析了所有这些基线模型的效率，包括延迟和吞吐量。我们将预填充长度在 [1024, 2048] 范围内变化，服务批量大小在 [1, 2, 4, 8, 16, 32, 64, 128, 256] 范围内变化，这使得瓶颈逐渐从内存访问转向计算。为了更接近真实的工业服务场景，我们使用 FP8 量化进行服务。延迟指标是在解码 256 个词元的 5 次独立运行中取平均值。我们在单个 GPU 上进行效率分析实验。

A.2.2 分析

实验结果。表 5 展示了各种变体在稠密模型和 MoE 模型上的性能。观察结果与我们的内部实验相似，即 PLT-2 取得了与普通循环 Transformer 非常相似的性能。图 4 进一步展示了效率分析，我们也得到了类似的观察结果，即 PLT-2 表现出明显优于普通循环 Transformer 的效率。


图 4 对超过10亿激活参数的普通 Transformer、PLT 和循环 Transformer 的推理效率分析，包括延迟和吞吐量。我们在推理过程中基于 VLLM [23] 使用 FP8 量化。

表 5 在基线、L=2 的普通循环 Transformer 和我们提出的 L=2 的 PLT 之间，1.2B 稠密模型和 1B/7B MoE 模型的性能评估。评估的基准包括：MMLU、Hellaswag (Hella.)、ARC-Challenge (ARC-C)、ARC-Easy (ARC-E)、PIQA、Winogrande (Wino.) 和 Commonsense QA (Comm.)。

A.3 消融研究

A.3.1 训练/评估设置

稠密模型设置。本节的训练设置主要遵循 A.2.1 节的设置，除了以下两处修改：
* 我们改用 GQA，有 16 个查询头和 4 个键值头，并将隐藏维度设置为 1536，MLP 隐藏大小设置为 8192。
* 我们将训练词元数量改为 100B，峰值学习率为 3e-4，使用余弦学习率调度器，预热 2000 步。

MoE 模型设置。本节的训练设置主要遵循 A.2.1 节的设置。不同的设置也如下所示：
* 我们将层数改为 12，隐藏维度改为 1536。
* 对于训练动态，我们也设置训练词元数量为 100B，峰值学习率为 3e-4，使用余弦学习率调度器，预热 2000 步。

评估设置。评估设置主要遵循 A.2.1 节，只是我们使用一块 H800 GPU进行评估。

A.3.2 分析

实验结果。表 A.3.1 展示了我们在 OLMo2 上的消融研究，其观察结果与我们的内部实验相似。

表 6 PLT 的组件分析。与普通循环 Transformer 相比，我们出于推理速度和性能的考虑，引入了两个额外的组件。Lat. 是延迟（ms）的缩写。