這六大方法，如何讓 Transformer 輕松應(yīng)對(duì)高難度長(zhǎng)文本序列？

本文作者： MrBear

2020-06-29 16:01

導(dǎo)語(yǔ)：面對(duì)長(zhǎng)文本，我們需要效率更高的 Transformer！

眾所周知，多頭注意力機(jī)制 (Multi-Head Self-Attention) 的計(jì)算開銷很大。在處理長(zhǎng)度為 n 的序列時(shí)，其 O(n2) 的時(shí)間復(fù)雜度會(huì)使得原始的 Transformer 模型難以處理長(zhǎng)文本序列。在過去的兩年里，已經(jīng)出現(xiàn)了多種有效的方法來(lái)應(yīng)對(duì)多頭注意力機(jī)制的復(fù)雜度問題，本文將重點(diǎn)討論在模型規(guī)模方面很有發(fā)展前景的方法。

一、密集型多頭注意力的時(shí)空復(fù)雜度

多頭自注意力機(jī)制擴(kuò)展到長(zhǎng)文本序列的能力很差，原因有二：

第一，計(jì)算注意力矩陣所要求的每秒浮點(diǎn)運(yùn)算數(shù)（FLOPs）與序列長(zhǎng)度的平方成正比，導(dǎo)致單個(gè)序列上的自注意力運(yùn)算的計(jì)算復(fù)雜度為，其中 h 是注意力頭的數(shù)量，d 是鍵向量和查詢向量的維數(shù)，n 是序列的長(zhǎng)度。

另一件糟糕的事是：點(diǎn)積自注意力運(yùn)算的空間復(fù)雜度也與序列長(zhǎng)度的平方成正比。計(jì)算注意力矩陣的空間復(fù)雜度為，其中 hdn 是存儲(chǔ)鍵和查詢所需的內(nèi)存的階，而是指存儲(chǔ)每個(gè)注意力頭產(chǎn)生的標(biāo)量注意力值所需內(nèi)存的階。

我們用 BERT-Base 中的一些具體數(shù)字來(lái)解釋一下復(fù)雜度到底有多高。BERT-Base 序列輸入的最大長(zhǎng)度為 512，768 個(gè)的隱藏維度和 12 個(gè)注意力頭，這意味著每個(gè)注意力頭有 64 維（768/12）。在這種設(shè)定下，需要 393,216 個(gè)浮點(diǎn)數(shù)（約為 1.5MB）（12 個(gè)注意力頭* 64 注意力頭的維度* 512 序列長(zhǎng)度）來(lái)存儲(chǔ)鍵和值，而存儲(chǔ)所有注意力頭得到的標(biāo)量注意力值所需的內(nèi)存將達(dá)到 3,145,728 個(gè)浮點(diǎn)數(shù)（12 * 512 * 512）或約 12MB 的設(shè)備內(nèi)存，這里所需的內(nèi)存幾乎是將鍵存儲(chǔ)在長(zhǎng)度為 512 個(gè)詞的上下文時(shí)的 10 倍。

為了進(jìn)行梯度運(yùn)算，必須在訓(xùn)練過程中緩存激活值，除非使用諸如梯度檢查點(diǎn)之類的重計(jì)算策略，因此對(duì)于每個(gè)示例來(lái)說，存儲(chǔ)所有 12 層 BERT base 的注意力矩陣就需要大概 150MB 的內(nèi)存。當(dāng)序列長(zhǎng)度為 1024 時(shí)，內(nèi)存需求則變?yōu)榧s 600MB，而序列長(zhǎng)度為 2048 時(shí)，對(duì)于每個(gè)示例的注意矩陣而言，內(nèi)存需求就已達(dá)到約 2.4GB。這意味著在訓(xùn)練時(shí)可以使用的批處理規(guī)模較小，并行性較差，并會(huì)進(jìn)一步阻礙我們訓(xùn)練該模型處理長(zhǎng)上下文的能力。

二、稀疏 Transformer

Rewon Child，Scott Gray，Alec Radford 和 Ilya Sutskever 在論文 "Generating Long Sequences with Sparse Transformers" 中提到，可以通過因式分解解決多頭自注意力的高時(shí)間和空間復(fù)雜度的問題。（論文地址：https://arxiv.org/abs/1904.10509）

1、分解后的注意力

