機器之心報導
參與：思、肖清、一鳴
在 Transformer 完全採用注意力機制之後，注意力機制有有了哪些改變？哈希算法、Head 之間的信息交流都需要考慮，顯存佔用、表徵能力都不能忽視。
注意力機制是非常優美而神奇的機制，在神經網絡「信息過載」的今天，讓 NN 學會只關注特定的部分，無疑會大幅度提升任務的效果與效率。藉助注意力機制，神經機器翻譯、預訓練語言模型等任務獲得了前所未有的提升。
但與此同時，注意力機制也面臨著重重問題，首先就是參數量太大，這有點類似於全連接層，注意力機制需要考慮兩兩之間的所有連接。我們可以看到，完全用注意力機制的模型，參數量輕輕鬆鬆破個億，而卷積這類參數共享的運算，參數量一般也就幾百萬。
如果只是參數量大也就算了，能完整利用這些參數，正是說明模型的表徵能力非常強。但問題在於，Transformer 採用的 Multi-head Attention 並沒有充分利用參數的表徵能力。舉個例子，Transformer 中每一個注意力 Head 都是相互獨立的，它們之間沒有信息交流，因此谷歌最近提出的 Talking-Head 就旨在解決這個問題。
本文從原 Multi-head Attention 出發，探索 Reformer 如何用哈希算法大量降低顯存需求，探索 Talking-Head 如何強化全注意力機制的表徵能力。
多頭注意力：開始的地方
Transformer 因在大型預訓練語言模型中的優秀性能而被世人所熟知。這一類模型已廣泛應用於多種預訓練語言模型中，如 BERT、GPT-2 等。當然，在廣泛應用之餘，人們也在思考 Transformer 存在的缺陷，並進行彌補。
論文：Attention Is All You Need論文連結：https://arxiv.org/abs/1706.03762 (https://arxiv.org/abs/1706.03762)提到 Transformer，人們總會想到這類模型最突出的幾個特徵：點乘注意力、多頭注意力、殘差連接、Mask 機制等。在最初介紹 Transformer 的論文中，Waswani 等引入了這些優秀的特性，使 Transformer 獨樹一幟。
優秀的特性，問題的根源
在 Transformer 中，首先引入的是注意力點乘機制。它將輸入分為查詢（Query）、鍵（Key）和值（Value），通過以下公式計算。其相當於根據序列元素之間的相似性，確定每一個元素都應該關注哪些信息。
而另一方面，僅計算一次注意力，不足以捕捉所有的文本結構特徵。因此，Transformer 的提出者想到類似於 CNN 中的 multi-channel 的方法：讓點乘注意力機制「重複」多次。這樣一來，Transformer 的模型複雜度顯著上升，捕捉了更多的語義和語言特徵，這種方法便是我們熟悉的多頭注意力。
圖 2：（左）點乘注意力；（右）多頭注意力，有多個注意力層並行。
多頭注意力的公式如下所示，每一個 Head 都是一個注意力層。為了提高差異性，每一個 Head 的輸入張量都會先經過它特有的線性變換，相當於希望該 Head 注意特定的某個方面。
儘管 Transformer 因為這些特性而產生了很好的性能，但其固有的缺陷也隨之而生。特別需要注意的就是多頭注意力機制了。在設計之初，為了讓不同的注意力頭能夠捕捉語言模型中的不同特徵，不論是 BERT 還是 GPT-2 的設計者都將增加了注意力頭的數量。
但是，由於 Query、Key 和 Value 在訓練中會進行多次線性變換，且每個頭的注意力單獨計算，因此產生了大量的冗餘參數，也需要極大地顯存。
怎樣才能降低計算複雜度呢？研究者們在這個問題上思考了很多方法，本文中可以看到採用哈希算法、參數壓縮等的方法。這些方法有的是強化多頭注意力機制。有的則是利用哈希算法的高效性，替換掉原本的架構，使模型性能更加優秀。
哈希改造下的注意力機制
Transformer 因為 Multi-Head Attention 而成功，也因為它而受到重重限制。在 2017 年 Transformer 提出以後，18 年到 19 年已經有很多研究在改進這種注意力機制，這裡介紹的是一種「哈希大法」，提出該方法的 Reformer 已經被接收為 ICLR 2020 的 Oral 論文。
ICLR 2020 程序主席評論道：「Reformer 這篇論文提出的新型注意力機制有效減少序列長度的複雜度，同時新機制也減少了存儲需求。實驗證明它們非常有效，該論文可能對研究社區會產生顯著影響。儘管有一些次要的觀點，但所有評審者都一致同意接收該論文。」
論文：Reformer: The Efficient Transformer論文地址：https://openreview.net/pdf?id=rkgNKkHtvB
注意力機制到底哪不對？
我們先回到注意力機制，其中非常重要的運算是 Query 與 Key 這兩個張量之間的矩陣乘法，其代表著餘弦相似性。具體而言，如果 Query 表示翻譯領域的中文序列，Key 表示英文序列，那麼矩陣乘法則表示在翻譯成某個中文詞時，需要注意哪些英文詞。
但問題在於，Query 與 Key 的張量維度都是 [batch size, length, d_model]，乘完後的維度是 [batch size, length, length]。要是序列長度 length 不長也就罷了，但如果上下文跨度很廣，需要很長的序列，也就是說 length 遠遠大於隱藏層大小 d_model。那可就麻煩了，研究者表示要是序列長度達到 64K，注意力機制計算一個樣本也需要 16GB 的顯存（Float32）。
你可能會說，我不一次性計算 Query 中的所有序列，而每次只算一步，例如一個中文詞，計算它與所有英文詞之間的相似性。這樣不就沒問題了麼？這是沒問題的，但並沒有實質上的改變。
哈希來幫忙
現在，假定序列長度還是 64K，對於 Query 序列中的每一個 q，正常的注意力機制需要計算它和 Key 序列 64K 個元素之間的相似度，並通過 Softmax 將相似度歸一化為概率。反正都是要計算概率，且一般只有概率最高的一些元素真正對 q 有很大的貢獻，那麼為什麼不直接找出這些元素？
之前超大的矩陣乘法會計算 Query 序列所有元素與 Key 序列所有元素之間的相似度，現在如果不通過矩陣乘法，只找每個 Query 序列「最相近」的 32 個或 64 個元素，那麼顯存與計算豈不是成千倍地減少？
這就是 Reformer 最核心的思想，完成查找「最相近」元素的算法即局部敏感哈希算法（Locality sensitive hashing）。
簡單而言，局部敏感哈希算法（LSH）在輸入數據彼此類似時，它們有很大概率映射後的哈希是一樣的；而當輸入數據彼此不同，它們映射後的哈希值相等概率極小。
LSH 算法根據局部敏感哈希函數族將類似的數據映射到相同的桶（Bucket）。
LSH 注意力機制
如下圖 a 所示，傳統注意力機制需要計算 q 與 k 之間的所有相關性，但實際上真正起作用的注意力分配是非常稀疏的。對於圖 b 來說，Query 和 Key 會先通過哈希桶排序，相似的 q 與 k 會落在相同的桶內，因此直接在桶內執行注意力計算就能逼近完整的注意力機制。
然而圖 b 只是理想情況，真正用於注意力機制還有很多困難要克服，比如說，一個桶內包含很多 q 而沒有一個 k。為此，研究者對 b 執行了一系列轉換，並使排序後的注意力權重，來自相同桶的 q 和 v 都集中在對角線上，就如同圖 c 一樣。注意 c 圖中 Q=K，表示 Transformer 中的自注意力機制。
排好序後，d 圖就可以分成不同的批量，並完成並行計算。這樣不僅計算量與參數量大大減少，還能加速訓練。
最後，整個 LSH 注意力機制就如上左圖所示。對於自注意力機制，先根據 LSH 分成不同的桶，然後對桶排序並分割為不同的批量，這些批量能夠並行處理。最後只要在相同的桶內執行注意力機制就行了，整個過程會節省大量存儲與計算資源。
Talking-Heads Attention
近日，來自 Google 的研究團隊提出一種「交談注意力機制」（Talking-Heads Attention），在 softmax 操作前後引入對多頭注意力之間的線性映射，以此增加多個注意力機制間的信息交流。這樣的操作雖然增加了模型的計算複雜度，卻能夠在多項語言處理問題上取得更好的效果。
論文：Talking-Heads Attention論文地址：https://arxiv.org/abs/2003.02436einsum 表示法
作者在論文的偽代碼中使用大寫字母表示張量，而小寫字母表示其對應維度大小。同時作者在張量的計算中使用了 einsum 表示法，也就是愛因斯坦求和約定。它在 numpy、tensorflow、pytorch 等 Python 擴展庫中均有實現。使用 einsum 表示法能夠僅使用一個函數，就可以優雅地實現如點積、外積、轉置、矩陣-向量乘法、矩陣-矩陣乘法等運算。
einsum 解釋：https://rockt.github.io/2018/04/30/einsum
限於本文篇幅，有關 einsum 表示法的概念，感興趣的小夥伴可參考上面這篇博客。這裡舉個慄子，兩個矩陣的乘法運算使用 einsum 表示法可寫為：
Y[a, c] = einsum(X[a, b], W[b, c])
於是，前面介紹的多頭注意力機制使用 einsum 表示法可改寫為如下形式：
同時，einsum 表示法還支持大於兩個矩陣作為輸入的運算。於是，以上偽代碼可進一步精簡為如下極簡模式：
交談注意力機制
不同於多頭注意力機制，交談注意力機制在 softmax 運算的前後增加了兩個可學習的線性映射 P_l 與 P_w，這使得多個注意力 logit 和注意力 weight 彼此之間能夠相互傳遞信息。
於是，交談注意力機制出現了三個不同的維度：h_k、h 和 h_v，其中 h_k 表示 Query 和 Key 的注意力頭數量，h 表示 logit 和 weight 的注意力頭數量，h_v 則表示值的注意力頭數量。其對應偽代碼表示如下，注釋中標出了每個 einsum 運算所對應的計算量。
Talking-Head 是 Multi-Head 的延伸
當假設注意力機制的輸入與輸出維度相同時，可得到計算多頭注意力機制所需進行的標量乘法運算數目為：
h·(dk +dv)·(n·dX +m·dM +n·m)
而交談注意力機制所需的標量乘法運算數目為：
(dk ·hk +dv ·hv)·(n·dX +m·dM +n·m)+n·m·h·(hk +hv)
可以看到上式中第一項與多頭注意力機制的計算量類似，而第二項是由交談注意力中的線性映射導致的。假如 h<d_k 且 h<d_v，那麼交談注意力機制的線性映射計算量，將分別小於對應多頭注意力機制中的 n·m·dk ·hk 與 n·m·dv ·hv。
作者指出，由於交談注意力機制引入了額外的小尺寸張量運算，在實際運行過程中很可能導致速度比 Transformer 更慢。
最後，多頭注意力機制與交談注意力機制都可看作一種「通用雙線性多頭注意力機制」（GBMA, i.e. general bilinear multihead attention）的特殊形式。我們可以從以下偽代碼看出，GBMA 具有兩個三維的參數張量。
從以下偽代碼不難看出，多頭注意力機制在數學上等效於，使用兩個因子的乘積分別表示 GBMA 中的各參數張量。
而交談注意力機制在數學上等效於，使用三個因子的乘積分別表示 GBMA 中的各參數張量。
這裡的 GBMA 僅作為理論研究探討，由於其計算量較大可能並不具備實用性。
「注意力」這個想法真的非常優雅，所以 2020 年剩下的是，我們如何才能更高效地實現「注意力」？