Transformers 長程上下文綜述

• 2020 年 9 月 1 日
• AI

字幕組雙語原文：Transformers 長程上下文綜述

英語原文：A Survey of Long-Term Context in Transformers

翻譯：雷鋒字幕組（wiige）

多頭自注意力開銷很大，這已經不是什麼秘密了——序列長度的O(n²)複雜度意味著不太可能讓經典的Transformer模型處理超長文本序列。在過去的兩年里，NLP社區已經開發出了名副其實的縫合怪方法來應付這種複雜性，但這篇文章將重點介紹在大規模下更有前景的方法。

1. 稀疏Transformers

2. 自適應Transformers

3. Transformer-XL

4. 壓縮Transformers

5. Reformer

6. 尋覓式 Transformer

密集多頭注意力的時間和空間複雜度

多頭注意力對序列長度的伸縮性很差，原因有二: 首先是計算注意力矩陣所需的FLOPs按序列長度的平方暴增，導致對單個序列的自注意力操作的計算複雜度為O(hdn²)，其中h是注意力頭數，d是key和query的維度，n是文本序列的長度。

第二點是，點積自注意力操作的空間複雜度也按序列長度的平方增大。 計算注意力矩陣的空間複雜度為O(hdn+hn²)——第一項是存儲key和query所需的記憶體，第二項是每個頭得出的注意力標量值。

讓我們向BERT-Base代入具體的數字來了解哪些項的開銷佔大頭。 BERT-Base使用的序列長度為512，隱藏層大小為768，12個頭，這意味著每個頭的維度為64(768 / 12)。 在這種設置下，需要393216 floats(約1.5MB)(頭12個*頭大小64*序列長度512)來存儲鍵和值，而存儲所有頭的注意力標量所需的空間為3145,728 floats (12*512*512)即約12MB的記憶體——幾乎10倍於鍵存儲的空間開銷，這還僅僅是僅僅512個上下文詞符。

由於在訓練過程中必須快取激活，以便進行梯度計算（除非使用梯度檢查點這種激活再計算策略），因此僅存儲這些12層的BERT-Base注意力矩陣，每個樣本就要消耗約150MB的記憶體。 當序列長度為1024時，這就上升為約600MB，而序列長度為2048時，每個樣本僅存儲注意力矩陣就已經消耗了約2.4GB的記憶體。這意味著訓練時的批處理尺寸(batch size)更小，並行性更差，進一步阻礙了模型利用長距離上下文的能力。

稀疏Transformers

“Generating Long Sequences with Sparse Transformers”   由 Rewon Child, Scott Gray, Alec Radford 和 Ilya Sutskever 通過離散化(factorization)方法解決了自注意力的時間和空間複雜度中的O(n²)項問題。

離散注意力

在一個典型的自注意力操作中，輸入序列中的每一項都會關注到輸入序列中的所有其他項，從而形成如下的注意力模式:

典型的自回歸設置下的自注意力連接模式。深藍色方塊代表 “查詢(query)”，淺藍色方塊代表 “鍵(key)”

經典自注意力帶來的益處是，其高連通性使得資訊在詞符之間很容易傳遞——只需要一層注意力就可以聚合任意兩個詞符的資訊。 但如果我們放寬這一約束，保證兩層內資訊可以在任意兩個詞符之間傳遞即可的話，就可以大幅降低按序列長度增長的複雜度。稀疏Transformer通過利用固定注意力模式的自定義核來實現這一目標。

固定的稀疏Transformer。深藍色方塊代表查詢，中淺藍色方塊代表奇數層注意力鍵索引，最淺藍色的方塊代表偶數層關注的關鍵索引。

一半的頭只關注短的局部上下文，而另一半則預分配均勻分布在整個序列中的索引。

通過這些聚合索引來傳遞資訊，網路就能夠從較遠的詞符獲取資訊並利用長距離上下文，同時將時間和空間複雜度降低到O(n√n)。最重要的一點是，只需要兩層隱藏層，任意詞符就都可以整合來自其他詞符的資訊。

