Transformer 核心机制 —— 代码对照笔记

tags:

• transformer
• self-attention
• deep-learning
• NLP
• GPT

created: 2026-06-02 source: "[[self-attention.py]] 代码逐段分析" paper: "[[Attention Is All You Need]]"

基于 self-attention.py（从零构建 GPT 风格语言模型）逐段分析，对照 Attention Is All You Need 论文中的核心设计。

一、超参数设定

batch_size = 32      # B: 每次并行处理的序列数量
block_size = 8       # T: 上下文的最大长度 (时间步)
max_iters = 5000     # 训练迭代次数
eval_interval = 500  # 评估间隔
learning_rate = 1e-3 # 学习率
n_embd = 32          # C: 模型内部的思考维度 (嵌入维度)
n_head = 4           # 多头注意力的头数
n_layer = 6          # Transformer Block 的层数
dropout = 0.2        # Dropout 比率

二、自注意力机制（Scaled Dot-Product Attention）

论文公式： $\text{Attention}(Q, K, V) = \text{softmax}\!\left(\frac{QK^T}{\sqrt{d_k}}\right) V$

2.1 单个注意力头 Head

class Head(nn.Module):
    def __init__(self, head_size):
        super().__init__()
        self.key   = nn.Linear(n_embd, head_size, bias=False)
        self.query = nn.Linear(n_embd, head_size, bias=False)
        self.value = nn.Linear(n_embd, head_size, bias=False)
        self.register_buffer('tril', torch.tril(torch.ones(block_size, block_size)))
        self.dropout = nn.Dropout(dropout)

    def forward(self, x):
        B, T, C = x.shape
        k = self.key(x)    # (B, T, head_size)
        q = self.query(x)  # (B, T, head_size)

        # 缩放点积注意力
        wei = q @ k.transpose(-2, -1) * (C ** -0.5)  # (B, T, T)
        # ↑ 除以 sqrt(d_k) 防止点积方差过大

        # 因果掩码：防止"偷看未来"
        wei = wei.masked_fill(self.tril[:T, :T] == 0, float('-inf'))

        # Softmax 归一化 → 得到注意力权重
        wei = F.softmax(wei, dim=-1)
        wei = self.dropout(wei)

        # 加权聚合 Value
        v = self.value(x)
        out = wei @ v   # (B, T, head_size)
        return out

🔑 关键细节拆解

① Query / Key / Value 三个线性投影

• self.key、self.query、self.value 是三个独立的线性层（无偏置）
• 输入 x 形状 (B, T, n_embd) → 分别投影到 head_size 维度
• 这是论文中的 $W^Q, W^K, W^V$ 矩阵

② 缩放因子 $\frac{1}{\sqrt{d_k}}$

wei = q @ k.transpose(-2, -1) * (C ** -0.5)

• 为什么需要缩放？ 当 $d_k$ 很大时，$QK^T$ 的点积值会很大
• 不做缩放 → Softmax 梯度会进入极小区间 → 梯度消失
• 除以 $\sqrt{d_k}$ 让方差回归 1，保持 Softmax 输出的概率分布"不极端"

③ 因果掩码（Causal Mask）

wei = wei.masked_fill(self.tril[:T, :T] == 0, float('-inf'))

• 使用下三角矩阵 tril，把未来位置的注意力分数设为 -inf
• Softmax 后，未来位置的权重变为 0
• 这是 GPT 风格 Decoder-only 的核心：只能看到过去，不能看到未来

论文原文：masked self-attention in the decoder prevents positions from attending to subsequent positions.

④ 注意力 Dropout

wei = self.dropout(wei)

• 在注意力权重矩阵上施加 Dropout
• 随机丢弃部分注意力连接，增加鲁棒性

三、多头注意力（Multi-Head Attention）

论文思想： 用多组 $W^Q, W^K, W^V$ 并行计算，让模型从不同"子空间"关注信息。

class MultiHeadAttention(nn.Module):
    def __init__(self, num_heads, head_size):
        super().__init__()
        self.heads = nn.ModuleList([Head(head_size) for _ in range(num_heads)])
        self.proj   = nn.Linear(n_embd, n_embd)
        self.dropout = nn.Dropout(dropout)

    def forward(self, x):
        # 每个头独立计算注意力，沿最后一维拼接
        out = torch.cat([h(x) for h in self.heads], dim=-1)
        # 线性投影回 n_embd 维度
        out = self.proj(out)
        return out

