← 返回博客

模型学习笔记 · 第二章

注意力机制：从工程角度理解信息如何流动

串起 Q/K/V、自注意力与交叉注意力，再看 MHA、MQA、GQA、MLA 的存储取舍。

Muxuan Liu · 模型学习笔记

从信息流动理解注意力

我想从工程计算的角度理解注意力机制。相比一开始就把它解释成“学习语言依赖”，我更愿意先问：当前位置怎样读取其他位置的信息，这个过程又怎样用矩阵运算实现？这是我阅读模型的一种视角，并不意味着长距离依赖只是设计中偶然出现的副作用。

从这个角度，我把注意力理解为一种信息路由：当前位置提出 query，各位置提供 key，二者比较得到权重，再按权重汇总 value。“我需要什么”和“我这里有什么”只是帮助理解的比喻，Q/K/V 实际上是学习得到的向量。

语言关系是这种计算的一种应用，图片和其他结构化输入也能使用它。这让我倾向于先理解通用的信息交互，再讨论任务中的语义。单个注意力权重不等于已经证明了一种语法关系，也不是模型决策的完整解释。

注意力与序列长度

最初，我把注意力与机器翻译中输入、输出 token 数量不同的问题联系在一起。对照计算过程后，我发现这个解释并不准确：RNN 编码器—解码器本来就能处理不等长序列，不能把这个能力归因于自注意力。

自注意力的 Q、K、V 来自同一序列，输出长度跟 query 的长度相同；交叉注意力则从一个序列取得 Q，从另一个序列取得 K/V，因而 query 与 key 的长度可以不同。交叉注意力并不只用于跨语言任务。

在编码器—解码器 Transformer 中，交叉注意力让目标端的每个 query 动态读取源序列的信息；输出多少个 token 则由解码过程和停止规则决定，多头注意力本身不决定输出长度。早期注意力缓解的是把整个输入压进固定向量的瓶颈。只有解码器的模型也不一定包含独立交叉注意力模块。

沿着张量形状看计算

单头 Q 的形状是 [Tq, Dh]，K 是 [Tk, Dh]，V 是 [Tk, Dv]。QKᵀ 得到 [Tq, Tk] 的分数，在 key 方向做 softmax，再乘以 V，输出便是 [Tq, Dv]。每个 query 得到一个汇总结果，输出长度随 query 而定。这个形状检查比只凭翻译的直觉推测机制更可靠。

例如，若有 6 个 query、10 个 key/value，权重矩阵就是 6×10，输出仍有 6 个位置；这不是把 10 个输入自动“翻译成”6 个 token。把查询长度和生成长度分开，是我修正原来理解的关键。

除以 √Dh 是为了控制点积分数的尺度。缩放点积公式属于 2017 年 Transformer，不能直接归给更早的 Bahdanau 注意力。多头是在多个子空间并行计算，再拼接并投影回统一表示。不同头可能学到不同模式，但“某个头负责动词，另一个负责时间”只能作为示意，不是预先分配或保证出现的职责。

Attention(Q, K, V) = softmax(QKᵀ / √Dh + mask) V MultiHead = Concat(head₁, …, headₕ) Wₒ

并行训练不等于并行生成所有 token

训练时目标序列已知，配合因果掩码可以同时计算多个位置。普通自回归生成时，下一个 token 仍依赖前面已生成的内容，因此要逐步解码。这是两种不同的计算阶段。

KV cache 保存已有前缀的 key 和 value，供后续重复使用。这里依赖因果注意力以及不变的前缀和配置：追加新 token 不会改变旧位置的因果状态。缓存减少重复计算，但新 query 仍需要与缓存中的位置交互。

MHA、MQA、GQA、MLA：缓存什么？

MHA 为各头保留独立 K/V，MQA 让所有 query 头共享一组 K/V，GQA 则让 h 个 query 头对应 g 个 K/V 头。g=h 时就是 MHA，g=1 时就是 MQA。固定 batch、长度、头维度和精度时，每层常规 KV 缓存元素数约为 2·B·T·g·Dh，实际速度还取决于内核和数据布局。

2024 年 DeepSeek-V2 描述的 MLA 通过潜在表示压缩 K/V，并另行处理位置相关部分。它不是简单地把 GQA 的 g 再缩小。下表比较的是结构思路，不保证模型质量相同，也不是通用的速度排名。

结构Query 头K/V 表示
MHAhh 组独立的头
MQAh1 组共享头
GQAhg 组共享头（1 ≤ g ≤ h）
MLA多个压缩潜在表示与位置相关部分

代码中不能省略的维度交换

示例允许 query 与 key 的序列长度不同。注意力加权后的张量是 [B, H, Tq, Dh]，必须先交换为 [B, Tq, H, Dh]，才能合并多头维度。直接 view 会把头与序列的布局混在一起。

这里的布尔 mask 约定 True 表示允许关注。代码会拒绝没有任何可用 key 的行，避免对整行 −∞ 做 softmax 而出现未定义结果。不同接口的掩码约定可能不同。这个版本用于理解计算，不是带完整 KV cache 或融合内核的生产实现。

展开 PyTorch 示例
import math
import torch
from torch import nn

class MultiHeadAttention(nn.Module):
    """Teaching example. Mask: True = allowed, [Tq,Tk] or [B,1,Tq,Tk]."""
    def __init__(self, d_model, num_heads, dropout=0.0):
        super().__init__()
        if d_model < 1 or num_heads < 1 or d_model % num_heads:
            raise ValueError('d_model must be divisible by num_heads')
        self.h = num_heads
        self.d = d_model // num_heads
        self.q_proj = nn.Linear(d_model, d_model)
        self.k_proj = nn.Linear(d_model, d_model)
        self.v_proj = nn.Linear(d_model, d_model)
        self.out_proj = nn.Linear(d_model, d_model)
        self.dropout = nn.Dropout(dropout)

    def forward(self, q, k, v, allowed=None):
        if any(x.ndim != 3 for x in (q, k, v)):
            raise ValueError('Expected [batch, sequence, d_model] inputs')
        if any(x.shape[0] != q.shape[0] for x in (k, v)):
            raise ValueError('Query, key and value batch sizes must match')
        if any(x.shape[-1] != self.h * self.d for x in (q, k, v)):
            raise ValueError('Input widths must equal d_model')
        if k.shape[1] == 0:
            raise ValueError('At least one key/value is required')
        if k.shape[1] != v.shape[1]:
            raise ValueError('Key and value lengths must match')
        b, tq, width = q.shape
        def split(x, projection):
            return projection(x).reshape(b, -1, self.h, self.d).transpose(1, 2)
        Q = split(q, self.q_proj)  # [B,H,Tq,Dh]
        K = split(k, self.k_proj)  # [B,H,Tk,Dh]
        V = split(v, self.v_proj)  # [B,H,Tk,Dh]
        scores = Q @ K.transpose(-2, -1) / math.sqrt(self.d)
        if allowed is not None:
            if allowed.dtype != torch.bool:
                raise ValueError('Use a boolean mask: True = allowed')
            if not allowed.any(dim=-1).all():
                raise ValueError('Each query must have at least one allowed key')
            scores = scores.masked_fill(~allowed, float('-inf'))
        weights = self.dropout(scores.softmax(dim=-1))
        out = weights @ V  # [B,H,Tq,Dh]
        out = out.transpose(1, 2).contiguous().reshape(b, tq, width)
        return self.out_proj(out)

参考文献