編譯 | Mr Bear
編輯 | 叢末
眾所周知，多頭注意力機制 (Multi-Head Self-Attention) 的計算開銷很大。在處理長度為 n 的序列時，其 的時間複雜度會使得原始的 Transformer 模型難以處理長文本序列。在過去的兩年里，已經出現了多種有效的方法來應對多頭注意力機制的複雜度問題，本文將重點討論在模型規模方面很有發展前景的方法。
1
密集型多頭注意力的時空複雜度
多頭自注意力機制擴展到長文本序列的能力很差，原因有二：
第一，計算注意力矩陣所要求的每秒浮點運算數（FLOPs）與序列長度的平方成正比，導致單個序列上的自注意力運算的計算複雜度為，其中 h 是注意力頭的數量，d 是鍵向量和查詢向量的維數，n 是序列的長度。
另一件糟糕的事是：點積自注意力運算的空間複雜度也與序列長度的平方成正比。計算注意力矩陣的空間複雜度為，其中 hdn 是存儲鍵和查詢所需的內存的階，而 是指存儲每個注意力頭產生的標量注意力值所需內存的階。
我們用 BERT-Base 中的一些具體數字來解釋一下複雜度到底有多高。BERT-Base 序列輸入的最大長度為 512，768 個的隱藏維度和 12 個注意力頭，這意味著每個注意力頭有 64 維（768/12）。在這種設定下，需要 393,216 個浮點數（約為 1.5MB）（12 個注意力頭* 64 注意力頭的維度* 512 序列長度）來存儲鍵和值，而存儲所有注意力頭得到的標量注意力值所需的內存將達到 3,145,728 個浮點數（12 * 512 * 512）或約 12MB 的設備內存，這裡所需的內存幾乎是將鍵存儲在長度為 512 個詞的上下文時的 10 倍。
為了進行梯度運算，必須在訓練過程中緩存激活值，除非使用諸如梯度檢查點之類的重計算策略，因此對於每個示例來說，存儲所有 12 層 BERT base 的注意力矩陣就需要大概 150MB 的內存。當序列長度為 1024 時，內存需求則變為約 600MB，而序列長度為 2048 時，對於每個示例的注意矩陣而言，內存需求就已達到約 2.4GB。這意味著在訓練時可以使用的批處理規模較小，並行性較差，並會進一步阻礙我們訓練該模型處理長上下文的能力。
2
稀疏 Transformer
Rewon Child，Scott Gray，Alec Radford 和 Ilya Sutskever 在論文 "Generating Long Sequences with Sparse Transformers" 中提到，可以通過因式分解解決多頭自注意力 的高時間和空間複雜度的問題。
論文地址：https://arxiv.org/abs/1904.10509
1、分解後的注意力
在典型的自注意力機制中，其輸入序列中的每一項都要計算與輸入序列中的其它項之間形成的注意力，從而得到如下所示的注意力模式：
圖 1：以一種自回歸的形式組織的傳統自注意力機制的連接模式。深藍色方塊代表「查詢」向量，而淺藍色方塊代表「鍵」向量。
傳統自我注意力機制的一大好處是，高度的連通性使詞例 (token) 之間的信息容易流通，僅需一層即可聚合來自任何兩個詞例的信息。但是，如果放寬此約束，並確保任意兩個詞例之間的信息只能在經過兩層之後才可以流通，則可以極大地降低長序列帶來的複雜性。稀疏 Transformer 通過編寫利用固定注意力模式的自定義核來實現此目標。
圖 2：稀疏 Transformer 的固定注意力變體。最深的藍色方塊代表「查詢」向量，較淺的藍色方塊代表被奇數層注意的「鍵」向量索引，最淺的藍色方塊代表被偶數層注意的「鍵」向量索引。
一半的注意力頭只關注較短的局部上下文中的項，而剩餘的一半注意力頭關注的是預先指定好的在整個序列中均勻分布的索引。通過根據這些聚合索引確定信息流動路徑，網絡仍然能夠使得相距較遠的詞例間傳遞信息，並使用長距離上下文，同時將時間和內存複雜性降低到 。重要的是，它只需要兩層就可以讓任意詞例考慮來自任何其它詞例的信息。
2、實驗結果
值得注意的是，分解後的注意力結構似乎不會對語言建模的性能產生負面影響，反而驚人地令每個字符需要的比特數比密集的注意力機制（原始 Transformer）在 Enwiki8 語料環境下要少一些，並且可以在包含多達 12,228 個詞例的上下文中得到高效的注意力。
可以想像的到，稀疏 transformer 之所以起作用，部分原因是它學到的的注意力模式與實際學習的密集注意力模式並沒有什麼不同。在 Kevin Clark 等人發表的文章"What Does BERT Look At? An Analysis of BERT’s Attention"中，作者探索了密集注意力學習到的模式，想要找到注意力在transformer 中起什麼作用。他們發現關注緊密相連的前面的詞例（類似於稀疏注意力機制中的局部注意力模式）以及關注特定聚合詞例（如 [SEP] 和句號）的注意力頭有重要作用。因此，可能編碼在稀疏 transformer 注意模式中的歸納偏置是有積極作用的。
論文地址：https://arxiv.org/abs/1906.04341
圖 3：BERT 學到的注意力模式示意圖，線的深度表明了注意力權重的強度（其中有一些注意力權重太小，以至於線變成了透明的）
要在自己的項目中使用固定注意力kernel，請查看OpenAI的blocksparse庫和作者作為開源發布的附帶示例。
3
自適應跨度 Transformer
Sainbayar Sukhbaatar 等人在論文"Adaptive Attention Span in Transformers"中選擇了另一種辦法來解決這個問題。和前一篇"What Does Bert Look At?"的類似，他們也觀察了 Transformer 中起作用的組件。
論文地址：https://arxiv.org/abs/1905.07799
值得注意的是，儘管密集注意力機制讓每個注意力頭可以關注整個上下文，但很多注意力頭只考慮局部上下文，而其它注意力頭則考慮了整個可用序列。他們建議通過使用一種自注意力的變體來利用這種觀察結果，這種變體允許模型選擇它的上下文大小。
1、自適應掩模
Adaptive Span Transformer 通過實現這一點的方式是：通過對序列進行掩模運算，使得學習到的每個注意力頭的上下文之外的詞例的貢獻迅速降為零。mask（M）與 softmax 操作的分對數相乘，從而將某些詞例對當前隱藏狀態 x 的貢獻歸零，超參數 R 控制最小跨度的大小。
圖 4："Adaptive Attention Span in Transformers" 採用的 soft-masking 函數
他們應用了一個學習到的z值的 懲罰項，以鼓勵模型僅在有益的情況下使用額外的上下文。
2、對注意力的思考和實驗結果
由於這些限制，大多數注意力頭只關注少於100 字符的上下文，而只有少數（主要是在網絡的後幾層）選擇添加一個 懲罰項，從而學習大於 1000 個字符的上下文。
圖 5：注意力跨度隨層數變化的示意圖
除了採用巧妙的緩存策略，這種對長距離上下文的懲罰項使得跨度自適應 Transformer 可以使用使用高達 8k 個字符的注意力跨度，同時仍然保持模型的總體計算開銷較低。此外，模型在基準對比測試中的性能仍然很高，在 Enwiki8 數據集上達到每字符占用 0.98 比特，在 text8 數據集上則達到每字符 1.07 比特。
然而，可變跨度大小的注意力在並行性方面並不理想，因為我們通常希望使用密集、大小一致的矩陣，從而獲得最佳性能。雖然這種方法可以顯著減少在預測時計算前向傳播所需的每秒浮點運算次數，但作者只給出了模糊的性能估計，聲明自適應跨度的實現使我們可以以與 2,048 個上下文詞例的上下文大小固定的模型相近的速率處理長達 8,192 個詞例的上下文。
Facebook AI 研究院也開源了他們的成果——代碼和預訓練模型請參閱以下連結：github.com/Facebook Research/adaptive-spans
4
Transformer-XL
相比於降低密集型注意力操作複雜度的方式，Zihang Dai 等人受 RNN 的啟發，在 transformer 的自注意機制之外引入了一種循環機制。他們的論文 "Transformer-XL: Attentive Language Models Beyond a Fixed-Length Context" （論文地址：https://arxiv.org/abs/1901.02860）介紹了兩個新概念：（1）一種新的組件，它將先前「段落」的隱藏狀態作為循環段落層的輸入；（2）使這種策略容易實現的相對位置編碼方案。
1、段落循環
標準 transformer 的上下文大小是固定的，要想處理長的輸入需要將輸入分成塊（段落），並分別處理每個塊（段落）。
然而，這種方法存在一個限制：前面段落的信息不能流向當前的詞例。這種段與段之間的獨立性在某種程度上有益於高效地對段落進行批處理，但從長距離一致性的角度出發，這又變成了一個主要限制因素。
圖 6：在 Transformer XL 的論文中，在一個自回歸 transformer 中的詞例注意力結構示意圖
Transformer XL 通過強制地使段落被串行處理來解決這一限制。處理完成第一個段之後，先前段的激活值將作為上下文傳遞給後續段的注意力，因此始終有 512 個緊鄰的字符的上下文被記錄。這意味著跨度為 N 個上下文大小 * L 層的信息可以傳播到給定的詞例。假設上下文大小為 640，並且模型有 16 層，那麼 Transformer XL 理論上至多可以考慮 10,240 個字符的信號。
圖7：Transformer-XL 中的詞例注意力模式，其中端的長度為 4
為了避免存儲先前所有段的激活值，作者阻止了梯度流向前面的段。
2、引入相對位置
Transformer-XL 還引入了一種新的位置編碼方案，即「相對位置編碼」。這種方案不再簡單地將內容和絕對位置嵌入的總和作為網絡輸入，而是對每層的注意力操進行分解，其中一部分基於內容執行注意力，另一部分基於相對位置執行注意力。因此，塊中的第 512 個詞例會注意到第 511 個詞例，在這裡會採用 相對位置 -1 相應的嵌入。
為了使相對位置編碼易於處理，他們將生成注意力權重的操作和鍵和生成查詢向量、鍵向量的操作分離。對於典型的密集注意力機制，進行 softmax 計算之前的注意力可以被分解如下：
在上面的公式中， 表示基於內容的注意力在位置 i 處的詞例的嵌入 ， 是詞例 j 的位置編碼嵌入。其中，每一項的含義如下：
a）基於鍵內容的查詢內容「尋址」
b）基於鍵位置的查詢內容位置偏置
c）根據鍵的內容的查詢位置偏置
d）根據鍵的位置的查詢位置偏置
使用相對位置嵌入時，公式被改寫為：
對於包含在查詢的位置中的項（c 和 d），我們使用兩個新的可學習參數 u 和 v 替代了 矩陣。這些向量可以被理解為兩個偏置，它們並不依賴於查詢的特徵。c 促使模型對某些項的注意力程度比對其它項更高，而 d 促使模型對某些相對位置的注意力程度比對其它位置更高。這種替代的方式受到了下面這一事實的啟發：查詢相對於自己的相對位置始終保持不變。
3、對注意力的思考和實驗結果
為了使 Transformer-XL 模型能夠使用長距離的上下文，每層中至少有一個注意力頭必須利用其注意力跨度內的全部上下文。平均注意權重圖顯示，每一層都有一些注意力頭，注意到先前許多的位置。
圖 8：在前面 640 個詞例上的平均注意力，其中每行對應於一個注意力頭，每列對應於一個相對位置。圖中共有 160 個注意力頭，每十個注意力頭同屬於同一層。顏色越深代表注意力值越大。
另外，Transformer-XL 的原始論文測量了有效上下文長度對困惑度（perplexity，交叉熵的指數形式）的影響。作者發現，增加上下文長度（上下文長度高達九百個詞例）會得到更好的困惑度分數（預測樣本更準確），這進一步證明了循環機制不僅理論上可行，而且實際上也十分有效。
如果想要開始將 Transformer-XL用於自己的項目：
原始碼請參見 Kimi Young 的 github：https://github.com/kimiyoung/transformer-xl
或查看HuggingFace 的實現：https://huggingface.co/transformers/model_doc/transformerxl.html
5
壓縮 Transformer
接下來我們將介紹壓縮 Transformer（Compressive Transformer），它是建立在 Transformer-XL 的架構之上的。壓縮 Transformer 使用了一個壓縮損失對 Transformer-XL 進行了拓展，考慮了更長的序列長度。在論文「Compressive Transformers for Long-Range Sequence Modelling」中，來自 DeepMind 的 Jack 等人詳細描述了一種模型架構，該架構能夠處理與整本書一樣長的序列。
論文地址：https://arxiv.org/abs/1911.05507
1、壓縮的 Transformer 注意力
圖 9：壓縮 Transformer 保留了過去的激活值的細粒度內存，它們隨後被壓縮到粒度較粗的壓縮內存中。上圖中的模型有一個三層的結構：（1）一個長度為 3 的序列（2）一個長度為 6 的內存（3）長度為 6 的壓縮內存。高亮顯示的內存在每層中被一種壓縮函數 f_c 壓縮到單個壓縮內存中，而不是直接被接下來的序列忽略。在本例中，壓縮比為 c=3。
遵循著 Transformer-XL 的設定，序列可以注意到一組存儲下來的之前段的激活值。另外，在相同的多頭注意里操作中，當前段中的詞例可以注意到存儲在「壓縮內存」中的第二組狀態。
在每一個時間步上，最早的一個壓縮內存會被丟棄，壓縮內存的索引會回退一位。接著，常規的內存段中最早的 n 個狀態會被壓縮，並且被轉移到壓縮內存中新開的槽中。
圖 10：將過去存儲的內存逐漸壓縮到壓縮內存中。
DeepMind 的研究團隊嘗試使用了各種各樣的壓縮操作（包括平均池化、最大池化、學習到的卷積操作等對比基線），但是他們最終決定訓練一個輔助網絡，該網絡被用於重建基於內容的被壓縮的內存的注意力矩陣。換而言之，他們學習了一種函數 ，該函數通過最小化壓縮內存上的注意力 與被壓縮的常規內存中狀態的注意力之間的差異，將最早的 n 個內存狀態壓縮到了單個壓縮後的內存狀態中：
由於讓狀態容易被壓縮與減小語言模型的損失是相矛盾的，他們選擇在一個獨立的優化循環中更新壓縮網絡，而不是同時訓練這種壓縮操作和主要的語言模型。
2、實驗結果
在實驗中，他們設置壓縮內存的大小為 512，正規內存的大小為 512，滑動窗口大小為 512，壓縮率為 2（這意味著最早的 2 個內存狀態會在壓縮步驟中被壓縮到單個狀態中）。在這樣的實驗環境下，他們取得了目前最好的實驗結果——在 WikiText-103 數據集上取得了 17.1 的困惑度。
由於通過利用更長的序列所獲得的收益往往是符合長尾分布的，他們專門關注了由詞頻刻畫的離散特徵的困惑度，並指出最稀有的詞例帶來的收益是最顯著的：
圖 11：通過詞頻刻畫的離散特徵桶（bucket）的困惑度
儘管他們的原始碼至今還沒有公開，DeepMind 已經開源了它們研發壓縮 Transformer 時所使用的數據集 PG-19。PG-19 是 Gutenberg 計劃的衍生品，旨在進一步研究長距離注意力。
數據集地址：https://github.com/deepmind/pg19
6
哈希注意力 Reformer
接下來，我們繼續關注由 Nikita Kitaev 等人完成的工作「Reformer：The Efficient Transformer」。Reformer 對於長度增長的序列採用了另一種策略，他們選擇通過局部敏感哈希技術將每個詞例的注意力範圍變窄，而不是引入循環機制或壓縮內存。
論文地址：https://arxiv.org/abs/2001.04451
局部敏感哈希，是一類將高維向量映射到一組離散值（桶/聚類簇）上的方法。在大多數情況下，它被用作一種近似的最近鄰搜索方法。
Reformer 的作者針對注意力操作的鍵向量和查詢向量使用了一種單一投影，他們使用了一種基於隨機旋轉的局部敏感哈希方法，將共享的鍵/查詢聚集到至多幾百個詞例的桶。哈希方法的示意圖如下：
圖12：角度局部敏感哈希算法示意圖。圖中顯示了三輪哈希，每輪哈希中將空間分成了 4 個桶。在下面一排示意圖的三輪哈希中中，向量都被映射到了同一個桶中，因為它們在輸入中也相近；而在上面一排的第一次哈希和第三次哈希中，向量映射到了不同的桶中。
它們計算了每個桶中的注意力矩陣，並對相應的值進行加權。由於它們只能注意到給定的中的元素，當選取的桶大小合適時，這樣做可以將注意力操作的整體空間複雜度從 降低到 。由於這裡的分桶（bucketing）處理是隨機並且基於隨機旋轉的，他們計算了若干次哈希，從而保證具有相近的「鍵-查詢」嵌入的詞例最終很有可能被分在同一個桶中。
他們額外地使用了 "The Reversible Residual Network: Backpropagation Without Storing Activations" 中提出的技術將訓練時的內存消耗維持在可控範圍內。逆殘差層使用了非常巧妙的架構，使模型可以容易地根據層的輸出重構層的輸入，並以額外的計算換取網絡深度不變的內存複雜性。
論文地址：https://arxiv.org/abs/1707.04585
儘管 Reformer 論文中給出的在 Enwiki-8 數據集上 1.05 比特每字符的性能比本文中給出的其它一些模型稍弱，但是 Reformer 仍然提出了一種令人耳目一新的機制引入了長距離上下文，人們期待看到 Reformer 進一步的拓展。
如果你有興趣更深入地探索 Reformer 架構，請參閱博文「 A Deep Dive into the Reformer」：
https://www.pragmatic.ml/reformer-deep-dive/Google/jax
Github 代碼倉庫中的 Reformer 的開源實現連結如下：
https://github.com/google/trax/tree/master/trax/models/reformer
一種 Phil Wang 訓練的 PyTorch 版的實現連結如下：
https://github.com/lucidrains/reformer-pytorch
7
路由Transformer
Aurko Roy 等人於 ICLR 2020 上發表的工作 "Efficient Content-Based Sparse Attention with 路由Transformers"與前文提到的 Reformer 有一些異曲同工之妙。它們將該問題建模為了一個路由問題，目的是讓模型學會選擇詞例的稀疏聚類 （將其作為內容 x 的函數）。
論文地址：https://openreview.net/profile?email=grangier%40google.com
在下面的圖表中，作者描述了他們的方法。他們並不只關注局部元素或每 n 個元素來增加稀疏性，而是學習了通過下圖 c 中的顏色表示的需要關注的聚類簇。重要的是，這些簇是關於每個鍵和查詢的內容的函數，而不僅僅與它們的絕對或相對位置相關。
圖 13：這裡的圖形顯示了路由 Transformer 與局部注意力機制和 Child 等人於 2019 年提出的跨步注意力機制的對比結果。每一行代表輸出，每一列代表輸入。對於局部注意力機制和跨步注意力機制來說，著色的方塊代表每一個輸出行注意到的元素。對於路由注意力機制來說，不同的顏色代表輸出詞例的聚類中的成員。
1、路由注意力
確保每個鍵和查詢向量都具有單位大小後，他們使用了一種公共的隨機權重矩陣對鍵和查詢的值進行投影，投影的尺寸為 ，其中D_K 是鍵和查詢的隱藏維度。
隨後，R 中的向量會根據一組 k-均值聚類中心 C 被聚類到 k 個簇中。K-均值聚類中心是通過使用每個批（batch）中的 K-均值更新學習到的，它與梯度下降過程相獨立。
在給定的聚類簇 C_i 中，他們使用了一種標準的加權求和方法，計算了一組新的上下文嵌入，其中每個注意力值 A_i 是使用標準的點乘自注意力計算而來的。
由於密集型注意力機制中的注意力模式往往由少數的關鍵元素，並且分配聚類的過程需要將具有高注意力權重的鍵和查詢分到同一個聚類簇中，作者認為這樣做保留了關鍵的信息。這說明 應用的密集型操作計算開銷過高。
最終，他們選擇了大約 個聚類，因此他們稀疏的基於內容的注意力機制的整體複雜度為 。為了讓整個過程易於並行化計算，並且可以處理統一大小的矩陣，作者使用了最接近每個聚類中心的前 k 個項來代替真正的 k-均值聚類。
除了基於內容的路由注意力機制，路由Transformer 還在一個大小為 256 的內容窗口上執行了局部注意力。
2、實驗結果
路由Transformer 在計算效率方面取得了提升，從而使模型在 Wikitext-103 上（一種單詞級別的語言建模對比基準）取得了在困惑度指標上性能的提升，他們的性能大大超過了上文介紹的 Tranformer-XL 模型。
表 1：在 Wikitext-103 數據集上的語言建模結果。Local Transformer 指的是由 Vaswani 等人於 2017 年提出的 Transformer，它用到了相對位置編碼和局部注意力。表中給出了測試集上的困惑度。
在 Enwiki-8 數據集上，路由Transformer 的性能也非常出色（儘管其性能稍微落後於自適應跨度 Transformer）。
表 2：在 Enwiki-8 數據集上的語言建模結果。Local Transformer 指的是由 Vaswani 等人於 2017 年提出的 Transformer，它用到了相對位置編碼和局部注意力。表中給出了測試集上的比特每位元組（bpc）。
路由Transformer 的原始碼連結如下：
https://storage.googleapis.com/routing_transformers_iclr20/routing_transformers_iclr20.zip
8
其它處理 Transformer 長文本序列方法
如果你想進一步了解其它處理 Transformer 長文本序列方法，可以參閱下列文章：
Efficient Content-Based Sparse Attention with Routing Transformers，https://openreview.net/forum?id=B1gjs6EtDr
Adaptively Sparse Transformers，https://arxiv.org/abs/1909.00015
BP-Transformer: Modelling Long-Range Context via Binary Partitioning，https://arxiv.org/abs/1911.04070v1
Scaling Laws for Neural Language Models，https://arxiv.org/abs/2001.08361
via https://www.pragmatic.ml/a-survey-of-methods-for-incorporating-long-term-context/
招 聘