Self-Attention：初步理解

Self-Attention 的基本结构与计算

Attention（注意力）实际上就是权重的另一种应用的称呼，其具体结构与初始输入的 content \(\vec{x_{1}}, \vec{x_{2}}, \cdots, \vec{x_{n}} \in \mathcal{X}\) 紧密相关。其中， \(\vec{x_{1}}, \vec{x_{2}}, \cdots, \vec{x_{n}}\) 为维度相同（设为 \(d\)，即 \(\vec{x_{i}} \in \mathbb{R}^{d}\) for \(\forall 1 \leq i \leq n\)）的向量。所谓 word embedding，实质是用低维的向量表示物体，但是，表示时需要注意，对于任意两种不同物体的 embedding，若两物体本身有着相似的属性（这个定义可以比较抽象，例如，绿巨人与钢铁侠、在地理上相近的两个物体、相似的声音等等都能称作具有某种相似的属性，具体需要看模型的任务和目的是什么），那么它们的 embedding 向量经过某种计算出来的结果，或 “距离” 需要很近。反之，如果两件物体风马牛不相及，或者在模型中我们极力希望将它们分开，那么它们的 embedding 相计算出的 “距离” 应当很远。

例如，在NLP任务中每个 \(\vec{x_{i}}\) 代表了一个 word embedding（原论文中每个word embedding 的维度 = 512，i.e., \(d = 512\)）。我们的实际任务是，对于每一个 \(\vec{x_{i}}\)，分别计算其对应的 attention \(A_{i}\)，具体计算方法如下：

对于每一个 word embedding \(\vec{x_{i}} \in \mathbb{R}^{d}\)，分别计算

• query： \(\vec{q_{i}} = \vec{x_{i}} W^{Q} \in \mathbb{R}^{d}\)
• key：\(\vec{k_{i}} = \vec{x_{i}} W^{K} \in \mathbb{R}^{d}\)
• value：\(\vec{v_{i}} = \vec{x_{i}} W^{V} \in \mathbb{R}^{d}\)

其中，\(W^{Q}, W^{K}, W^{V}\) 分别为 \(d \times d\) 的参数方阵，那么 \(\vec{q_{i}}, \vec{k_{i}}, \vec{v_{i}}\) 皆为 \(d\) 维行向量。对于 \(1 \leq i \leq n\)，可以合并写为矩阵形式，i.e.，

\[X_{n\times d} = \begin{pmatrix} —— ~ \vec{x_{1}} ~ —— \\ —— ~ \vec{x_{2}} ~ —— \\ \vdots \\ —— ~ \vec{x_{n}} ~ —— \\ \end{pmatrix} ~\\ ~\\ Q_{n\times d} = X W^{Q} = \begin{pmatrix} —— ~ \vec{x_{1}} ~ —— \\ —— ~ \vec{x_{2}} ~ —— \\ \vdots \\ —— ~ \vec{x_{n}} ~ —— \\ \end{pmatrix} \begin{pmatrix} & \Big| & \Big| & & \Big| \\ & \vec{w^{Q}_{1}}, & \vec{w^{Q}_{2}}, & \cdots, &\vec{w^{Q}_{d}}\\ & \Big| & \Big| & & \Big| \\ \end{pmatrix} = \begin{pmatrix} —— ~ \vec{q_{1}} ~ —— \\ —— ~ \vec{q_{2}} ~ —— \\ \vdots \\ —— ~ \vec{q_{n}} ~ —— \\ \end{pmatrix} ~\\ ~\\ K_{n\times d} = X W^{K} = \begin{pmatrix} —— ~ \vec{x_{1}} ~ —— \\ —— ~ \vec{x_{2}} ~ —— \\ \vdots \\ —— ~ \vec{x_{n}} ~ —— \\ \end{pmatrix} \begin{pmatrix} & \Big| & \Big| & & \Big| \\ & \vec{w^{K}_{1}}, & \vec{w^{K}_{2}}, & \cdots, &\vec{w^{K}_{d}}\\ & \Big| & \Big| & & \Big| \\ \end{pmatrix} = \begin{pmatrix} —— ~ \vec{k_{1}} ~ —— \\ —— ~ \vec{k_{2}} ~ —— \\ \vdots \\ —— ~ \vec{k_{n}} ~ —— \\ \end{pmatrix} ~\\ ~\\ V_{n\times d} = X W^{V} = \begin{pmatrix} —— ~ \vec{x_{1}} ~ —— \\ —— ~ \vec{x_{2}} ~ —— \\ \vdots \\ —— ~ \vec{x_{n}} ~ —— \\ \end{pmatrix} \begin{pmatrix} & \Big| & \Big| & & \Big| \\ & \vec{w^{V}_{1}}, & \vec{w^{V}_{2}}, & \cdots, &\vec{w^{V}_{d}}\\ & \Big| & \Big| & & \Big| \\ \end{pmatrix} = \begin{pmatrix} —— ~ \vec{v_{1}} ~ —— \\ —— ~ \vec{v_{2}} ~ —— \\ \vdots \\ —— ~ \vec{v_{n}} ~ —— \\ \end{pmatrix} \]

