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