Transformer 是近期 NLP 領(lǐng)域里最熱門的模型之一，但因?yàn)樗懔ο倪^(guò)大，對(duì)于個(gè)人研究者來(lái)說(shuō)一直不太友好。近日一篇入選 ICLR 2020 的研究提出了「Reformer」，把跑 Transformer 模型的硬件要求壓縮到了只需一塊 GPU，同時(shí)效果不變。

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

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

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

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

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

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

引言

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

諸如此類的大型 Transformer 模型頻頻出現(xiàn)，到底是客觀上必須要求如此多的資源，還是僅僅是因?yàn)樘幚硇什粔蚋?

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

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

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

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

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

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

研究者提出了一種 Reformer 模型來(lái)解決剛才說(shuō)的那些問(wèn)題：

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

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

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

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

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

首先回顧一下點(diǎn)乘注意力機(jī)制，如下所示：

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

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

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

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

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

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

最終，對(duì) attention 進(jìn)行哈希處理的流程如下：

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

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

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

怎么使用這種新型 attention

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

可逆 Transformer

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

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

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

，并遵循以下方程：

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

可逆 Transformer b 不需要在每一層中激活存儲(chǔ)，于是無(wú)需使用 nl 項(xiàng)。

分塊

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

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

實(shí)驗(yàn)

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

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

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

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

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

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

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

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

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

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

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

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

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