🧩 Transformer 总结文档

李萨如
35 阅读7分钟

🧩 Transformer 总结文档

整理近期对《Attention Is All You Need》学习的一份系统回顾文档，帮助后续继续深入和工程实践。

1. 总体视角：Transformer 在干什么？

Transformer 是一种 序列建模结构，典型目标：

  • 给定一个序列 x1, x2, ..., xT，学习它的概率分布。
  • 在位置 t，根据前面的 token 预测“下一个 token”。

典型应用形态：

  • Encoder–Decoder：机器翻译等，src → tgt。
  • Encoder-only：BERT，一般用于理解（分类、检索、问答）。
  • Decoder-only：GPT，一般用于生成（大模型、聊天、写代码）。

一个典型 Encoder–Decoder 流程：

src_ids
  → src_embeddings (+ 位置编码)
  → N 层 Encoder Block
  → enc_out: [T_src, d_model]

tgt_prefix_ids
  → tgt_embeddings (+ 位置编码)
  → N 层 Decoder Block (Masked Self-Attn + Enc-Dec Attn)
  → 线性层 + softmax
  → 在每个位置预测“下一个 token”

2. 输入侧：词向量 Embedding 与位置编码

2.1 词向量（Token Embedding）

  • 词表大小：vocab_size
  • 隐藏维度：d_model

定义一个可学习矩阵：

E: [vocab_size, d_model]
E[token_id] = 该 token 的向量
  • 实现：查表操作等价于 one_hot(token) @ E。
  • 训练：E 和其它权重一起通过梯度下降更新。
  • 效果：语义相近的词，向量方向会更接近（点积/余弦相似度更大）。

2.2 位置编码（Positional Encoding）

Self-Attention 对序列顺序本身不敏感，必须显式注入位置信息。

统一做法是：

x_emb[pos]   = token_emb[pos]        # [d_model]
pos_emb[pos] = position_encoding(pos)   # [d_model]
x_with_pos[pos] = x_emb[pos] + pos_emb[pos]

常见几类：

  • 固定正弦/余弦位置编码（sinusoidal）
    每一维是不同频率的正弦/余弦函数：

    PE[pos][2i]   = sin( pos / 10000^(2i/d_model) )
PE[pos][2i+1] = cos( pos / 10000^(2i/d_model) )
    • 不需要训练，任意长度都能计算。
    • 用一组“波形”编码绝对位置和一定的相对位置信息。

  • 可学习位置 Embedding
    类似词向量，直接定义：

    PosEmbed: [max_len, d_model]
pos_emb[pos] = PosEmbed[pos]   # 可训练

  • 相对位置编码 / RoPE（旋转位置编码）（进阶）

    • 不只编码“我是第几位”，更关注“i 和 j 的相对距离”。
    • RoPE 将 Q/K 的每对维度视为二维向量，对不同位置施加不同角度的旋转：
      点积中自然包含 (pos_i - pos_j) 的信息。
    • LLaMA 等现代大模型普遍使用。

3. 单头 Attention：Scaled Dot-Product Attention

对单层、单头，忽略 batch：

  • 输入：X: [T, d_model]
  • 参数矩阵： 
    W_Q: [d_model, d_k]
W_K: [d_model, d_k]
W_V: [d_model, d_v]

计算流程：

Q = X @ W_Q           # [T, d_k]
K = X @ W_K           # [T, d_k]
V = X @ W_V           # [T, d_v]

scores  = Q @ K^T     # [T, T]，所有位置两两点积
scores  = scores / sqrt(d_k)    # 缩放，控制数值范围

weights = softmax(scores, dim=-1)   # 按行 softmax，得到注意力权重 [T, T]

O = weights @ V        # [T, d_v]，对 V 做加权平均

