大模型教我成为大模型算法工程师之day11：Transformer 架构

Day 11: Transformer 架构

摘要：2017年 Google 的那篇《Attention Is All You Need》横空出世，彻底改变了深度学习的格局。Transformer 抛弃了 RNN 的循环结构，完全依赖注意力机制，实现了并行计算和长距离依赖的完美解决。本文将拆解 Transformer 的核心组件：Self-Attention、Multi-Head Attention、位置编码以及 Encoder-Decoder 架构。

---

1. 为什么抛弃 RNN？

虽然 LSTM 解决了梯度消失，但它有两个硬伤：

1. 无法并行：RNN 必须等 $t-1$ 算完才能算 $t$，在大规模数据上训练极慢。
2. 长距离衰减：即使有 LSTM，隔了 1000 个词后，信息传递依然会损耗。

Transformer 的口号是：不要循环，只要注意力 (Attention)。
它可以一次性把整句话输入进去（并行），而且不管两个词隔多远，它们之间的距离都是 1（直接 Attention）。

---

2. Self-Attention (自注意力机制)

这是 Transformer 的灵魂。它的作用是：让每个词都能看到整句话，并根据相关性聚焦于重要的词。

2.1 Q, K, V 的比喻

对于每个词向量 $x$，我们通过三个权重矩阵生成三个向量：

• Query (Q)：我在找什么？（查询）
• Key (K)：我是什么？（索引）
• Value (V)：我的内容是什么？（实际信息）

思考：为什么是 3 个？能合并吗？

• 灵感来自数据库查询：Query 查 Key，返回 Value。
• 能不能 K=V？ 可以，早期 Attention 就是这么做的。但这样表达能力会受限，就像强制要求一本书的“目录信息”必须等于“正文内容”，不够灵活。
• Attention 是过程性的吗？
  • 单层 Attention 是一次性的聚合。
  • 但 Transformer 通过堆叠 N 层模拟了人类的“反复思考”过程：第一层看到局部关系，第二层基于第一层的理解看到更深的关系，层层递进。

通俗例子：图书管借书

• Query：我想找一本关于“深度学习”的书。
• Key：每本书封面上贴的标签（“机器学习”、“神经网络”、“做菜指南”）。
• Value：书里的具体内容。
• 过程：拿你的 Q 去和每本书的 K 算相似度（点积），相似度高（Attention Score 大）的书，你就多读一点它的 V。

2.2 计算公式

$$\text{Attention}(Q,K,V)=\text{softmax}(\frac{Q{K}^T}{\sqrt{d_k}})V$$

1. $Q{K}^T$：计算查询 Q 和所有键 K 的相似度（分数）。
2. $\sqrt{d_k}$：缩放因子。防止点积结果太大导致 Softmax 梯度消失。
3. Softmax：把分数归一化成概率（和为1）。
4. $\times V$：用概率对 V 进行加权求和。

---

3. Multi-Head Attention (多头注意力)

为什么要有多个头？
这就好比虽然是同一句话，但我们可以从不同角度去理解。

• Head 1 关注语法结构（主谓宾关系）。
• Head 2 关注指代关系（“it” 指的是前面的 “dog”）。
• Head 3 关注情感倾向。

实现：把大的 Q, K, V 拆成 $h$ 份（比如 8 个头），分别做 Self-Attention，最后把结果拼接起来，再过一个线性层融合。

---

4. 位置编码 (Positional Encoding)

Transformer 是并行输入的，它把 “I love you” 作为一个集合扔进去，如果不加位置信息，网络就分不清 “I love you” 和 “You love I”。

• RNN：天生有顺序（第1步、第2步）。
• Transformer：天生无序。
• 解决：人为给每个词向量加上一个位置向量 (PE)。
  $$P{E}_{(pos,2i)}=\sin (pos/10000^{2i/d_{model}})$$
  $$P{E}_{(pos,2i+1)}=\cos (pos/10000^{2i/d_{model}})$$
  这种正弦/余弦编码允许模型学习到相对位置关系。

---

5. 整体架构 (Encoder-Decoder)

5.1 Encoder (编码器)

• 作用：理解输入序列，提取特征。
• 结构：由 N 个 Block 堆叠而成，每个 Block 包含：
  1. Multi-Head Self-Attention
  2. Feed Forward Network (FFN)
  3. 残差连接 (Add) & 层归一化 (Norm)：每个子层后都有 LayerNorm(x + Sublayer(x))。

5.2 Decoder (解码器)

