深度学习理论-直观理解 Attention输出特征与输入特征维度一致，但这个特征已经是经过注意力加权的特征，相关性更高的

本文首先介绍 Attention 的原始公式，然后以 Self-Attention 为例，简化后逐步分析 Attention 计算结果表达的含义

Attention

Attention 公式如下：

$Attention = softmax(\frac{Q \cdot K^T}{\sqrt{d_k}}) \cdot V$

其中 softmax 作用是归一化，公式如下：

$softmax(x) = \frac{e^x}{\sum_{i=1}^n{e^{x_i}}}$

我们将 $\frac{Q \cdot K^T}{\sqrt{d_k}}$ 称为 attention score，归一化后 $softmax(\frac{Q \cdot K^T}{\sqrt{d_k}})$ 称为 attention weight

在 Self-Attention 中，输入为 $X$ ，乘以不同的权重矩阵，就得到了不同的 $Q$ 、 $K$ 、 $V$

$Q = X \cdot W_q$

$K = X \cdot W_k$

$V = X \cdot W_v$

为了方便理解，我们先做简化，把权重矩阵 $W_q, W_k, W_v$ 和缩放因子 $\sqrt{d_k}$ 都假设为 1

简化后，Self-Attention 长这样

$softmax(X\cdot X^T) \cdot X$

首先来看 $X \cdot X^T$ 的含义，我们先复习一下，向量内积表示的是两个向量的相关性

假设 $X$ 是 $n \times d$ 的矩阵， $n$ 是输入的数量， $d$ 是特征的维度， $X \cdot X^T$ 是 $n \times n$ 的矩阵，表示输入的每个元素，与其它元素的相关性

softmax 就是做归一化，使得权重的和为 1，表达的含义跟 score 一致，相关性高的权重也高，通过非线性函数 $e^x$ 后，变成了概率分布

用相关性矩阵乘以输入向量，得到了 $n \times d$ 的矩阵，跟输入 $X$ 的尺度一致，含义也一致，依然表示 $n$ 个输入变量的 $d$ 维特征，但这个特征已经是经过注意力加权的特征，相关性更高的元素响应更高。

看到这里，相信你对 attention 机制已经有了直观的理解。下面就把之前简化的细枝末节加回来。

权重矩阵 $W_q, W_k, W_v$ 都是可训练的参数，具有以下作用

dk 作为缩放因子有如下两个作用：

5.1 防止数值过大，避免梯度消失或梯度爆炸

$QK^T$ 是两个向量的点积，其值范围随着维度 $d_k$ 增大而增大。
如果 不除以 $\sqrt{d_k}$ ，那么较大维度时，点积的数值会变得非常大，导致 softmax 结果变得极端（接近 0 或 1），从而导致梯度消失，影响训练稳定性。
除以 dkdk 后，使得点积值的范围保持在适当区间，从而让 softmax 更平滑。

5.2 保持不同维度下的数值稳定性

假设 $Q$ $K$ $V$ 的维度分别为 $N \times d$ ， $M \times d$ ， $M \times d$

时间复杂度: $O(NMd)$
- 计算 Q, K, V: $O(Nd^2)$ $O(Md^2)$ $O(Md^2)$ （线性变换）
- 计算 Attention-Score $QK^T$ : $O(NMd)$
- 计算 Softmax: $O(NM)$
- 计算加权求和 $softmax(QK^T)V$ : $O(NMd)$
- 总体上，主要瓶颈是 $QK^T$ 的计算和加权求和，因此时间复杂度为 $O(NMd)$ 。
空间复杂度: $O(NM)$
- 由于需要存储 $QK^T$ （一个 $N×M$ 的矩阵），因此空间复杂度是 $O(NM)$ 。

对于 Self-Attention，由于 $N=M$ ，时间复杂度为 $O(N^2d)$ ，空间复杂度为 $O(N^2)$

如果想屏蔽某些特征，应该如何做？mask 是怎样实现的？

Google T5 不除以 $\sqrt{d_k}$ 为什么也能够收敛？