稀疏Transformers的實驗結果

重要的是，離散化注意力結構似乎不會對語言模型的性能產生負面影響，它的每字元比特數比enwiki8上的密集注意力(令人驚訝地)更高，且可以在高達12228個詞符的上下文上產生有效的注意力。  

可以看到，稀疏Transformer的注意結構之所以有效，部分原因是因為這種注意力模式與真正習得的密集注意力並沒有什麼不同。 Kevin Clark, Urvashi Khandelwal, Omer Levy, 和 Christopher D. Manning在他們的 “What Does BERT Look At? An Analysis of BERT』s Attention”  一文中探究了密集注意力所習得的模式，試圖弄明白注意力在Transformer模型中承載了什麼功能。他們發現注意力頭傾向於關注緊接在前的詞符（類似於稀疏注意力中的局部注意力模式），以及如[SEP]和句號的特定詞符。所以，也許稀疏Transformer的注意力模式中包含的歸納偏差是有用而非有害的。

例子: BERT習得的注意力模式

如果想在自己的項目中採用固定注意力核，可以查看OpenAI的blockparse庫以及作者發布的配套示例。

自適應窗口Transformers

Sainbayar Sukhbaatar, Edouard Grave, Piotr Bojanowski,  和 Armand Joulin  在他們的文章 “Adaptive Attention Span in Transformers” 中對複雜性問題採取了不同的方法.   他們與 “What Does Bert Look At?” 的作者得到了同樣的觀察，並指出雖然密集注意力允許每個頭關注完整的上下文，但許多注意力頭只關注局部上下文，而剩下的頭才會關注整個文本序列。 他們建議使用某種自注意力來利用這一現象，以此讓模型可以選擇其上下文尺寸。

自適應遮罩(Adaptive Masking)

自適應窗口Transformer通過對序列進行遮罩，使每頭習得的上下文之外的詞符貢獻迅速消失。 遮罩(M)與softmax操作的對數相乘，使某些詞符對當前隱狀態x的貢獻歸零，其中超參數R控制最小窗口(span)大小。

“Adaptive Attention Span in Transformers”所用的軟遮罩函數。圖片來自論文所附的Facebook AI博文

為鼓勵模型只在有益的地方使用額外上下文，對習得的z值施加了ℓ1懲罰。

注意力工作細探與實驗結果

在這些限制條件下，大部分的頭都會選擇關注<100個字元的上下文，只有少數的頭（主要是在網路的後幾層）會選擇頂著ℓ1ℓ1的懲罰來關注>1000個字元的上下文。

除了巧妙的快取機制，這種對長上下文的懲罰使得自適應窗口高達8k個字元，同時仍然保持模型的整體計算成本可控。此外，它在基準線上的性能仍然很高——在enwiki8上達到了0.98比特/字元，text8數據集上達到了1.07比特/字元。

然而，可變窗口大小在便於並行方面並不理想，我們通常希望密集的、尺寸規整的矩陣來實現最佳性能。雖然這種方法可以大幅減少預測時前向計算所需的flops數量，但作者只提供了模糊的性能估計，並表示自適應跨度的實現能夠以2048個上下文詞符的固定上下文大小模型的速度最高處理8192個詞符  。 

Facebook AI Research還開源了他們的工作——程式碼和預訓練的模型可以在github.com/facebookresearch/adaptive-spans上獲得。

Transformer-XL

比起讓密集注意力操作的成本更低廉，Zihang Dai, Zhilin Yang, Yiming Yang, Jaime Carbonell, Quoc V. Le和 Ruslan Salakhutdinov 則從RNNs中獲得靈感，並在Transformer中引入了除自關注機制外的循環機制。 他們的文章”Transformer-XL: Attentive Language Models Beyond a Fixed-Length Context”，引入了兩個新穎的概念——第一個是將之前 “段(segments)”的隱狀態作為輸入回饋給當前段層的組件，另一個是促進這一策略的相對位置編碼方案。

段循環

