機器之心報導
機器之心編輯部
Transformer 是近期 NLP 領域裡最熱門的模型之一，但因為算力消耗過大，對於個人研究者來說一直不太友好。近日一篇入選 ICLR 2020 的研究提出了「Reformer」，把跑 Transformer 模型的硬體要求壓縮到了只需一塊 GPU，同時效果不變。
大型的 Transformer 往往可以在許多任務上實現 sota，但訓練這些模型的成本很高，尤其是在序列較長的時候。在 ICLR 的入選論文中，我們發現了一篇由谷歌和伯克利研究者發表的優質論文。文章介紹了兩種提高 Transformer 效率的技術，最終的 Reformer 模型和 Transformer 模型在性能上表現相似，並且在長序列中擁有更高的存儲效率和更快的速度。論文最終獲得了「8，8，6」的高分。
在最開始，文章提出了將點乘注意力（dot-product attention）替換為一個使用局部敏感哈希（locality-sensitive hashing）的點乘注意力，將複雜度從 O(L2 ) 變為 O(L log L)，此處 L 指序列的長度。
此外，研究者使用可逆殘差（reversible residual layers）代替標準殘差（standard residuals），這使得存儲在訓練過程中僅激活一次，而不是 n 次（此處 n 指層數）。最終的 Reformer 模型和 Transformer 模型在性能上表現相同，同時在長序列中擁有更高的存儲效率和更快的速度。
這篇論文在評審過程中收穫了「一致通過」，並被認為將產生重大影響，也經過了幾位外部評審的詳細審查，最終獲得了「8，8，6」的高分。
論文地址：https://openreview.net/forum?id=rkgNKkHtvB代碼：https://github.com/google/trax/blob/master/trax/models/research/reformer.py
引言
Transformer 架構被廣泛用於自然語言處理中，並且在許多任務中實現了 sota。為了獲得這些結果，研究者不得不開始訓練更大的 Transformer 模型。在最大的配置中，參數數量已經超過了 0.5B/層，層數多達 64。
諸如此類的大型 Transformer 模型頻頻出現，到底是客觀上必須要求如此多的資源，還是僅僅是因為處理效率不夠高？
可以參考下面這些數據：
0.5B 的參數佔據了 2GB 的內存，嵌入大小為 1024、批處理大小為 8 的 64K token 的激活要用 64K×1K×8 = 0.5B 浮點數，需要另外 2GB 的內存。
如果說每層的內存佔用只有這麼一些的話，部署 Transformer 會比實際中更容易，但是事情並不是這樣的。以上的估計只包括了每層的內存佔用情況和輸入的激活損失，並沒有考慮 Transformer 上的內存佔用問題：
由於激活需要被存儲並用於反向傳播，有著 N 層的模型的大小比單層大了 N 倍；由於中間的全連接層的深度 d_ff 通常遠大於注意力激活層的深度 d_model，因此需要佔用很大的內存；在長度為 L 的序列上的 attention 的計算和時間複雜度是 O(L2)，所以即使是一個有 64K 字符的序列就會耗盡 GPU 的內存。研究者提出了一種 Reformer 模型來解決剛才說的那些問題：
可逆層（Reversible layer），這個東西最早是 Gomez 等人引入的，在整個模型中啟用單個副本，所以 N factor 就消失了；在前饋層（feed-forward layer）分開激活和分塊處理，消除 d_ff factor，節省前饋層的內存；基於局部敏感哈希（locality-sensitive hashing，LSH）的近似注意力計算，讓注意力層的 O(L2) 因子替代 O(L) 因子，實現在長序列上的操作。局部敏感哈希注意力（LSH Attention）
Transformer 中的多頭注意力層是造成內存佔用大的主要原因，因此研究者從這裡入手解決問題。
首先回顧一下點乘注意力機制，如下所示：
在多頭注意力中，多個注意力層平行計算併疊加。每個注意力層會線性地投影 queries、keys 和 values h 次。
在計算中可以發現，這種注意力機制帶來的內存佔用是很大的。回到公式 1，假設 Q、K、V 都有 [batch size, length, d_model] 這樣的 shape。主要的問題就在於 QK^T，因為它的 shape 是 [batch size, length, length]。如果實驗中序列的長度是 64k，在批大小為 1 的情況下，這就是一個 64K × 64K 的矩陣了，如果是 32 位浮點計算就需要 16GB 的內存。因此，序列越長，Transformer 性能就越受到影響。
如果要減少內存佔用的話，在這裡就需要讓 Q 和 K 保持一致。這是很容易的，只要從同樣的線性層 A 提取即可，並單獨分離一個給 V。QK 共享不會對 Transformer 的性能造成影響，即使對 K 加入額外的正則長度。
對於局部敏感哈希注意力而言，需要 Q=K，以及 V，它們的 shape 都是 [batch size,length,d_model]，而重點關注的是 QK^T，有著 [batch size,length,length] 的 shape。進一步來說，對於每個 q_i，實際需要關注的是它們在 key 的接近值。例如，如果 K 是 64K，對於每個 q_i，只需要考慮一小部分，如 32 個到 64 個最接近的 keys。
這樣一來就需要找到最近鄰的值，這就需要局部敏感哈希（LSH）了，它能夠快速在高維空間中找到最近鄰。一個局部敏感哈希算法可以將每個向量 x 轉換為 hash h(x)，和這個 x 靠近的哈希更有可能有著相同的哈希值，而距離遠的則不會。在這裡，研究者希望最近的向量最可能得到相同的哈希值，或者 hash-bucket 大小相似的更有可能相同。
圖 1：研究中使用的局部敏感哈希算法。這種算法使用隨機旋轉的方法，對投影的點建立分塊，建立的規則依據對給定軸的投影進行比較。在本圖中，兩個點 x、y 由於三次隨機旋轉投影中的兩次都不靠近，所以不太可能有相同的哈希值。而另一個例子中他們投影后都在同一個。
最終，對 attention 進行哈希處理的流程如下：
圖 2：簡化的局部敏感哈希注意力，展示了 hash-bucketing、排序和分塊步驟，並最終實現注意力機制。
不同注意力類型的複雜度對比結果見下表 1：
表 1：Scaled Dot-Product、Memory-Efficient 與 LSH 注意力的內存和複雜度對比。l 表示長度，b 表示批量大小，n_h 表示 head 數量，n_c 表示 LSH 塊數量，n_r 表示哈希重複次數。
怎麼使用這種新型 attention
在一個大型 Transformer 中，通常設置 d_ff = 4K、n_l = 16，所以，如果 n_l = 16，那內存佔用就會達到 16GB。在論文中，研究者首先通過可逆層來解決 n_l 問題，然後展示了如何利用分塊來解決 d_ff 問題。
可逆 Transformer
研究者在 Transformer 上應用了 RevNet 思想，將注意力和前饋層結合在 RevNet 塊內
常規的殘差層執行一個作用於單個輸入並產生單個輸出的
函數，其形式為 y = x + F (x)，可逆層作用於成對的輸入/輸出：
，並遵循以下方程：
在上面的公式中，F 成為注意力層，而 G 成為前饋層。
可逆 Transformer b 不需要在每一層中激活存儲，於是無需使用 nl 項。
分塊
比較厚的層仍然會佔用大量內存。前饋層的計算在序列中是完全獨立的，所以可以分塊：
一般這一層會通過執行所有位置的操作來進行批處理，但是每次進行一塊的處理方法會減少內存佔用，反向計算（reverse computation）和反向過程（backward pass）也會被分塊。
實驗
在實驗部分，研究者逐個分析上述每種技術，以確定哪種組合會對性能產生影響。首先，他們證明了可逆的層和共享的查詢-鍵空間對性能沒有影響。接下來，他們開始分析哈希注意力以及整個 Reformer 模型。
研究者在 imagenet64 和 enwik8-64K 任務上進行了實驗，其中，後者是 enwik8 的一個變體，被分為 2 個 16 = 64K token 的子序列。研究者使用 3 層的模型進行控制變量實驗，以便與常規 transformer 進行比較。所有的實驗都有 d_model = 1024、d_ff = 4096、n_heads = 8。這些模型在每塊 GPU 上進行批大小為一個序列的訓練，總共有 8 塊 GPU 並行。
研究者首先考慮了共享 QK 注意力對於常規 Transformer 模型的影響。共享 QK 注意力使得
，並且防止 token 注意到自身（除非沒有其他可用的語境）。在下圖 3 的左半部分，研究者繪製了常規和共享 QK 注意力的困惑度曲線。
共享的查詢-鍵空間並不比常規注意力表現差；實際上，對於 enwik8 來說，前者甚至訓練得稍快一些。換句話說，採用共享 QK 注意力並不會造成準確率的損失。
圖 3：在 enwik8 和 imagenet64 訓練中，共享查詢-鍵空間（左）和可逆性（右）對於性能的影響。
可逆層又會產生什麼影響呢？如上圖 3 右所示，研究者對比了常規 Transformer 和文中提到的可逆 Transformer。它們擁有相同的參數量，學習曲線也幾乎一樣。結果表明，可逆 Transformer 節省內存的同時也不以犧牲準確率為代價。
如下圖 4 所證，LSH 注意力是全注意力的近似值，它的準確率隨著哈希值的增加而提升。當哈希值為 8 時，LSH 注意力幾乎等同於全注意力。一般而言，模型的計算開銷隨哈希值的增加而增大，所以研究者可以根據自身計算預算調整哈希值。
圖 4：在 imagenet64 上 LSH 注意力性能基於哈希值的變化曲線圖。
如下表 2 所示，研究者可以在評估的時候增加哈希值，從而使得結果更加準確。
如下圖 5 右所示，研究者描述出不同注意力類型的速度和序列長度的變化曲線圖，同時保持 token 總數量不變。結果顯示，常規注意力隨著序列長度的增加而速度減緩，而 LSH 注意力速度保持平穩。
圖 5 左：在 enwik8 上 LSH 注意力隨層數增加的性能變化曲線；圖 5 右：全注意力和 LSH 注意力的評估速度呈現出不同的曲線變化。
此外，為了驗證 Reformer 的確可以在單核心上擬合大模型，並能夠在長序列上快速訓練，研究者在 enwik8 和 imagenet64 上訓練了多達 20 層的大型 Reformer。如上圖 5 所示，這些模型擬合內存和訓練。