论文公式： $\text{MultiHead}(Q,K,V) = \text{Concat}(\text{head}1, ..., \text{head}h) W^O$

🔑 关键细节

• 这里 head_size = n_embd // n_head = 32 // 4 = 8
• 每个头处理 8 维子空间，拼起来恢复 32 维

四、前馈神经网络（Position-wise Feed-Forward）

论文公式： $\text{FFN}(x) = \max(0, xW1 + b1)W2 + b2$

class FeedForward(nn.Module):
    def __init__(self, n_embd):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(n_embd, 4 * n_embd),  # 扩张到 4 倍
            nn.ReLU(),                       # 非线性激活
            nn.Linear(4 * n_embd, n_embd),  # 压缩回原维度
            nn.Dropout(dropout),
        )

    def forward(self, x):
        return self.net(x)

🔑 关键细节

• 扩张比 4 倍： 论文中的 $d{ff} = 4 \times d{\text{model}}$，这里 4 * n_embd
• 逐位置（Position-wise）： 对序列中每个位置独立做相同的线性变换
• ReLU 激活： 论文原始用的是 ReLU（后来 GPT 系列改用 GELU）
• Dropout： 在 FFN 输出后施加，防止过拟合

五、残差连接（Residual Connections）

论文原文： We employ a residual connection around each of the two sub-layers.

class Block(nn.Module):
    def forward(self, x):
        x = x + self.sa(self.ln1(x))    # 残差：注意力子层
        x = x + self.ffwd(self.ln2(x))  # 残差：前馈子层
        return x

🔑 关键细节

• x = x + f(x) 就是残差连接的标准写法
• 为什么需要残差？
• 没有残差 → 深层网络梯度消失 → 无法训练
• 有残差 → 梯度可以走"短路"直接回传 → 堆再多层也不怕
• 论文中把残差连接写为：$\text{LayerNorm}(x + \text{Sublayer}(x))$
• ⚠️ 注意：这里的实现是先 LN 再子层（Pre-LN），和论文原始的 Post-LN 顺序不同

Pre-LN vs Post-LN

当前代码用的是 Pre-LN，这是 GPT-2 之后的业界标准做法。

六、层归一化（Layer Normalization）

论文原文： We employ layer normalization after each sub-layer.

self.ln1 = nn.LayerNorm(n_embd)  # 注意力子层前
self.ln2 = nn.LayerNorm(n_embd)  # 前馈子层前

self.blocks.append(nn.LayerNorm(n_embd))

🔑 Batch Norm vs Layer Norm

LN 的数学形式

$$\text{LN}(x) = \gamma \cdot \frac{x - \mu}{\sigma} + \beta$$

• $\mu, \sigma$：在特征维度上计算均值和标准差
• $\gamma, \beta$：可学习的缩放和平移参数
• 输入 (B, T, n_embd) → 在最后一维 $n_{\text{embd}}$ 上归一化

七、Dropout —— 防止过拟合

论文原文： We apply dropout to the output of each sub-layer, before it is added to the sub-layer input and normalized. In addition, we apply dropout to the sums of the embeddings and the positional encodings.

dropout = 0.2  # 每次随机丢弃 20% 的神经元

代码中使用 Dropout 的位置

Dropout 的工作原理

• 训练时： 以概率 $p$ 随机将神经元输出置零
• 推理时： 所有神经元参与，但输出乘以 $(1-p)$（PyTorch 自动处理）
• model.eval() 会自动关闭 Dropout，model.train() 会重新打开

训练时相当于训练了 $2^n$ 个不同子网络的集成，推理时等价于它们的平均。

八、位置编码（Positional Encoding）

论文原文： Since our model contains no recurrence and no convolution, we must inject some information about the relative or absolute position of the tokens in the sequence.

self.position_embedding_table = nn.Embedding(block_size, n_embd)

pos_emb = self.position_embedding_table(torch.arange(T, device=device))
x = tok_emb + pos_emb  # 词义 + 位置 → 融合

🔑 关键细节

• 论文原文使用正弦/余弦固定编码（Sinusoidal）
• 这里的代码使用的是可学习的位置嵌入（Learnable Position Embedding）
• 这是 GPT 系列的做法，实践中效果相当甚至更好
• 原理：让模型知道 "第 0 个位置" 和 "第 3 个位置" 是不同的

九、Transformer Block 完整组装