在典型的自注意力機(jī)制中，其輸入序列中的每一項(xiàng)都要計(jì)算與輸入序列中的其它項(xiàng)之間形成的注意力，從而得到如下所示的注意力模式：

圖 1：以一種自回歸的形式組織的傳統(tǒng)自注意力機(jī)制的連接模式。深藍(lán)色方塊代表「查詢」向量，而淺藍(lán)色方塊代表「鍵」向量。

傳統(tǒng)自我注意力機(jī)制的一大好處是，高度的連通性使詞例 (token) 之間的信息容易流通，僅需一層即可聚合來(lái)自任何兩個(gè)詞例的信息。但是，如果放寬此約束，并確保任意兩個(gè)詞例之間的信息只能在經(jīng)過兩層之后才可以流通，則可以極大地降低長(zhǎng)序列帶來(lái)的復(fù)雜性。稀疏 Transformer 通過編寫利用固定注意力模式的自定義核來(lái)實(shí)現(xiàn)此目標(biāo)。

圖 2：稀疏 Transformer 的固定注意力變體。最深的藍(lán)色方塊代表「查詢」向量，較淺的藍(lán)色方塊代表被奇數(shù)層注意的「鍵」向量索引，最淺的藍(lán)色方塊代表被偶數(shù)層注意的「鍵」向量索引。

一半的注意力頭只關(guān)注較短的局部上下文中的項(xiàng)，而剩余的一半注意力頭關(guān)注的是預(yù)先指定好的在整個(gè)序列中均勻分布的索引。通過根據(jù)這些聚合索引確定信息流動(dòng)路徑，網(wǎng)絡(luò)仍然能夠使得相距較遠(yuǎn)的詞例間傳遞信息，并使用長(zhǎng)距離上下文，同時(shí)將時(shí)間和內(nèi)存復(fù)雜性降低到。重要的是，它只需要兩層就可以讓任意詞例考慮來(lái)自任何其它詞例的信息。

2、實(shí)驗(yàn)結(jié)果

值得注意的是，分解后的注意力結(jié)構(gòu)似乎不會(huì)對(duì)語(yǔ)言建模的性能產(chǎn)生負(fù)面影響，反而驚人地令每個(gè)字符需要的比特?cái)?shù)比密集的注意力機(jī)制（原始 Transformer）在 Enwiki8 語(yǔ)料環(huán)境下要少一些，并且可以在包含多達(dá) 12,228 個(gè)詞例的上下文中得到高效的注意力。

可以想象的到，稀疏 transformer 之所以起作用，部分原因是它學(xué)到的的注意力模式與實(shí)際學(xué)習(xí)的密集注意力模式并沒有什么不同。在 Kevin Clark 等人發(fā)表的文章"What Does BERT Look At? An Analysis of BERT’s Attention"（論文地址：https://arxiv.org/abs/1906.04341）中，作者探索了密集注意力學(xué)習(xí)到的模式，想要找到注意力在transformer 中起什么作用。他們發(fā)現(xiàn)關(guān)注緊密相連的前面的詞例（類似于稀疏注意力機(jī)制中的局部注意力模式）以及關(guān)注特定聚合詞例（如 [SEP] 和句號(hào)）的注意力頭有重要作用。因此，可能編碼在稀疏 transformer 注意模式中的歸納偏置是有積極作用的。

圖 3：BERT 學(xué)到的注意力模式示意圖，線的深度表明了注意力權(quán)重的強(qiáng)度（其中有一些注意力權(quán)重太小，以至于線變成了透明的）

要在自己的項(xiàng)目中使用固定注意力kernel，請(qǐng)查看OpenAI的blocksparse庫(kù)和作者作為開源發(fā)布的附帶示例。

三、自適應(yīng)跨度 Transformer

Sainbayar Sukhbaatar 等人在論文"Adaptive Attention Span in Transformers"（論文地址：https://arxiv.org/abs/1905.07799）中選擇了另一種辦法來(lái)解決這個(gè)問題。和前一篇"What Does Bert Look At?"（論文地址：https://arxiv.org/abs/1906.04341）的類似，他們也觀察了 Transformer 中起作用的組件。

值得注意的是，盡管密集注意力機(jī)制讓每個(gè)注意力頭可以關(guān)注整個(gè)上下文，但很多注意力頭只考慮局部上下文，而其它注意力頭則考慮了整個(gè)可用序列。他們建議通過使用一種自注意力的變體來(lái)利用這種觀察結(jié)果，這種變體允許模型選擇它的上下文大小。

