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$，导致内积越大，也就是其中一向量越能 “代表” 另一向量，通俗的解释即： “注意力在此处更集中”。