Looped Transformer：参数共享与重复计算#

整理于 2026-09-08。本文采用 Pre-Norm、decoder-only 的示意结构；原始 Transformer采用 Encoder–Decoder 与 Post-LN，两者的归一化顺序和模块组成不同。循环机制参考 Universal Transformers 与 Scaling up Test-Time Compute with Latent Reasoning。本文只记录已讨论的结构，不展开具体模型的训练与退出策略。

Looped Transformer 让同一组 Transformer Block 反复处理内部状态，把独立参数量和计算深度部分分开：不新增一套层参数，也可以多执行一轮计算。

这次理解的转折是：先知道了“参数共享、表示在变化”，但还不知道一个 Block 里面是什么。把 Block 拆开后，Loop 就不再是一个黑盒套另一个黑盒。

一句话版本：

同一组 Block 反复计算；参数共享，hidden states 持续更新。

1. 先分清参数和状态#

最初的理解是“参数共享，embedding 在变化”。方向对，但这里应该把 embedding 换成 hidden states（隐藏状态），避免把参数表和计算结果混在一起。

这里的“不修改参数”说的是一次前向计算内部，不是说这些参数不参与训练。训练时优化器仍会更新它们。

以 decoder-only 模型为例，忽略位置编码等细节，数据流为：

token IDs
    ↓
Embedding 查表
    ↓
H₀ → Block 1 → H₁ → Block 2 → H₂ → … → Block L → Hₗ
                                                        ↓
                                                   最后的 Norm
                                                        ↓
                                            LM Head：投影到词表
                                                        ↓
                                               下一个 token 的 logits

Embedding 和 LM Head 在这组 Block 外面。 经过 Block 后，传给下一层的是刚算出的 hidden states，不是重新对原始 token IDs 查一次表。

忽略 batch，记：

$$H \in \mathbb{R}^{T\times D}$$

$T$ 是 token 数，$D$ 是每个 token 的表示维度；每一行对应一个 token 当前的表示。例如：

10 个 token，每个 token 512 维：

H：[10, 512] → 一个 Block → H′：[10, 512]

形状不变，不代表数值不变。 Block 内部可以升维、分头，输出时仍回到 $T\times D$，所以能继续接下一个 Block。

2. 一个 Transformer Block 里面有什么#

这里采用 Pre-Norm Block。这是理解结构用的具体版本，不把它当成所有 Transformer 变体的唯一模板。

先记住两个子层：

Attention 汇集不同 token 的信息；FFN 分别加工每个 token 的表示。两者各配一组 Norm 和残差连接。

2.1 完整计算只有两行#

$$U = H + \operatorname{Attention}\bigl(\operatorname{Norm}_1(H)\bigr)$$

$$H' = U + \operatorname{FFN}\bigl(\operatorname{Norm}_2(U)\bigr)$$

对应前向示意如下。省略 dropout；attn、ffn、ln1、ln2 是 Block 初始化时创建的模块。

def forward(self, x, mask=None):
    x = x + self.attn(self.ln1(x), mask)
    x = x + self.ffn(self.ln2(x))
    return x

第二行使用第一行更新后的 x。 两个 Norm 各有自己的参数；残差分支加回的是原始的 $H$ 或 $U$，不是归一化后的向量。

2.2 Attention：从允许访问的位置汇集信息#

先得到 $X=\operatorname{Norm}_1(H)$，再投影出：

$$Q=XW_Q,\qquad K=XW_K,\qquad V=XW_V$$

对其中一个 attention head：

$$A=\operatorname{softmax}\left(\frac{QK^\top}{\sqrt{d_h}}+M\right),\qquad Z=AV$$

$d_h$ 是这个 head 的 Q/K 维度，$M$ 是加性 mask。因果语言模型中，允许访问的位置为 0，未来位置为 $-\infty$；softmax 沿 key 位置计算。多个 head 的结果拼接后，再通过输出投影 $W_O$ 回到 $T\times D$。

在 $AV$ 中，一个位置的输出可以汇集其他允许访问位置的 V，因此 Attention 在这里负责 token 之间的信息交换。

难点在于分清两类“权重”：

注意力矩阵 $A$ 不是一张固定的参数表。 同一套投影矩阵遇到不同的 hidden states，会算出不同的 Q/K/V，注意力分布也可能变化。

2.3 FFN：逐 token 做非线性加工#

FFN（Feed-Forward Network）在这个版本中是一个小型 MLP：

$$\operatorname{FFN}(x)=\operatorname{GELU}(xW_{\mathrm{up}}+b_{\mathrm{up}})W_{\mathrm{down}}+b_{\mathrm{down}}$$

先升维，经过非线性函数，再降回原维度。已有代码默认中间维度为 $4D$，因此示例中的形状是：

[10, 512] → Linear → [10, 2048] → GELU → [10, 2048] → Linear → [10, 512]

token 数没有改变，变化的是每个 token 内部的特征维度。 所有位置使用同一套 FFN 参数，但分别计算；第 $i$ 个位置的 FFN 输出只取决于这个位置送入 FFN 的向量。

从矩阵运算看，Attention 的 $AV$ 在组合不同行，而 FFN 的 $xW$ 在每一行内部变换特征。不过，FFN 的输入已经经过 Attention，不直接访问其他 token，不等于没有上下文信息。

