基础 第 3 讲 / 共 17 讲

第 3 讲：Transformer 架构与 Hyperparameters

1. Self-Attention 机制

Self-Attention 是 Transformer 的核心组件。它允许序列中的每个位置"关注"序列中的其他位置， 从而捕获长距离依赖关系。

1.1 Query, Key, Value

Self-Attention 的核心思想借鉴了信息检索：给定一个 Query（查询）， 在所有 Key（键）中找到最相关的，然后返回对应 Value（值）的加权组合。

具体地，对于输入序列 X ∈ ℝT×d：

# Linear projections
Q = X @ W_Q  # (T, d) @ (d, d_k) = (T, d_k)
K = X @ W_K  # (T, d) @ (d, d_k) = (T, d_k)
V = X @ W_V  # (T, d) @ (d, d_v) = (T, d_v)

其中 WQ、WK、WV 是可学习的投影矩阵。 直觉上：Q 代表"我在找什么"，K 代表"我提供什么"，V 代表"我的实际内容"。

1.2 Scaled Dot-Product Attention

计算 attention 的完整公式：

Attention(Q, K, V) = softmax(QKT / √dk) · V

逐步分解：

  1. 计算相似度：QKT ∈ ℝT×T，每个元素表示两个位置之间的相关性
  2. 缩放：除以 √dk，防止 dot product 在高维空间中值过大导致 softmax 饱和
  3. Causal Mask：在 decoder 中，将未来位置的 attention score 设为 -∞，确保每个位置只能看到自己和之前的内容
  4. Softmax：归一化为概率分布
  5. 加权求和：用 attention weights 对 V 加权求和
import torch
import torch.nn.functional as F

def scaled_dot_product_attention(Q, K, V, mask=None):
    """
    Q: (B, T, d_k)
    K: (B, T, d_k)
    V: (B, T, d_v)
    mask: (T, T) causal mask
    """
    d_k = Q.shape[-1]

    # Step 1-2: Compute scaled attention scores
    scores = (Q @ K.transpose(-2, -1)) / (d_k ** 0.5)  # (B, T, T)

    # Step 3: Apply causal mask
    if mask is not None:
        scores = scores.masked_fill(mask == 0, float('-inf'))

    # Step 4: Softmax
    attn_weights = F.softmax(scores, dim=-1)  # (B, T, T)

    # Step 5: Weighted sum
    output = attn_weights @ V  # (B, T, d_v)

    return output

为什么要缩放？

假设 Q 和 K 的每个元素独立服从 N(0,1)，那么 QKT 的每个元素的方差为 dk。当 dk 很大时（如 128），dot product 的值会很大，softmax 会接近 one-hot 分布，梯度几乎为零。除以 √dk 将方差归一化为 1，缓解这个问题。

2. Multi-Head Attention

单一的 attention 只能捕获一种"关注模式"。Multi-Head Attention 通过并行运行多个 attention head，让模型同时关注不同的表示子空间。

MultiHead(X) = Concat(head1, ..., headh) · WO
where headi = Attention(X · WQi, X · WKi, X · WVi)

关键参数关系：

class MultiHeadAttention(nn.Module):
    def __init__(self, d_model: int, n_heads: int):
        super().__init__()
        self.n_heads = n_heads
        self.d_k = d_model // n_heads

        # Combined QKV projection for efficiency
        self.qkv_proj = nn.Linear(d_model, 3 * d_model, bias=False)
        self.out_proj = nn.Linear(d_model, d_model, bias=False)

    def forward(self, x):
        B, T, d = x.shape

        # Project to Q, K, V
        qkv = self.qkv_proj(x)  # (B, T, 3*d)
        Q, K, V = qkv.chunk(3, dim=-1)  # each (B, T, d)

        # Reshape for multi-head: (B, T, d) -> (B, n_heads, T, d_k)
        Q = Q.view(B, T, self.n_heads, self.d_k).transpose(1, 2)
        K = K.view(B, T, self.n_heads, self.d_k).transpose(1, 2)
        V = V.view(B, T, self.n_heads, self.d_k).transpose(1, 2)

        # Scaled dot-product attention
        attn = scaled_dot_product_attention(Q, K, V, causal_mask)

        # Reshape back: (B, n_heads, T, d_k) -> (B, T, d)
        attn = attn.transpose(1, 2).contiguous().view(B, T, d)

        return self.out_proj(attn)

Grouped Query Attention (GQA)