如上所示，\(\vec{w^{Q}_{i}}, \vec{w^{K}_{i}}, \vec{w^{V}_{i}}\) 为 \(d \times 1\) 的列向量 for \(\forall 1 \leq i \leq d\)。

现在，对于 word embedding \(\vec{x_{i}}\)，已求得其对应的\(\vec{q_{i}}, \vec{k_{i}}, \vec{v_{i}}\)，因此 \(\vec{x_{i}}\) 的 attention 记作：

\[A_{i}(q_{i}, K, V) = \sum\limits^{n}_{i=1} \frac{\exp(q_{i}k_{i}^{T})}{\sum\limits^{n}_{j=1} \exp(q_{j}k_{j}^{T})} v_{i} \]

其中，\(q_{i}k_{i}^{T}\) 与 \(q_{j}k_{j}^{T}\) 代表了 query 与 key 的内积，结果为标量。则 \(A_{i}(q_{i}, K, V)\) 的维度与最后乘上的 value \(v_{i}\) 相同，即为 \(1 \times d\) 的行向量。由于一共有 \(n\) 个 word embedding （\(1 \leq i \leq n\)），对应地，最终也应有 \(n\) 个维度为 \(1 \times d\) 的attention。写作矩阵形式为：

\[A(X) = A(Q, K, V) = \mbox{softmax} \big( \frac{QK^{T}}{\sqrt{d}} \big) V \]

\(A(X)\) 即为 \(n \times d\) 的矩阵，softmax 定义为：

\[\mbox{softmax}(z_{i}) = \frac{e^{z_{i}}}{\sum\limits^{n}_{j=1}e^{z_{j}}} \]

注意，最终式中除以\(\sqrt{d}\) 的原因是，维度 \(d\) 的增大会导致整个向量的方差增大，因此更容易出现极端值（即非常大与非常小的值），使 softmax 的梯度变得极小。

从 Nadaraya–Watson Kernel Regression 到 Attention

Attention 其实就是 Nadaraya–Watson Kernel Regression 在 Deep Learning 中的应用，核心思想完全一致，实际上这种思想在机器学习中随处可见，尤其在非参估计（Non-parametric estimation)中。

线性回归及其衍生（e.g. Lasso, Ridge and etc.）存在的一个缺陷是，如果我们不知道independent variables 与 dependent variables 之间联系的参数形式，那么就无法建立模型并对参数进行估计。因此，Kernel Regression 所解决的便是在没有模型假设的情况下对一个新的 test point \(\vec{x}\) 进行 label 的预测。

一个顺应逻辑的想法是，将新的 test point \(\vec{x}\) 的 local neighborhood \(X\) 中所包含的全部 observed data （or training data）的 label 的平均值视为 estimate \(\hat{y}\)，即：

\[\hat{y} = f(\vec{x}) = \mbox{average estimate } y \mbox{ of observed data in a local neighborhood } X \mbox{ of } \vec{x} \]

也就是说，对于新的 test data \(\vec{x}\), 它的 label 可以被估计为邻域中所有已知数据的 label 的平均值。当然，我们对于邻域的选择是灵活的，并且 “平均值” 也只是其中一种估计法。总得来说，我们有 Kernel Regression 的一般式：

\[\hat{y} = \hat{f_{n}}(\vec{x}) = \sum\limits^{\infty}_{i=1} w_{i}(\vec{x}) y_{i} \]

其中，\(w_{i}(\vec{x})\) 为突显 local observation 的权重，定义为：

\[w_{i}(x) = \frac{K_{h}(x, x_{i})}{\sum\limits^{n}_{j=1} K_h(x, x_{j})} \]

