注意力机制

一切来源于一句朴素的直觉："注意力放在哪里，计算就倾向于哪里。"

如果你接触过深度学习，多半听过一个说法：Transformer 时代的基石是注意力机制（Attention Mechanism）。无论是 GPT、BERT，还是机器翻译、图像生成，几乎所有大模型都离不开它。这篇文章不堆砌术语，而是从"为什么要注意力"讲起，一步步拆解它的动机、数学和实现。

一、为什么要注意力：RNN 的瓶颈

在注意力出现之前，处理序列数据的主力是循环神经网络（RNN/LSTM）。RNN 把上一时刻的隐藏状态 h_{t-1} 传入下一时刻，像一条流水线一样逐词处理句子：

"我 喜欢 深度学习"  →  逐词读入 → 最终隐藏状态

这种方式有两个致命的弱点：

1. 长距离依赖难：信息要经过很多步才能从序列开头传到结尾，中间经过非线性变换，早期的信息很容易"遗忘"。
2. 无法并行：t 时刻必须等 t-1 时刻算完，GPU 的并行能力被浪费。

注意力机制正是来解这两个问题的：它让序列中任意两个位置之间可以直接建立联系，路径长度是 1，且计算天然可并行。

二、核心思想：Query / Key / Value

注意力机制的直觉，用一个"查资料"的场景就能说清：

你（Query，问题）要在一堆资料里找答案。每份资料上有个标题（Key，索引），以及正文（Value，内容）。你先扫一遍标题，找出与自己问题最相关的几份，然后主要阅读这几份的正文。

映射到神经网络里：

• Q（Query）：你当前在关注什么（比如要预测的词）。
• K（Key）：序列里每个位置"是什么"（用来被匹配）。
• V（Value）：序列里每个位置"携带什么信息"（真正被取用的内容）。

注意力要做的就一件事：用 Q 去匹配每一个 K，算出一个"相关度分数"，再按分数加权求和所有的 V。

分数越高的位置，在输出里占的比重越大——这就是"注意力"。

三、缩放点积注意力（Scaled Dot-Product Attention）

现代 Transformer 用的是最常用的形式，公式如下：

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

拆成四步，一步步看：

第 1 步：算相关度

用 Q 与所有 K 做点积，得到一个分数矩阵：

$$\text{score}(q, k_i) = q \cdot k_i$$

点积越大，说明两者方向越接近、越相关。

第 2 步：缩放 1/√d_k

这就是"缩放"二字的由来。为什么要除？

当向量维度 d_k 很大时，点积的值会变得很大，Softmax 的梯度会趋近于 0（处于饱和区），训练变得困难。除以 √d_k 可以让点积的方差保持稳定，让 Softmax 始终落在梯度合理的区间。

第 3 步：Softmax 归一化

对每一行做 Softmax，把分数变成"和为 1 的概率分布"：

$$\alpha_i = \frac{\exp(\text{score}_i)}{\sum_j \exp(\text{score}_j)}$$

分数最高的位置拿到最大的权重，这就是"注意力集中"的体现。

第 4 步：加权求和

用归一化后的权重对 V 加权求和，得到最终的输出向量：

$$\text{output} = \sum_i \alpha_i v_i$$

四、自注意力（Self-Attention）：关注自己

如果 Q、K、V 全部来自同一个序列本身，就称为自注意力。此时每个词都能看到句子里的其他所有词，并决定自己该"多关注谁"。

以 "The animal didn't cross the street because it was too tired" 为例，编码 "it" 时，自注意力会把较大的权重分给 "animal"，从而知道这里的 "it" 指的是动物，而不是别的名词。

自注意力让每个 token 的表示都融合了全局上下文，这正是"长距离依赖"问题的终极解法——任意两词之间只需一步就能互相看见。

五、多头注意力（Multi-Head Attention）

只算一次注意力，相当于所有人只从一个角度去"看"问题。多头注意力则是把 Q、K、V 切分成多个子空间，各自独立算注意力，最后拼接起来：

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

其中每个 head：

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

不同 head 能学到不同的关注模式：有的 head 关注语法关系，有的关注相邻词，有的关注长距离指代。多个视角叠加，表达能力更强。

六、别忘了位置：位置编码

注意力公式本身是"无序"的——它对序列位置不敏感，把顺序打乱结果不变。但语言是有顺序的，"你打了我"和"我打了你"完全不同。

所以 Transformer 会给每个 token 加上位置编码（Positional Encoding），把位置信息注入输入向量，让模型能区分"第 1 个词"和"第 5 个词"。