2.4 Norm 和残差：先调整尺度，再把更新加回状态#

这里的 Norm 使用 LayerNorm，对每个 token 自己的 $D$ 个特征归一化：

$$\operatorname{LN}(h)=\gamma\odot\frac{h-\mu(h)}{\sqrt{\sigma^2(h)+\epsilon}}+\beta$$

$\gamma,\beta$ 是可学习参数；均值和方差由当前输入计算。Norm 不改变 $T\times D$ 的形状。“Pre-Norm”就是把 Norm 放在 Attention / FFN 前面。

残差连接则是：

$$\text{新状态}=\text{旧状态}+\text{本次算出的更新}$$

所以完整 Block 可以读成：

当前状态 H
  → 归一化后做 Attention，把结果加回 H，得到 U
  → 归一化后做 FFN，把结果加回 U，得到 H′

3. 普通堆叠和循环有什么区别#

3.1 普通堆叠：结构相同，参数不同#

把完整 Block 命名为 $B_\theta$。这里 $\theta$ 包括 Attention 投影、FFN 和两组 Norm 的可学习参数，而不包括每次计算产生的 Q/K/V 或 hidden states。

普通的三个 Block 是：

$$H_0\xrightarrow{B_{\theta_1}}H_1\xrightarrow{B_{\theta_2}}H_2\xrightarrow{B_{\theta_3}}H_3$$

它们可以长得完全一样，但各有自己的参数。“相同结构的层”不等于“同一层对象被反复调用”。

3.2 单 Block 循环：重复调用同一套参数#

最简单的循环是：

$$H_0\xrightarrow{B_\theta}H_1\xrightarrow{B_\theta}H_2\xrightarrow{B_\theta}H_3$$

也就是：

$$\boxed{H_{r+1}=B_\theta(H_r)}$$

$r$ 表示计算轮次，不是 token 的位置。共享的是完整 Block 的参数，更新的是它的输入输出状态。

3.3 一组 Block 循环：组内独立，跨轮共享#

“8 层循环 4 次”不是全模型只有一个 Block 的参数，而是：

第 1 轮：B1 → B2 → … → B8
第 2 轮：B1 → B2 → … → B8
第 3 轮：B1 → B2 → … → B8
第 4 轮：B1 → B2 → … → B8

B1 到 B8 各有自己的参数；每轮复用同一组 B1 到 B8。
上一轮的输出状态，接着作为下一轮的输入状态。

把这一组整体记为 $F_\Theta=B_{\theta_8}\circ\cdots\circ B_{\theta_1}$，其中 $\Theta=(\theta_1,\ldots,\theta_8)$，则仍可写成：

$$H_{r+1}=F_\Theta(H_r)$$

结构伪代码如下，省略位置编码、缓存等实现细节；shared_blocks 在循环外创建：

h = embedding(token_ids)

for _ in range(num_loops):
    for block in shared_blocks:
        h = block(h, mask)

logits = lm_head(final_norm(h))

这里既没有在每轮重新创建一组 Block，也没有把 h 重置回初始 embedding。

4. 同样的参数，为什么还能继续计算#

复用计算规则，不是复用上轮的计算结果。 一般来说：

$$B_\theta(B_\theta(H))\neq B_\theta(H)$$

把完整 Block 的更新抽象成：

$$H_{r+1}=H_r+\Delta_\theta(H_r)$$

第一轮计算 $\Delta_\theta(H_0)$，第二轮计算 $\Delta_\theta(H_1)$。输入已经不同，产生的更新也可以不同。

联系刚拆开的 Block 就更具体了：

参数 W_Q / W_K / W_V 不变
        ↓
本轮 hidden states 改变
        ↓
重新计算 Q / K / V，注意力分布可以改变
        ↓
FFN 也对新的输入重新计算
        ↓
得到下一轮 hidden states

可以用迭代求解器作类比：同一套更新规则，作用在不断变化的当前解上。这个类比说明“再次计算可以不同”，不保证模型一定收敛，也不保证每多一轮答案就更好。

它更不意味着每轮必定有“理解题目、制定计划、检查答案”这样的人类可读分工。结构本身只规定了参数复用与状态更新。

5. 用 8 层 × 4 轮算一笔账#

假设每个 Block 的参数量都是 $P$，对相同输入执行一次的计算量都是 $C$。暂时忽略 embedding、输出头和额外模块，只比较这些 Block：

和普通 8 层相比，循环版本没有增加这组 Block 的独立参数，但多执行了计算，计算链更深。

和普通 32 层相比，循环版本的层执行次数相同，但独立参数只有四分之一。可以把它理解为一个受跨层权重共享约束的 32 层展开网络，而不是拥有 32 套独立参数的网络。

因此：

8 层循环 4 次 = 32 次层计算，不等于自动获得普通 32 层模型的全部能力。

这个账本只回答参数量和计算深度。它不直接推出总显存、实际延迟或精度；“参数少了”不等于“计算免费了”。训练策略、自适应停止和跨轮 KV cache 优化留到后续单独讨论，这一版不展开。

6. 要点回顾#

Looped Transformer = Block 参数共享 + hidden states 逐轮更新。 先用 Attention 汇集信息，再用 FFN 逐 token 加工；重复的是完整计算过程。收益是用较少的独立参数支持更深的计算，代价是仍需执行那些计算，而效果不能只靠循环次数保证。