Skip to main content
Version: Next

Transformer 逐层推导

2017 年《Attention Is All You Need》提出的 Transformer 是全部现代 LLM 的骨架。本篇按"一个 token 流过模型的完整路径"逐层拆解，每层给出数学与 PyTorch 代码，全部是现代 LLM（LLaMA 系）采用的当前最佳实践：Pre-Norm、RMSNorm、RoPE、SwiGLU。

全局结构​

token ids ──→ Embedding ──→ ┌────────── Transformer Block ×N ──────────┐
                            │  x = x + Attention(RMSNorm(x))           │
                            │  x = x + FFN(RMSNorm(x))                 │
                            └──────────────────────────────────────────┘
                                   ──→ RMSNorm ──→ Linear ──→ logits

输入 (B, T) 整数序列，输出 (B, T, vocab_size) 的 logits。

1. Embedding：id → 向量​

x = W_emb[token_ids]        W_emb: (vocab_size, d_model)

本质是一次查表，W_emb 是可训练参数。训练结束后，语义相近的 token 在嵌入空间中距离相近。

tok_emb = nn.Embedding(vocab_size, d_model)
x = tok_emb(idx)                       # (B, T) → (B, T, d_model)

LLaMA 等模型不再使用任何加性位置编码（不在这里加位置信息），位置完全交给注意力里的 RoPE 处理。

2. RMSNorm：更省更稳的归一化​

LayerNorm 做"减均值除方差 + 仿射"；RMSNorm 发现减均值并不重要，只做均方根缩放：

RMSNorm(x) = x / sqrt(mean(x²) + ε) · g        g: 可学习增益向量

省掉均值计算与对应梯度，速度快 ~10-40%，效果无损。

class RMSNorm(nn.Module):
    def __init__(self, dim, eps=1e-6):
        super().__init__()
        self.eps = eps
        self.weight = nn.Parameter(torch.ones(dim))

    def forward(self, x):
        norm = x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)
        return norm * self.weight

Pre-Norm vs Post-Norm：现代 LLM 一律 Pre-Norm（先归一化再进子层，残差支路保持原始 x），深层训练更稳，不需要精心调 warmup。

3. 因果多头注意力（核心）​

单头注意力的数学​

Q = xW_q,  K = xW_k,  V = xW_v                每个 (T, d_k)
scores = QK^T / sqrt(d_k)                      (T, T)
scores = scores + causal_mask                  上三角置 -inf
A = softmax(scores)                            每行和为 1
out = A V                                      (T, d_k)

逐步理解：

  • Q·K 打分：第 i 个 token 的 query 与第 j 个 token 的 key 内积 → "i 应该多关注 j"。
  • 除以 √d_k：d_k 大时内积方差随之变大，softmax 会饱和到 one-hot、梯度消失。缩放让分数方差保持在 O(1)。
  • 因果掩码：scores[i][j] = -inf (j > i)，softmax 后变 0——token 只能看自己及左边，这是自回归合法性的保证，也是"GPT 式 Decoder-only"与 BERT 式双向编码器的分水岭。
  • 加权和 V：按注意力权重聚合各 token 的信息。
class CausalSelfAttention(nn.Module):
    def __init__(self, d_model, n_head):
        super().__init__()
        assert d_model % n_head == 0
        self.n_head = n_head
        self.d_head = d_model // n_head
        self.qkv = nn.Linear(d_model, 3 * d_model, bias=False)
        self.proj = nn.Linear(d_model, d_model, bias=False)

    def forward(self, x):
        B, T, C = x.shape
        q, k, v = self.qkv(x).chunk(3, dim=-1)
        # (B, T, C) → (B, n_head, T, d_head)
        q = q.view(B, T, self.n_head, self.d_head).transpose(1, 2)
        k = k.view(B, T, self.n_head, self.d_head).transpose(1, 2)
        v = v.view(B, T, self.n_head, self.d_head).transpose(1, 2)

        scores = q @ k.transpose(-2, -1) / math.sqrt(self.d_head)
        mask = torch.triu(torch.ones(T, T, device=x.device, dtype=torch.bool), 1)
        scores = scores.masked_fill(mask, float("-inf"))
        attn = torch.softmax(scores, dim=-1)
        out = attn @ v                         # (B, n_head, T, d_head)

        out = out.transpose(1, 2).contiguous().view(B, T, C)
        return self.proj(out)

为什么是多头？​

一个头只能在一个子空间里度量相关性。把 d_model 切成 n_head 份，各头独立做注意力再拼接——有的头学语法搭配，有的头学指代，有的头学位置邻近。维度分配：每头 d_head = d_model / n_head，常见 d_head=64/128。

