在现代自然语言处理（NLP）和计算机视觉（CV）的任务中，Self-Attention（自注意力机制）是不可或缺的核心技术。无论是 Transformer 架构，还是各种基于 Attention 的模型，都离不开这一机制。本文将深入解析 Self-Attention 的计算流程，详细说明中间的维度变化，帮助你更好地理解这个重要的概念。
什么是 Self-Attention？
Self-Attention 是一种机制，它允许一个序列中的每个元素与其他元素进行交互，基于它们之间的关系生成加权表示。在 NLP 中，Self-Attention 用于捕捉句子中单词之间的依赖关系，甚至可以捕捉远程的上下文信息。
Self-Attention 计算的步骤
Self-Attention 的计算过程可以通过几个关键的矩阵变换来描述。我们从输入矩阵开始，经过一系列的计算步骤，最终得到 Self-Attention 输出。具体的步骤如下：
1. 输入矩阵 X 和权重矩阵
假设输入矩阵 X 由 n 个 token 组成，每个 token 的表示为一个 d_{model} 维的向量。输入矩阵的形状为：
X \in \mathbb{R}^{n \times d_{model}}
接下来，我们需要通过三个训练得出的权重矩阵将输入矩阵 X 转换为 Query（Q）、Key（K）和 Value（V）矩阵：
Q = XW^Q
K = XW^K
V = XW^V
其中，W^Q, W^K, W^V 是学习得到的权重矩阵，维度分别为 (d_{model}, d_k) 和 (d_{model}, d_v)，通常我们取 d_k = d_v = d_{model} / h，其中 h 是注意力头的数量。
2. 计算 Attention Score（打分）
计算 Self-Attention 的核心是计算一个表示注意力分布的矩阵，通常使用 Query 和 Key 的点积来衡量它们之间的相似度：
\text{Attention Score} = QK^T
这里，Q 的维度是 (n, d_k)，K^T 的维度是 (d_k, n)，所以 QK^T 得到的矩阵维度是 (n, n)。这个矩阵的每个元素表示一个 token 与其它 token 之间的相似度。
3. 缩放与 Softmax
由于点积值会随 $d_k$ 增加而变得较大，我们需要对点积结果进行缩放，公式如下：
\text{Scaled Attention Score} = \frac{QK^T}{\sqrt{d_k}}
接下来，使用 Softmax 函数对每一行进行归一化，确保每一行的值之和为 1。Softmax 的操作是对每行的元素进行计算，因此 Softmax 是沿着行方向进行的。
\text{Attention Weights} = \text{softmax}\left(\frac{QK^T}{\sqrt{d_k}}\right)
这时，我们得到了一个维度为 $(n, n)$ 的矩阵，其中每行的值代表了当前 token 对其它 token 的关注程度。每一行的和为 1，表示当前 token 会将自己的注意力分配给其它 token。
4. 计算加权输出
最终，利用上述得到的 Attention Weights 与 Value 矩阵 $V$ 相乘，得到每个 token 的加权表示。公式如下：
\text{Output} = \text{Attention Weights} \times V
由于 Attention Weights 的维度是 (n, n)，而 V 的维度是 (n, d_v)，所以最终的输出维度是 (n, d_v)。
5. 多头注意力（Multi-Head Attention）
在多头注意力机制中，我们并行计算多个 Attention Heads。每个头独立地执行上述的计算，并生成不同的加权输出。最终，将这些输出拼接起来，并通过一个线性变换映射回原始维度：
\text{Multi-Head Output} = \text{Concat}(head_1, head_2, ..., head_h) W^O
其中，W^O 是一个线性变换矩阵，通常将拼接后的结果映射回 d_{model} 维度。
Self-Attention 的维度总结
通过上述步骤，我们可以看到 Self-Attention 的各个计算环节和对应的维度变化：
1.输入矩阵 X 的维度是 (n, d_{model})。
2.经过权重矩阵变换后，得到 Q、K 和 V，它们的维度分别是 (n, d_k)、(n, d_k) 和 (n, d_v)。
3.点积计算后，得到 Attention Score，维度是 (n, n)。
4.通过 Softmax 归一化后，得到 Attention Weights，维度仍为 (n, n)。
5.最终输出的维度是 (n, d_v)，并通过线性变换调整为原始维度 (n, d_{model})。
Softmax 是对行还是对列？
在 Self-Attention 的计算中，Softmax 是对每一行进行的。具体来说，Softmax 会将每一行的打分归一化，使得每一行的注意力权重和为 1。这确保了每个 token 对其它 token 的注意力分布是有效的。
结论
Self-Attention 是 Transformer 和其他现代神经网络架构的核心机制。通过理解其计算过程和维度变化，我们能够更好地掌握其在 NLP 和计算机视觉任务中的应用。希望通过本文的解析，你能对 Self-Attention 的计算流程有更深入的了解，进而帮助你在相关领域的实践和研究中更加得心应手。
引用资料：
- Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L., Gomez, A. A., Kaiser, Ł., & Polosukhin, I. (2017). Attention is all you need. NeurIPS.
- Radford, A., Narasimhan, K., Salimans, T., & Sutskever, I. (2018). Improving Language Understanding by Generative Pre-Training. OpenAI.
这篇文章为你深入解析了 Self-Attention 的计算过程和细节，带你了解中间的维度变化和 softmax 的具体应用。如果你觉得有用，记得分享给朋友，也欢迎提出问题或者讨论！
