Transformer又又又被挑戰(zhàn)了！

這次的挑戰(zhàn)者來自大名鼎鼎的谷歌DeepMind，并且一口氣推出了兩種新架構(gòu)，——Hawk和Griffin。

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

這種將門控線性RNN與局部注意力混合在一起的模型新架構(gòu)的表現(xiàn)相當亮眼。

首先，同為線性RNN架構(gòu)的Griffin，憑借著1/2的訓(xùn)練數(shù)據(jù)，在所有評測中全面優(yōu)于之前大火的Mamba。

更重要的是，Griffin將模型成功擴展到了14B，做到了Mamba想做卻沒能做的事。

其次，面對基于Transformer架構(gòu)的模型，Griffin則憑借著1/6的訓(xùn)練數(shù)據(jù)，打平甚至超越了同等參數(shù)量的Llama 2！

同時，模型能夠利用很長的上下文來改進其預(yù)測，表明線性RNN的外推能力可以遠遠超出它們訓(xùn)練的序列長度。

此外，團隊還證明了這種組合構(gòu)架保留了Transformer在合成任務(wù)上的許多功能，例如從長上下文中復(fù)制和檢索token。

文章共同一作興奮地發(fā)推表示，Griffin這種新的模型構(gòu)架效率非常高——它集合了線性RNN的所有的效率優(yōu)勢和Transformer的表現(xiàn)力（expressiveness）和規(guī)?；膬?yōu)勢。

推上給的配圖也很有趣，還記得之前Mamba配圖用自己的蟒蛇挑戰(zhàn)變形金剛嗎？這次Griffin（獅鷲）直接霸氣占據(jù)C位，在賽道上領(lǐng)先。

小編莫名覺得Griffin對Mamba的傷害比較大，畢竟老鷹抓蛇......

不過這只是玩笑，因為我們可以發(fā)現(xiàn)Mamba的作者Albert Gu也在這篇文章的作者之中——所以有可能天下的線性RNN都是一家人。

說回正題，雖然現(xiàn)在Transformer稱霸江湖，但它平方級別的計算和存儲開銷對科研和工業(yè)界都造成了很大的壓力（盡管喂飽了老黃......）。

大家在拼命做優(yōu)化（比如Mamba另一位作者Tri Dao開發(fā)的FlashAttention系列）之余，也想另辟蹊徑，于是就誕生了這些挑戰(zhàn)Transformer的架構(gòu)。

——如果真成了，就能跟「Attention Is All You Need」一樣名留青史。

最近大火的幾個著名研究都與線性RNN相關(guān)，比如RWKV、Mamba，以及我們今天的Hawk和Griffin。

循環(huán)神經(jīng)網(wǎng)絡(luò)（RNN）在處理長序列數(shù)據(jù)時表現(xiàn)出色，因為它的推理開銷是線性的，相比于Transformer有天然的計算和存儲優(yōu)勢。

不過RNN系列在記憶和選擇性提取信息方面相對于Transformer又有著原理上的劣勢，所以在當前的這些任務(wù)上表現(xiàn)力不行。

另外，由于結(jié)構(gòu)問題，訓(xùn)練大規(guī)模的RNN是非常困難的。

為此，研究人員提出了Hawk，一種采用了門控線性循環(huán)的RNN；以及Griffin，一個將門控線性反饋與局部注意力機制相結(jié)合的混合模型。

首先，研究人員提出了RG-LRU層，這是一個新穎的門控線性循環(huán)層，并圍繞它設(shè)計了一個新的循環(huán)塊來替代MQA。

接著研究人員以這個循環(huán)塊為基礎(chǔ)，構(gòu)建了兩個新模型：Hawk（將MLP與循環(huán)塊混合使用），以及Griffin（將MLP與循環(huán)塊和局部注意力混合使用）。

具體來說：

1. Hawk和Griffin模型在訓(xùn)練FLOPs和保留損失方面，展現(xiàn)出了與Transformer模型類似的冪律縮放關(guān)系，即使參數(shù)量達到7B以上也是如此（圖1(a)）。

2. 在所有模型規(guī)模上，Griffin的保留損失略低于強大的Transformer基線。

