革命新架構掀翻Transformer！無限上下文處理，2萬億token碾壓Llama 2

新智元報道

編輯：桃子 好睏

【新智元導讀】Transformer王座即將被取而代之！Meta、USC、CMU和UCSD聯合提出了革命性新架構Megalodon，能夠處理無限上下文，在2萬億token訓練任務中，性能超越Llama2-7B實現了非凡的效率。

繼Mamba之後，又一敢於挑戰Transformer的架構誕生了！

來自Meta、南加州大學（USC）、CMU和UCSD的研究人員提出了全新的神經網絡架構——Megalodon（巨齒鯊）。

這是專爲有效處理「無限上下文」長度的LLM預訓練，以及推理而設計的架構。

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

我們都知道，Transformer架構個在處理長上下文時，會受到二次複雜度，以及長度外推能力弱的限制。

儘管已有次二次方解決方案（諸如線性注意力，狀態空間模型），但它們在預訓練效率，甚至下游任務的準確率上，通常還不及Transformer。

Megalodon的出現，就是爲了解決無限處理上下文的難題。

同時，它可以同時實現高效訓練（減少通信和計算量），以及高效推理（保持恆定的KV緩存）。

值得一提的是，在與Llama 2的直接比較中，Megalodon在處理70億參數和2萬億訓練token的任務上，不僅訓練更高效，而且準確率也超過了Transformer。

具體來說，Megalodon的訓練損失爲1.70，位於Llama2-7B（1.75）和 13B（1.67）之間。

這一改變範式的創新代表着AI領域的巨大飛躍，Megalodon開啓了計算效率和性能的新時代。

GPT-3發佈以來最大里程碑

網友表示，先是谷歌，又是Meta，無限上下文離我們更進一步，LLM將會釋放出無限潛力。

還有人認爲「無限上下文長度，絕對是遊戲規則的改變者」！

更有甚者，初創公司CEO稱，「這是自GPT-3發佈以來最大的里程碑，但卻沒有任何動靜？！

Megalodon就相當於是AGI的基礎」。

「Meta的Megalodon是一項突破性進展，對AGI具有重要意義。它的無限上下文長度模擬了人類的認知，實現了無縫任務切換」。

論文作者Hao Zhang表示，這是一種全新替代Transformer的架構。

論文作者Beidi Chen稱，「注意力雖好，但你不需要完整的注意力機制」！

普林斯頓助理教授Tri Dao表示，「將SSM/RNN/EMA與注意力相結合是獲得更高質量、更長上下文和更快推理的方法！Griffin、Jamba、Zamba和現在的Megalodon都是很好的例子」。

革命性架構，訓練更穩定

那麼，Megalodon架構採用了怎樣的設計，才能取得如此優異的表現？

據介紹，它基於MEGA架構進行了改進，並新增了多個技術組件。

首先，複雜指數移動平均（CEMA）組件是一種全新技術，擴展了MEGA中使用的多維阻尼指數移動平均方法到複數域，可以增強模型處理複雜數據的能力。

其次，研究人員提出了一種創新的歸一化技術——「時間步歸一化層」。

它將傳統的組歸一化技術擴展到自迴歸序列建模任務中，允許模型在處理序列數據時，進行有效的歸一化。

以往，「層歸一化」（Layer Normalization）與Transformer相結合性能，雖令人印象深刻。

但很明顯，層歸一化並不能直接減少時間步長或順序維度的內部協變量偏移。

另外，「組歸一化」（Group Normalization）雖比「層歸一化」在CV任務中獲得改進，但它卻無法直接應用於Transformer的自迴歸序列建模，因未來信息會通過時間步維度的均值和方差泄漏。

如下圖所示，c展示了Megalodon架構中，層標準化和時間步標準化的方法。

最後，研究人員爲了增強大規模LLM預訓練的穩定性，提出了將歸一化注意力，和帶有兩跳殘差的預歸一化相結合的配置。

這種配置可以優化模型的學習過程，提高訓練的穩定性。

下圖3中，a是Megalodon的完整框架草圖。

中間和右邊兩張圖分別介紹了，預歸一化和帶有兩跳殘差預歸一化的配置。

2T token訓練，性能超越Llama2-7B

在具體實驗評估中，研究人員將Megalodon擴展到70億參數規模，並將其應用於2萬億token的大規模LLM預訓練中。

此外，作者還在中/小參數規模的序列建模基準上進行了實驗，包括Long Range Arena (LRA) 、Speech Commands上的原始語音分類、ImageNet-1K上的圖像分類，以及WikiText-103和PG19上的語言建模。

結果顯示，在這些任務中，Megalodon在各種數據模式下的表現明顯優於所有最先進的基線模型。

數據學習效率

通過訓練損失圖以及多個benchmark的結果可以看出，Megalodon比Transformer在7B參數下有更好的數據學習效率。

計算效率

針對不同的4K和32K上下文長度，Megalodon這一架構的預訓練的計算效率也是非常強的。

學術基準上短上下文評估

具體來說，研究人員在短上下文（4K token）的標準學術基準上，對Megalodon與Llama 2，以及開源基礎模型進行了比較。

在相同的2萬億token訓練後，Megalodon-7B的表現明顯優於Llama2-7B。

長上下文評估

針對不同長上下文困惑度，證明了Megalodon可以利用很長的上下文進行下一個token預測的能力。

圖5顯示了，驗證數據集在4K到2M各種上下文長度下的困惑度（PPL）。

在Scroll數據集中的長上下文QA任務中，Megalodon在NaQA上獲得最佳F1，並與Llama 2 Long相競爭。

中等規模基準評估

在Long Range Arena（LRA）的測試中，新架構顯著縮小了分塊注意力和全注意力之間的性能差距。

其他評測集，如原始語音分類、ImageNet-1K、WikiText-103和PG-19的結果如下：

一些感想

這裡quote一下這項研究原作者的一些感悟和經歷：

通過這個項目，研究者們也體會到了在大模型時代做新的模型架構時要注意的問題。總結來說：

對於兩個不同模型架構的比較必須要在數據完全相同的條件下才有說服力。當數據不同的時候，哪怕不同的比例很小（<10%)，最後的結果也可能有明顯的差別。包括training loss和下游任務的結果，都受到訓練數據的很大影響。

對於不同的架構，一定要在模型得到充分訓練的條件下的比較纔有意義。例如對於7B大小的模型，2T的訓練數據幾乎是基本要求。有的模型可能在數據少的時候表現的很好，但是數據規模增大後反而落後其他模型。因此，對於大模型架構的比較，結果有說服力的前提是充分的訓練。

對於架構差別很大的模型，傳統的基於flops的scaling law的比較意義在降低。原因是兩個不同架構的模型，即使有相同的flops，他們的實際速度可能差幾倍。這個和架構算法本身是不是適合在最先進的GPU上計算有很大的關係。因此，真正貼合實際的比較方法是像本文中那樣分成數據學習效率和計算效率兩個方面。但是這樣在實際中對於研究員的工程能力有很高的要求。在大模型時代，新算法的開發已經和系統等方面高度結合在一起。

參考資料：

https://arxiv.org/abs/2404.08801

https://zhuanlan.zhihu.com/p/692682649