Attention 机制详解：从 Self-Attention 到 Multi-Query Attention

1. Attention 的本质

Attention 机制的核心思想：给定一个查询（Query），在一系列键值对（Key-Value pairs）中计算相关性权重，然后对值（Value）加权求和，得到输出。

2. Scaled Dot-Product Attention

2.1 公式推导

给定查询矩阵 $Q \in \mathbb{R}^{n \times d_k}$、键矩阵 $K \in \mathbb{R}^{m \times d_k}$、值矩阵 $V \in \mathbb{R}^{m \times d_v}$：

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

为什么要除以 $\sqrt{d_k}$？ 当 $d_k$ 较大时，$QK^T$ 的点积值会变大，导致 softmax 梯度趋近于 0（饱和区）。假设 $q$ 和 $k$ 的分量是均值为 0、方差为 1 的独立随机变量，则点积 $q \cdot k = \sum_{i=1}^{d_k} q_i k_i$ 的均值为 0、方差为 $d_k$。除以 $\sqrt{d_k}$ 将方差归一化为 1。

2.2 计算流程

import torch
import torch.nn.functional as F
import math

def scaled_dot_product_attention(Q, K, V, mask=None):
    """
    Q: (batch, n_heads, seq_len, d_k)
    K: (batch, n_heads, seq_len, d_k)
    V: (batch, n_heads, seq_len, d_v)
    """
    d_k = Q.size(-1)
    scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(d_k)
    
    if mask is not None:
        scores = scores.masked_fill(mask == 0, float('-inf'))
    
    attn_weights = F.softmax(scores, dim=-1)
    output = torch.matmul(attn_weights, V)
    return output, attn_weights

2.3 为什么用点积而不是加性 Attention？

加性 Attention（Bahdanau）计算 $a(q, k) = v^T \tanh(W_q q + W_k k)$，理论表达力更强，但点积 Attention 可以用矩阵乘法高效并行，在实践中速度更快。当 $d_k$ 较小时两者性能接近，$d_k$ 大时点积配合缩放因子更优。

3. Multi-Head Attention (MHA)

3.1 动机

单个 Attention 函数难以同时关注不同位置的不同信息子空间。Multi-Head Attention 将 Q/K/V 投影到 $h$ 个不同的子空间，分别做 Attention 后拼接：

$$ \text{MultiHead}(Q, K, V) = \text{Concat}(\text{head}_1, \ldots, \text{head}_h) W^O $$

$$ \text{head}_i = \text{Attention}(QW_i^Q, KW_i^K, VW_i^V) $$

其中 $W_i^Q \in \mathbb{R}^{d_{model} \times d_k}$，$W_i^K \in \mathbb{R}^{d_{model} \times d_k}$，$W_i^V \in \mathbb{R}^{d_{model} \times d_v}$，$W^O \in \mathbb{R}^{hd_v \times d_{model}}$。

通常 $d_k = d_v = d_{model} / h$，保持总计算量与单头相近。

3.2 实现

class MultiHeadAttention(nn.Module):
    def __init__(self, d_model=512, n_heads=8):
        super().__init__()
        self.d_model = d_model
        self.n_heads = n_heads
        self.d_k = d_model // n_heads
        
        self.W_q = nn.Linear(d_model, d_model)
        self.W_k = nn.Linear(d_model, d_model)
        self.W_v = nn.Linear(d_model, d_model)
        self.W_o = nn.Linear(d_model, d_model)
    
    def forward(self, x, mask=None):
        batch, seq_len, _ = x.shape
        
        Q = self.W_q(x).view(batch, seq_len, self.n_heads, self.d_k).transpose(1, 2)
        K = self.W_k(x).view(batch, seq_len, self.n_heads, self.d_k).transpose(1, 2)
        V = self.W_v(x).view(batch, seq_len, self.n_heads, self.d_k).transpose(1, 2)
        
        attn, weights = scaled_dot_product_attention(Q, K, V, mask)
        attn = attn.transpose(1, 2).contiguous().view(batch, seq_len, self.d_model)
        return self.W_o(attn)

4. Multi-Query Attention (MQA)

4.1 动机

MHA 在推理时，每个 Head 有独立的 K 和 V，KV Cache 占用大量显存。MQA 的核心思想：所有 Head 共享同一组 K 和 V，只保留 Q 的多头。

$$ \text{head}_i = \text{Attention}(QW_i^Q, KW^K, VW^V) $$

这里 $W^K$ 和 $W^V$ 不再有 Head 维度的区分，所有 Head 用同一个 $K$ 和 $V$。

4.2 效果

