记录

论文笔记：Attention Is All You Need

科研 #论文笔记#Transformer#深度学习

第一遍读的时候，卡在”为什么需要多头”和”位置编码为什么用正弦函数”上。这次重读把推导写下来。

核心问题

RNN 的顺序依赖让并行化变得困难：第 t 步必须等第 t-1 步算完。卷积可以并行，但要堆很多层才能覆盖长距离依赖。这篇论文的出发点就是去掉循环，只用注意力。

自注意力怎么算

对输入序列每个位置生成三组向量：Query、Key、Value。

$$ \mathrm{Attention}(Q, K, V) = \mathrm{softmax}!\left(\frac{QK^\top}{\sqrt{d_k}}\right)V $$

几个关键点：

  • $QK^\top$ 得到的是任意两个位置之间的相关性分数矩阵，形状 (seq_len, seq_len)——这一步没有循环，整段序列可以并行算。
  • 除以 $\sqrt{d_k}$ 是缩放。$d_k$ 越大，点积的方差越大，softmax 会推向极端值，梯度就接近 0。缩放让分布保持在梯度好用的区间。
  • 乘 $V$ 是加权求和：每个位置的输出是所有位置 value 的加权组合，权重就是注意力分数。
import torch
import torch.nn.functional as F

def scaled_dot_product_attention(q, k, v, mask=None):
    d_k = q.size(-1)
    scores = q @ k.transpose(-2, -1) / d_k ** 0.5
    if mask is not None:
        scores = scores.masked_fill(mask == 0, float('-inf'))
    weights = F.softmax(scores, dim=-1)
    return weights @ v, weights

多头的作用

不是”为了更强”，而是给了模型在不同子空间里看不同关系的能力。单头注意力被 softmax 归一化后，一个位置的所有注意力权重加起来是 1，只能表达一种”关注模式”；切成 h 个头之后，每个头可以各自关注语法依赖、共指、位置邻近等不同关系，最后拼接再做线性变换。

实现上，多头不是循环 h 次，而是把 d_model 拆成 h × d_k 一次性做批量矩阵乘：

# x: (batch, seq_len, d_model)
qkv = self.qkv_proj(x)                       # -> (batch, seq_len, 3*d_model)
q, k, v = qkv.chunk(3, dim=-1)
# reshape 成 (batch, n_heads, seq_len, d_k)，然后一次 attention

位置编码

注意力本身是置换等变的：把输入顺序打乱，输出也只是跟着打乱，模型完全不知道谁在前谁在后。所以必须显式注入位置信息。

$$ PE_{(pos, 2i)} = \sin!\left(\frac{pos}{10000^{2i/d_{model}}}\right), \quad PE_{(pos, 2i+1)} = \cos!\left(\frac{pos}{10000^{2i/d_{model}}}\right) $$

用正弦而不是可学习的位置嵌入，好处是可以外推到训练时没见过的更长序列——后面的位置由同一条曲线给出。代价是实际外推效果有限（后来的 RoPE 就是冲这个来的）。

还没搞懂的

  • Layer Norm 放在残差之后（Post-LN）会导致训练不稳定，需要 warmup；后来的 Pre-LN 改法为什么就能稳，我目前只停留在”梯度路径更短”的直觉上，没有推。
  • 编码器里的 mask 和解码器里的 causal mask 我经常混淆，下次单独整理一篇。

参考

  • Vaswani et al., Attention Is All You Need, NeurIPS 2017. arXiv:1706.03762