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\)，導致內積越大，也就是其中一向量越能「代表」另一向量，通俗的解釋即：「注意力在此處更集中」。