上一篇我们理解了为什么需要注意力——它让序列里任意两个 token 之间的信息传递路径从 O(N)O(N) 步降到 O(1)O(1) 步。但「让每个词直接看见所有其他词」具体是怎么做到的?今天这篇就进入自注意力机制的核心运算本身。

我们要回答的问题是:给你一个句子,自注意力如何算出每个词的新表示——一个「看过上下文」的表示?

答案藏在三个角色——Query、Key、Value——和一套精巧的加权平均运算里。

检索的类比:Query 问,Key 答,Value 的内容被取出

想象你去图书馆找一本书。你心里有一个 查询(Query):「关于深度学习训练技巧的书」。图书馆里的每本书都贴着一个 标签(Key),比如「计算机视觉」「优化算法」「Python 入门」。管理员会拿你的 Query 去匹配所有 Key,找到最相关的那几本,然后从这些书里取出 实际内容(Value)给你。

自注意力机制借用了完全相同的逻辑。对于序列里的每个词,我们给它分配三个向量:

  • Query(查询):这个词想问什么?它代表当前词在「寻找」什么样的上下文。
  • Key(键):这个词能提供什么?它代表当前词在「被匹配」时的特征。
  • Value(值):这个词真正要贡献的内容是什么?

这三个向量不是凭空来的——它们来自同一个输入向量与三个不同的权重矩阵相乘。假设输入是词嵌入 xx,则:

q=xWQ,k=xWK,v=xWVq = xW^Q, \quad k = xW^K, \quad v = xW^V

其中 WQW^QWKW^KWVW^V 是模型在训练中学到的参数矩阵。同一个词,用不同的投影得到三个不同的角色,这就是自注意力的第一个巧妙设计:让每个词同时扮演提问者、应答者和内容提供者 1

一个词怎么聚合其他词的信息:五步走

现在,假设我们有一个长度为 nn 的句子,每个词 ii 都已经算好了自己的 qiq_ikik_iviv_i。我们要计算词 ii 的「上下文表示」——即它该从其他词身上吸收多少信息。

第一步:算匹配分数。 拿词 ii 的 Query 向量 qiq_i,去和句子中每个词 jj 的 Key 向量 kjk_j 做点积:

scoreij=qikj=m=1dkqi(m)kj(m)\text{score}_{ij} = q_i \cdot k_j = \sum_{m=1}^{d_k} q_i^{(m)} k_j^{(m)}

点积的含义是「相似度」:两个向量方向越一致,点积越大。这一步回答了「词 ii 和词 jj 有多相关?」。词 ii 和自己的分数通常最高,但有时也会出现其他词分数更高——比如在「it」指向「animal」时。

第二步:缩放。 把分数除以 dk\sqrt{d_k},其中 dkd_k 是 Key 向量的维度。为什么要除?我们稍后解释。

第三步:Softmax 归一化。 对缩放后的分数做 softmax,让它们变成一组和为 1 的正数——注意力权重:

αij=exp(scoreij/dk)j=1nexp(scoreij/dk)\alpha_{ij} = \frac{\exp(\text{score}_{ij} / \sqrt{d_k})}{\sum_{j'=1}^{n} \exp(\text{score}_{ij'} / \sqrt{d_k})}

αij\alpha_{ij} 就是词 ii 分配给词 jj 的注意力权重。值越大,词 ii 就越「关注」词 jj。Softmax 保证权重非负且和为 1,让这变成一个概率分布——相当于一个「软性选择」:不是只选一个词,而是所有词都选,但权重不同。

第四步:加权求和。 用每个 αij\alpha_{ij} 去乘对应的 Value 向量 vjv_j,然后求和:

zi=j=1nαijvjz_i = \sum_{j=1}^{n} \alpha_{ij} v_j

ziz_i 就是词 ii 经过自注意力之后的新表示。它的一大特点是:每个位置的 Value 贡献多少,由 Query 和 Key 的匹配程度决定。如果词 ii 认为词 jj 很重要(αij\alpha_{ij} 很大),vjv_j 就会大量涌入 ziz_i;如果认为词 jj 无关,αij\alpha_{ij} 接近于 0,vjv_j 几乎被忽略。

自注意力中 QKV 的创建与运算流程

上面这张图来自 Jay Alammar 的经典讲解,展示了从输入向量 x1x_1 通过三个矩阵投影得到 q1q_1k1k_1v1v_1,然后计算注意力分数的过程 1

为什么需要缩放?一个数学上的理由

这可能是初学者最容易被忽略的细节。假设 qqkk 的每个分量都是均值为 0、方差为 1 的独立随机变量(训练初期权重随机初始化时大致如此),那么它们的点积 qk=i=1dkqikiq \cdot k = \sum_{i=1}^{d_k} q_i k_i 的方差就是 dkd_k 2

dk=64d_k = 64 时,点积的方差是 64,标准差是 8。这意味着大部分点积会落在 [24,24][-24, 24] 的范围内(均值 ± 3 个标准差)。把这些大数送进 softmax——softmax 的指数函数 exp(x)\exp(x) 会把大数变得极大——结果就是:最大的那个分数会「压倒」所有其他分数,注意力权重几乎变成 one-hot 分布(一个接近 1,其余接近 0)。

更致命的是,在 one-hot 的区域,softmax 的梯度接近于 0,模型几乎学不动了 2

除以 dk\sqrt{d_k} 后,方差被缩回到 1,分数回到合理的范围,softmax 的梯度也变得温和。所以缩放不是可有可无的工程技巧,而是让训练稳定进行的必要条件。

矩阵形式:一次算出所有词

实际实现时,不会一个词一个词地算,而是把所有词的 Query、Key、Value 分别拼成矩阵 QQKKVV,一次算完:

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

  • QQ 的形状是 (n×dk)(n \times d_k),每行是一个词的 Query
  • KK 的形状是 (n×dk)(n \times d_k),每行是一个词的 Key
  • QKQK^\top 得到 (n×n)(n \times n) 的分数矩阵,第 ii 行第 jj 列就是 qikjq_i \cdot k_j
  • Softmax 沿着每一行(每个 query 对应的所有 key)做归一化
  • 最后乘上 VV(形状 (n×dv)(n \times d_v)),得到输出矩阵 (n×dv)(n \times d_v)

这意味着一次矩阵乘法就能算出所有词对所有词的注意力权重和新的表示。这正是 GPU 最擅长的操作——整个自注意力层可以表示为几个矩阵乘法,并行度极高 3

矩阵形式的自注意力运算

上图展示了 XX 矩阵乘以 WQW^Q 得到 QQ 矩阵,然后 QKQK^\top 得到分数矩阵,再经过 softmax 和 VV 矩阵相乘的完整流程 1

这一运算的意义

自注意力层的输出 ziz_i 是一个 上下文感知的表示:它不只是词 ii 本身的含义,还包括了从整个句子中聚合来的、与词 ii 最相关的信息。这种「动态聚合」是 Transformer 的核心能力。

下一篇我们会进入多头注意力——为什么用八组不同的 QKV 投影比只用一组更好,以及「多头」如何让模型同时关注不同层面的关系(比如语法关系和语义关系)。