本文介紹的是ICLR2020入選 Oral 論文《Reformer: The Efficient Transformer》，作者來自UC 伯克利和谷歌大腦。

作者 | 李光明

編輯 | 叢 末

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

Transformer是NLP中廣為應用的成熟技術，在許多任務中取得了驕人的成績，尤其是長序列文本上表現突出，但卻極其耗費算力和內存資源，Transformer網絡一層的參數量約0.5B，需要2G的內存空間，單層網絡的Transformer在單台機器上尚可滿足，但鑒於以下考慮，整個Transformer網絡所需要的資源是驚人的：

• 一個N層的網絡需要的內存資源要多於一層所需內存的N倍，因為同時需要存儲激活結果，在反向傳播時使用。

• Transformer前饋全連接神經網絡的寬度（神經單元數）要比attention激活的寬度（可理解為embedding的size）多，需要更多的內存消耗。

• 對一個長度為L的序列，Attention層的複雜度是，這對長序列文本處理是無法接受的。

針對上述問題，這篇文章通過下面幾項技術解決上面提到的幾個問題：

• 使用可逆殘差層取代標準殘差層，在訓練階段只需要存儲一層的激活結果而不是N層（N是網絡層數）（消除了網絡中N的倍數）。

• 分離前饋全連接層的激活部分，分區塊進行處理，消除對內存的消耗。

• 使用局部敏感哈希（Local-Sensitive Hashing, LSH）技術把計算attention部分的複雜度（主要來自於點乘）從降至（其中L代表序列長度）。

1 LSH Attention

Transformer的注意力計算公式（encoder和decoder相同）如下：

不妨假設Q, K, V的 shape 都是[, , ]，計算複雜度主要來源於，其 shape為[, , ]，處理64k長度的序列時，即使 為1，一個64k * 64k的 float矩陣需要16G內存，這通常是不切實際的。

其實，矩陣並不需要全部存儲在內存中，只需要計算每個query 的

進一步分析，其實並不需要計算和全部keys的結果，因為點乘的最大值會在softmax之後得到最大的權重，從而在attention的結果中影響最大，因此只需要關心與相近的key即可（得到較大的點乘值），這樣可以大幅提升效率。LSH是尋找最近鄰的有效手段，它會賦予每個向量一個Hash值，並且相近的向量會大機率獲得相同的Hash值，相距較遠的向量則不同。

上圖是LSH的一個簡單示意圖，在示意圖的上部分，x和y不屬於近鄰，所以在三次隨意旋轉後，有兩次投影都不一樣；而在示意圖的下部分，x和y相距很近，在三次的隨意旋轉後，三次都投影都一樣，這就是LSH的基本原理。

為計算LSH attention，首先重寫的一次attention計算如下：

其中表示需要關注的key集合，是分區函數（可以理解成softmax的分母），為了書寫方便，這裡省去了。為了批處理方便，把公式進一步修改為：

其中是比更大的集合（可以理解成全集），從公式可以看出，如果不屬於的都被置為正無窮（趨近於0)，相當於對應的 key 被 mask 掉了。LSH 的作用就是生成，僅限與在同一個 Hash 桶中的 key 才會參與 attention 計算，滿足：

下圖的分別表示傳統的係數attention矩陣，以及根據Hash分桶排序後的attention矩陣，從圖中可以看出Hash值分桶可能發生不均勻的情況（跨多個桶的批處理是比較困難的），而且一個桶內的queries和keys數量也不一定相等，甚至有可能存在桶中只有queries而沒有keys的情況。

為了解決這個問題，首先通過確保。接下來根據桶號和桶內序列位置對 queries 進行雙層排序，排序後可以得到一個新的排序，排序後的 attention 矩陣，同一個桶的將聚集在對角線附近，如下圖c所示（Q=K，文中的設定，實驗階段會驗證此設定不會造成太大影響），接下來可以定義一種批處理方法，使得排序後的 m 個連續的塊相互關聯，如下圖d所示，關聯規則是後面的塊只需要往前看一個塊，並對設置如下：

其中（認為桶的長度超過平均桶長度兩倍的機率很小）

LSH attention的整個處理流程如下圖所示：

單個Hash函數仍然存在較小機率導致相似的item不能分到相同的桶里，可以通過多輪Hash函數避免這個問題，多輪Hash使用不同的Hash函數，有：

其中

多輪Hash過程可以並行。

2 可逆Transformer

