單個(gè)GPU上可運(yùn)行的Transformer，谷歌&伯克利最新研究開源 | ICLR 2020
魚羊 發(fā)自 凹非寺
量子位 報(bào)道 | 公眾號(hào) QbitAI
Transformer自誕生以來，就在NLP領(lǐng)域刷新一個(gè)又一個(gè)紀(jì)錄，稱作當(dāng)下最流行的深度學(xué)習(xí)框架亦不為過。
不過，拿下SOTA并不意味著十全十美。
比如，在長序列訓(xùn)練上，Transformer就存在計(jì)算量巨大、訓(xùn)練成本高的問題。
其對(duì)內(nèi)存的要求從GB級(jí)別到TB級(jí)別不等。這意味著，模型只能處理簡短的文本，生成簡短的音樂。
此外，許多大型Transformer模型在經(jīng)過模型并行訓(xùn)練之后，無法在單個(gè)GPU上進(jìn)行微調(diào)。
現(xiàn)在，谷歌和UC伯克利推出了一個(gè)更高效的Transformer模型——Reformer。
在長度為L的序列上，將復(fù)雜度從 O(L2)降低到了O(L logL)。
并且，模型訓(xùn)練后，可以僅使用16GB內(nèi)存的單個(gè)GPU運(yùn)行。
論文已被ICLR 2020接收。模型也已開源。（地址見文末）
Reformer
Reformer主要引入了以下技術(shù)：
- 可逆層。使用可逆殘差層替代標(biāo)準(zhǔn)殘差，這樣以來，模型在訓(xùn)練過程中僅存儲(chǔ)一次activations，而非N次（N為層數(shù)）。
- 在前饋層內(nèi)拆分activations，并進(jìn)行分段處理，消除dff因子，節(jié)省前饋層內(nèi)存。
- 將點(diǎn)乘注意力替換為局部敏感哈希（locality sensitive hashing）注意力，將復(fù)雜度從 O(L2) 降低到 O(L logL)，使其能對(duì)長序列進(jìn)行操作。
局部敏感哈希（LSH）注意力
Transformer中使用的標(biāo)準(zhǔn)注意力是點(diǎn)乘注意力。
其中，Q矩陣由一組query的注意力函數(shù)組成，key打包為矩陣K，value打包為矩陣V，dk為query和key的維度。
在softmax(QKT)中，softmax受最大元素控制，因此對(duì)于每個(gè)query（qi），只需要關(guān)注K中最接近qi的key。這樣效率會(huì)高得多。
那么如何在key中尋找最近鄰居呢？
局部敏感哈希就可以解決在高維空間中快速找到最近鄰居的問題。
局部敏感哈希指的是，如果鄰近的向量很可能獲得相同的哈希值，而遠(yuǎn)距離的向量沒可能，則給每個(gè)向量x分配哈希值h(x)。
在這項(xiàng)研究中，實(shí)際上僅需要求鄰近向量以高概率獲得相同的哈希，并且哈希桶的大小高概率相似。
△LSH注意力的簡化描述
只要近似值是可接受的，注意力的復(fù)雜度就會(huì)序列長度的平方降低至線性。
可逆的Transformer
那么問題來了，從下面這張表格中可以看出，在某些情況下，內(nèi)存復(fù)雜度和時(shí)間復(fù)雜度還是很高。
于是，研究人員引入了可逆層和分段處理，來進(jìn)一步降低成本。
可逆Transformer無需在每個(gè)層中存儲(chǔ)activations。
這樣一來，整個(gè)網(wǎng)絡(luò)中activations占用的內(nèi)存就與層數(shù)無關(guān)了。
實(shí)驗(yàn)結(jié)果
研究人員在enwik8和imagenet64數(shù)據(jù)集上對(duì)20層的Reformer模型進(jìn)行了訓(xùn)練。
實(shí)驗(yàn)表明，Reformer能達(dá)到與Transformer相同的性能，并且內(nèi)存效率更高，模型在長序列任務(wù)上訓(xùn)練更快。
還可以僅在單個(gè)GPU上進(jìn)行微調(diào)。
應(yīng)用
Reformer將Transformer的建模能力和高效處理長序列任務(wù)的架構(gòu)相結(jié)合，即使對(duì)于具有大量層的模型，內(nèi)存成本也會(huì)大幅降低。
研究人員認(rèn)為，這將有助于大型的、參數(shù)豐富的Transformer模型的推廣。
并且，Reformer的長序列處理能力，將把Transformer的功能引入更廣闊的領(lǐng)域，比如時(shí)間序列預(yù)測(cè)，音樂、圖像和視頻生成。
傳送門
論文地址：https://arxiv.org/abs/2001.04451
GitHub：https://github.com/google/trax/tree/master/trax/models/reformer
— 完 —
- 標(biāo)準(zhǔn)化3D生成質(zhì)量榜單來了！首創(chuàng)層次化評(píng)價(jià)體系，告別“誰的demo更吸睛”主觀評(píng)估2025-08-16
- 阿里閃電入局Agent Infra！智能體新基建亮相WAIC，“超級(jí)大腦”開箱即用2025-07-31
- 世紀(jì)華通謝斐：在“三大平衡”中領(lǐng)跑，實(shí)現(xiàn)游戲行業(yè)更高質(zhì)量的發(fā)展2025-08-01
- 1.5B參數(shù)撬動(dòng)“吉卜力級(jí)”全能體驗(yàn)，國產(chǎn)開源之光多模態(tài)統(tǒng)一模型，來了2025-07-30
相關(guān)閱讀
Transformer挑戰(zhàn)者出現(xiàn)！斯坦福CMU聯(lián)合團(tuán)隊(duì)，模型代碼都開源
來自FlashAttention作者
只用一張圖+相機(jī)走位，AI就能腦補(bǔ)周圍環(huán)境，來自華人團(tuán)隊(duì) | CVPR2022
AI看一眼門口，就能腦補(bǔ)出房間內(nèi)部畫面
谷歌大模型研究陷重大爭議：訓(xùn)練數(shù)據(jù)之外完全無法泛化？網(wǎng)友：AGI奇點(diǎn)推遲了
網(wǎng)友找出論文中更多關(guān)鍵卻被忽略的細(xì)節(jié)，比如只做了GPT-2規(guī)模的試驗(yàn)等
