自注意力（Self-Attention）
本文包含较多数学公式。若公式未正常渲染，请刷新页面后重试。
transformer模型的编码器和解码器层都有自注意力层。
以语言翻译为例：
The animal didn't cross the street because it was too tired
这句话中的it是什么？指的是animal还是street？对人类来说简单但是对机器来说不简单。
当transformer模型处理单词"it"时，自注意力机制使其能够将"it"与"animal"关联起来。
自注意力机制是如何具体实现的呢？
以单个token的视角来看
以标准transformer的编码器为例，
输入词 (经tokenizer分词转id后) 会被embedding转化为512维向量 (例如，9832 -> [0.23,0.11,...]) ，对于整个句子，每个单词都会经过编码器。
第一步，对于每个输入的向量，都会创建一个查询向量（Q）、一个键向量（K）和一个值向量（V），这些向量是通过将输入向量与训练过程中训练的三个矩阵（Wq,Wk,Wv)相乘而创建的。以标准transformer为例，512维输入，QKV向量为64维。
例如，将x1乘以WQ权重矩阵得到与该词关联的"查询"向量q1。
第二步，计算自注意力的第二步是计算一个分数。同样以语言翻译为例，该分数决定了在编码某个位置的单词时，对输入句子其他部分的关注程度，因此，需要将输入句子中的每个单词与该单词进行评分。第一个得分是q1和k1的点积。第二个得分是q1和k2的点积。
第三步和第四步是将分数除以键向量维度的平方根，方便获得更稳定的梯度；然后将结果通过softmax操作。Softmax对分数进行归一化，使它们都为正且总和为1。
第五步是将每个值向量乘以softmax分数
第六步是求和加权后的值向量。这将产生第一个单词的自注意力层的输出。
一般的，（为了方便理解）
计算第 1 个 token 的自注意力输出
- q1点乘k1
- 除以K维度的平方根
- 拓展到其它token,对于每一个token都有q1* ki , i 从1到n
- softmax这个结果
- vi乘以softmax后的结果
- 所有结果相加求和，既得该token的自注意力输出
\sum_{i=1}^{n}
V_i \cdot
\operatorname{Softmax}
\left(
\frac{q_1 \cdot k_i}{\sqrt{d_k}}
\right)
这里的写法是为了逐项展示"每个 (V_i) 都会乘上一个注意力权重"。
实际计算时，Softmax 会同时作用于(\frac{q_1k_1^T}{\sqrt{d_k}},\ldots,\frac{q_1k_n^T}{\sqrt{d_k}}) 这一整组分数，而不是分别对单个标量计算。
我最开始是从"第 1 个 token 对所有 Value 加权求和"这个直觉出发的,
下面把这个直觉展开成严格的矩阵形式。
二、输入矩阵
假设输入序列中有 (n) 个 token：
X=
\begin{bmatrix}
x_1\\
x_2\\
\vdots\\
x_n
\end{bmatrix}
其中，每个 token 的输入向量维度为：
x_i\in\mathbb{R}^{1\times d_{\text{model}}}
假设：
d_{\text{model}}=512
那么：
x_i\in\mathbb{R}^{1\times512}
整个输入矩阵为：
X\in\mathbb{R}^{n\times512}
设每个注意力头中的 Query、Key 和 Value 维度为：
d_q=d_k=d_v=64
则投影矩阵的维度为：
W_Q\in\mathbb{R}^{512\times64}
W_K\in\mathbb{R}^{512\times64}
W_V\in\mathbb{R}^{512\times64}
三、计算 Query、Key 和 Value
通过线性投影得到：
Q=XW_Q
K=XW_K
V=XW_V
其中：
Q=
\begin{bmatrix}
q_1\\
q_2\\
\vdots\\
q_n
\end{bmatrix}
\in\mathbb{R}^{n\times64}
K=
\begin{bmatrix}
k_1\\
k_2\\
\vdots\\
k_n
\end{bmatrix}
\in\mathbb{R}^{n\times64}
V=
\begin{bmatrix}
v_1\\
v_2\\
\vdots\\
v_n
\end{bmatrix}
\in\mathbb{R}^{n\times64}
对于第 (i) 个 token：
q_i=x_iW_Q
k_i=x_iW_K
v_i=x_iW_V
并且：
q_i,k_i,v_i\in\mathbb{R}^{1\times64}
四、计算 (QK^T)
由于：
Q\in\mathbb{R}^{n\times64}
而：
K^T=
\begin{bmatrix}
k_1^T & k_2^T & \cdots & k_n^T
\end{bmatrix}
\in\mathbb{R}^{64\times n}
所以：
QK^T\in\mathbb{R}^{n\times n}
展开后：
QK^T
=
\begin{bmatrix}
q_1k_1^T & q_1k_2^T & \cdots & q_1k_n^T\\
q_2k_1^T & q_2k_2^T & \cdots & q_2k_n^T\\
\vdots & \vdots & \ddots & \vdots\\
q_nk_1^T & q_nk_2^T & \cdots & q_nk_n^T
\end{bmatrix}
矩阵中的第 ((i,j)) 个元素为：
q_ik_j^T
它表示：
第 (i) 个 token 的 Query 与第 (j) 个 token 的 Key 之间的匹配分数。
因此，(QK^T) 的第 (i) 行表示：
第 (i) 个 token 对所有 token 的原始注意力分数。
五、计算缩放后的原始注意力分数
定义原始注意力分数矩阵：
S=
\frac{QK^T}{\sqrt{d_k}}
其中：
S\in\mathbb{R}^{n\times n}
除以 键向量维度的平方根 的作用是缩放点积，防止点积数值随着维度增大而过大，从而使 Softmax 过早进入饱和区域。
矩阵 (S) 可以写成：
S=
\begin{bmatrix}
S_{11} & S_{12} & \cdots & S_{1n}\\
S_{21} & S_{22} & \cdots & S_{2n}\\
\vdots & \vdots & \ddots & \vdots\\
S_{n1} & S_{n2} & \cdots & S_{nn}
\end{bmatrix}
其中：
S_{ij}
=
\frac{q_ik_j^T}{\sqrt{d_k}}
对于第 (1) 个 token：
S_1=
\begin{bmatrix}
S_{11} & S_{12} & \cdots & S_{1n}
\end{bmatrix}
即：
S_{11}
=
\frac{q_1k_1^T}{\sqrt{d_k}}
S_{12}
=
\frac{q_1k_2^T}{\sqrt{d_k}}
\cdots
S_{1n}
=
\frac{q_1k_n^T}{\sqrt{d_k}}
一般地，令：
S_i=
\begin{bmatrix}
S_{i1} & S_{i2} & \cdots & S_{in}
\end{bmatrix}
则：
S=
\begin{bmatrix}
S_1\\
S_2\\
\vdots\\
S_n
\end{bmatrix}
六、逐行计算 Softmax
Softmax 沿着矩阵 (S) 的每一行计算。
对于第 (i) 个 token：
a_i
=
\operatorname{softmax}(S_i)
其中：
a_i=
\begin{bmatrix}
a_{i1} & a_{i2} & \cdots & a_{in}
\end{bmatrix}
第 (j) 个注意力权重为：
a_{ij}
=
\frac{e^{S_{ij}}}
{\displaystyle\sum_{t=1}^{n}e^{S_{it}}}
因此：
a_{ij}>0
并且：
\sum_{j=1}^{n}a_{ij}=1
将所有 token 的注意力权重组合起来，得到注意力矩阵：
A=\operatorname{softmax}(S)
即：
A
=
\operatorname{softmax}
\left(
\frac{QK^T}{\sqrt{d_k}}
\right)
其中：
A=
\begin{bmatrix}
a_{11} & a_{12} & \cdots & a_{1n}\\
a_{21} & a_{22} & \cdots & a_{2n}\\
\vdots & \vdots & \ddots & \vdots\\
a_{n1} & a_{n2} & \cdots & a_{nn}
\end{bmatrix}
\in\mathbb{R}^{n\times n}
矩阵 (A) 的第 (i) 行：
a_i=
\begin{bmatrix}
a_{i1} & a_{i2} & \cdots & a_{in}
\end{bmatrix}
表示：
第 (i) 个 token 对所有 token 的注意力权重。
七、计算最终输出
注意力层的输出为：
Z=AV
其中：
A\in\mathbb{R}^{n\times n}
V\in\mathbb{R}^{n\times d_v}
所以：
Z\in\mathbb{R}^{n\times d_v}
展开后：
Z
=
\begin{bmatrix}
a_1V\\
a_2V\\
\vdots\\
a_nV
\end{bmatrix}
=
\begin{bmatrix}
z_1\\
z_2\\
\vdots\\
z_n
\end{bmatrix}
对于第 (i) 个 token：
z_i=a_iV
由于：
a_i=
\begin{bmatrix}
a_{i1} & a_{i2} & \cdots & a_{in}
\end{bmatrix}
而：
V=
\begin{bmatrix}
v_1\\
v_2\\
\vdots\\
v_n
\end{bmatrix}
所以：
z_i
=
\begin{bmatrix}
a_{i1} & a_{i2} & \cdots & a_{in}
\end{bmatrix}
\begin{bmatrix}
v_1\\
v_2\\
\vdots\\
v_n
\end{bmatrix}
展开后：
z_i
=
a_{i1}v_1
+
a_{i2}v_2
+
\cdots
+
a_{in}v_n
也可以写成：
z_i
=
\sum_{j=1}^{n}a_{ij}v_j
其中：
z_i\in\mathbb{R}^{1\times d_v}
它的含义是：
第 (i) 个 token 根据自己的注意力权重，从所有 token 的 Value 表示中提取信息，得到结合上下文后的新表示 (z_i)。
因此：
Z
=
AV
=
\begin{bmatrix}
z_1\\
z_2\\
\vdots\\
z_n
\end{bmatrix}
\in\mathbb{R}^{n\times d_v}
综上，完整自注意力计算公式为：
\boxed{
\operatorname{Attention}(Q,K,V)
=
\operatorname{softmax}
\left(
\frac{QK^T}{\sqrt{d_k}}
\right)V
}
多头注意力（Multi-Head Attention）
原始transfromer采用了一种名为多头注意力的机制进一步提升性能：通过多头扩展了观察不同位置的能力和给注意力层提供了多个表示子空间。
一、多头注意力的基本过程
多头注意力的计算过程可以概括为：
- 使用多组不同的 (W_Q)、(W_K)、(W_V) 并行计算；
- 每个注意力头分别得到一个输出Z^r；
- 将所有注意力头的输出拼接起来；
- 通过输出投影矩阵 (W^O) 进行融合。
不同的注意力头拥有不同的参数，因此可以用不同的方式判断一个 token 应该关注哪些 token。
二、第 (r) 个注意力头
对于第 (r) 个注意力头：
Q^{(r)}=XW_Q^{(r)}
K^{(r)}=XW_K^{(r)}
V^{(r)}=XW_V^{(r)}
第 (r) 个头的注意力矩阵为：
A^{(r)}
=
\operatorname{softmax}
\left(
\frac{
Q^{(r)}K^{(r)T}
}{
\sqrt{d_k^{(r)}}
}
\right)
第 (r) 个头的输出为：
Z^{(r)}
=
A^{(r)}V^{(r)}
完整地写为：
Z^{(r)}
=
\operatorname{softmax}
\left(
\frac{
Q^{(r)}K^{(r)T}
}{
\sqrt{d_k^{(r)}}
}
\right)
V^{(r)}
对于第 (i) 个 token：
z_i^{(r)}
=
a_i^{(r)}V^{(r)}
展开后：
z_i^{(r)}
=
\sum_{j=1}^{n}
a_{ij}^{(r)}v_j^{(r)}
因此：
第 (r) 个注意力头以自己的方式判断第 (i) 个 token 应该关注哪些 token，再从该头的 Value 表示中汇总一份上下文信息。
三、每个注意力头的维度
假设：
d_{\text{model}}=512
注意力头数为：
h=8
通常令每个注意力头的维度为：
d_k=d_v=\frac{d_{\text{model}}}{h}
因此：
d_k=d_v=\frac{512}{8}=64
对于每个注意力头：
Q^{(r)}\in\mathbb{R}^{n\times64}
K^{(r)}\in\mathbb{R}^{n\times64}
V^{(r)}\in\mathbb{R}^{n\times64}
注意力矩阵为：
A^{(r)}\in\mathbb{R}^{n\times n}
每个头的输出为：
Z^{(r)}\in\mathbb{R}^{n\times64}
四、拼接所有注意力头
将 (8) 个注意力头的输出沿特征维度拼接：
Z_{\text{cat}}
=
\operatorname{Concat}
\left(
Z^{(1)},
Z^{(2)},
\ldots,
Z^{(8)}
\right)
由于每个头的输出维度为：
Z^{(r)}\in\mathbb{R}^{n\times64}
所以拼接后的维度为：
Z_{\text{cat}}
\in
\mathbb{R}^{n\times(8\cdot64)}
即：
Z_{\text{cat}}
\in
\mathbb{R}^{n\times512}
一般地，对于 (h) 个注意力头：
Z_{\text{cat}}
=
\operatorname{Concat}
\left(
Z^{(1)},
Z^{(2)},
\ldots,
Z^{(h)}
\right)
并且：
Z_{\text{cat}}
\in
\mathbb{R}^{n\times(hd_v)}
五、通过 (W^O) 融合多头信息
拼接完成后，通过输出投影矩阵 (W^O)：
Y=Z_{\text{cat}}W^O
在当前例子中：
W^O\in\mathbb{R}^{512\times512}
因此：
Y\in\mathbb{R}^{n\times512}
一般地：
W^O
\in
\mathbb{R}^{(hd_v)\times d_{\text{model}}}
对于第 (i) 个 token，先将所有注意力头的结果拼接：
z_{i,\text{cat}}
=
\begin{bmatrix}
z_i^{(1)}
\;\Vert\;
z_i^{(2)}
\;\Vert\;
\cdots
\;\Vert\;
z_i^{(h)}
\end{bmatrix}
其中，符号 || 表示向量拼接。
随后计算：
y_i
=
z_{i,\text{cat}}W^O
即：
y_i
=
\begin{bmatrix}
z_i^{(1)}
\;\Vert\;
z_i^{(2)}
\;\Vert\;
\cdots
\;\Vert\;
z_i^{(h)}
\end{bmatrix}
W^O
(W^O) 的作用是：
让模型根据最终任务，通过训练自适应地学习如何融合不同注意力头提取的信息。
六、多头注意力的完整公式
定义第 (r) 个注意力头：
\operatorname{head}_r
=
\operatorname{Attention}
\left(
XW_Q^{(r)},
XW_K^{(r)},
XW_V^{(r)}
\right)
即：
\operatorname{head}_r
=
\operatorname{softmax}
\left(
\frac{
XW_Q^{(r)}
\left(XW_K^{(r)}\right)^T
}{
\sqrt{d_k}
}
\right)
XW_V^{(r)}
最终的多头注意力输出为：
\boxed{
\operatorname{MultiHead}(X)
=
\operatorname{Concat}
\left(
\operatorname{head}_1,
\operatorname{head}_2,
\ldots,
\operatorname{head}_h
\right)
W^O
}