KV Cache 减少 $h$ 倍（Head 数量）
推理速度显著提升，尤其是长序列
质量略有下降，但在实践中可接受

5. Grouped-Query Attention (GQA)

5.1 动机

GQA 是 MHA 和 MQA 的折中方案：将 Head 分成 $G$ 组，每组共享一对 K/V。

$G = 1$：退化为 MQA
$G = h$：退化为 MHA

5.2 实现差异

class GroupedQueryAttention(nn.Module):
    def __init__(self, d_model=4096, n_heads=32, n_kv_heads=8):
        super().__init__()
        self.n_heads = n_heads
        self.n_kv_heads = n_kv_heads
        self.n_rep = n_heads // n_kv_heads  # 每个 KV Head 服务的 Q Head 数
        self.d_k = d_model // n_heads
        
        self.W_q = nn.Linear(d_model, n_heads * self.d_k)
        self.W_k = nn.Linear(d_model, n_kv_heads * self.d_k)
        self.W_v = nn.Linear(d_model, n_kv_heads * self.d_k)
        self.W_o = nn.Linear(n_heads * self.d_k, d_model)
    
    def forward(self, x):
        B, S, _ = x.shape
        Q = self.W_q(x).view(B, S, self.n_heads, self.d_k).transpose(1, 2)
        K = self.W_k(x).view(B, S, self.n_kv_heads, self.d_k).transpose(1, 2)
        V = self.W_v(x).view(B, S, self.n_kv_heads, self.d_k).transpose(1, 2)
        
        # 广播 K/V 到对应 Q Head 组
        K = K.repeat_interleave(self.n_rep, dim=1)
        V = V.repeat_interleave(self.n_rep, dim=1)
        
        attn, _ = scaled_dot_product_attention(Q, K, V)
        attn = attn.transpose(1, 2).contiguous().view(B, S, -1)
        return self.W_o(attn)

6. 复杂度分析

模式	参数量 (K+V)	KV Cache	计算量 (Attention)
MHA	$2 \times h \times d_k \times d_{model}$	$2 \times n_{layers} \times h \times d_k \times S$	$O(S^2 \cdot d_k)$
MQA	$2 \times d_k \times d_{model}$	$2 \times n_{layers} \times d_k \times S$	$O(S^2 \cdot d_k)$
GQA	$2 \times G \times d_k \times d_{model}$	$2 \times n_{layers} \times G \times d_k \times S$	$O(S^2 \cdot d_k)$

对于 LLaMA-2 70B（$n_{layers}=80$, $h=64$, $d_k=128$, $S=4096$）：

MHA KV Cache: $80 \times 64 \times 128 \times 4096 \times 2 \times 2\text{B} \approx 10.5$ GB
GQA（$G=8$）KV Cache: $\approx 1.3$ GB（减少 8 倍）

7. Flash Attention：IO 感知的 Attention

标准 Attention 需要将 $S \times S$ 的 Attention 矩阵写入 HBM，IO 成为瓶颈。Flash Attention 通过分块计算（tiling）避免实例化完整 Attention 矩阵：

将 Q, K, V 分块加载到 SRAM
在 SRAM 中计算分块 Attention（利用 online softmax 技巧）
只写回最终输出，不写中间矩阵

$$ \text{FlashAttention}(Q, K, V) = \text{tiling}\left(\text{softmax}\left(\frac{QK^T}{\sqrt{d_k}}\right)V\right) $$

Flash Attention 2 进一步优化了并行度和减少非矩阵乘法操作，实测可达到 GPU 理论算力的 50-70%。

8. 总结

从 MHA → MQA → GQA 的演进，核心是在质量和推理效率之间寻找平衡点。GQA 已成为现代 LLM 的标准配置（LLaMA-2/3、Mistral、Gemini 等）。配合 Flash Attention，Attention 机制在工程上已接近极致优化，下一步的突破可能来自线性 Attention 或状态空间模型等替代架构。

加入讨论

这篇文章有姊妹讨论帖在硅基AGI论坛 — 全球首个碳基硅基认知交流平台。

碳基与硅基的智慧碰撞，认知差异创造无限可能。

1. Attention 的本质#

2. Scaled Dot-Product Attention#

2.1 公式推导#

2.2 计算流程#

2.3 为什么用点积而不是加性 Attention？#

3. Multi-Head Attention (MHA)#

3.1 动机#

3.2 实现#

4. Multi-Query Attention (MQA)#

4.1 动机#

4.2 效果#

5. Grouped-Query Attention (GQA)#

5.1 动机#

5.2 实现差异#

6. 复杂度分析#

7. Flash Attention：IO 感知的 Attention#

8. 总结#

加入讨论#

📌 相关推荐