使用具有固定上下文大小的標準變換器，處理長輸入時需要將輸入分割成塊（或稱段）分別處理。然而，這種方法有一個局限，即前一段的資訊與當前段的詞符無關。 這種段獨立性在某種程度上是有益的，因為它讓我們能分批地有效處理每一段，但如果目標保持是長序列的一致性，這就會成為一個很大的限制。

帶密集注意力的自回歸Transformer的詞符注意力結構來自Transformer-XL。

Transformer-XL通過強制進行分段串聯來克服這一限制。 在第一段之後，後續段中的詞符將始終保持512個詞符的上下文尺寸，之前段的激活作為上下文傳遞給後續段的注意力操作。 這意味著來自NN上下文大小*LL層外的資訊可以傳播到一個給定的令牌。 假設上下文尺寸為640，模型有16層，理論上Transformer-XL可以納入10,240個詞符的資訊。

Transformer-XL原文中的詞符注意力模式。

為了避免存儲所有段的激活，作者阻止了梯度流向之前的段。

納入相對位置

Transformer-XL還引入了一種新穎的位置編碼方案，稱為 “相對位置編碼”。 不是簡單地將內容和絕對位置嵌入的和輸入網路，而是將每層的注意力操作分解為基於內容的部分和基於相對位置的部分——如果一個段中的第512個詞符要關注第511個詞符，它的相對位置嵌入就是-1。

為了使相對位置編碼的使用變得簡單，他們將從鍵和查詢中產生注意力權重的操作進行了分解。對於一個典型的密集注意力操作，pre-softmax注意力權重可以如下分解:

上式中，Exi是詞符在位置i的基於內容的嵌入，Uj是詞符j的位置嵌入。

(a) 將查詢的內容與鍵的內容相關聯

(b) 將查詢的內容與鍵的位置相關聯

(c) 將查詢的位置與鍵的內容相關聯

(d) 將查詢的位置與鍵的位置相關聯

當使用相對位置嵌入時，作者將式子改成如下形式:

在b)和d)中，Uj已經被它的相對位置嵌入Ri-j替換.

對於包含查詢位置的項，我們將矩陣UiUi替換為兩個新學習參數uu和vv。這些向量可以理解為兩個不依賴於查詢具體細節的偏置——cc鼓勵多關注一些項，dd鼓勵多關注相對位置。 作出這種替換是因為對自身的查詢其相對位置保持不變。

注意力工作細探與實驗結果

要想讓Transformer-XL模型利用這樣的長程上下文，每一層至少有一個頭要利用其注意力跨度的全部上下文。 平均注意力權重圖顯示，每一層都有頭主要關注先前的位置。

Transformer-XL論文中的平均注意力權重圖

此外Transformer-XL論文還測量了有效上下文長度對困惑度的影響，並發現增大上下文長度至~900個詞符時會產生更好的困惑度得分——這進一步證明了循環機制在實踐中的有效性，而非僅僅停留於理論。

Kimi Young的github提供了源程式碼，HuggingFace的實現可以幫助你搭建自己的Transformer-XL項目。

壓縮Transformers

列表上的下一個模型是Transfromer，它建立在Transformer-XL架構的基礎上，但通過壓縮損失擴展了該方法，以便納入更長的文本序列。在來自DeepMind的Jack W. Rae、Anna Potapenko、Siddhant M. Jayakumar和Timothy P. Lillicrap的文章 “Compressive Transformers for Long-Range Sequence Modelling “中，詳細介紹了一種能夠關注與全書長度的序列模型架構。

壓縮Transformer注意力

” Compressive Transformers For Long Range Sequence Modeling “中的壓縮記憶(Compressed Memory)示意圖

按照Transformer-XL的做法，當前序列可以關注已存儲的前段的激活。此外，在同樣的多頭注意操作中，當前段中的詞符可以關注存儲在 “壓縮記憶 “中的第二組狀態。

在每一步中，最早的壓縮記憶被丟棄，壓縮記憶向後移動一個索引。 然後，來自正常記憶段中最早的nn個狀態進行壓縮，然後轉移到壓縮記憶新空出來的槽中。