class Block(nn.Module):
    """ Transformer 块：通信 (Attention) + 思考 (FeedForward) """

    def __init__(self, n_embd, n_head):
        super().__init__()
        head_size = n_embd // n_head
        self.sa   = MultiHeadAttention(n_head, head_size)  # 通信
        self.ffwd = FeedForward(n_embd)                    # 思考
        self.ln1  = nn.LayerNorm(n_embd)                   # 归一化 1
        self.ln2  = nn.LayerNorm(n_embd)                   # 归一化 2

    def forward(self, x):
        x = x + self.sa(self.ln1(x))    # Pre-LN + 注意力 + 残差
        x = x + self.ffwd(self.ln2(x))  # Pre-LN + FFN + 残差
        return x

一张图理解 Block

输入 x
  │
  ├─→ LayerNorm → Multi-Head Attention ─→ (+) ─┐
  │                                    ↑        │
  │                              残差连接 (x)    │
  │                                             ▼
  └─→ LayerNorm → Feed-Forward ────────→ (+) ─→ 输出
                                       ↑
                                 残差连接 (x)

十、多层堆叠

self.blocks = nn.Sequential(
    *[Block(n_embd, n_head=n_head) for _ in range(n_layer)]
)
self.blocks.append(nn.LayerNorm(n_embd))  # 最终 LayerNorm

• 论文中 Encoder 和 Decoder 各堆 6 层（$N=6$）
• 这里 n_layer=6，每个 Block 里注意力 + FFN 各一个子层 → 共 12 个子层
• 最后的 LayerNorm 是 GPT-2 风格，确保输出归一化后再做预测

十一、训练相关

11.1 优化器 —— AdamW

optimizer = torch.optim.AdamW(model.parameters(), lr=learning_rate)

• AdamW = Adam + 解耦的权重衰减（Decoupled Weight Decay）
• Transformer 训练的标配优化器
• 论文原版用的是 Adam，AdamW 是后来的改进

11.2 学习率 Warmup（论文中有，本代码未实现）

论文使用公式：$lr = d{\text{model}}^{-0.5} \cdot \min(\text{step\num}^{-0.5}, \text{step\num} \cdot \text{warmup\steps}^{-1.5})$

• 特征：先线性增长，再按步数平方根衰减
• 本代码因为用了 Pre-LN，不需要 warmup 也能稳定训练

11.3 训练/评估模式切换

model.eval()   # 评估模式：关闭 Dropout
model.train()  # 训练模式：开启 Dropout

十二、架构总览图

                    BigramLanguageModel
                           │
              ┌────────────┼────────────┐
              ▼            ▼            ▼
        Token Embed   Position Embed   输入 idx
              │            │
              └───── (+) ──┘
                      │
              ┌───────▼────────┐
              │   Block × 6    │
              │  ┌──────────┐  │
              │  │  LN → MHA│  │  ← 通信（词与词交互）
              │  │  + 残差   │  │
              │  │  LN → FFN│  │  ← 思考（逐位置加工）
              │  │  + 残差   │  │
              │  └──────────┘  │
              └───────┬────────┘
                      │
                 Final LN
                      │
                 lm_head (Linear)
                      │
                    logits

十三、关键 Q&A

Q1: 为什么 Attention 叫"通信"，FFN 叫"思考"？

Attention 让序列中不同位置的词相互"交流"信息（跨位置操作）；FFN 在每个位置上独立地对融合后的信息做非线性变换（不跨位置）。所以一个是"开会讨论"，一个是"各自消化"。

Q2: 残差连接为什么能让深层网络可训练？

反向传播时，梯度可以走残差支路（恒等映射）直接回传，不会因为经过太多层而衰减到零。数学上：$\frac{\partial}{\partial x}(x + f(x)) = 1 + \frac{\partial f}{\partial x}$，那个 $+1$ 保证了梯度不会消失。

Q3: Pre-LN 和 Post-LN 到底差在哪？

Post-LN：LN(x + Sublayer(x)) —— 残差加完再归一化，梯度路径经过 LN 可能被压缩，需要 warmup。 Pre-LN：x + Sublayer(LN(x)) —— 先归一化再做子层、再加回残差，梯度走残差路径时不经过 LN，训练更稳定。

Q4: Dropout 加在不同位置有什么不同效果？

参考资料

• [[Attention Is All You Need]] (Vaswani et al., 2017)
• self-attention.py —— Andrej Karpathy 风格的从零实现
• GPT-2 论文：Language Models are Unsupervised Multitask Learners