1、自適應(yīng)掩模

Adaptive Span Transformer 通過實(shí)現(xiàn)這一點(diǎn)的方式是：通過對(duì)序列進(jìn)行掩模運(yùn)算，使得學(xué)習(xí)到的每個(gè)注意力頭的上下文之外的詞例的貢獻(xiàn)迅速降為零。mask（M）與 softmax 操作的分對(duì)數(shù)相乘，從而將某些詞例對(duì)當(dāng)前隱藏狀態(tài) x 的貢獻(xiàn)歸零，超參數(shù) R 控制最小跨度的大小。

圖 4："Adaptive Attention Span in Transformers" 采用的 soft-masking 函數(shù)

他們應(yīng)用了一個(gè)學(xué)習(xí)到的z值的懲罰項(xiàng)，以鼓勵(lì)模型僅在有益的情況下使用額外的上下文。

2、對(duì)注意力的思考和實(shí)驗(yàn)結(jié)果

由于這些限制，大多數(shù)注意力頭只關(guān)注少于100 字符的上下文，而只有少數(shù)（主要是在網(wǎng)絡(luò)的后幾層）選擇添加一個(gè) 懲罰項(xiàng)，從而學(xué)習(xí)大于 1000 個(gè)字符的上下文。

圖 5：注意力跨度隨層數(shù)變化的示意圖

除了采用巧妙的緩存策略，這種對(duì)長(zhǎng)距離上下文的懲罰項(xiàng)使得跨度自適應(yīng) Transformer 可以使用使用高達(dá) 8k 個(gè)字符的注意力跨度，同時(shí)仍然保持模型的總體計(jì)算開銷較低。此外，模型在基準(zhǔn)對(duì)比測(cè)試中的性能仍然很高，在 Enwiki8 數(shù)據(jù)集上達(dá)到每字符占用 0.98 比特，在 text8 數(shù)據(jù)集上則達(dá)到每字符 1.07 比特。

然而，可變跨度大小的注意力在并行性方面并不理想，因?yàn)槲覀兺ǔＯＭ褂妹芗⒋笮∫恢碌木仃嚕瑥亩@得最佳性能。雖然這種方法可以顯著減少在預(yù)測(cè)時(shí)計(jì)算前向傳播所需的每秒浮點(diǎn)運(yùn)算次數(shù)，但作者只給出了模糊的性能估計(jì)，聲明自適應(yīng)跨度的實(shí)現(xiàn)使我們可以以與 2,048 個(gè)上下文詞例的上下文大小固定的模型相近的速率處理長(zhǎng)達(dá) 8,192 個(gè)詞例的上下文。

Facebook AI 研究院也開源了他們的成果——代碼和預(yù)訓(xùn)練模型請(qǐng)參閱以下鏈接：github.com/Facebook Research/adaptive-spans

四、Transformer-XL

相比于降低密集型注意力操作復(fù)雜度的方式，Zihang Dai 等人受 RNN 的啟發(fā)，在 transformer 的自注意機(jī)制之外引入了一種循環(huán)機(jī)制。他們的論文 "Transformer-XL: Attentive Language Models Beyond a Fixed-Length Context" （論文地址：https://arxiv.org/abs/1901.02860）介紹了兩個(gè)新概念：（1）一種新的組件，它將先前「段落」的隱藏狀態(tài)作為循環(huán)段落層的輸入；（2）使這種策略容易實(shí)現(xiàn)的相對(duì)位置編碼方案。

1、段落循環(huán)

標(biāo)準(zhǔn) transformer 的上下文大小是固定的，要想處理長(zhǎng)的輸入需要將輸入分成塊（段落），并分別處理每個(gè)塊（段落）。

然而，這種方法存在一個(gè)限制：前面段落的信息不能流向當(dāng)前的詞例。這種段與段之間的獨(dú)立性在某種程度上有益于高效地對(duì)段落進(jìn)行批處理，但從長(zhǎng)距離一致性的角度出發(fā)，這又變成了一個(gè)主要限制因素。

圖 6：在 Transformer XL 的論文中，在一個(gè)自回歸 transformer 中的詞例注意力結(jié)構(gòu)示意圖