• 作用：根据 Encoder 的输出生成目标序列。
• 区别：
  1. Masked Self-Attention：解码时不能“偷看”后面的词（未来的信息）。所以要把 $Q{K}^T$ 矩阵的上三角部分 Mask 掉（设为负无穷）。
  2. Cross Attention：Decoder 的 Q 来自自己，但 K 和 V 来自 Encoder 的输出。这意味着：“我在生成翻译时，要去查阅原文的信息”。

---

6. PyTorch 维度操作扫盲 (必读)

写 Transformer 代码时，维度变换是最容易晕的地方。

1. reshape vs view：
   • 都用于改变形状。view 要求内存连续，reshape 更鲁棒（不连续时会自动复制）。推荐无脑用 reshape。
2. permute / transpose：
   • 维度换位。如 [Batch, Seq, Head, Dim] -> [Batch, Head, Seq, Dim]。
   • 坑：换位后内存顺序乱了，如果你马上调 view 会报错。必须先调 .contiguous()。
3. unsqueeze / squeeze：
   • unsqueeze(1)：增加维度（插队）。[32, 100] -> [32, 1, 100]。常用于广播。
   • squeeze(1)：压缩维度（挤掉）。如果是 1 就去掉。
4. einsum (爱因斯坦求和)：
   • 神器。不用思考怎么转置怎么乘。
   • torch.einsum("nqhd,nkhd->nhqk", Q, K)
   • 规则：输入有 d，输出没 d -> 对 d 维度求和（点积）。

---

7. 代码实践：Self-Attention 实现

import torch
import torch.nn as nn
import math

class SelfAttention(nn.Module):
    def __init__(self, embed_size, heads):
        super(SelfAttention, self).__init__()
        self.embed_size = embed_size
        self.heads = heads
        self.head_dim = embed_size // heads

        assert (
            self.head_dim * heads == embed_size
        ), "Embed size needs to be divisible by heads"

        # 定义 Q, K, V 的线性变换层
        self.values = nn.Linear(self.head_dim, self.head_dim, bias=False)
        self.keys = nn.Linear(self.head_dim, self.head_dim, bias=False)
        self.queries = nn.Linear(self.head_dim, self.head_dim, bias=False)
        
        # 最后融合的线性层
        self.fc_out = nn.Linear(heads * self.head_dim, embed_size)

    def forward(self, values, keys, query, mask):
        # 这里的 values, keys, query 其实输入通常都是同一个 x (Self-Attention)
        N = query.shape[0] # Batch Size
        
        # 1. 拆分 Head: [N, seq_len, embed_size] -> [N, seq_len, heads, head_dim]
        # reshape 后 permute 变成 [N, heads, seq_len, head_dim] 方便并行计算
        values = values.reshape(N, -1, self.heads, self.head_dim)
        keys = keys.reshape(N, -1, self.heads, self.head_dim)
        query = query.reshape(N, -1, self.heads, self.head_dim)
        
        # 2. 线性投影 Q, K, V
        values = self.values(values)
        keys = self.keys(keys)
        queries = self.queries(query)

        # 3. 计算 Energy (QK^T)
        # queries shape: [N, heads, query_len, head_dim]
        # keys shape:    [N, heads, key_len, head_dim]
        # energy shape:  [N, heads, query_len, key_len]
        energy = torch.einsum("nqhd,nkhd->nhqk", [queries, keys])

        # 4. 缩放与Mask
        if mask is not None:
            # mask 为 0 的位置填成负无穷，Softmax 后变成 0
            energy = energy.masked_fill(mask == 0, float("-1e20"))

        attention = torch.softmax(energy / (self.embed_size ** (1 / 2)), dim=3)

        # 5. 加权求和: Attention * V
        # out shape: [N, heads, query_len, head_dim]
        out = torch.einsum("nhqk,nkhd->nqhd", [attention, values])
        
        # 6. 拼接 Heads 并通过最后线性层
        out = out.reshape(N, -1, self.heads * self.head_dim)
        out = self.fc_out(out)
        
        return out

---

7. 总结

• Self-Attention 解决了长距离依赖，实现了全并行训练。
• Multi-Head 增强了模型捕捉多种特征关系的能力。
• Positional Encoding 补足了序列位置信息。
• ResNet + LayerNorm 保证了深层网络的稳定训练。

Transformer 不仅统一了 NLP，后来通过 ViT 也统一了 CV，是目前 AI 领域绝对的基石架构。

---

参考资料

• Attention Is All You Need (Original Paper)
• The Illustrated Transformer (Jay Alammar)

如果有更多问题，欢迎关注公众号【算法最TOP】进一步讨论！