可逆殘差網絡的思路是在反向傳播梯度時，僅依賴模型參數，就可以利用後續網絡層恢復任何給定層的激活結果，從而節省內存空間。標準的殘差層是從一個輸入到一個輸出 的映射，而可逆層的輸入輸出都是成對的: ，計算公式如下：

逆向過程通過減殘差實現：

將可逆殘差網絡的思想應用到 Transformer 中，充分結合 Attention 層和 Feed-Forward 層，上式中的變成 Attention 層，變成 Feed-Forward 層，於是有：

3 分塊

Transformer 前饋網絡的中間向量維度甚至更高維度，依然非常占用內存；其實，序列中不同位置的文本可以完全獨立計算，因此可以拆分成多個區塊並行計算：

前述BP過程的可逆操作也可以分塊並行。

4 實驗

實驗階段會使用 imagenet64 和 enwik8-64K 兩個數據集，依次驗證可逆層和共享 Q-K（Q = K）不會對模型效果造成影響，進而分析驗證 Hash Attention 和整個 Reformer 架構的結果。下圖的困惑度曲線（資訊理論中衡量機率分布預測的指標，此處用來評價模型好壞）展示了共享 Q-K 對比獨立的 Q-K （左）和可逆層（右）對比正常殘差層在兩個數據集上的效果，從圖中可知共享 Q-K 和可逆層的使用並沒有導致模型準度的降低。

下圖展示了LSH Attention的效果，隨著hash函數量的增加，LSH Attention的精度損失隨之減少，當hash函數量達到8個時，基本與full attention的效果持平。

上圖同時說明了hash函數量（可視為一個超參數）會影響Reformer整體效果，同樣也會影響Reformer的性能，下圖展示了hash函數量對性能的影響：

隨著序列長度的不斷增加，full attention模型變得越來越慢，而Reformer基本變化不大，提速效果非常明顯。

ICLR 2020 系列論文解讀

0、ICLR 2020 會議動態報道

疫情嚴重，ICLR2020 將舉辦虛擬會議，非洲首次 AI 國際頂會就此泡湯

疫情影響，ICLR 突然改為線上模式，2020年將成為頂會變革之年嗎？

火爆的圖機器學習，ICLR 2020上有哪些研究趨勢？

1、直播

回放 | 華為諾亞方舟ICLR滿分論文：基於強化學習的因果發現

2、Oral

01. Oral | 一種鏡像生成式機器翻譯模型：MGNMT

02. Oral | 額外高斯先驗目標，緩解負多樣性無知

03. Oral | 引入額外門控運算，LSTM稍做修改，性能便堪比Transformer-XL

04. Oral | 並行蒙卡樹搜索，性能無損，線性加速，勇闖「消消樂」1000關！

05. Oral | 元強化學習迎來一盆冷水：不比元Q學習好多少

06. Oral | 用群卷積建立深度、等變的膠囊網絡

07. Oral | 谷歌推出分布式強化學習框架SEED，性能「完爆」IMPALA，可擴展數千台機器，還很便宜

3、Spotlight

01. Spotlight | 模型參數這麼多，泛化能力為什麼還能這麼強？

02. Spotlight | 公平與精確同樣重要！CMU提出學習公平表征方法，實現算法公平

03. Spotlight | 組合泛化能力太差？用深度學習融合組合求解器試試

04. Spotlight | 加速NAS，僅用0.1秒完成搜索

05. Spotlight | 華盛頓大學：圖像分類中對可實現攻擊的防禦（視頻解讀）

06. Spotlight | 超越傳統，基於圖神經網絡的歸納矩陣補全

4、Poster

01. Poster | 華為諾亞：巧妙思想，NAS與「對抗」結合，速率提高11倍

02. Poster | 拋開卷積，多頭自注意力能夠表達任何卷積操作

03. Poster | NAS 太難了，搜索結果堪比隨機採樣！華為給出 6 條建議

04. Poster | 清華提 NExT 框架，用「神經元執行樹」學習可解釋性

05. Poster | 谷歌最新研究：用「復合散度」量化模型合成泛化能力

06. Poster | 完勝 BERT，谷歌最佳 NLP 預訓練模型開源，單卡訓練僅需 4 天

07. Poster | FSNet：利用卷積核概要進行深度卷積神經網絡的壓縮

08. Poster | "同步平均教學"框架為無監督學習提供更魯棒的偽標籤

09. Poster | 快速神經網絡自適應技術