Transformer XL 通過強(qiáng)制地使段落被串行處理來(lái)解決這一限制。處理完成第一個(gè)段之后，先前段的激活值將作為上下文傳遞給后續(xù)段的注意力，因此始終有 512 個(gè)緊鄰的字符的上下文被記錄。這意味著跨度為 N 個(gè)上下文大小 * L 層的信息可以傳播到給定的詞例。假設(shè)上下文大小為 640，并且模型有 16 層，那么 Transformer XL 理論上至多可以考慮 10,240 個(gè)字符的信號(hào)。

圖7：Transformer-XL 中的詞例注意力模式，其中端的長(zhǎng)度為 4

為了避免存儲(chǔ)先前所有段的激活值，作者阻止了梯度流向前面的段。

2、引入相對(duì)位置

Transformer-XL 還引入了一種新的位置編碼方案，即“相對(duì)位置編碼”。這種方案不再簡(jiǎn)單地將內(nèi)容和絕對(duì)位置嵌入的總和作為網(wǎng)絡(luò)輸入，而是對(duì)每層的注意力操進(jìn)行分解，其中一部分基于內(nèi)容執(zhí)行注意力，另一部分基于相對(duì)位置執(zhí)行注意力。因此，塊中的第 512 個(gè)詞例會(huì)注意到第 511 個(gè)詞例，在這里會(huì)采用相對(duì)位置 -1 相應(yīng)的嵌入。

為了使相對(duì)位置編碼易于處理，他們將生成注意力權(quán)重的操作和鍵和生成查詢向量、鍵向量的操作分離。對(duì)于典型的密集注意力機(jī)制，進(jìn)行 softmax 計(jì)算之前的注意力可以被分解如下：

在上面的公式中，表示基于內(nèi)容的注意力在位置 i 處的詞例的嵌入，是詞例 j 的位置編碼嵌入。其中，每一項(xiàng)的含義如下：

a）基于鍵內(nèi)容的查詢內(nèi)容“尋址”

b）基于鍵位置的查詢內(nèi)容位置偏置

c）根據(jù)鍵的內(nèi)容的查詢位置偏置

d）根據(jù)鍵的位置的查詢位置偏置

使用相對(duì)位置嵌入時(shí)，公式被改寫為：

對(duì)于包含在查詢的位置中的項(xiàng)（c 和 d），我們使用兩個(gè)新的可學(xué)習(xí)參數(shù) u 和 v 替代了矩陣。這些向量可以被理解為兩個(gè)偏置，它們并不依賴于查詢的特征。c 促使模型對(duì)某些項(xiàng)的注意力程度比對(duì)其它項(xiàng)更高，而 d 促使模型對(duì)某些相對(duì)位置的注意力程度比對(duì)其它位置更高。這種替代的方式受到了下面這一事實(shí)的啟發(fā)：查詢相對(duì)于自己的相對(duì)位置始終保持不變。

3、對(duì)注意力的思考和實(shí)驗(yàn)結(jié)果

為了使 Transformer-XL 模型能夠使用長(zhǎng)距離的上下文，每層中至少有一個(gè)注意力頭必須利用其注意力跨度內(nèi)的全部上下文。平均注意權(quán)重圖顯示，每一層都有一些注意力頭，注意到先前許多的位置。

圖 8：在前面 640 個(gè)詞例上的平均注意力，其中每行對(duì)應(yīng)于一個(gè)注意力頭，每列對(duì)應(yīng)于一個(gè)相對(duì)位置。圖中共有 160 個(gè)注意力頭，每十個(gè)注意力頭同屬于同一層。顏色越深代表注意力值越大。

另外，Transformer-XL 的原始論文測(cè)量了有效上下文長(zhǎng)度對(duì)困惑度（perplexity，交叉熵的指數(shù)形式）的影響。作者發(fā)現(xiàn)，增加上下文長(zhǎng)度（上下文長(zhǎng)度高達(dá)九百個(gè)詞例）會(huì)得到更好的困惑度分?jǐn)?shù)（預(yù)測(cè)樣本更準(zhǔn)確），這進(jìn)一步證明了循環(huán)機(jī)制不僅理論上可行，而且實(shí)際上也十分有效。

如果想要開始將 Transformer-XL用于自己的項(xiàng)目，源代碼請(qǐng)參見 Kimi Young 的 github：https://github.com/kimiyoung/transformer-xl，或查看HuggingFace 的實(shí)現(xiàn)（https://huggingface.co/transformers/model_doc/transformerxl.html）。

五、壓縮 Transformer