3. 對Hawk和Griffin模型在不同規(guī)模下使用300B token進行了過訓(xùn)練。在一系列下游任務(wù)中，Hawk的表現(xiàn)超越了2倍token訓(xùn)練的Mamba模型，而Griffin則可以和6倍token訓(xùn)練的Llama-2相媲美。

4. 在TPU-v3上，Hawk和Griffin達到了與Transformers相當?shù)挠?xùn)練效率。通過在Pallas中設(shè)計的RG-LRU層內(nèi)核，研究人員在最大限度地減少內(nèi)存?zhèn)鬏數(shù)耐瑫r，克服了由于對角線RNN層的內(nèi)存限制帶來的挑戰(zhàn)。

5. 在推理階段，Hawk和Griffin的吞吐量顯著高于MQA Transformers（圖1(b)），并且在處理長序列時延遲更低。

6. Griffin在處理訓(xùn)練期間未見過的更長序列時，表現(xiàn)優(yōu)于Transformers，同時還能高效地從訓(xùn)練數(shù)據(jù)中學(xué)習復(fù)制和檢索任務(wù)。然而，如果不進行微調(diào)，直接使用預(yù)訓(xùn)練模型進行復(fù)制和精確檢索任務(wù)的評估，Hawk和Griffin的表現(xiàn)則不如Transformers。

模型架構(gòu)

所有的模型都包含以下三大核心部分：(1) 一個殘差塊，(2) 一個MLP塊，以及 (3)一個 時間混合塊。

殘差塊和MLP塊在所有模型中保持不變，而時間混合塊則有三種不同的實現(xiàn)方式：

1. 全局多查詢注意力（MQA）；

2. 局部（滑動窗口）MQA；

3. 研究人員提出的循環(huán)塊。

在循環(huán)塊中，研究人員受線性循環(huán)單元的啟發(fā)，提出了一種新型的循環(huán)層——真實門控線性循環(huán)單元（RG-LRU）。

如圖2(a)所示，模型的全局結(jié)構(gòu)由殘差塊定義，其設(shè)計靈感來源于預(yù)歸一化的Transformers架構(gòu)。

首先將輸入序列進行嵌入處理，然后讓它通過??個這樣的殘差塊（這里的??代表模型的深度），接著應(yīng)用RMSNorm生成最終激活。

為了計算token的概率，研究人員在最后加入了一個線性層和一個softmax函數(shù)。值得一提的是，這個線性層的權(quán)重是與輸入嵌入層共享的。

像Transformer一樣高效擴展

縮放研究評估了三種不同的模型，并展示了它們在參數(shù)量達到并超過7B時的性能表現(xiàn)。

1. MQA Transformer基準模型

使用了殘差模式和門控MLP塊，并結(jié)合了MQA和RoPE技術(shù)。

2. Hawk模型

Hawk遵循了與MQA Transformer相同的殘差模式和MLP塊設(shè)計，但在時間混合部分采用了一種新的循環(huán)塊（包含RG-LRU層），而不是MQA。通過將循環(huán)塊的寬度擴大約約3/4，其參數(shù)量也達到了與多頭注意力塊大致相當?shù)乃?，從而可以匹配模型維度??的設(shè)置。

3. Griffin模型

Griffin的一個顯著優(yōu)點是，相比于全局注意力機制，它通過固定的狀態(tài)大小來總結(jié)序列，而不是像MQA那樣讓KV緩存隨著序列長度的增加而線性增長。由于局部注意力具有相同的特性，因此將循環(huán)塊與局部注意力相結(jié)合，可以保持這一優(yōu)點。結(jié)果顯示，這種組合非常有效，因為局部注意力能夠精準模擬近期的信息，而循環(huán)層則能夠在長序列中傳遞信息。

Griffin使用了與Transformer基線相同的殘差模式和MLP塊，但與上述兩種模型不同的是，它采用了循環(huán)塊和MQA塊的混合。具體而言，研究人員設(shè)計了一種分層結(jié)構(gòu)——先交替使用兩個殘差塊和一個循環(huán)塊，然后再使用一個基于局部注意力機制的殘差塊。通常，局部注意力的窗口大小為1024個token。

大規(guī)模并行訓(xùn)練

