Transformer又又又被挑戰了！

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

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

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

首先，同為線性RNN架構的Griffin，憑借著1/2的訓練數據，在所有評測中全面優于之前大火的Mamba。

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

其次，面對基于Transformer架構的模型，Griffin則憑借著1/6的訓練數據，打平甚至超越了同等參數量的Llama 2！

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

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

文章共同一作興奮地發推表示，Griffin這種新的模型構架效率非常高——它集合了線性RNN的所有的效率優勢和Transformer的表現力（expressiveness）和規模化的優勢。

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

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

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

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

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

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

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

循環神經網絡（RNN）在處理長序列數據時表現出色，因為它的推理開銷是線性的，相比于Transformer有天然的計算和存儲優勢。

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

另外，由于結構問題，訓練大規模的RNN是非常困難的。

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

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

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

具體來說：

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

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

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

4. 在TPU-v3上，Hawk和Griffin達到了與Transformers相當的訓練效率。通過在Pallas中設計的RG-LRU層內核，研究人員在最大限度地減少內存傳輸的同時，克服了由于對角線RNN層的內存限制帶來的挑戰。

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

6. Griffin在處理訓練期間未見過的更長序列時，表現優于Transformers，同時還能高效地從訓練數據中學習復制和檢索任務。然而，如果不進行微調，直接使用預訓練模型進行復制和精確檢索任務的評估，Hawk和Griffin的表現則不如Transformers。

模型架構

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

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

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

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

3. 研究人員提出的循環塊。

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

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

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

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

像Transformer一樣高效擴展

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

1. MQA Transformer基準模型

使用了殘差模式和門控MLP塊，并結合了MQA和RoPE技術。

2. Hawk模型

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

3. Griffin模型

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

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

大規模并行訓練

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

因此，在訓練期間，作者使用模型并行性跨設備對大型模型進行分片。由于不同訓練設備之間的通信成本很高，因此有效地對模型進行分片對于大規模快速訓練至關重要。

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

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

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

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

另外優化器（比如Adam）的狀態參數可能會消耗大量內存，超過模型參數本身的大小。

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

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

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

為什么不使用卷積或者關聯掃描？

線性RNN模型的優勢之一是高并行化，源于它計算的關聯性，可以通過卷積或前綴和算法（關聯掃描）在設備上高效執行。

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

根據經驗，在TPU-v3上，關聯掃描明顯慢于原生Jax線性掃描。作者推測，并行前綴-求和算法的樹重組的隨機訪問性質不太適合TPU架構，導致內存傳輸速度更慢。

推理速度

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

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

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

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

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

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

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

當序列長度和預填充長度（影響KV緩存的大小）增加時，這一點尤其明顯。Griffin實現了與Hawk相似的延遲，展示了線性RNN和局部注意力的出色兼容性。

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

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

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

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

網友評論

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

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

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

還有通過陰謀論推導出模型不行的：

「DeepMind有什么動機來發表這樣的研究？如果他們想要在與OpenAI的競爭中占據優勢，那么如果他們發現一些很棒的新架構，他們會將其保密。」

「這是否意味著這些結果現在是好的，但還不夠好，不足以在提供競爭優勢方面具有革命性？」

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

作者介紹

Albert Gu

Albert Gu也是之前爆火的Mamba架構的作者

2015年，Albert Gu在卡內基梅隆大學（CMU）取得了計算機科學和數學雙學士學位。

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

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

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

他的研究方向包括：

用于機器學習的結構化表示，包括結構化線性；

代數和嵌入，序列模型的分析和設計；

關于長上下文，非歐幾里得表示學習。

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

Soham De

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

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

Samuel L. Smith

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

George-Cristian Muraru

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

Aleksandar Botev

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