模型内部 · LayerNorm 与残差

LayerNorm 与残差

前置 LN · 第 1 块的 ln1

图中数值
玩具规模
特征
8 GPT-2 768
LayerNorm
25 每块 2 个 + lnf
每个的参数
16 GPT-2 1,536
ε
1e-5

全部步骤

  1. 残差流

    残差流径直穿过每一个块：每个子层读取一份归一化后的副本，再把结果加回去，所以信息只会被累加，而这条直通路径让训练保持稳定。GPT-2 在每个子层之前归一化（前置 LN）；2017 年的 Transformer 则是在每次相加之后归一化。点击 attn 或 mlp 打开它们。

    GPT-2：前置 LN，12 个块

  2. 减去均值

    LayerNorm 一次处理一个词元，作用在它的各个特征上。首先减去该词元特征的均值，让它们以 0 为中心。（视觉模型中常见的 BatchNorm 则是对每个特征在整个批次上归一化；那样一个词元的结果会依赖批次里的其他序列，所以 Transformer 用 LayerNorm。）

    在 768 个特征上求 μ

  3. 除以标准差

    然后全部除以标准差，这样无论残差流长到多大，每个词元的特征都有相同的离散度。这让每个子层的输入保持在它训练时见过的范围内。

    σ = √(var + ε)

  4. 缩放与平移

    最后，每个特征乘以 γ、加上 β，两者都是学到的。结果 X 就是注意力层收到的输入。

    X = γ ⊙ x̂ + β

代码

class Block(nn.Module):
    def forward(self, x):
        x = x + self.attn(self.ln_1(x))   # read a normalised copy, add back
        x = x + self.mlp(self.ln_2(x))
        return x

# what ln_1 computes for each token's 768 features (F.layer_norm):
mu = x.mean(-1, keepdim=True)
var = x.var(-1, keepdim=True, unbiased=False)
x_hat = (x - mu) / torch.sqrt(var + 1e-5)
y = self.weight * x_hat + self.bias      # γ ⊙ x̂ + β

延伸阅读