哈希算法、愛因斯坦求和約定，這是2020年的注意力機制

機器之心報導

參與：思、肖清、一鳴

在 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.02436

einsum 表示法

作者在論文的偽代碼中使用大寫字母表示張量，而小寫字母表示其對應維度大小。同時作者在張量的計算中使用了 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 表示法可改寫為如下形式：