魚羊 發自 凹非寺

量子位 報導 | 公眾號 QbitAI

Transformer自誕生以來，就在NLP領域刷新一個又一個紀錄，稱作當下最流行的深度學習框架亦不為過。

不過，拿下SOTA並不意味著十全十美。

比如，在長序列訓練上，Transformer就存在計算量巨大、訓練成本高的問題。

其對內存的要求從GB級別到TB級別不等。這意味著，模型只能處理簡短的文本，生成簡短的音樂。

此外，許多大型Transformer模型在經過模型並行訓練之後，無法在單個GPU上進行微調。

現在，谷歌和UC伯克利推出了一個更高效的Transformer模型——Reformer。

在長度為L的序列上，將複雜度從 O(L2)降低到了O(L logL)。

並且，模型訓練後，可以僅使用16GB內存的單個GPU運行。

論文已被ICLR 2020接收。模型也已開源。（地址見文末）

Reformer

Reformer主要引入了以下技術：

• 可逆層。使用可逆殘差層替代標準殘差，這樣以來，模型在訓練過程中僅存儲一次activations，而非N次（N為層數）。
• 在前饋層內拆分activations，並進行分段處理，消除dff因子，節省前饋層內存。
• 將點乘注意力替換為局部敏感哈希（locality sensitive hashing）注意力，將複雜度從 O(L2) 降低到 O(L logL)，使其能對長序列進行操作。

局部敏感哈希（LSH）注意力

Transformer中使用的標準注意力是點乘注意力。

其中，Q矩陣由一組query的注意力函數組成，key打包為矩陣K，value打包為矩陣V，dk為query和key的維度。

在softmax(QKT)中，softmax受最大元素控制，因此對於每個query（qi），只需要關注K中最接近qi的key。這樣效率會高得多。

那麼如何在key中尋找最近鄰居呢？

局部敏感哈希就可以解決在高維空間中快速找到最近鄰居的問題。

局部敏感哈希指的是，如果鄰近的向量很可能獲得相同的哈希值，而遠距離的向量沒可能，則給每個向量x分配哈希值h(x)。

在這項研究中，實際上僅需要求鄰近向量以高機率獲得相同的哈希，並且哈希桶的大小高機率相似。

△LSH注意力的簡化描述

只要近似值是可接受的，注意力的複雜度就會序列長度的平方降低至線性。

可逆的Transformer

那麼問題來了，從下面這張表格中可以看出，在某些情況下，內存複雜度和時間複雜度還是很高。

於是，研究人員引入了可逆層和分段處理，來進一步降低成本。

可逆Transformer無需在每個層中存儲activations。

這樣一來，整個網絡中activations占用的內存就與層數無關了。

實驗結果

研究人員在enwik8和imagenet64數據集上對20層的Reformer模型進行了訓練。

實驗表明，Reformer能達到與Transformer相同的性能，並且內存效率更高，模型在長序列任務上訓練更快。

還可以僅在單個GPU上進行微調。

應用

Reformer將Transformer的建模能力和高效處理長序列任務的架構相結合，即使對於具有大量層的模型，內存成本也會大幅降低。

研究人員認為，這將有助於大型的、參數豐富的Transformer模型的推廣。

並且，Reformer的長序列處理能力，將把Transformer的功能引入更廣闊的領域，比如時間序列預測，音樂、圖像和視頻生成。

傳送門

論文地址：https://arxiv.org/abs/2001.04451

GitHub：https://github.com/google/trax/tree/master/trax/models/reformer

— 完 —

量子位 QbitAI · 頭條號簽約

關注我們，第一時間獲知前沿科技動態