注意力机制

直觉理解

读一句话时，理解某个词往往要参考句子里的其他词。「它」指代什么，要回头去找前面的名词；英语 bank 指「银行」还是「河岸」，要看周围的词。注意力机制让模型在处理每个位置时，动态地决定该关注序列中哪些位置、关注多少，再把这些位置的信息按权重汇总过来。

可以把它想象成一次「软」的查字典：

• 查询（query，$q$）：当前位置想找什么。
• 键（key，$k$）：每个位置用来被匹配的标签。
• 值（value，$v$）：每个位置真正提供的内容。

用查询和所有键逐一比较相似度，得到一组权重；再用这组权重对所有值做加权平均。与普通字典只返回一个精确匹配不同，注意力会按相似程度从所有条目里各取一部分。

Bahdanau 等人于 2014 年公开的神经机器翻译论文，是现代序列注意力的代表性起点：让解码器在生成每个词时回看源句的不同部分。这不意味着所有神经网络注意力思想都始于该工作。2017 年的 Attention Is All You Need 去掉循环与卷积，用注意力、逐位置前馈层等组件搭建了 Transformer。

定义与公式

缩放点积注意力

把 $n$ 个查询、$m$ 个键和值分别排成矩阵 $Q\in {\mathbb{R}}^{n\times d_k}$、$K\in {\mathbb{R}}^{m\times d_k}$、$V\in {\mathbb{R}}^{m\times d_v}$：

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

逐步来看：

1. $Q{K}^{\top }$ 是一个 $n\times m$ 的矩阵，第 $(i,j)$ 个元素是第 $i$ 个查询与第 $j$ 个键的点积，衡量两者的相似度。
2. 除以 $\sqrt{d_k}$ 做缩放。
3. 对每一行做 softmax，得到第 $i$ 个查询对所有位置的注意力权重，每行和为 1。
4. 用权重对 $V$ 的行加权求和，得到每个查询的输出。

为什么要除以 $\sqrt{d_k}$

假设 $q$ 和 $k$ 的各分量相互独立、均值为 0、方差为 1，那么点积 $q\cdot k={\sum }_{i=1}^{d_k}{q}_i{k}_i$ 的均值为 0、方差为 $d_k$。维度越高，点积的数值就越大。

softmax 对输入的尺度很敏感：输入差距很大时，输出会接近 one-hot，几乎所有权重集中在一个位置，其他位置的梯度接近 0，训练变得困难。除以 $\sqrt{d_k}$ 把点积的方差拉回 1，让 softmax 工作在梯度正常的区间。

自注意力与交叉注意力

$Q$、$K$、$V$ 都由输入经过线性变换得到：$Q=X_q{W}^Q$，$K=X_{kv}{W}^K$，$V=X_{kv}{W}^V$。根据 $X_q$ 和 $X_{kv}$ 是不是同一个序列，分成两种：

• 自注意力（self-attention）：$X_q=X_{kv}$，序列内部的每个位置互相关注。Transformer 的编码器和解码器都用它来建模上下文。
• 交叉注意力（cross-attention）：查询来自一个序列，键和值来自另一个序列。例如在翻译模型中，解码器的查询去关注编码器输出的源句表示；在多模态模型中，文本去关注图像特征。

因果掩码

语言模型按从左到右的顺序生成文本。输入位置 $i$ 的表示可以看到 $x_1,\dots ,x_i$，用来预测 $x_{i+1}$；预测 $x_t$ 时则只能以 $x_{<t}$ 为条件，不能看到目标 $x_t$ 本身及后续内容。做法是在 softmax 之前，把所有 $j>i$ 位置的分数设为 $-\infty$，并让标签相对输入错开一位：

