Self-Attention 到 GQA

什么是 Attention 机制

Attention（注意力机制）是现代深度学习中最重要的创新之一，它使模型能够在处理序列数据时动态地关注不同位置的信息。

在 Transformer 架构中，Attention 机制的核心思想是：让每个位置的 token 根据其与其他 token 的相关性，对它们进行加权求和。

值得先说清楚它替代了什么。RNN 处理序列时，第 $n$ 个位置的信息要传到第 $m$ 个位置，必须经过 $|m - n|$ 步递推，路径越长信号衰减越严重，而且这些步骤天然串行、无法并行。Attention 把这条路径压到常数：任意两个位置之间只隔一次点积。代价是计算量从 $O(N)$ 变成 $O(N^2)$，显存也随序列长度平方增长。这笔交易是后面 FlashAttention、GQA、KV cache 等一整条优化线索的起点。

标准 Self-Attention 计算流程

给定输入序列，Self-Attention 通过以下步骤计算输出：

1. 线性变换：将输入 $X$ 分别投影为 Query ($Q$)、Key ($K$)、Value ($V$)

   $$Q = XW_Q, \quad K = XW_K, \quad V = XW_V$$

2. 计算注意力分数：通过 Query 和 Key 的点积计算相似度

   $$S = \frac{QK^T}{\sqrt{D}}$$

   其中 $D$ 是 head dimension，除以 $\sqrt{D}$ 用于数值稳定性（scaled dot-product）

3. 应用 Softmax：将分数转换为概率分布

   $$A = \text{softmax}(S)$$

4. 加权求和：用注意力权重对 Value 进行加权

   $$O = AV$$

直观理解：

• $Q$（Query）：我想要什么信息？
• $K$（Key）：我提供什么信息？
• $V$（Value）：我实际包含的信息是什么？
• 注意力分数：衡量"我需要的"和"你提供的"有多匹配

PyTorch 参考实现

import torch
import torch.nn.functional as F

def self_attention(X, W_q, W_k, W_v):
    """
    X: (batch, seq_len, d_model)
    W_q, W_k, W_v: (d_model, d_head)
    """
    Q = X @ W_q  # (batch, seq_len, d_head)
    K = X @ W_k
    V = X @ W_v

    # 计算注意力分数
    scores = Q @ K.transpose(-2, -1)  # (batch, seq_len, seq_len)
    scores = scores / (K.shape[-1] ** 0.5)  # 缩放

    # Softmax 归一化
    attn_weights = F.softmax(scores, dim=-1)

    # 加权求和
    output = attn_weights @ V  # (batch, seq_len, d_head)

    return output, attn_weights