魚羊 發自 凹非寺
量子位 報道 | 公衆号 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
作者系網易新聞·網易号“各有态度”簽約作者
— 完 —