DeepMind部落格上的一個gif很好地說明了這個過程:

將過去的記憶逐漸壓縮成壓縮記憶——圖片由 DeepMind  部落格提供

DeepMind團隊嘗試了多種壓縮操作（包括平均池化、最大池化和學習卷積等等），但最終決定訓練一個二級網路來重建縮記憶中基於內容的注意力矩陣。 

換句話說，他們學習了一個函數fc，通過最小化壓縮記憶的注意力（C-1=fc(Mold)）和正被壓縮的正常記憶狀態的注意力之差，將n個最早的記憶狀態壓縮為一個壓縮記憶狀態。:

他們沒有將這種壓縮操作與主語言模型聯合訓練，而是選擇在一個單獨的優化循環中更新壓縮網路。因為讓注意力狀態容易被壓縮，對降低語言模型的損失會起到反作用。

實驗結果

在他們的實驗中，壓縮記憶的尺寸為512，正常記憶尺寸為512，窗口尺寸為512，壓縮率為2——意味著在壓縮步驟中，最早的2個記憶狀態將被壓縮為1個狀態。在這個設置下，他們在WikiText-103上實現了17.1的SOTA困惑度。

由於利用較長序列長度的收益通常是長尾的，他們特別研究了不同詞頻下的困惑度，發現在最稀有的詞符上的收益尤其顯著:

按詞頻分類的困惑度。圖片由DeepMind部落格提供

雖然他們還沒有公開源碼，但DeepMind已經開放了PG-19的源碼，這是他們在研究壓縮Transformer時開發的數據集。PG-19是Project Gutenberg的衍生品，旨在進一步研究長期注意力。

Reformer

接下來是Nikita Kitaev, Łukasz Kaiser, Anselm Levskaya 的文章”Reformer: The Efficient Transformer”.  Reformer在增大序列長度方面採取了不同的策略——他們沒有引入循環機制或壓縮記憶，而是選擇通過使用局部敏感哈希(LSH)技術來縮小每個詞符的注意力範圍.  

局部敏感哈希是一系列將高維向量映射到一組離散值（桶/簇）的方法。它常常用在近似最近鄰搜索中。

Reformer的作者對注意力操作的鍵和查詢使用了同一個投影，並使用基於隨機旋轉的局部敏感哈希方法將共享的鍵/查詢分到最多幾百個詞符的不同桶中。下圖是該方法的一個示例:

Reformer論文中的角度化LSH。該圖展示了一個有3次旋轉的4桶哈希設置。下方圖中的三個向量已經映射到同一個哈希桶中，因為其的原始向量就很接近，而上方的第一個圖和最後一個展示了向量被映射到不同的哈希桶中。

它們計算每個桶內的注意力矩陣，然後取對應值的加權和。 由於只關注一個給定的桶內的元素，所以如果桶的大小選擇得當，這可以將注意力操作的整體複雜度從O(n2)降低到O(nlogn)。由於分桶是隨機的，而且旋轉映射向量也是隨機的，所以他們計算了多個哈希值，以確保具有相似的共享鍵-查詢嵌入的詞符最終以高概率落入同一個桶中。

他們還採用了 ” The Reversible Residual Network: Backpropagation Without Storing Activations   “中介紹的技術來控制訓練時的空間消耗。可逆殘差層使用巧妙的架構結構，允許從層輸出中輕鬆重構層輸入，並在網路深度不變的情況下以空間複雜度降低了計算量。

通過局部敏感的哈希技巧來降低計算成本，以及可逆殘差來降低記憶體消耗，Reformer架構能夠在單個加速器上處理長達64000個詞符的序列。

雖然在enwiki-8上得出的1.05比特/字元的得分落後於我們在這篇文章中所研究的其他一些模型，但Reformer是一個令人耳目一新的獨特機制，我很期待看到這種方法如何繼續擴大規模。