接下來(lái)我們將介紹壓縮 Transformer（Compressive Transformer），它是建立在 Transformer-XL 的架構(gòu)之上的。壓縮 Transformer 使用了一個(gè)壓縮損失對(duì) Transformer-XL 進(jìn)行了拓展，考慮了更長(zhǎng)的序列長(zhǎng)度。在論文“Compressive Transformers for Long-Range Sequence Modelling”（論文地址：https://arxiv.org/abs/1911.05507）中，來(lái)自 DeepMind 的 Jack 等人詳細(xì)描述了一種模型架構(gòu)，該架構(gòu)能夠處理與整本書一樣長(zhǎng)的序列。

1、壓縮的 Transformer 注意力

圖 9：壓縮 Transformer 保留了過去的激活值的細(xì)粒度內(nèi)存，它們隨后被壓縮到粒度較粗的壓縮內(nèi)存中。上圖中的模型有一個(gè)三層的結(jié)構(gòu)：（1）一個(gè)長(zhǎng)度為 3 的序列（2）一個(gè)長(zhǎng)度為 6 的內(nèi)存（3）長(zhǎng)度為 6 的壓縮內(nèi)存。高亮顯示的內(nèi)存在每層中被一種壓縮函數(shù) f_c 壓縮到單個(gè)壓縮內(nèi)存中，而不是直接被接下來(lái)的序列忽略。在本例中，壓縮比為 c=3。

遵循著 Transformer-XL 的設(shè)定，序列可以注意到一組存儲(chǔ)下來(lái)的之前段的激活值。另外，在相同的多頭注意里操作中，當(dāng)前段中的詞例可以注意到存儲(chǔ)在「壓縮內(nèi)存」中的第二組狀態(tài)。

在每一個(gè)時(shí)間步上，最早的一個(gè)壓縮內(nèi)存會(huì)被丟棄，壓縮內(nèi)存的索引會(huì)回退一位。接著，常規(guī)的內(nèi)存段中最早的 n 個(gè)狀態(tài)會(huì)被壓縮，并且被轉(zhuǎn)移到壓縮內(nèi)存中新開的槽中。

圖 10：將過去存儲(chǔ)的內(nèi)存逐漸壓縮到壓縮內(nèi)存中。

DeepMind 的研究團(tuán)隊(duì)嘗試使用了各種各樣的壓縮操作（包括平均池化、最大池化、學(xué)習(xí)到的卷積操作等對(duì)比基線），但是他們最終決定訓(xùn)練一個(gè)輔助網(wǎng)絡(luò)，該網(wǎng)絡(luò)被用于重建基于內(nèi)容的被壓縮的內(nèi)存的注意力矩陣。換而言之，他們學(xué)習(xí)了一種函數(shù) ，該函數(shù)通過最小化壓縮內(nèi)存上的注意力與被壓縮的常規(guī)內(nèi)存中狀態(tài)的注意力之間的差異，將最早的 n 個(gè)內(nèi)存狀態(tài)壓縮到了單個(gè)壓縮后的內(nèi)存狀態(tài)中：

由于讓狀態(tài)容易被壓縮與減小語(yǔ)言模型的損失是相矛盾的，他們選擇在一個(gè)獨(dú)立的優(yōu)化循環(huán)中更新壓縮網(wǎng)絡(luò)，而不是同時(shí)訓(xùn)練這種壓縮操作和主要的語(yǔ)言模型。

2、實(shí)驗(yàn)結(jié)果

在實(shí)驗(yàn)中，他們?cè)O(shè)置壓縮內(nèi)存的大小為 512，正規(guī)內(nèi)存的大小為 512，滑動(dòng)窗口大小為 512，壓縮率為 2（這意味著最早的 2 個(gè)內(nèi)存狀態(tài)會(huì)在壓縮步驟中被壓縮到單個(gè)狀態(tài)中）。在這樣的實(shí)驗(yàn)環(huán)境下，他們?nèi)〉昧四壳白詈玫膶?shí)驗(yàn)結(jié)果——在 WikiText-103 數(shù)據(jù)集上取得了 17.1 的困惑度。

由于通過利用更長(zhǎng)的序列所獲得的收益往往是符合長(zhǎng)尾分布的，他們專門關(guān)注了由詞頻刻畫的離散特征的困惑度，并指出最稀有的詞例帶來(lái)的收益是最顯著的：

圖 11：通過詞頻刻畫的離散特征桶（bucket）的困惑度

