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