$$\mathrm{softmax}\left(\frac{Q{K}^{\top }}{\sqrt{d_k}}+M\right),M_{ij}=\{\begin{array}{ll}0& j\le i\\ -\infty & j>i\end{array}$$

$e^{-\infty }=0$，这些位置的权重就变成 0。这个下三角形状的掩码叫因果掩码（causal mask）。有了它，一条长度为 $n$ 的序列在一次前向计算里就能同时训练 $n$ 个「预测下一个 token」的任务。

另一种常见的是填充掩码：批次里的序列长短不一，短的要补齐，补齐的位置也要用掩码屏蔽掉。

多头注意力

单头注意力在每个查询处产生一组注意力权重。多头注意力用 $h$ 组可学习的投影，把完整输入映射到不同子空间；不是先把原始输入切成互不相干的 $h$ 份。常见配置取每头 $d_k=d_v=d_{\text{model}}/h$，各头独立做注意力，最后拼接输出并做一次线性变换：

$$\mathrm{MultiHead}(X_q,X_{kv})=\mathrm{Concat}({\text{head}}_1,\dots ,{\text{head}}_h)W^O$$ $${\text{head}}_i=\mathrm{Attention}(X_q{W}_i^Q,X_{kv}{W}_i^K,X_{kv}{W}_i^V)$$

不同的头可以学到不同的关系，比如有的关注相邻的词，有的关注句法上的依存，有的关注指代。由于每个头的维度变小了，总计算量与单头全维度注意力相当。

计算复杂度

对于全序列的稠密自注意力，设 $h$ 为头数、每头维度为 $d_k$、$d=h{d}_k$：$Q{K}^{\top }$ 与对 $V$ 加权求和的总计算量为 $O(n^{2}d)$；朴素实现显式保存各头的分数或权重，空间为 $O(h{n}^2)$，并不是 $O(n^{2}d)$。Q/K/V 和输出投影另需 $O(n{d}^2)$ 计算。FlashAttention 可以避免存下完整权重矩阵，但不改变稠密注意力计算量的平方阶。使用 KV 缓存后，单个新 token 的注意力计算随已有上下文长度线性增长。围绕这些瓶颈有大量的改进工作：

• 从精确计算的实现入手，减少显存读写，如 FlashAttention，见注意力加速内核。
• 只计算部分位置对，见稀疏注意力与局部注意力。
• 让多个查询头共享键和值，减少推理时的 KV 缓存，如 MQA、GQA，见KV缓存压缩与注意力变体。
• 换成亚二次复杂度的结构，见高效序列架构。

代码示例

手写缩放点积注意力和多头注意力，并与 PyTorch 内置实现对照：

import math
import torch
import torch.nn as nn
import torch.nn.functional as F
 
def scaled_dot_product_attention(q, k, v, mask=None):
    # q: (..., n_q, d_k)  k: (..., n_k, d_k)  v: (..., n_k, d_v)
    d_k = q.size(-1)
    scores = q @ k.transpose(-2, -1) / math.sqrt(d_k)
    if mask is not None:
        scores = scores.masked_fill(~mask, float("-inf"))
    weights = scores.softmax(dim=-1)
    return weights @ v, weights
 
class MultiHeadAttention(nn.Module):
    def __init__(self, d_model, n_heads):
        super().__init__()
        assert d_model % n_heads == 0
        self.h, self.d_k = n_heads, 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 split(self, x):                   # (B, n, d_model) -> (B, h, n, d_k)
        B, n, _ = x.shape
        return x.view(B, n, self.h, self.d_k).transpose(1, 2)
 
    def forward(self, x_q, x_kv, mask=None):
        q, k, v = self.split(self.w_q(x_q)), self.split(self.w_k(x_kv)), self.split(self.w_v(x_kv))
        out, _ = scaled_dot_product_attention(q, k, v, mask)
        B, _, n, _ = out.shape
        return self.w_o(out.transpose(1, 2).reshape(B, n, -1))   # 拼接各头再投影
 
torch.manual_seed(0)
x = torch.randn(2, 5, 16)                 # (batch, 序列长度, d_model)
causal = torch.tril(torch.ones(5, 5, dtype=torch.bool))   # 因果掩码：只看自己和之前
 
mha = MultiHeadAttention(d_model=16, n_heads=4)
print(mha(x, x, causal).shape)            # 自注意力：torch.Size([2, 5, 16])
 
memory = torch.randn(2, 7, 16)            # 例如编码器的输出
print(mha(x, memory).shape)               # 交叉注意力：torch.Size([2, 5, 16])
 
# 与 PyTorch 内置实现对照
q = k = v = torch.randn(2, 4, 5, 8)
ours, _ = scaled_dot_product_attention(q, k, v, causal)
ref = F.scaled_dot_product_attention(q, k, v, attn_mask=causal)
print(torch.allclose(ours, ref, atol=1e-6))   # True

注意掩码的约定：这里和 F.scaled_dot_product_attention 一样，布尔掩码中 True 表示「允许关注」；而 nn.MultiheadAttention 的布尔掩码中 True 表示「屏蔽」，正好相反，使用时要看清文档。

在 AI 中的作用

注意力是 Transformer 的核心运算，因而也是几乎所有大语言模型、视觉 Transformer 和多模态模型的核心。它能在一步之内连接任意两个位置，不像循环网络那样要逐步传递信息，既擅长捕捉长距离依赖，又便于在 GPU 上并行计算。推理时，已经算过的键和值可以缓存下来重复使用，这就是 KV 缓存，它是大模型推理优化的中心话题之一。

参考资料

• Attention Is All You Need，第 3.2 节与表 1
• FlashAttention：精确注意力的 IO 与存储优化

相关

• Transformer：注意力如何与前馈层、残差连接、归一化组成完整的模型。
• Attention Is All You Need：提出缩放点积注意力、多头注意力与 Transformer 的论文。
• 线性代数：矩阵乘法与点积。
• 激活函数：softmax 的定义。
• 语言模型：因果掩码与自回归生成。
• KV缓存压缩与注意力变体：推理时如何缓存和压缩键值。
• 长上下文：注意力的平方复杂度与长序列。