LLaMA 2/3 采用了 GQA：多个 Query head 共享同一组 Key 和 Value head。例如，32 个 Q head 可能只有 8 个 KV head。这减少了 KV cache 的大小，在推理时尤其有用。GQA 是 Multi-Head Attention（每个 head 独立 KV）和 Multi-Query Attention（所有 head 共享一组 KV）之间的折中方案。

3. Positional Encoding

Self-Attention 本身是置换不变（permutation-invariant）的—— 它不关心 token 的顺序。因此需要显式注入位置信息。

3.1 绝对位置编码

原始 Transformer 使用正弦/余弦函数生成固定的位置编码，或者使用可学习的 position embedding。 缺点是：无法泛化到训练时未见过的序列长度。

3.2 RoPE（Rotary Position Embedding）

RoPE 是目前 LLM 中最主流的位置编码方法（LLaMA, Mistral 等均采用）。 其核心思想是：通过旋转 Query 和 Key 向量来注入位置信息， 使得两个位置之间的 attention score 只依赖于它们的相对距离。

对于位置 m 处的向量 x，RoPE 将相邻的两个维度视为一个 2D 向量，并旋转角度 mθ：

RoPE(xm, m) = R(m) · xm
其中 R(m) 是分块对角旋转矩阵
θi = 10000-2i/d 
def apply_rope(x, positions):
    """
    x: (B, n_heads, T, d_k)
    positions: (T,) position indices
    """
    d_k = x.shape[-1]

    # Compute rotation angles
    freqs = 1.0 / (10000.0 ** (torch.arange(0, d_k, 2).float() / d_k))
    angles = positions.unsqueeze(-1) * freqs.unsqueeze(0)  # (T, d_k/2)

    cos = angles.cos()  # (T, d_k/2)
    sin = angles.sin()  # (T, d_k/2)

    # Split into pairs and rotate
    x1 = x[..., ::2]   # Even dimensions
    x2 = x[..., 1::2]  # Odd dimensions

    rotated = torch.stack([
        x1 * cos - x2 * sin,
        x1 * sin + x2 * cos,
    ], dim=-1).flatten(-2)

    return rotated

RoPE 的数学美妙之处在于：

⟨RoPE(q, m), RoPE(k, n)⟩ = ⟨q, R(n-m) · k⟩

即两个位置的 attention score 只取决于相对距离 (n-m)，而非绝对位置。 这使得模型可以更好地泛化到不同长度的序列。

4. Feed-Forward Network

Transformer 中每个 Attention 层后都跟着一个 Feed-Forward Network（FFN）。 标准 FFN 是一个两层的 MLP：

FFN(x) = W2 · σ(W1 · x + b1) + b2

4.1 SwiGLU 激活函数

现代 LLM（如 LLaMA）通常使用 SwiGLU 替代传统的 ReLU/GELU。 SwiGLU 结合了 Swish 激活函数和 Gated Linear Unit (GLU) 的思想：

SwiGLU(x) = (W1x ⊙ Swish(Wgatex)) · W2
where Swish(x) = x · σ(x), σ is sigmoid 
class SwiGLUFFN(nn.Module):
    def __init__(self, d_model: int, d_ff: int):
        super().__init__()
        self.w1 = nn.Linear(d_model, d_ff, bias=False)
        self.w_gate = nn.Linear(d_model, d_ff, bias=False)
        self.w2 = nn.Linear(d_ff, d_model, bias=False)

    def forward(self, x):
        # Gate branch with Swish activation
        gate = F.silu(self.w_gate(x))  # silu = swish = x * sigmoid(x)
        # Up projection
        up = self.w1(x)
        # Element-wise product and down projection
        return self.w2(gate * up)

参数量变化

注意 SwiGLU 有三个权重矩阵（W1, Wgate, W2），而标准 FFN 只有两个。为了保持总参数量不变，使用 SwiGLU 时通常将 d_ff 设为 (2/3) × 4d，而非标准的 4d。例如 LLaMA-7B 的 d_model=4096, d_ff=11008 ≈ (2/3) × 4 × 4096。

5. Layer Normalization

5.1 Pre-Norm vs Post-Norm

原始 Transformer 采用 Post-Norm 结构（先 sublayer 再 norm）， 但现代 LLM 几乎全部使用 Pre-Norm（先 norm 再 sublayer）： 

# Post-Norm (original Transformer)
x = LayerNorm(x + Attention(x))
x = LayerNorm(x + FFN(x))

# Pre-Norm (modern LLMs)
x = x + Attention(RMSNorm(x))
x = x + FFN(RMSNorm(x))

Pre-Norm 的优势是训练更稳定，不需要 learning rate warmup 也能训练。 直觉上，Pre-Norm 中的 residual connection 直接连接输入和输出， 梯度可以更顺畅地从输出层流回输入层。

