Transformer 与注意力

自注意力怎么计算

输入序列的每个 token 的向量，分别乘三个可学习的权重矩阵得到 Query、Key、Value：

Q = X·W_Q    K = X·W_K    V = X·W_V

Attention(Q, K, V) = softmax( Q·Kᵀ / √d_k ) · V

逐步理解：

Q·Kᵀ —— 每个 token 的 query 和所有 token 的 key 做点积，得到 n×n 的相关性矩阵。第 i 行第 j 列表示「token i 应该多关注 token j」
/ √d_k —— 缩放
softmax —— 按行归一化成概率分布，权重和为 1
· V —— 用这组权重对所有 token 的 value 做加权求和，得到该位置的新表示

直觉理解：query 是「我在找什么」，key 是「我能提供什么」，value 是「我的实际内容」。

为什么要除以 √d_k

几乎必问。两层回答：

数学层面 —— 假设 q 和 k 的各分量是独立的、均值 0 方差 1 的随机变量，那么点积 q·k = Σ qᵢkᵢ 共 d_k 项，其方差为 d_k（独立变量和的方差等于方差之和）。d_k = 64 时标准差是 8，d_k = 512 时是 22.6 —— 维度越高，点积的数值波动越大。除以 √d_k 正好把方差拉回 1。

后果层面 —— 如果不缩放，点积的量级随维度增长，送进 softmax 后会落在饱和区：最大的那个值的概率趋近 1，其余趋近 0，变成近似 one-hot 的「硬选择」。而 softmax 在饱和区的梯度趋近于 0，反向传播时梯度消失，模型训练不起来。

提示

追问「为什么是 √d_k 而不是 d_k」—— 因为要抵消的是标准差（√方差），除以 d_k 会过度缩放，让分布太平坦、注意力失去区分度。

多头注意力（MHA）

把 d_model（如 512）拆成 h 个头（如 8 个，每头 d_k = 64），每个头独立做注意力，最后把结果拼接再过一个线性层 W_O。

为什么要多头：单个注意力头的输出是一个加权平均，表达的是一种相关性模式。多头让模型在不同子空间并行捕捉不同类型的关系 —— 实验上观察到有的头学语法依赖、有的头学共指消解、有的头关注相邻位置。类比 CNN 里的多个卷积核。

注意总计算量与单头基本相同，因为维度被切分了，不是复制了 8 倍。

位置编码

自注意力本身是排列不变的 —— 打乱输入顺序，每个 token 的输出不变（只是跟着换位置）。所以必须显式注入位置信息。

方案	做法	特点
正弦绝对位置编码	原始论文，用不同频率的 sin/cos 生成固定向量加到 embedding 上	无参数，理论上能外推，实际外推效果一般
可学习绝对位置	BERT / GPT-2，把位置当作可训练参数	简单有效，但无法超出训练长度
RoPE（旋转位置编码）	LLaMA / Qwen 等主流模型采用。把 Q、K 向量按位置旋转一个角度	旋转后的点积只依赖相对位置差，兼具绝对与相对编码优点，外推性好
ALiBi	直接在注意力分数上按距离加线性惩罚偏置	实现极简，长度外推能力强

RoPE 为什么好：核心性质是 <f(q,m), f(k,n)> = g(q, k, m-n) —— 经过位置 m 和 n 的旋转后，内积只与 m-n（相对距离）有关。这让模型天然理解相对位置，且配合 NTK 插值 / YaRN 等方法可以在不重训的情况下扩展上下文窗口。

KV Cache 与注意力变体

KV Cache 的动机：自回归生成时，每生成一个新 token 都要对全部历史做注意力。但历史 token 的 K、V 是不变的，没必要重算。缓存起来，每步只算新 token 的 Q 和它与缓存 K、V 的注意力 —— 把每步的复杂度从 O(n²) 降到 O(n)。

代价是显存。KV Cache 大小的估算公式：

2 (K和V) × batch × seq_len × num_layers × num_kv_heads × head_dim × dtype_bytes

举例：70B 模型（80 层、64 头、head_dim 128）、fp16、batch=1、序列 4096： 2 × 1 × 4096 × 80 × 64 × 128 × 2 ≈ 10.7 GB —— 仅 KV Cache 就要 10GB，这是长上下文推理的主要瓶颈。

于是有了 KV 头共享：

变体	KV 头数	效果
MHA（多头注意力）	= Q 头数	质量最好，显存最大
MQA（多查询注意力）	1 组，所有 Q 头共享	显存降到 1/h，质量有损
GQA（分组查询注意力）	分成 g 组，组内共享	LLaMA-2/3 采用，质量与显存的折中
MLA（多头潜在注意力）	DeepSeek 提出，把 KV 压缩到低维潜空间再还原	比 GQA 省更多且质量更好

其他架构要点

FFN（前馈网络） —— 每层注意力之后接一个两层 MLP，中间维度通常是 4 × d_model。注意力负责「token 之间交换信息」，FFN 负责「对每个 token 单独做非线性变换」。模型参数量的大头在 FFN（约 2/3）。现代模型用 SwiGLU 替代 ReLU，效果更好。

残差连接 + 层归一化 —— 残差让梯度能直接回传到浅层，使几十上百层可训。归一化位置有两种：Post-LN（原论文，LN(x + Sublayer(x))）训练不稳定，需要 warmup；Pre-LN（现代主流，x + Sublayer(LN(x))）训练稳定得多。LLaMA 等用 RMSNorm 替代 LayerNorm —— 去掉了减均值和偏置项，只做缩放，计算更省且效果不降。

Flash Attention —— 注意不是近似算法，计算结果与标准注意力完全相同。它的优化点是 IO：标准实现要把 n×n 的注意力矩阵写进 HBM 再读回，而 Flash Attention 用分块（tiling）在 SRAM 里增量计算 softmax，避免了中间矩阵的落盘。所以它省的是显存带宽和显存占用，不是 FLOPs，但实际速度快 2～4 倍。这个「省 IO 不省算力」的区分是考点。

MoE（混合专家） —— 把 FFN 替换成 N 个专家网络 + 一个路由器，每个 token 只激活 top-k 个专家（如 8 选 2）。于是总参数量很大但单次推理的激活参数量小，用更少计算换更大容量。难点是负载均衡（需要辅助损失防止所有 token 都挤向少数专家）和专家并行的通信开销。

自注意力怎么计算​

为什么要除以 √d_k​

多头注意力（MHA）​

位置编码​

KV Cache 与注意力变体​

其他架构要点​