Transformer 是近期 NLP 領(lǐng)域里最熱門的模型之一，但因為算力消耗過大，對于個人研究者來說一直不太友好。近日一篇入選 ICLR 2020 的研究提出了「Reformer」，把跑 Transformer 模型的硬件要求壓縮到了只需一塊 GPU，同時效果不變。

大型的 Transformer 往往可以在許多任務(wù)上實現(xiàn) sota，但訓練這些模型的成本很高，尤其是在序列較長的時候。在 ICLR 的入選論文中，我們發(fā)現(xiàn)了一篇由谷歌和伯克利研究者發(fā)表的優(yōu)質(zhì)論文。文章介紹了兩種提高 Transformer 效率的技術(shù)，最終的 Reformer 模型和 Transformer 模型在性能上表現(xiàn)相似，并且在長序列中擁有更高的存儲效率和更快的速度。論文最終獲得了「8，8，6」的高分。

在最開始，文章提出了將點乘注意力(dot-product attention)替換為一個使用局部敏感哈希(locality-sensitive hashing)的點乘注意力，將復(fù)雜度從 O(L2 ) 變?yōu)?O(L log L)，此處 L 指序列的長度。

此外，研究者使用可逆殘差(reversible residual layers)代替標準殘差(standard residuals)，這使得存儲在訓練過程中僅激活一次，而不是 n 次(此處 n 指層數(shù))。最終的 Reformer 模型和 Transformer 模型在性能上表現(xiàn)相同，同時在長序列中擁有更高的存儲效率和更快的速度。

這篇論文在評審過程中收獲了「一致通過」，并被認為將產(chǎn)生重大影響，也經(jīng)過了幾位外部評審的詳細審查，最終獲得了「8，8，6」的高分。

論文地址：https://openreview.net/forum?id=rkgNKkHtvB

代碼：https://github.com/google/trax/blob/master/trax/models/research/reformer.py

引言

Transformer 架構(gòu)被廣泛用于自然語言處理中，并且在許多任務(wù)中實現(xiàn)了 sota。為了獲得這些結(jié)果，研究者不得不開始訓練更大的 Transformer 模型。在最大的配置中，參數(shù)數(shù)量已經(jīng)超過了 0.5B/層，層數(shù)多達 64。

諸如此類的大型 Transformer 模型頻頻出現(xiàn)，到底是客觀上必須要求如此多的資源，還是僅僅是因為處理效率不夠高?

可以參考下面這些數(shù)據(jù)：

0.5B 的參數(shù)占據(jù)了 2GB 的內(nèi)存，嵌入大小為 1024、批處理大小為 8 的 64K token 的激活要用 64K×1K×8 = 0.5B 浮點數(shù)，需要另外 2GB 的內(nèi)存。

如果說每層的內(nèi)存占用只有這么一些的話，部署 Transformer 會比實際中更容易，但是事情并不是這樣的。以上的估計只包括了每層的內(nèi)存占用情況和輸入的激活損失，并沒有考慮 Transformer 上的內(nèi)存占用問題：

由于激活需要被存儲并用于反向傳播，有著 N 層的模型的大小比單層大了 N 倍;

由于中間的全連接層的深度 d_ff 通常遠大于注意力激活層的深度 d_model，因此需要占用很大的內(nèi)存;

在長度為 L 的序列上的 attention 的計算和時間復(fù)雜度是 O(L2)，所以即使是一個有 64K 字符的序列就會耗盡 GPU 的內(nèi)存。

研究者提出了一種 Reformer 模型來解決剛才說的那些問題：

可逆層(Reversible layer)，這個東西最早是 Gomez 等人引入的，在整個模型中啟用單個副本，所以 N factor 就消失了;

在前饋層(feed-forward layer)分開激活和分塊處理，消除 d_ff factor，節(jié)省前饋層的內(nèi)存;