隨著模型大小的增加，我們無法在訓(xùn)練期間將模型安裝在單個設(shè)備上，即使每個設(shè)備的批大小為1。

因此，在訓(xùn)練期間，作者使用模型并行性跨設(shè)備對大型模型進行分片。由于不同訓(xùn)練設(shè)備之間的通信成本很高，因此有效地對模型進行分片對于大規(guī)?？焖儆?xùn)練至關(guān)重要。

對于gated-MLP，這里使用Megatron式的分片（即張量并行），這需要在前向和后向傳遞中執(zhí)行all-reduce操作。相同的策略也可以應(yīng)用于注意力塊中的線性層，將多頭分配給不同的設(shè)備。

循環(huán)塊包含每個分支的兩個線性層。所以也可以高效地應(yīng)用Megatron的分片方式。而Conv1D層跨通道獨立運行，可以在設(shè)備之間拆分其參數(shù)，并不會產(chǎn)生任何通信開銷。

為了避免額外的跨設(shè)備通信，研究人員對RG-LRU中的gate使用block-diagonal weights（下面的公式1和2），而不是密集矩陣。

對于本文中的所有實驗，作者使用16個塊作為recurrence gate和 input gate。recurrence的對角線結(jié)構(gòu)具有與Conv1D相同的優(yōu)勢，允許在沒有任何通信的情況下進行參數(shù)分片和計算。使用這種策略，循環(huán)區(qū)塊的通信要求與MLP區(qū)塊的通信要求相同。

另外優(yōu)化器（比如Adam）的狀態(tài)參數(shù)可能會消耗大量內(nèi)存，超過模型參數(shù)本身的大小。

為了解決這個問題，研究人員采用了ZeRO，還使用bfloat16表示模型參數(shù)和激活，從而最大限度地減少任何數(shù)據(jù)傳輸開銷。

上圖以2K序列長度的MQA作為基線，顯示了Griffin和MQA的訓(xùn)練耗時，這里測試了不同的模型大小和序列長度。

隨著序列長度的增加，實驗中會按比例降低批大小，以便每個批的token總數(shù)保持不變。

為什么不使用卷積或者關(guān)聯(lián)掃描？

線性RNN模型的優(yōu)勢之一是高并行化，源于它計算的關(guān)聯(lián)性，可以通過卷積或前綴和算法（關(guān)聯(lián)掃描）在設(shè)備上高效執(zhí)行。

然而，RG-LRU的門控機制與卷積視圖不兼容。雖然原則上仍然可以使用關(guān)聯(lián)掃描減少所需的FLOP數(shù)，但不會減少內(nèi)存開銷，這是在實踐中的主要瓶頸。

根據(jù)經(jīng)驗，在TPU-v3上，關(guān)聯(lián)掃描明顯慢于原生Jax線性掃描。作者推測，并行前綴-求和算法的樹重組的隨機訪問性質(zhì)不太適合TPU架構(gòu)，導(dǎo)致內(nèi)存?zhèn)鬏斔俣雀?/span>

推理速度

在評估推理速度時，需要考慮兩個主要指標。

第一個是延遲，它衡量在特定批大小下生成指定數(shù)量的token所花費的時間。第二個是吞吐量，衡量采樣指定數(shù)量的token時，每秒可以在單個設(shè)備上生成的最大token數(shù)。

由于吞吐量與token采樣、批大小以及延遲有關(guān)，因此可以通過減少延遲或減少內(nèi)存使用量來提高吞吐量，從而允許在設(shè)備上使用更大的批大小。

對于需要快速響應(yīng)時間的實時應(yīng)用程序，可以考慮延遲。而在考慮其他語言應(yīng)用程序時，吞吐量很重要，例如人類反饋強化學(xué)習（RLHF）或評分語言模型輸出。

這里測試了大小為1B參數(shù)的模型的推理結(jié)果，基線為MQA（Transformer）。MQA在推理過程中比文獻中經(jīng)常使用的標準MHA要快得多。

參與比較的模型是：MQA Transformer，Hawk，和Griffin。

圖4比較了批大小為16的模型的延遲，預(yù)填充為空，或4096個token。對于長序列，Hawk和Griffin的采樣延遲比MQA Transformer更快。