关键点：

  • scores[i][j] = Q[i]·K[j]
    表示第 i 个 token 认为第 j 个 token 的相关性。
  • softmax 使每行变成一组和为 1 的权重分布。
  • 输出 O[i] = Σ_j weights[i][j] * V[j]
    即：第 i 个 token 从所有 token 的信息中，按权重聚合得到新向量。

缩放 1/sqrt(d_k) 的原因：

  • 维度 d_k 越大，点积 Q·K 的典型大小会变大；
  • 不缩放直接给 softmax，会导致分布极度“尖锐”（接近 one-hot），梯度差等问题；
  • 除以 sqrt(d_k) 将分数拉回稳定范围，让 softmax 的行为在不同 d_k 下比较一致。

4. 多头注意力（Multi-Head Attention）

单头只在一个子空间里学习模式，多头让模型在多个子空间并行学习。

假定：

d_model  = 8
num_heads = 2
d_k = d_v = d_model / num_heads = 4

对 head h：

W_Q^h: [d_model, d_k]
W_K^h: [d_model, d_k]
W_V^h: [d_model, d_v]

Q_h = X @ W_Q^h              # [T, 4]
K_h = X @ W_K^h
V_h = X @ W_V^h

O_h = Attention(Q_h, K_h, V_h)  # [T, 4]

然后在特征维度上拼接所有 head：

O_concat = concat(O_1, O_2, ..., O_H, dim=-1)   # [T, H*d_v] = [T, 8]

再通过输出线性层：

W_O: [H*d_v, d_model] = [8, 8]
O   = O_concat @ W_O          # [T, 8] -> [T, 8]

解释：

  • 每个 head 学习不同的依赖模式（语法、语义、局部/全局等）。
  • W_O 可以看作一个“多视角信息混合器”，把各 head 输出综合成统一的 d_model 表示。

5. FFN：前馈网络（按位置独立的两层 MLP）

FFN 对每个位置独立应用相同的 MLP：

W1: [d_model, d_ff]   (通常 d_ff = 4 * d_model)
b1: [d_ff]
W2: [d_ff, d_model]
b2: [d_model]

FFN(x) = W2( activation( x @ W1 + b1 ) ) + b2   # [d_model] -> [d_ff] -> [d_model]

特点：

  • 不在 token 之间做交互，只对每个 token 的特征维做非线性变换。
  • Attention 做“信息在不同位置间怎么流动”，FFN 做“每个位置自身特征如何变换”。
  • 在 Transformer 中，FFN 的参数量通常比 Attention 还多（尤其是大模型里）。

6. 残差连接（Residual）与 LayerNorm

6.1 残差连接：y = x + F(x)

  • 避免深层网络退化：
    若 F(x) ≈ 0，则 y ≈ x，多加几层不会让表达变差。
  • 更易优化：
    网络学的是“在原有表示上加一点增量”，而不是“重造一个新表示”。
  • 梯度更稳：
    反向传播时梯度可以通过 +x 这条“直通路径”有效传到前层。

在 Transformer 中，每个子层（Attention 或 FFN）都包在：

sub_out = SubLayer(x)    # Attention 或 FFN 的输出
x_res  = x + sub_out     # 残差连接
y      = LayerNorm(x_res)

6.2 LayerNorm（Layer Normalization）

对每个 token 的向量 x: [d_model]，做：

  1. 计算该向量所有维度的均值和方差：

    mean = average(x[i] over i)
var  = average( (x[i] - mean)^2 over i )

  2. 归一化：

    x_norm[i] = (x[i] - mean) / sqrt(var + eps)

  3. 再乘以可学习尺度 gamma[i]，加偏置 beta[i]：

    y[i] = gamma[i] * x_norm[i] + beta[i]

特点：

  • 对每个 token 单独归一化，不跨样本、不跨时间步。
  • 不依赖 batch 大小，适合自回归模型和变长序列。
  • 让每层输出的数值分布更稳定，有利于深层训练。

7. Encoder 与 Decoder Block

7.1 Encoder Block

