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

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

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

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

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

更重要的是，Griffin將模型成功擴(kuò)展到了14B，做到了Mamba想做卻沒(méi)能做的事。

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

同時(shí)，模型能夠利用很長(zhǎng)的上下文來(lái)改進(jìn)其預(yù)測(cè)，表明線性RNN的外推能力可以遠(yuǎn)遠(yuǎn)超出它們訓(xùn)練的序列長(zhǎng)度。

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

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

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

小編莫名覺(jué)得Griffin對(duì)Mamba的傷害比較大，畢竟老鷹抓蛇......

不過(guò)這只是玩笑，因?yàn)槲覀兛梢园l(fā)現(xiàn)Mamba的作者Albert Gu也在這篇文章的作者之中——所以有可能天下的線性RNN都是一家人。

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

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

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

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

循環(huán)神經(jīng)網(wǎng)絡(luò)（RNN）在處理長(zhǎng)序列數(shù)據(jù)時(shí)表現(xiàn)出色，因?yàn)樗耐评黹_(kāi)銷(xiāo)是線性的，相比于Transformer有天然的計(jì)算和存儲(chǔ)優(yōu)勢(shì)。

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

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

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

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

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

具體來(lái)說(shuō)：

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

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

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

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

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

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

模型架構(gòu)

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

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

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

2. 局部（滑動(dòng)窗口）MQA；

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

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

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

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

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

像Transformer一樣高效擴(kuò)展

縮放研究評(píng)估了三種不同的模型，并展示了它們?cè)趨?shù)量達(dá)到并超過(guò)7B時(shí)的性能表現(xiàn)。

1. MQA Transformer基準(zhǔn)模型

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

2. Hawk模型

Hawk遵循了與MQA Transformer相同的殘差模式和MLP塊設(shè)計(jì)，但在時(shí)間混合部分采用了一種新的循環(huán)塊（包含RG-LRU層），而不是MQA。通過(guò)將循環(huán)塊的寬度擴(kuò)大約約3/4，其參數(shù)量也達(dá)到了與多頭注意力塊大致相當(dāng)?shù)乃剑瑥亩梢云ヅ淠Ｐ途S度??的設(shè)置。

3. Griffin模型

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

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

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

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

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

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

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

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

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

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

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

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

隨著序列長(zhǎng)度的增加，實(shí)驗(yàn)中會(huì)按比例降低批大小，以便每個(gè)批的token總數(shù)保持不變。

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

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

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

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

推理速度

在評(píng)估推理速度時(shí)，需要考慮兩個(gè)主要指標(biāo)。

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

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

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

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

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

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

當(dāng)序列長(zhǎng)度和預(yù)填充長(zhǎng)度（影響KV緩存的大小）增加時(shí)，這一點(diǎn)尤其明顯。Griffin實(shí)現(xiàn)了與Hawk相似的延遲，展示了線性RNN和局部注意力的出色兼容性。

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

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

這在一定程度上是由于重復(fù)模型具有較低的延遲，但也主要是因?yàn)镚riffin和Hawk可以在單個(gè)設(shè)備上容納比MQA Transformer更大的批大小，因?yàn)樗鼈兊木彺娓　?/span>

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

網(wǎng)友評(píng)論

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

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

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

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

「DeepMind有什么動(dòng)機(jī)來(lái)發(fā)表這樣的研究？如果他們想要在與OpenAI的競(jìng)爭(zhēng)中占據(jù)優(yōu)勢(shì)，那么如果他們發(fā)現(xiàn)一些很棒的新架構(gòu)，他們會(huì)將其保密。」

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

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

作者介紹

Albert Gu

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

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

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

2011Albert Gu曾在Facebook做過(guò)軟件工程實(shí)習(xí)生，2015年又在英國(guó)倫敦的Jump Trading做算法交易實(shí)習(xí)生。

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

他的研究方向包括：

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

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

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

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

Soham De

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

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

Samuel L. Smith

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

George-Cristian Muraru

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

Aleksandar Botev

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