盡管他們的源代碼至今還沒有公開，DeepMind 已經(jīng)開源了它們研發(fā)壓縮 Transformer 時(shí)所使用的數(shù)據(jù)集 PG-19（https://github.com/deepmind/pg19）。PG-19 是 Gutenberg 計(jì)劃的衍生品，旨在進(jìn)一步研究長(zhǎng)距離注意力。

六、哈希注意力 Reformer

接下來(lái)，我們繼續(xù)關(guān)注由 Nikita Kitaev 等人完成的工作「Reformer：The Efficient Transformer」（論文地址：https://arxiv.org/abs/2001.04451）。Reformer 對(duì)于長(zhǎng)度增長(zhǎng)的序列采用了另一種策略，他們選擇通過局部敏感哈希技術(shù)將每個(gè)詞例的注意力范圍變窄，而不是引入循環(huán)機(jī)制或壓縮內(nèi)存。

局部敏感哈希，是一類將高維向量映射到一組離散值（桶/聚類簇）上的方法。在大多數(shù)情況下，它被用作一種近似的最近鄰搜索方法。

Reformer 的作者針對(duì)注意力操作的鍵向量和查詢向量使用了一種單一投影，他們使用了一種基于隨機(jī)旋轉(zhuǎn)的局部敏感哈希方法，將共享的鍵/查詢聚集到至多幾百個(gè)詞例的桶。哈希方法的示意圖如下：

圖12：角度局部敏感哈希算法示意圖。圖中顯示了三輪哈希，每輪哈希中將空間分成了 4 個(gè)桶。在下面一排示意圖的三輪哈希中中，向量都被映射到了同一個(gè)桶中，因?yàn)樗鼈冊(cè)谳斎胫幸蚕嘟欢谏厦嬉慌诺牡谝淮喂：偷谌喂Ｖ校蛄坑成涞搅瞬煌耐爸小?/p>

它們計(jì)算了每個(gè)桶中的注意力矩陣，并對(duì)相應(yīng)的值進(jìn)行加權(quán)。由于它們只能注意到給定的中的元素，當(dāng)選取的桶大小合適時(shí)，這樣做可以將注意力操作的整體空間復(fù)雜度從降低到。由于這里的分桶（bucketing）處理是隨機(jī)并且基于隨機(jī)旋轉(zhuǎn)的，他們計(jì)算了若干次哈希，從而保證具有相近的「鍵-查詢」嵌入的詞例最終很有可能被分在同一個(gè)桶中。

他們額外地使用了 "The Reversible Residual Network: Backpropagation Without Storing Activations"（論文地址：https://arxiv.org/abs/1707.04585）中提出的技術(shù)將訓(xùn)練時(shí)的內(nèi)存消耗維持在可控范圍內(nèi)。逆殘差層使用了非常巧妙的架構(gòu)，使模型可以容易地根據(jù)層的輸出重構(gòu)層的輸入，并以額外的計(jì)算換取網(wǎng)絡(luò)深度不變的內(nèi)存復(fù)雜性。

盡管 Reformer 論文中給出的在 Enwiki-8 數(shù)據(jù)集上 1.05 比特每字符的性能比本文中給出的其它一些模型稍弱，但是 Reformer 仍然提出了一種令人耳目一新的機(jī)制引入了長(zhǎng)距離上下文，人們期待看到 Reformer 進(jìn)一步的拓展。

如果你有興趣更深入地探索 Reformer 架構(gòu)，請(qǐng)參閱博文「 A Deep Dive into the Reformer」（https://www.pragmatic.ml/reformer-deep-dive/）。Google/jax Github 代碼倉(cāng)庫(kù)中的 Reformer 的開源實(shí)現(xiàn)鏈接如下：https://github.com/google/trax/tree/master/trax/models/reformer。一種 Phil Wang 訓(xùn)練的 PyTorch 版的實(shí)現(xiàn)鏈接如下：https://github.com/lucidrains/reformer-pytorch

七、路由Transformer

Aurko Roy 等人于 ICLR 2020 上發(fā)表的工作 "Efficient Content-Based Sparse Attention with 路由Transformers"（https://openreview.net/profile?email=grangier%40google.com）與前文提到的 Reformer 有一些異曲同工之妙。它們將該問題建模為了一個(gè)路由問題，目的是讓模型學(xué)會(huì)選擇詞例的稀疏聚類（將其作為內(nèi)容 x 的函數(shù)）。

