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_{ntimes d} = begin{pmatrix} —— ~ vec{x_{1}} ~ —— \ —— ~ vec{x_{2}} ~ —— \ vdots \ —— ~ vec{x_{n}} ~ —— \ end{pmatrix} ~\ ~\ Q_{ntimes 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_{ntimes 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_{ntimes 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) = sumlimits^{n}_{i=1} frac{exp(q_{i}k_{i}^{T})}{sumlimits^{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}}}{sumlimits^{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}) = sumlimits^{infty}_{i=1} w_{i}(vec{x}) y_{i} ]

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

[w_{i}(x) = frac{K_{h}(x, x_{i})}{sumlimits^{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)，导致内积越大，也就是其中一向量越能 “代表” 另一向量，通俗的解释即： “注意力在此处更集中”。