输入 X: [T_src, d_model]（已加位置编码）：

1) Multi-Head Self-Attention on X
   attn_out = MHA(X, X, X)    # Q,K,V 都来自 X，自注意力
   x1 = LayerNorm(X + attn_out)

2) FFN
   ffn_out = FFN(x1)
   Y = LayerNorm(x1 + ffn_out)    # 这一层 Encoder Block 输出
  • 多层 Encoder 堆叠：X^0 → X^1 → ... → X^N = enc_out。

7.2 Decoder Block（Encoder–Decoder Transformer）

输入 x_dec^{l-1}: [T_tgt, d_model]，以及 enc_out: [T_src, d_model]：

1) Masked Self-Attention（目标序列内部，只看自己和过去）
   attn1 = MHA_masked(x_dec^{l-1}, x_dec^{l-1}, x_dec^{l-1})
   x1 = LayerNorm(x_dec^{l-1} + attn1)

2) Encoder–Decoder Attention（跨序列，Q 从 Decoder，K/V 从 Encoder）
   attn2 = MHA(x1 (Q), enc_out (K), enc_out (V))
   x2 = LayerNorm(x1 + attn2)

3) FFN
   ffn_out = FFN(x2)
   x_dec^l = LayerNorm(x2 + ffn_out)
  • 多层 Decoder 堆叠：x_dec^0 → x_dec^1 → ... → x_dec^N。
  • 顶部再接一个线性层 + softmax 预测下一个 token。

8. Masked Self-Attention 与训练/推理

8.1 因果 Mask（Causal Mask）

目标序列长度 T_tgt 时，自注意力的 scores 为 [T_tgt, T_tgt]：

  • 训练时，每个位置 t 只能使用 <= t 的信息预测下一个 token。
  • 构造下三角 Mask：
mask[i][j] = 1 if j <= i else 0

在 softmax 前，对 mask=0 的位置置为极小值（如 -1e9），让其权重接近 0。

8.2 训练：Teacher Forcing + Cross-Entropy

以目标序列：

tgt = [y1, y2, ..., yT]

构造：

  • Decoder 输入：[<bos>, y1, y2, ..., y_{T-1}]
  • 训练标签： [y1, y2, ..., yT]

在每个位置 t，Decoder 输出对“下一个 token”的分布 probs_t：

  • loss 形式： 
    loss_t = -log( probs_t[ label_t ] )
loss_sentence = average_t(loss_t)
loss_batch = average_over_batch(loss_sentence)
  • 一次前向计算所有位置的 loss，一次 backward 更新所有参数。

8.3 推理：逐步生成直到 <eos>

推理时（以翻译为例）：

  1. Encoder：src → enc_out。
  2. 初始化目标侧序列：dec_seq = [<bos>]。
  3. 循环：
    • 用当前 dec_seq 作为 Decoder 输入（带因果 mask）。
    • 取最后一个位置的输出，通过 softmax 得到 probs_next。
    • 选一个 token（贪心/beam/top-k/top-p），记为 y_next。
    • 如果 y_next == <eos>，停止；否则 append 到 dec_seq，继续。

结束条件：

  • 生成 <eos>，或
  • 达到最大长度上限（防止异常不过结束）。

9. 几种典型结构对比

  • Encoder–Decoder Transformer（原始论文）

    • N 层 Encoder + N 层 Decoder
    • Decoder 中既有 Masked Self-Attn，又有 Encoder–Decoder Attn

  • Encoder-only（BERT 等）

    • 只有 Encoder 堆栈
    • Self-Attn 不带因果 mask（双向）
    • 用于理解任务：分类、问答、检索等

  • Decoder-only（GPT 系列）

    • 只有 Decoder 堆栈
    • 每层都是 Masked Self-Attention + FFN
    • 训练目标：自回归 next-token prediction
    • 用于文本/代码生成、大模型对话等
avatar
文章
阅读
粉丝