選自towardsdatascience
作者：Raimi Karim
機器之心編譯
參與：Panda W、杜偉
BERT 及其多種變體已經在多種語言理解任務上取得了非常出色的表現，這些架構全都基於 Transformer，而 Transformer 又使用了一種名為「自注意力」的方法。本文將通過圖示和代碼對自注意力機制進行透徹的解讀。當然，在閱讀本文之前，你可能也想了解什麽是注意力機制。沒有問題，同一位作者機器學習工程師 Raimi Karim 之前已經通過類似的方式解讀過了：《圖解神經機器翻譯中的注意力機制》。
BERT、RoBERTa、ALBERT、SpanBERT、DistilBERT、SesameBERT、SemBERT、MobileBERT、TinyBERT 和 CamemBERT 有什麽共同點？別說「BERT」，那不是我想要的答案。
答案：自注意力（self-attention）。
我們要探討的不僅是名字裡面帶有「BERT」的架構，而是「基於 Transformer」的架構。基於 Transformer 的架構主要用在語言理解任務的建模中；這種架構沒有使用神經網絡中的循環（recurrence）來學習輸入和輸出之間的全局依賴關係，而是完全依靠自注意力。
但自注意力背後的數學原理是怎樣的呢？
這就是本文所要探討的主題。本文的主要內容是帶你縱覽自注意力模塊中所涉及的數學運算。你在讀完本文之後，應該就有能力從頭開始編寫自注意力模塊代碼了。
本文的目標不是提供自注意力模塊中不同數值運算和數學運算背後的直觀理解和解讀，也不是為了說明 Transformer 中使用自注意力的原因和方式（這方面的解釋網上能找到很多）。本文也不會刻意說明注意力（attention）和自注意力的差別。
自注意力是什麽？
你可能會想自注意力與注意力可能有相似之處，你的想法是對的！它們的概念基本一樣，也有很多共同的數學運算。
自注意力模塊有 n 個輸入，並會返回 n 個輸出。這個模塊中發生了什麽？用門外漢的話來說，自注意力機制讓每個輸入都會彼此互動（自），然後找到它們應該更加關注的輸入（注意力）。自注意力模塊的輸出是這些互動的聚合和注意力分數。
圖示
下面將按照以下步驟通過圖示來說明自注意力：
1. 準備輸入
2. 初始化權重
3. 推導鍵（key）、查詢（query）和值（value）
4. 計算輸入 1 的注意力分數
5. 計算 softmax
6. 將分數與值相乘
7. 對加權的值求和，得到輸出 1
8. 為輸入 2 和 3 重複 4-7 步驟
備注：在實踐中，這些數學運算是經過向量化的，即所有輸入會一起經歷這樣的數學運算。我們會在後面的代碼章節看到這一點。
第一步：準備輸入
圖 1.1：準備輸入。
為了方便說明，我們先來 3 個輸入，每個輸入的維度為 4.
第二步：初始化權重
每個輸入必須有 3 個表征（見下圖）。這些表征被稱為鍵（key，橙色）、查詢（query，紅色）和值（value，紫色）。在此示例中，我們設這些表征的維度為 3。因為每個輸入的維度為 4，所以這意味著每組權重的形狀為 4×3。
備注：後面我們會看到，值的維度也就是輸出的維度。
圖 1.2：為每個輸入推導鍵、查詢和值。
為了得到這些表征，每個輸入（綠色）都要與一組鍵的權重、一組查詢的權重、一組值的權重相乘。在這個示例中，我們按如下方式初始化這三個權重：
鍵的權重：
查詢的權重：
值的權重：
備注：在神經網絡設置中，這些權重通常是較小的數值，初始化也是使用合適的隨機分布來實現，比如高斯分布、Xavier 分布、Kaiming 分布。
第三步：推導鍵、查詢和值
現在我們有三組權重了，我們來實際求取每個輸入的鍵、查詢和值的表征：
輸入 1 的鍵表征：
使用同樣一組權重求取輸入 2 的鍵表征：
使用同樣一組權重求取輸入 3 的鍵表征：
向量化以上運算能實現更快的速度：
圖 1.3a：推導每個輸入的鍵表征。
通過類似的方式，我們求取每個輸入的值表征：
圖 1.3b：推導每個輸入的值表征。
最後還有查詢表征：
圖 1.3c：推導每個輸入的查詢表征。
備注：在實踐中，可能還會為矩陣乘法的積添加一個偏置向量。
第四步：計算輸入 1 的注意力分數
圖 1.4：根據查詢 1 計算注意力分數（藍色）。
為了求取注意力分數，我們首先求輸入 1 的查詢（紅色）與所有鍵（橙色，包括其自身的鍵）的點積。因為有 3 個鍵表征（因為輸入有 3 個），所以會得到 3 個注意力分數（藍色）。
注意這裡僅使用了輸入 1 的查詢。後面我們會為其它查詢重複同一步驟。
備注：上面的運算也被稱為點積注意力（dot product attention），這是眾多評分函數中的一個，其它評分函數還包括擴展式點積和 additive/concat，請參閱《圖解神經機器翻譯中的注意力機制》。
第五步：計算 softmax
圖 1.5：對注意力分數（藍色）執行 softmax。
對這些注意力分數（藍色）進行 softmax:
第六步：將分數與值相乘
圖 1.6：通過將值（紫色）與分數（藍色）相乘，推導加權的值表征（黃色）。
每個輸入的經過 softmax 的注意力分數（藍色）與其對應的值（紫色）相乘。這會得到 3 個對齊向量（alignment vector，黃色）。在本教程中，我們稱之為加權值（weighted value）。
第七步：對加權的值求和，得到輸出 1
圖 1.7：對所有加權值（黃色）求和得到輸出 1（深綠色）。
將所有加權值（黃色）按元素求和：
所得到的向量 [2.0, 7.0, 1.5]（深綠色）是輸出 1，這是基於輸入 1 的查詢表征與所有其它鍵（包括其自身的）的互動而得到的。
第八步：為輸入 2 和 3 重複 4-7 步驟
現在已經完成了對輸出 1 的求解，我們再為輸出 2 和輸出 3 重複步驟 4-7。我相信現在你完全能自己完成這些計算了。
圖 1.8：為輸入 2 和 3 重複之前的步驟。
備注：因為使用了點積評分函數，所以查詢和鍵的維度必須總是一致。但是，值的維度可能不同於查詢和鍵。由此造成的結果是所得輸出的維度與值的維度一致。
代碼
下面是用 PyTorch 寫的代碼。PyTorch 是一個使用 Python 的常用深度學習框架。為了在代碼中享受 @ 算子、.T 和 None 索引方法的 API 的便利，請確保你使用的是 Python 3.6 或更新版本以及 PyTorch 1.3.1。請跟隨以下步驟，直接將代碼複製到 Python/IPython REPL 或 Jupyter Notebook 中。
第一步：準備輸入
第二步：初始化權重
第三步：推導鍵、查詢和值
第四步：計算注意力分數
第五步：計算 softmax
第六步：將分數與值相乘
第七步：對加權值求和
備注：PyTorch 已經為這些運算提供了一個 API，即 nn.MultiheadAttention。但是，這個 API 需要你輸入鍵、查詢和值的 PyTorch 張量。此外，這個模塊的輸出會經過一次線性變換。
擴展用於 Transformer
那麽，接下來又如何呢？Transformer！實際上我們正處於深度學習研究和高計算資源齊頭並進的激動人心的時代。Transformer 源自論文《Attention Is All You Need》https://arxiv.org/abs/1706.03762，最早是為神經機器翻譯而生的。在此基礎上，研究者繼續向前——組合、分拆、添加和擴展，最後將 Transformer 擴展到了更多語言任務上。
這裡我簡要說說我們可以怎樣將自注意力擴展用於 Transformer 架構：
在自注意力模塊內：維度、偏置；
自注意力模塊的輸入：嵌入模塊、位置編碼、截斷、掩碼；
添加更多自注意模塊：多頭、層堆疊；
自注意力模塊之間的模塊：線性變換、層歸一化
香港九龍灣馬來街興發大廈15樓D室