當序列長度和預(yù)填充長度（影響KV緩存的大?。┰黾訒r，這一點尤其明顯。Griffin實現(xiàn)了與Hawk相似的延遲，展示了線性RNN和局部注意力的出色兼容性。

隨后，研究人員比較了相同模型在空提示下對512、1024、2048和4196個token進行采樣時的最大吞吐量（token/秒）。

如上圖所示，我們可以看到Griffin和Hawk的吞吐量都明顯高于MQA Transformer基線。

這在一定程度上是由于重復(fù)模型具有較低的延遲，但也主要是因為Griffin和Hawk可以在單個設(shè)備上容納比MQA Transformer更大的批大小，因為它們的緩存更小。

Hawk實現(xiàn)了比Griffin更高的吞吐量，因為當批大小較大時，本地注意力緩存的大小最終會與參數(shù)的大小相當。

網(wǎng)友評論

雖然論文的效果顯著，但也有網(wǎng)友并不買賬，質(zhì)疑本文的模型與Mamba的比較并不「公平」：

「這篇論文說明了LLM研究中的一個巨大問題。他們聲稱在更少的token上表現(xiàn)優(yōu)于Mamba。然而，直到第3.2節(jié)，他們才承認使用了與Mamba完全不同的數(shù)據(jù)集進行訓(xùn)練。」

「由于數(shù)據(jù)實際上是最重要的東西，因此性能的比較是沒有用的。完全沒用。無法獲得任何科學(xué)結(jié)論或見解......」

還有通過陰謀論推導(dǎo)出模型不行的：

「DeepMind有什么動機來發(fā)表這樣的研究？如果他們想要在與OpenAI的競爭中占據(jù)優(yōu)勢，那么如果他們發(fā)現(xiàn)一些很棒的新架構(gòu)，他們會將其保密?！?/span>

「這是否意味著這些結(jié)果現(xiàn)在是好的，但還不夠好，不足以在提供競爭優(yōu)勢方面具有革命性？」

大浪淘沙，結(jié)果如何，讓我們期待時間的檢驗。

作者介紹

Albert Gu

Albert Gu也是之前爆火的Mamba架構(gòu)的作者

2015年，Albert Gu在卡內(nèi)基梅隆大學(xué)（CMU）取得了計算機科學(xué)和數(shù)學(xué)雙學(xué)士學(xué)位。

隨后就讀于斯坦福大學(xué)，專業(yè)是計算機科學(xué)，目前在CMU擔任Assistant Professor。

2011Albert Gu曾在Facebook做過軟件工程實習生，2015年又在英國倫敦的Jump Trading做算法交易實習生。

2019年，Albert Gu到DeepMind實習，作為研究科學(xué)家實習生。

他的研究方向包括：

用于機器學(xué)習的結(jié)構(gòu)化表示，包括結(jié)構(gòu)化線性；

代數(shù)和嵌入，序列模型的分析和設(shè)計；

關(guān)于長上下文，非歐幾里得表示學(xué)習。

Albert Gu近年來在ICML、ICLR、NeurlPS等AI頂級會議上發(fā)表了多篇高被引文章：

Soham De

論文共一Soham De是谷歌DeepMind的研究科學(xué)家，工作重點是深入理解并提升大規(guī)模深度學(xué)習的性能。

此前，他于2018年在美國馬里蘭大學(xué)取得博士學(xué)位，導(dǎo)師是Dana Nau教授和Tom Goldstein教授，期間主要研究用于機器學(xué)習問題的快速隨機優(yōu)化算法。

Samuel L. Smith

論文共一Samuel在劍橋大學(xué)一路本碩博，2016年獲得理論物理學(xué)博士學(xué)位，隨后在谷歌大腦和DeepMind工作。

George-Cristian Muraru

論文共一George-Cristian Muraru也是谷歌DeepMind的軟件工程師，在谷歌的工作之前，他曾在彭博干過程序員。

Aleksandar Botev

論文共一Aleksandar Botev畢業(yè)于英國倫敦大學(xué)學(xué)院，獲得機器學(xué)習碩士學(xué)位，深度學(xué)習博士學(xué)位，曾在OpenAI和谷歌DeepMind實習，于2019年加入DeepMind，擔任Reseach Scientist。