七、PyTorch 实现：从零写一个多头注意力

理论讲完，来点能跑的。下面是一个完整的 PyTorch 实现，包含单头注意力与多头注意力：

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


class ScaledDotProductAttention(nn.Module):
    """缩放点积注意力"""

    def __init__(self, d_k):
        super().__init__()
        self.d_k = d_k  # 每个 head 的维度，用于缩放

    def forward(self, q, k, v, mask=None):
        # q, k: (B, heads, L, d_k)  v: (B, heads, L, d_v)
        scores = torch.matmul(q, k.transpose(-2, -1))  # (B, heads, L, L)
        scores = scores / (self.d_k ** 0.5)            # 缩放，稳定梯度

        if mask is not None:
            scores = scores.masked_fill(mask == 0, float('-inf'))

        weights = F.softmax(scores, dim=-1)            # 归一化成权重
        out = torch.matmul(weights, v)                 # 加权求和
        return out


class MultiHeadAttention(nn.Module):
    """多头注意力：切分 d_model，并行算注意力后拼接"""

    def __init__(self, d_model, n_heads):
        super().__init__()
        assert d_model % n_heads == 0, "d_model 必须能被 n_heads 整除"
        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, bias=False)
        self.w_k = nn.Linear(d_model, d_model, bias=False)
        self.w_v = nn.Linear(d_model, d_model, bias=False)
        self.w_o = nn.Linear(d_model, d_model, bias=False)

        self.attention = ScaledDotProductAttention(self.d_k)

    def _split_heads(self, x):
        # (B, L, d_model) -> (B, n_heads, L, d_k)
        B, L, _ = x.size()
        return x.view(B, L, self.n_heads, self.d_k).transpose(1, 2)

    def forward(self, q, k, v):
        # q/k/v: (B, L, d_model)
        q, k, v = self.w_q(q), self.w_k(k), self.w_v(v)
        q, k, v = self._split_heads(q), self._split_heads(k), self._split_heads(v)

        out = self.attention(q, k, v)           # (B, n_heads, L, d_k)
        out = out.transpose(1, 2).contiguous()  # 拼回头维度
        B, L, _, _ = out.size()
        out = out.view(B, L, self.d_model)      # (B, L, d_model)
        return self.w_o(out)


# 快速验证
if __name__ == "__main__":
    torch.manual_seed(42)
    d_model, n_heads, L = 64, 8, 10
    x = torch.randn(2, L, d_model)             # 模拟 (batch, 序列长度, 特征维度)
    mha = MultiHeadAttention(d_model, n_heads)
    y = mha(x, x, x)                           # 自注意力：Q=K=V=x
    print("输入:", x.shape, "-> 输出:", y.shape)

运行结果应该是 输入: torch.Size([2, 10, 64]) -> 输出: torch.Size([2, 10, 64])，维度保持不变，方便直接接入 Transformer 的残差连接。

代码里值得注意的两个点：

• scores / d_k**0.5 就是公式里的"缩放"，别小看这一行，少了它深层模型很容易训不动。
• mask：解码器里做自回归时需要挡住"未来的词"（把分数设为 -inf，Softmax 后权重为 0），这是 GPT 等模型"只允许看过去"的关键。

八、注意力机制的应用

• 机器翻译：Attention 最早在 NMT（神经机器翻译）里大放异彩，解决了 RNN 翻译长句时的信息丢失。
• Transformer：纯注意力架构，"Attention is All You Need"。
• BERT / GPT：预训练语言模型全部建立在 Transformer 之上，自注意力是其核心组件。
• 计算机视觉：Vision Transformer（ViT）把图像切成 patch 当序列处理；目标检测里的 DETR 也用注意力建模目标关系。
• 扩散模型 / 多模态：Stable Diffusion 用交叉注意力把文本条件注入图像生成过程。

九、总结

最后用一张图收个尾，把整个流程串起来：

Q ─┐
K ─┤→  Q·Kᵀ →  /√d_k →  softmax →  × V →  输出
V ─┘

注意力机制用一个统一的模式解决了序列建模的两大痛点：

1. 任意两位置一步可达，长距离依赖不再是问题；
2. 计算可并行，大模型才训得动。

理解了 Q/K/V、缩放、Softmax、多头这些零件，你再去读任何一篇 Transformer 相关的论文，都会觉得熟悉——因为它们的核心，都是注意力。

相关主题可以继续阅读：Transformer 架构详解、位置编码的直觉、预训练语言模型 BERT/GPT。