基于局部敏感哈希(locality-sensitive hashing，LSH)的近似注意力計算，讓注意力層的 O(L2) 因子替代 O(L) 因子，實現(xiàn)在長序列上的操作。

局部敏感哈希注意力(LSH Attention)

Transformer 中的多頭注意力層是造成內(nèi)存占用大的主要原因，因此研究者從這里入手解決問題。

首先回顧一下點乘注意力機制，如下所示：

在多頭注意力中，多個注意力層平行計算并疊加。每個注意力層會線性地投影 queries、keys 和 values h 次。

在計算中可以發(fā)現(xiàn)，這種注意力機制帶來的內(nèi)存占用是很大的。回到公式 1，假設(shè) Q、K、V 都有 [batch size, length, d_model] 這樣的 shape。主要的問題就在于 QK^T，因為它的 shape 是 [batch size, length, length]。如果實驗中序列的長度是 64k，在批大小為 1 的情況下，這就是一個 64K × 64K 的矩陣了，如果是 32 位浮點計算就需要 16GB 的內(nèi)存。因此，序列越長，Transformer 性能就越受到影響。

如果要減少內(nèi)存占用的話，在這里就需要讓 Q 和 K 保持一致。這是很容易的，只要從同樣的線性層 A 提取即可，并單獨分離一個給 V。QK 共享不會對 Transformer 的性能造成影響，即使對 K 加入額外的正則長度。

對于局部敏感哈希注意力而言，需要 Q=K，以及 V，它們的 shape 都是 [batch size,length,d_model]，而重點關(guān)注的是 QK^T，有著 [batch size,length,length] 的 shape。進一步來說，對于每個 q_i，實際需要關(guān)注的是它們在 key 的接近值。例如，如果 K 是 64K，對于每個 q_i，只需要考慮一小部分，如 32 個到 64 個最接近的 keys。

這樣一來就需要找到最近鄰的值，這就需要局部敏感哈希(LSH)了，它能夠快速在高維空間中找到最近鄰。一個局部敏感哈希算法可以將每個向量 x 轉(zhuǎn)換為 hash h(x)，和這個 x 靠近的哈希更有可能有著相同的哈希值，而距離遠的則不會。在這里，研究者希望最近的向量最可能得到相同的哈希值，或者 hash-bucket 大小相似的更有可能相同。

圖 1：研究中使用的局部敏感哈希算法。這種算法使用隨機旋轉(zhuǎn)的方法，對投影的點建立分塊，建立的規(guī)則依據(jù)對給定軸的投影進行比較。在本圖中，兩個點 x、y 由于三次隨機旋轉(zhuǎn)投影中的兩次都不靠近，所以不太可能有相同的哈希值。而另一個例子中他們投影后都在同一個。

最終，對 attention 進行哈希處理的流程如下：

圖 2：簡化的局部敏感哈希注意力，展示了 hash-bucketing、排序和分塊步驟，并最終實現(xiàn)注意力機制。

不同注意力類型的復(fù)雜度對比結(jié)果見下表 1：

表 1：Scaled Dot-Product、Memory-Efficient 與 LSH 注意力的內(nèi)存和復(fù)雜度對比。l 表示長度，b 表示批量大小，n_h 表示 head 數(shù)量，n_c 表示 LSH 塊數(shù)量，n_r 表示哈希重復(fù)次數(shù)。

怎么使用這種新型 attention

在一個大型 Transformer 中，通常設(shè)置 d_ff = 4K、n_l = 16，所以，如果 n_l = 16，那內(nèi)存占用就會達到 16GB。在論文中，研究者首先通過可逆層來解決 n_l 問題，然后展示了如何利用分塊來解決 d_ff 問題。

可逆 Transformer

研究者在 Transformer 上應(yīng)用了 RevNet 思想，將注意力和前饋層結(jié)合在 RevNet 塊內(nèi)。

常規(guī)的殘差層執(zhí)行一個作用于單個輸入并產(chǎn)生單個輸出的