5.2 RMSNorm

RMSNorm 是 LayerNorm 的简化版本，去掉了均值中心化步骤，只做方差归一化：

LayerNorm(x) = γ ⊙ (x - μ) / √(σ² + ε) + β

RMSNorm(x) = γ ⊙ x / √(mean(x²) + ε) 
class RMSNorm(nn.Module):
    def __init__(self, d_model: int, eps: float = 1e-6):
        super().__init__()
        self.weight = nn.Parameter(torch.ones(d_model))
        self.eps = eps

    def forward(self, x):
        # RMS = root mean square
        rms = torch.sqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)
        return x / rms * self.weight

RMSNorm 比 LayerNorm 更快（少了一次 mean 计算和减法操作）， 而且在实践中效果几乎一样好。

6. Transformer Block 完整结构

将以上组件组合起来，一个现代 Transformer Block（如 LLaMA 风格）的完整结构如下： 

class TransformerBlock(nn.Module):
    def __init__(self, d_model, n_heads, d_ff):
        super().__init__()
        self.attn_norm = RMSNorm(d_model)
        self.attention = MultiHeadAttention(d_model, n_heads)
        self.ffn_norm = RMSNorm(d_model)
        self.ffn = SwiGLUFFN(d_model, d_ff)

    def forward(self, x, positions):
        # Self-attention with pre-norm and residual
        h = self.attn_norm(x)
        h = self.attention(h, positions)
        x = x + h

        # FFN with pre-norm and residual
        h = self.ffn_norm(x)
        h = self.ffn(h)
        x = x + h

        return x

class Transformer(nn.Module):
    def __init__(self, vocab_size, d_model, n_heads, n_layers, d_ff):
        super().__init__()
        self.embedding = nn.Embedding(vocab_size, d_model)
        self.layers = nn.ModuleList([
            TransformerBlock(d_model, n_heads, d_ff)
            for _ in range(n_layers)
        ])
        self.norm = RMSNorm(d_model)
        self.output = nn.Linear(d_model, vocab_size, bias=False)

    def forward(self, token_ids):
        x = self.embedding(token_ids)
        positions = torch.arange(token_ids.shape[1])

        for layer in self.layers:
            x = layer(x, positions)

        x = self.norm(x)
        logits = self.output(x)
        return logits

Weight Tying

一些模型（如 GPT-2）将 embedding 层和 output projection 层共享权重（weight tying），即 self.output.weight = self.embedding.weight。这减少了参数量（节约 vocab_size × d_model 个参数），但在大模型中收益递减。LLaMA 系列不使用 weight tying。

7. 超参数选择

Transformer 有多个关键超参数，它们之间存在经验性的比例关系：

超参数含义典型关系
d_model模型隐藏维度基础参数
n_headsAttention head 数量d_model / d_k（d_k 通常为 64 或 128）
n_layersTransformer 层数与 d_model 成正比
d_ffFFN 中间维度≈ (8/3) × d_model（SwiGLU 时）
vocab_size词表大小32K - 128K
context_length最大序列长度2K - 128K

7.1 参数量估算

对于一个 L 层的 Transformer，总参数量近似为：

N ≈ 12 × L × d² (不含 embedding)
N ≈ 12 × L × d² + V × d (含 embedding，V 为词表大小)

其中 12d² 来自每层的：QKV projection (3d²) + output projection (d²) + FFN up/down (8d²，SwiGLU 时)。

7.2 代表性模型配置

模型参数量d_modeln_headsn_layersd_ff
GPT-2 Small124M76812123072
GPT-2 XL1.5B160025486400
LLaMA 2 7B6.7B4096323211008
LLaMA 2 70B70B8192648028672
LLaMA 3 405B405B1638412812653248

8. 现代架构变体

相比原始 Transformer，现代 LLM 做了多项改进。以 LLaMA 系列为例：

改进点原始 TransformerLLaMA 风格
NormalizationPost-Norm + LayerNormPre-Norm + RMSNorm
位置编码Sinusoidal / LearnedRoPE
激活函数ReLU / GELUSwiGLU
AttentionMulti-HeadGrouped Query (GQA)
Bias有 bias无 bias（所有 Linear 层）
Weight Tying通常使用通常不使用

为什么去掉 Bias？

去掉 bias 的主要原因是简化模型，减少参数量和计算量。实验表明，在大规模预训练场景下，bias 对模型性能的影响可以忽略不计，而去掉 bias 后可以简化 Tensor Parallelism 等并行策略的实现。

9. 本讲总结