工程实现提示​

实际代码会用 F.scaled_dot_product_attention（PyTorch ≥2.0），它自动调度 FlashAttention 内核：分块计算、不物化 (T,T) 矩阵，显存从 O(T²) 降到 O(T)，速度快数倍：

out = F.scaled_dot_product_attention(q, k, v, is_causal=True)

4. RoPE：用旋转编码位置​

绝对位置编码（sinusoidal/可学习）是"加"在 embedding 上的，外推到更长上下文时表现差。RoPE（旋转位置编码，LLaMA/Qwen 标配）的思路：给 q、k 的每对维度按位置做旋转，使得内积自动携带相对位置信息。

数学：把 d_head 维向量两两分组，第 m 个位置、第 i 对维度旋转角度 m·θᵢ，其中 θᵢ = 10000^(-2i/d)：

(q′₂ᵢ,   q′₂ᵢ₊₁) = (q₂ᵢ cos mθᵢ - q₂ᵢ₊₁ sin mθᵢ,  q₂ᵢ sin mθᵢ + q₂ᵢ₊₁ cos mθᵢ)

由于 ⟨RoPE(q,m), RoPE(k,n)⟩ = f(q, k, m−n)，注意力分数只依赖相对距离——模型天然学到"隔多远"而不是"在第几个位置"。

def apply_rope(x, cos, sin):
    # x: (B, n_head, T, d_head); cos/sin: (T, d_head/2)
    x1, x2 = x[..., ::2], x[..., 1::2]
    out1 = x1 * cos - x2 * sin
    out2 = x1 * sin + x2 * cos
    return torch.stack([out1, out2], dim=-1).flatten(-2)

# 预计算
def precompute_rope(d_head, max_len, base=10000.0):
    inv_freq = base ** (-torch.arange(0, d_head, 2).float() / d_head)
    t = torch.arange(max_len).float()
    freqs = torch.outer(t, inv_freq)           # (T, d_head/2)
    return freqs.cos(), freqs.sin()

长上下文扩展（YaRN、NTK 缩放、位置插值）都是在 base 或 freqs 上做文章，本篇不展开。

5. SwiGLU：现代 FFN​

每个 Block 的第二半是前馈网络，承担"逐 token 的非线性知识加工"。原版是 GELU(xW₁)W₂；现代标准是门控变体 SwiGLU：

FFN(x) = (Swish(xW_gate) ⊙ xW_up) W_down
Swish(x) = x·sigmoid(x)

三个矩阵、中间维度约为 (8/3)d_model 再取整到 256 的倍数，总参数量与原版持平但效果更优。

class SwiGLU(nn.Module):
    def __init__(self, d_model, hidden):
        super().__init__()
        self.gate = nn.Linear(d_model, hidden, bias=False)
        self.up   = nn.Linear(d_model, hidden, bias=False)
        self.down = nn.Linear(hidden, d_model, bias=False)

    def forward(self, x):
        return self.down(F.silu(self.gate(x)) * self.up(x))

6. 残差连接与输出头​

每个子层都是 x = x + Sublayer(RMSNorm(x))。残差让梯度直通浅层，也让"删除某一层"近似等于恒等——这是深网络可训练的根本原因。

最后过一遍 RMSNorm 后接输出头：

logits = x_final W_out        W_out: (d_model, vocab_size)

很多实现让 W_out = W_emb^T（权重绑定），省 1/vocab 的参数且效果略好。

7. 参数量速算​

以 LLaMA-7B 为例（d_model=4096, n_layer=32, vocab=32000, hidden≈11008）：

Embedding:        32000 × 4096            ≈ 0.13B
每层注意力:       4 × 4096²               ≈ 67M
每层 FFN:         3 × 4096 × 11008        ≈ 135M
每层合计 ≈ 202M × 32 层 ≈ 6.5B，加 embedding ≈ 6.7B ✓

注意力与 FFN 参数量比约 1:2，这是通用配比。注意：每层 KV 投影在现代模型中常改用 GQA（多个 Q 头共享 K/V 头）以压缩推理显存，详见推理与部署。

层数、宽度与深度的权衡​

  • 深度（层数）更影响"推理/组合"能力，宽度（d_model）更影响"记忆容量"；
  • 深而窄的模型训练更不稳，宽而浅则表达效率低——主流比例大致 d_model : n_layer ∈ [64, 256]；
  • 上下文长度 T 只影响计算量与显存（注意力 O(T²)，已被 FlashAttention 缓解），不影响参数量。

所有零件就绪，下一篇 从零实现 mini-LLM 把它们拼成一个能训练、能生成文字的完整模型。