函數(shù)，其形式為 y = x + F (x)，可逆層作用于成對的輸入/輸出：

，并遵循以下方程：

在上面的公式中，F(xiàn) 成為注意力層，而 G 成為前饋層。

可逆 Transformer b 不需要在每一層中激活存儲，于是無需使用 nl 項。

分塊

比較厚的層仍然會占用大量內(nèi)存。前饋層的計算在序列中是完全獨立的，所以可以分塊：

一般這一層會通過執(zhí)行所有位置的操作來進行批處理，但是每次進行一塊的處理方法會減少內(nèi)存占用，反向計算(reverse computation)和反向過程(backward pass)也會被分塊。

實驗

在實驗部分，研究者逐個分析上述每種技術(shù)，以確定哪種組合會對性能產(chǎn)生影響。首先，他們證明了可逆的層和共享的查詢-鍵空間對性能沒有影響。接下來，他們開始分析哈希注意力以及整個 Reformer 模型。

研究者在 imagenet64 和 enwik8-64K 任務(wù)上進行了實驗，其中，后者是 enwik8 的一個變體，被分為 2 個 16 = 64K token 的子序列。研究者使用 3 層的模型進行控制變量實驗，以便與常規(guī) transformer 進行比較。所有的實驗都有 d_model = 1024、d_ff = 4096、n_heads = 8。這些模型在每塊 GPU 上進行批大小為一個序列的訓練，總共有 8 塊 GPU 并行。

研究者首先考慮了共享 QK 注意力對于常規(guī) Transformer 模型的影響。共享 QK 注意力使得

，并且防止 token 注意到自身(除非沒有其他可用的語境)。在下圖 3 的左半部分，研究者繪制了常規(guī)和共享 QK 注意力的困惑度曲線。

共享的查詢-鍵空間并不比常規(guī)注意力表現(xiàn)差;實際上，對于 enwik8 來說，前者甚至訓練得稍快一些。換句話說，采用共享 QK 注意力并不會造成準確率的損失。

圖 3：在 enwik8 和 imagenet64 訓練中，共享查詢-鍵空間(左)和可逆性(右)對于性能的影響。

可逆層又會產(chǎn)生什么影響呢?如上圖 3 右所示，研究者對比了常規(guī) Transformer 和文中提到的可逆 Transformer。它們擁有相同的參數(shù)量，學習曲線也幾乎一樣。結(jié)果表明，可逆 Transformer 節(jié)省內(nèi)存的同時也不以犧牲準確率為代價。

如下圖 4 所證，LSH 注意力是全注意力的近似值，它的準確率隨著哈希值的增加而提升。當哈希值為 8 時，LSH 注意力幾乎等同于全注意力。一般而言，模型的計算開銷隨哈希值的增加而增大，所以研究者可以根據(jù)自身計算預(yù)算調(diào)整哈希值。

圖 4：在 imagenet64 上 LSH 注意力性能基于哈希值的變化曲線圖。

如下表 2 所示，研究者可以在評估的時候增加哈希值，從而使得結(jié)果更加準確。

如下圖 5 右所示，研究者描述出不同注意力類型的速度和序列長度的變化曲線圖，同時保持 token 總數(shù)量不變。結(jié)果顯示，常規(guī)注意力隨著序列長度的增加而速度減緩，而 LSH 注意力速度保持平穩(wěn)。

圖 5 左：在 enwik8 上 LSH 注意力隨層數(shù)增加的性能變化曲線;圖 5 右：全注意力和 LSH 注意力的評估速度呈現(xiàn)出不同的曲線變化。

此外，為了驗證 Reformer 的確可以在單核心上擬合大模型，并能夠在長序列上快速訓練，研究者在 enwik8 和 imagenet64 上訓練了多達 20 層的大型 Reformer。如上圖 5 所示，這些模型擬合內(nèi)存和訓練。