0
| 本文作者： AI研習社-譯站 | 2020-09-01 15:20 |
字幕組雙語原文：Transformers 長程上下文綜述
英語原文：A Survey of Long-Term Context in Transformers
翻譯：雷鋒字幕組（wiige）
多頭自注意力開銷很大，這已經(jīng)不是什么秘密了——序列長度的O(n2)復雜度意味著不太可能讓經(jīng)典的Transformer模型處理超長文本序列。在過去的兩年里，NLP社區(qū)已經(jīng)開發(fā)出了名副其實的縫合怪方法來應付這種復雜性，但這篇文章將重點介紹在大規(guī)模下更有前景的方法。
稀疏Transformers
自適應Transformers
Transformer-XL
壓縮Transformers
Reformer
尋覓式 Transformer
多頭注意力對序列長度的伸縮性很差，原因有二: 首先是計算注意力矩陣所需的FLOPs按序列長度的平方暴增，導致對單個序列的自注意力操作的計算復雜度為O(hdn2)，其中h是注意力頭數(shù)，d是key和query的維度，n是文本序列的長度。
第二點是，點積自注意力操作的空間復雜度也按序列長度的平方增大。 計算注意力矩陣的空間復雜度為O(hdn+hn2)——第一項是存儲key和query所需的內存，第二項是每個頭得出的注意力標量值。
讓我們向BERT-Base代入具體的數(shù)字來了解哪些項的開銷占大頭。 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時，每個樣本僅存儲注意力矩陣就已經(jīng)消耗了約2.4GB的內存。這意味著訓練時的批處理尺寸(batch size)更小，并行性更差，進一步阻礙了模型利用長距離上下文的能力。
"Generating Long Sequences with Sparse Transformers" 由 Rewon Child, Scott Gray, Alec Radford 和 Ilya Sutskever 通過離散化(factorization)方法解決了自注意力的時間和空間復雜度中的O(n2)項問題。
在一個典型的自注意力操作中，輸入序列中的每一項都會關注到輸入序列中的所有其他項，從而形成如下的注意力模式:
典型的自回歸設置下的自注意力連接模式。深藍色方塊代表 "查詢(query)"，淺藍色方塊代表 "鍵(key)"
經(jīng)典自注意力帶來的益處是，其高連通性使得信息在詞符之間很容易傳遞——只需要一層注意力就可以聚合任意兩個詞符的信息。 但如果我們放寬這一約束，保證兩層內信息可以在任意兩個詞符之間傳遞即可的話，就可以大幅降低按序列長度增長的復雜度。稀疏Transformer通過利用固定注意力模式的自定義核來實現(xiàn)這一目標。
固定的稀疏Transformer。深藍色方塊代表查詢，中淺藍色方塊代表奇數(shù)層注意力鍵索引，最淺藍色的方塊代表偶數(shù)層關注的關鍵索引。
一半的頭只關注短的局部上下文，而另一半則預分配均勻分布在整個序列中的索引。
通過這些聚合索引來傳遞信息，網(wǎng)絡就能夠從較遠的詞符獲取信息并利用長距離上下文，同時將時間和空間復雜度降低到O(n√n)。最重要的一點是，只需要兩層隱藏層，任意詞符就都可以整合來自其他詞符的信息。
重要的是，離散化注意力結構似乎不會對語言模型的性能產(chǎn)生負面影響，它的每字符比特數(shù)比enwiki8上的密集注意力(令人驚訝地)更高，且可以在高達12228個詞符的上下文上產(chǎn)生有效的注意力。
可以看到，稀疏Transformer的注意結構之所以有效，部分原因是因為這種注意力模式與真正習得的密集注意力并沒有什么不同。 Kevin Clark, Urvashi Khandelwal, Omer Levy, 和 Christopher D. Manning在他們的 "What Does BERT Look At? An Analysis of BERT’s Attention" 一文中探究了密集注意力所習得的模式，試圖弄明白注意力在Transformer模型中承載了什么功能。他們發(fā)現(xiàn)注意力頭傾向于關注緊接在前的詞符（類似于稀疏注意力中的局部注意力模式），以及如[SEP]和句號的特定詞符。所以，也許稀疏Transformer的注意力模式中包含的歸納偏差是有用而非有害的。
例子: BERT習得的注意力模式
如果想在自己的項目中采用固定注意力核，可以查看OpenAI的blockparse庫以及作者發(fā)布的配套示例。
Sainbayar Sukhbaatar, Edouard Grave, Piotr Bojanowski, 和 Armand Joulin 在他們的文章 "Adaptive Attention Span in Transformers" 中對復雜性問題采取了不同的方法. 他們與 "What Does Bert Look At?" 的作者得到了同樣的觀察，并指出雖然密集注意力允許每個頭關注完整的上下文，但許多注意力頭只關注局部上下文，而剩下的頭才會關注整個文本序列。 他們建議使用某種自注意力來利用這一現(xiàn)象，以此讓模型可以選擇其上下文尺寸。
自適應窗口Transformer通過對序列進行遮罩，使每頭習得的上下文之外的詞符貢獻迅速消失。 遮罩(M)與softmax操作的對數(shù)相乘，使某些詞符對當前隱狀態(tài)x的貢獻歸零，其中超參數(shù)R控制最小窗口(span)大小。
"Adaptive Attention Span in Transformers"所用的軟遮罩函數(shù)。圖片來自論文所附的Facebook AI博文
為鼓勵模型只在有益的地方使用額外上下文，對習得的z值施加了?1懲罰。
在這些限制條件下，大部分的頭都會選擇關注<100個字符的上下文，只有少數(shù)的頭（主要是在網(wǎng)絡的后幾層）會選擇頂著?1?1的懲罰來關注>1000個字符的上下文。
除了巧妙的緩存機制，這種對長上下文的懲罰使得自適應窗口高達8k個字符，同時仍然保持模型的整體計算成本可控。此外，它在基線上的性能仍然很高——在enwiki8上達到了0.98比特/字符，text8數(shù)據(jù)集上達到了1.07比特/字符。
然而，可變窗口大小在便于并行方面并不理想，我們通常希望密集的、尺寸規(guī)整的矩陣來實現(xiàn)最佳性能。雖然這種方法可以大幅減少預測時前向計算所需的flops數(shù)量，但作者只提供了模糊的性能估計，并表示自適應跨度的實現(xiàn)能夠以2048個上下文詞符的固定上下文大小模型的速度最高處理8192個詞符 。
Facebook AI Research還開源了他們的工作——代碼和預訓練的模型可以在github.com/facebookresearch/adaptive-spans上獲得。
比起讓密集注意力操作的成本更低廉，Zihang Dai, Zhilin Yang, Yiming Yang, Jaime Carbonell, Quoc V. Le和 Ruslan Salakhutdinov 則從RNNs中獲得靈感，并在Transformer中引入了除自關注機制外的循環(huán)機制。 他們的文章"Transformer-XL: Attentive Language Models Beyond a Fixed-Length Context"，引入了兩個新穎的概念——第一個是將之前 "段(segments)"的隱狀態(tài)作為輸入反饋給當前段層的組件，另一個是促進這一策略的相對位置編碼方案。
使用具有固定上下文大小的標準變換器，處理長輸入時需要將輸入分割成塊（或稱段）分別處理。然而，這種方法有一個局限，即前一段的信息與當前段的詞符無關。 這種段獨立性在某種程度上是有益的，因為它讓我們能分批地有效處理每一段，但如果目標保持是長序列的一致性，這就會成為一個很大的限制。
帶密集注意力的自回歸Transformer的詞符注意力結構來自Transformer-XL。
Transformer-XL通過強制進行分段串聯(lián)來克服這一限制。 在第一段之后，后續(xù)段中的詞符將始終保持512個詞符的上下文尺寸，之前段的激活作為上下文傳遞給后續(xù)段的注意力操作。 這意味著來自NN上下文大小*LL層外的信息可以傳播到一個給定的令牌。 假設上下文尺寸為640，模型有16層，理論上Transformer-XL可以納入10,240個詞符的信息。
Transformer-XL原文中的詞符注意力模式。
為了避免存儲所有段的激活，作者阻止了梯度流向之前的段。
Transformer-XL還引入了一種新穎的位置編碼方案，稱為 "相對位置編碼"。 不是簡單地將內容和絕對位置嵌入的和輸入網(wǎng)絡，而是將每層的注意力操作分解為基于內容的部分和基于相對位置的部分——如果一個段中的第512個詞符要關注第511個詞符，它的相對位置嵌入就是-1。
為了使相對位置編碼的使用變得簡單，他們將從鍵和查詢中產(chǎn)生注意力權重的操作進行了分解。對于一個典型的密集注意力操作，pre-softmax注意力權重可以如下分解:
上式中，Exi是詞符在位置i的基于內容的嵌入，Uj是詞符j的位置嵌入。
(a) 將查詢的內容與鍵的內容相關聯(lián)
(b) 將查詢的內容與鍵的位置相關聯(lián)
(c) 將查詢的位置與鍵的內容相關聯(lián)
(d) 將查詢的位置與鍵的位置相關聯(lián)
當使用相對位置嵌入時，作者將式子改成如下形式:
在b)和d)中，Uj已經(jīng)被它的相對位置嵌入Ri-j替換.
對于包含查詢位置的項，我們將矩陣UiUi替換為兩個新學習參數(shù)uu和vv。這些向量可以理解為兩個不依賴于查詢具體細節(jié)的偏置——cc鼓勵多關注一些項，dd鼓勵多關注相對位置。 作出這種替換是因為對自身的查詢其相對位置保持不變。
要想讓Transformer-XL模型利用這樣的長程上下文，每一層至少有一個頭要利用其注意力跨度的全部上下文。 平均注意力權重圖顯示，每一層都有頭主要關注先前的位置。
Transformer-XL論文中的平均注意力權重圖
此外Transformer-XL論文還測量了有效上下文長度對困惑度的影響，并發(fā)現(xiàn)增大上下文長度至~900個詞符時會產(chǎn)生更好的困惑度得分——這進一步證明了循環(huán)機制在實踐中的有效性，而非僅僅停留于理論。
Kimi Young的github提供了源代碼，HuggingFace的實現(xiàn)可以幫助你搭建自己的Transformer-XL項目。
列表上的下一個模型是Transfromer，它建立在Transformer-XL架構的基礎上，但通過壓縮損失擴展了該方法，以便納入更長的文本序列。在來自DeepMind的Jack W. Rae、Anna Potapenko、Siddhant M. Jayakumar和Timothy P. Lillicrap的文章 "Compressive Transformers for Long-Range Sequence Modelling "中，詳細介紹了一種能夠關注與全書長度的序列模型架構。
" Compressive Transformers For Long Range Sequence Modeling "中的壓縮記憶(Compressed Memory)示意圖
按照Transformer-XL的做法，當前序列可以關注已存儲的前段的激活。此外，在同樣的多頭注意操作中，當前段中的詞符可以關注存儲在 "壓縮記憶 "中的第二組狀態(tài)。
在每一步中，最早的壓縮記憶被丟棄，壓縮記憶向后移動一個索引。 然后，來自正常記憶段中最早的nn個狀態(tài)進行壓縮，然后轉移到壓縮記憶新空出來的槽中。
DeepMind博客上的一個gif很好地說明了這個過程:
將過去的記憶逐漸壓縮成壓縮記憶——圖片由 DeepMind 博客提供
DeepMind團隊嘗試了多種壓縮操作（包括平均池化、最大池化和學習卷積等等），但最終決定訓練一個二級網(wǎng)絡來重建縮記憶中基于內容的注意力矩陣。
換句話說，他們學習了一個函數(shù)fc，通過最小化壓縮記憶的注意力（C-1=fc(Mold)）和正被壓縮的正常記憶狀態(tài)的注意力之差，將n個最早的記憶狀態(tài)壓縮為一個壓縮記憶狀態(tài)。:
他們沒有將這種壓縮操作與主語言模型聯(lián)合訓練，而是選擇在一個單獨的優(yōu)化循環(huán)中更新壓縮網(wǎng)絡。因為讓注意力狀態(tài)容易被壓縮，對降低語言模型的損失會起到反作用。
在他們的實驗中，壓縮記憶的尺寸為512，正常記憶尺寸為512，窗口尺寸為512，壓縮率為2——意味著在壓縮步驟中，最早的2個記憶狀態(tài)將被壓縮為1個狀態(tài)。在這個設置下，他們在WikiText-103上實現(xiàn)了17.1的SOTA困惑度。
由于利用較長序列長度的收益通常是長尾的，他們特別研究了不同詞頻下的困惑度，發(fā)現(xiàn)在最稀有的詞符上的收益尤其顯著:
按詞頻分類的困惑度。圖片由DeepMind博客提供
雖然他們還沒有公開源碼，但DeepMind已經(jīng)開放了PG-19的源碼，這是他們在研究壓縮Transformer時開發(fā)的數(shù)據(jù)集。PG-19是Project Gutenberg的衍生品，旨在進一步研究長期注意力。
接下來是Nikita Kitaev, ?ukasz Kaiser, Anselm Levskaya 的文章"Reformer: The Efficient Transformer". Reformer在增大序列長度方面采取了不同的策略——他們沒有引入循環(huán)機制或壓縮記憶，而是選擇通過使用局部敏感哈希(LSH)技術來縮小每個詞符的注意力范圍.
局部敏感哈希是一系列將高維向量映射到一組離散值（桶/簇）的方法。它常常用在近似最近鄰搜索中。
Reformer的作者對注意力操作的鍵和查詢使用了同一個投影，并使用基于隨機旋轉的局部敏感哈希方法將共享的鍵/查詢分到最多幾百個詞符的不同桶中。下圖是該方法的一個示例:
Reformer論文中的角度化LSH。該圖展示了一個有3次旋轉的4桶哈希設置。下方圖中的三個向量已經(jīng)映射到同一個哈希桶中，因為其的原始向量就很接近，而上方的第一個圖和最后一個展示了向量被映射到不同的哈希桶中。
它們計算每個桶內的注意力矩陣，然后取對應值的加權和。 由于只關注一個給定的桶內的元素，所以如果桶的大小選擇得當，這可以將注意力操作的整體復雜度從O(n2)降低到O(nlogn)。由于分桶是隨機的，而且旋轉映射向量也是隨機的，所以他們計算了多個哈希值，以確保具有相似的共享鍵-查詢嵌入的詞符最終以高概率落入同一個桶中。
他們還采用了 " The Reversible Residual Network: Backpropagation Without Storing Activations "中介紹的技術來控制訓練時的空間消耗?？赡鏆埐顚邮褂们擅畹募軜嫿Y構，允許從層輸出中輕松重構層輸入，并在網(wǎng)絡深度不變的情況下以空間復雜度降低了計算量。
通過局部敏感的哈希技巧來降低計算成本，以及可逆殘差來降低內存消耗，Reformer架構能夠在單個加速器上處理長達64000個詞符的序列。
雖然在enwiki-8上得出的1.05比特/字符的得分落后于我們在這篇文章中所研究的其他一些模型，但Reformer是一個令人耳目一新的獨特機制，我很期待看到這種方法如何繼續(xù)擴大規(guī)模。
如果你有興趣進一步探索Reformer架構，可以看看我最近關于這個主題的博文 "A Deep Dive into the Reformer"。 在google/jax Github倉庫中，有一個Reformer的開源實現(xiàn)示例。 Phil Wang也維護了一個PyTorch版本。
我原打算在這里就結束本文，但在Aran Komatsauzaki的建議下，我加入了最后一篇論文。Aurko Roy、Mohammad Taghi Saffar、David Grangier和Ashish Vaswani的第二篇ICLR 2020的論文 "Efficient Content-Based Sparse Attention with Routing Transformers "與前述的Reformer有一些異曲同工之處。他們將問題框定為路徑問題，旨在學習選擇稀疏詞符聚類簇Si和內容x的函數(shù)。
作者的在下圖中說明了他們的方法。 他們不是只關注局部元素或第n個元素來增加稀疏性，而是學習了（圖c中用彩色表示的）需要關注的聚類。重要的是這些簇是每個鍵和查詢的內容的函數(shù)，而不僅僅與絕對或相對位置相關。
Routing Attention與Local和Strided Attention的比較，來自Routing Transformer論文
在確保每個鍵和查詢向量都單位化后，他們使用尺寸為（Dk,Dk）的隨機正交權重的共享矩陣來投影鍵和查詢的值，其中Dk是鍵和查詢的隱藏維度。
然后，根據(jù)一組k-means中心將R中的向量分組到k個簇中，每個k-means中心通過每批使用k-means更新來習得，與梯度下降過程無關。
在一個給定的聚類簇Ci內，它們使用標準的加權求和方法計算了一組新的上下文嵌入，其中每個注意力值Ai都是使用標準的點積自注意力計算的。
由于密集注意力中的注意力模式通常由少數(shù)關鍵元素主導，而且聚類分配過程需要將具有高注意力權重的鍵和查詢歸入同一個聚類中，作者認為這個機制能夠有效告知我們X′i進行了高開銷的密集操作。
最后，他們選擇了一些接近√n的聚類簇，這樣基于稀疏內容的注意力機制的整體復雜度就變成了O(n√n)。 為了使整個過程易于并行化和統(tǒng)一矩陣尺寸，作者最接近每個中心點的top-k個項來代替真正的k-means聚類。
除了基于內容的路徑注意力， 路徑注意力Transformer還在尺寸為256的上下文上執(zhí)行局部注意力。
Routing Transformer在計算效率上的提升也使得其在Wikitext-103這個詞級語言建?；鶞蕼y試上的困惑度得到提升，它們的性能明顯超過了之前描述的Transformer-XL模型。
Wikitext-103測試集上的困惑度結果，來自Routing Transformer論文。
在enwiki-8上，Routing Transformer的表現(xiàn)也相當不錯，不過他們的結果略微落后于Adaptive Span Transformer。
enwiki-8的測試集的每個字符比特數(shù)結果。
我本來沒找到Routing Transformer的實現(xiàn)，但Aurko Roy很好心地給我了一份他們源碼的壓縮包，這份源碼是提交給ICLR審閱的版本。
如果你對在Transformer中加入長程上下文的其他方法感興趣，那你還可以讀一下:
Efficient Content-Based Sparse Attention with Routing Transformers
BP-Transformer: Modelling Long-Range Context via Binary Partitioning
雷鋒字幕組是一個由 AI 愛好者組成的翻譯團隊，匯聚五百多位志愿者的力量，分享最新的海外AI資訊，交流關于人工智能技術領域的行業(yè)變革與技術創(chuàng)新的見解。
團隊成員有大數(shù)據(jù)專家、算法工程師、圖像處理工程師、產(chǎn)品經(jīng)理、產(chǎn)品運營、IT咨詢人、在校師生；志愿者們來自IBM、AVL、Adobe、阿里、百度等知名企業(yè)，北大、清華、港大、中科院、南卡羅萊納大學、早稻田大學等海內外高校研究所。
如果，你也是位熱愛分享的AI愛好者。歡迎與雷鋒字幕組一起，學習新知，分享成長。
雷鋒網(wǎng)雷鋒網(wǎng)
雷峰網(wǎng)版權文章，未經(jīng)授權禁止轉載。詳情見轉載須知。