如果你有興趣進一步探索Reformer架構，可以看看我最近關於這個主題的博文 “A Deep Dive into the Reformer”。 在google/jax Github倉庫中，有一個Reformer的開源實現示例。 Phil Wang也維護了一個PyTorch版本。

路徑 Transformer

我原打算在這裡就結束本文，但在Aran Komatsauzaki的建議下，我加入了最後一篇論文。Aurko Roy、Mohammad Taghi Saffar、David Grangier和Ashish Vaswani的第二篇ICLR 2020的論文 “Efficient Content-Based Sparse Attention with Routing Transformers “與前述的Reformer有一些異曲同工之處。他們將問題框定為路徑問題，旨在學習選擇稀疏詞符聚類簇Si和內容x的函數。

作者的在下圖中說明了他們的方法。 他們不是只關注局部元素或第n個元素來增加稀疏性，而是學習了（圖c中用彩色表示的）需要關注的聚類。重要的是這些簇是每個鍵和查詢的內容的函數，而不僅僅與絕對或相對位置相關。

Routing Attention與Local和Strided Attention的比較，來自Routing Transformer論文

Routing Attention

在確保每個鍵和查詢向量都單位化後，他們使用尺寸為（Dk,Dk）的隨機正交權重的共享矩陣來投影鍵和查詢的值，其中Dk是鍵和查詢的隱藏維度。

然後，根據一組k-means中心將R中的向量分組到k個簇中，每個k-means中心通過每批使用k-means更新來習得，與梯度下降過程無關。

在一個給定的聚類簇Ci內，它們使用標準的加權求和方法計算了一組新的上下文嵌入，其中每個注意力值Ai都是使用標準的點積自注意力計算的。

由於密集注意力中的注意力模式通常由少數關鍵元素主導，而且聚類分配過程需要將具有高注意力權重的鍵和查詢歸入同一個聚類中，作者認為這個機制能夠有效告知我們X′i進行了高開銷的密集操作。

最後，他們選擇了一些接近√n的聚類簇，這樣基於稀疏內容的注意力機制的整體複雜度就變成了O(n√n)。 為了使整個過程易於並行化和統一矩陣尺寸，作者最接近每個中心點的top-k個項來代替真正的k-means聚類。

除了基於內容的路徑注意力， 路徑注意力Transformer還在尺寸為256的上下文上執行局部注意力。

實驗結果

Routing Transformer在計算效率上的提升也使得其在Wikitext-103這個詞級語言建模基準測試上的困惑度得到提升，它們的性能明顯超過了之前描述的Transformer-XL模型。

Wikitext-103測試集上的困惑度結果，來自Routing Transformer論文。

在enwiki-8上，Routing Transformer的表現也相當不錯，不過他們的結果略微落後於Adaptive Span Transformer。

enwiki-8的測試集的每個字元比特數結果。

我本來沒找到Routing Transformer的實現，但Aurko Roy很好心地給我了一份他們源碼的壓縮包，這份源碼是提交給ICLR審閱的版本。

其他在Transformer中的長程上下文方法

如果你對在Transformer中加入長程上下文的其他方法感興趣，那你還可以讀一下:

• Efficient Content-Based Sparse Attention with Routing Transformers

• Adaptively Sparse Transformers

• BP-Transformer: Modelling Long-Range Context via Binary Partitioning

• Scaling Laws for Neural Language Models

雷鋒字幕組是一個由 AI 愛好者組成的翻譯團隊，匯聚五百多位志願者的力量，分享最新的海外AI資訊，交流關於人工智慧技術領域的行業變革與技術創新的見解。

團隊成員有大數據專家、演算法工程師、影像處理工程師、產品經理、產品運營、IT諮詢人、在校師生；志願者們來自IBM、AVL、Adobe、阿里、百度等知名企業，北大、清華、港大、中科院、南卡羅萊納大學、早稻田大學等海內外高校研究所。

如果，你也是位熱愛分享的AI愛好者。歡迎與雷鋒字幕組一起，學習新知，分享成長。

雷鋒網雷鋒網