在下面的圖表中，作者描述了他們的方法。他們并不只關(guān)注局部元素或每 n 個(gè)元素來(lái)增加稀疏性，而是學(xué)習(xí)了通過下圖 c 中的顏色表示的需要關(guān)注的聚類簇。重要的是，這些簇是關(guān)于每個(gè)鍵和查詢的內(nèi)容的函數(shù)，而不僅僅與它們的絕對(duì)或相對(duì)位置相關(guān)。

圖 13：這里的圖形顯示了路由 Transformer 與局部注意力機(jī)制和 Child 等人于 2019 年提出的跨步注意力機(jī)制的對(duì)比結(jié)果。每一行代表輸出，每一列代表輸入。對(duì)于局部注意力機(jī)制和跨步注意力機(jī)制來(lái)說，著色的方塊代表每一個(gè)輸出行注意到的元素。對(duì)于路由注意力機(jī)制來(lái)說，不同的顏色代表輸出詞例的聚類中的成員。

1、路由注意力

確保每個(gè)鍵和查詢向量都具有單位大小后，他們使用了一種公共的隨機(jī)權(quán)重矩陣對(duì)鍵和查詢的值進(jìn)行投影，投影的尺寸為，其中D_K 是鍵和查詢的隱藏維度。

隨后，R 中的向量會(huì)根據(jù)一組 k-均值聚類中心 C 被聚類到 k 個(gè)簇中。K-均值聚類中心是通過使用每個(gè)批（batch）中的 K-均值更新學(xué)習(xí)到的，它與梯度下降過程相獨(dú)立。

在給定的聚類簇 C_i 中，他們使用了一種標(biāo)準(zhǔn)的加權(quán)求和方法，計(jì)算了一組新的上下文嵌入，其中每個(gè)注意力值 A_i 是使用標(biāo)準(zhǔn)的點(diǎn)乘自注意力計(jì)算而來(lái)的。

由于密集型注意力機(jī)制中的注意力模式往往由少數(shù)的關(guān)鍵元素，并且分配聚類的過程需要將具有高注意力權(quán)重的鍵和查詢分到同一個(gè)聚類簇中，作者認(rèn)為這樣做保留了關(guān)鍵的信息。這說明應(yīng)用的密集型操作計(jì)算開銷過高。

最終，他們選擇了大約個(gè)聚類，因此他們稀疏的基于內(nèi)容的注意力機(jī)制的整體復(fù)雜度為。為了讓整個(gè)過程易于并行化計(jì)算，并且可以處理統(tǒng)一大小的矩陣，作者使用了最接近每個(gè)聚類中心的前 k 個(gè)項(xiàng)來(lái)代替真正的 k-均值聚類。

除了基于內(nèi)容的路由注意力機(jī)制，路由Transformer 還在一個(gè)大小為 256 的內(nèi)容窗口上執(zhí)行了局部注意力。

2、實(shí)驗(yàn)結(jié)果

路由Transformer 在計(jì)算效率方面取得了提升，從而使模型在 Wikitext-103 上（一種單詞級(jí)別的語(yǔ)言建模對(duì)比基準(zhǔn)）取得了在困惑度指標(biāo)上性能的提升，他們的性能大大超過了上文介紹的 Tranformer-XL 模型。

圖 14：在 Wikitext-103 數(shù)據(jù)集上的語(yǔ)言建模結(jié)果。Local Transformer 指的是由 Vaswani 等人于 2017 年提出的 Transformer，它用到了相對(duì)位置編碼和局部注意力。表中給出了測(cè)試集上的困惑度。

在 Enwiki-8 數(shù)據(jù)集上，路由Transformer 的性能也非常出色（盡管其性能稍微落后于自適應(yīng)跨度 Transformer）。

圖 15：在 Enwiki-8 數(shù)據(jù)集上的語(yǔ)言建模結(jié)果。 Local Transformer 指的是由 Vaswani 等人于 2017 年提出的 Transformer，它用到了相對(duì)位置編碼和局部注意力。表中給出了測(cè)試集上的比特每字節(jié)（bpc）。

路由Transformer 的源代碼鏈接如下：https://storage.googleapis.com/routing_transformers_iclr20/routing_transformers_iclr20.zip 雷鋒網(wǎng)雷鋒網(wǎng)雷鋒網(wǎng)