对于 Kernel Regression 中 “核” （即kernel，或 localization function） 的选择，一般来说有：

• Gaussian Kernal: \(\quad K_{h}(x, x^{'}) = e^{-\frac{||x - x^{'}||^{2}}{h}}\)

• Box Kernel: \(\quad K_{h}(x, x^{'}) = \mathbb{I}_{\left\{ ||x-x^{'}|| \leq h \right\}}\)

• Triangle Kernel: \(\quad K_{h}(x, x^{'}) = \left[ 1 - \frac{||x - x^{'}||}{h} \right]_{+}\)

Kernel 的选择是灵活的，其本质只是衡量任意 observed data 对一个新数据点的预测值的贡献程度。因此通常满足：对于距待预测数据 \(\vec{x}\) 越近的 \(\vec{x_{i}}\)，所得到的函数结果 \(K_{h}(\vec{x}, \vec{x_{i}})\) 应越大。

到这里我们可以很清晰地发现，attention 就是一个运用了 exponential function 作为 kernel 的权重运算结果。因此，attention 的计算也可以形象地写为：

• 根据已知数据 \(x_{i}\) 与相应的 label \(y_{i}\) (\(1 \leq i \leq n\)) ，预测在 \(x\) 处的 label \(y\)。\(x\) 即为要查询的 query，\(x_{i}\) 即为 key，\(y_{i}\) 即为 value，满足：

\[\begin{align*} y = \sum \limits^{\infty}_{i=1} \alpha(x, x_{i})y_{i}\\ \alpha(x, x_{i}) = \frac{k(x, x_{i})}{\sum_{j} k(x, x_{j})} \end{align*} \]

同时，这也揭示了为什么它的名字叫做 “attention（注意力）”，这个注意力就像 Kernel Regression 我们取的 local neighborhood，代表了我们在预测 \(\vec{x}\) 的 label 时，注意力放在了结果权重大的 neighborhood 中，而对于 neighborhood 以外，权重相对很小，因此不需要过分关注。

Attention 结构的意义

现在我们知道：

\[A(X) = A(Q, K, V) = \mbox{softmax} \big( \frac{QK^{T}}{\sqrt{d}} \big) V \]

其中 \(Q = X W^{Q}, K = X W^{K}, V = X W^{V}\)。

我们知道，\(X W^{Q}\) （\(XW^{K}, XW^{V}\) 同理） 的本质是将 \(X\) 中的各行向量：\(\vec{x_{1}}, \vec{x_{2}}, \ldots, \vec{x_{n}}\) 变换到 \(W^{Q}\) 中以各列向量：\(\vec{w^{Q}_{1}}, \vec{w^{Q}_{2}}, \ldots, \vec{w^{Q}_{d}}\)为基所表示的向量空间中。所得新矩阵的第 \(m\) 列，为 \(X\) 在 \(W^{Q}\) 的第 m 个基（即 \(\vec{w^{Q}_{m}}\)）上的投影。 那么， 对于公式中分子 \(Q K^{T}\)，本质上是变换到两个向量空间中的 \(X\) 的矩阵相乘，

\[QK^{T} = XW^{Q} (W^{K})^{T} X^{T} \]

从实际意义上可以理解为：

\[X X^{T} = \begin{pmatrix} —— ~ \vec{x_{1}} ~ —— \\ —— ~ \vec{x_{2}} ~ —— \\ \vdots \\ —— ~ \vec{x_{n}} ~ —— \\ \end{pmatrix} \begin{pmatrix} & \Big| & \Big| & & \Big| \\ & \vec{x_{1}}^{T}, & \vec{x_{2}}^{T}, & \cdots, &\vec{x_{d}}^{T}\\ & \Big| & \Big| & & \Big| \\ \end{pmatrix} \]

以上的矩阵运算实际上是令 \(\vec{x_{1}}, \vec{x_{2}}, \ldots, \vec{x_{n}}\) 两两分别做内积（包括与自身），而向量内积：

\[a \cdot b = |a| \cdot |b| \cdot \cos \theta \]

其中 \(\theta\) 为向量 \(a, b\) 之间的夹角。因此，内积运算反映了两个向量相似度。当两个向量越相似，即夹角越小，i.e. \(\theta \rightarrow 0, \cos \theta \rightarrow 1\)，导致内积越大，也就是其中一向量越能 “代表” 另一向量，通俗的解释即： “注意力在此处更集中”。