太平洋科技要聞

首頁 > 科技要聞 > 科技> 正文

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

新智元整合編輯：龔震發(fā)布于：2024-04-27 23:07

繼Mamba之后，又一敢于挑戰(zhàn)Transformer的架構(gòu)誕生了！

來自Meta、南加州大學(xué)（USC）、CMU和UCSD的研究人員提出了全新的神經(jīng)網(wǎng)絡(luò)架構(gòu)——Megalodon（巨齒鯊）。

這是專為有效處理「無限上下文」長度的LLM預(yù)訓(xùn)練，以及推理而設(shè)計的架構(gòu)。

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

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

盡管已有次二次方解決方案（諸如線性注意力，狀態(tài)空間模型），但它們在預(yù)訓(xùn)練效率，甚至下游任務(wù)的準(zhǔn)確率上，通常還不及Transformer。

Megalodon的出現(xiàn)，就是為了解決無限處理上下文的難題。

同時，它可以同時實(shí)現(xiàn)高效訓(xùn)練（減少通信和計算量），以及高效推理（保持恒定的KV緩存）。

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

具體來說，Megalodon的訓(xùn)練損失為1.70，位于Llama2-7B（1.75）和 13B（1.67）之間。

這一改變范式的創(chuàng)新代表著AI領(lǐng)域的巨大飛躍，Megalodon開啟了計算效率和性能的新時代。

GPT-3發(fā)布以來最大里程碑

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

還有人認(rèn)為「無限上下文長度，絕對是游戲規(guī)則的改變者」！

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

Megalodon就相當(dāng)于是AGI的基礎(chǔ)」。

「Meta的Megalodon是一項突破性進(jìn)展，對AGI具有重要意義。它的無限上下文長度模擬了人類的認(rèn)知，實(shí)現(xiàn)了無縫任務(wù)切換」。

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

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

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

革命性架構(gòu)，訓(xùn)練更穩(wěn)定

那么，Megalodon架構(gòu)采用了怎樣的設(shè)計，才能取得如此優(yōu)異的表現(xiàn)？

據(jù)介紹，它基于MEGA架構(gòu)進(jìn)行了改進(jìn)，并新增了多個技術(shù)組件。

首先，復(fù)雜指數(shù)移動平均（CEMA）組件是一種全新技術(shù)，擴(kuò)展了MEGA中使用的多維阻尼指數(shù)移動平均方法到復(fù)數(shù)域，可以增強(qiáng)模型處理復(fù)雜數(shù)據(jù)的能力。

其次，研究人員提出了一種創(chuàng)新的歸一化技術(shù)——「時間步歸一化層」。

它將傳統(tǒng)的組歸一化技術(shù)擴(kuò)展到自回歸序列建模任務(wù)中，允許模型在處理序列數(shù)據(jù)時，進(jìn)行有效的歸一化。

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

但很明顯，層歸一化并不能直接減少時間步長或順序維度的內(nèi)部協(xié)變量偏移。

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

如下圖所示，c展示了Megalodon架構(gòu)中，層標(biāo)準(zhǔn)化和時間步標(biāo)準(zhǔn)化的方法。

最后，研究人員為了增強(qiáng)大規(guī)模LLM預(yù)訓(xùn)練的穩(wěn)定性，提出了將歸一化注意力，和帶有兩跳殘差的預(yù)歸一化相結(jié)合的配置。

這種配置可以優(yōu)化模型的學(xué)習(xí)過程，提高訓(xùn)練的穩(wěn)定性。

下圖3中，a是Megalodon的完整框架草圖。中間和右邊兩張圖分別介紹了，預(yù)歸一化和帶有兩跳殘差預(yù)歸一化的配置。

2T token訓(xùn)練，性能超越Llama2-7B

在具體實(shí)驗(yàn)評估中，研究人員將Megalodon擴(kuò)展到70億參數(shù)規(guī)模，并將其應(yīng)用于2萬億token的大規(guī)模LLM預(yù)訓(xùn)練中。

此外，作者還在中/小參數(shù)規(guī)模的序列建�；鶞�(zhǔn)上進(jìn)行了實(shí)驗(yàn)，包括Long Range Arena (LRA) 、Speech Commands上的原始語音分類、ImageNet-1K上的圖像分類，以及WikiText-103和PG19上的語言建模。

結(jié)果顯示，在這些任務(wù)中，Megalodon在各種數(shù)據(jù)模式下的表現(xiàn)明顯優(yōu)于所有最先進(jìn)的基線模型。

數(shù)據(jù)學(xué)習(xí)效率

通過訓(xùn)練損失圖以及多個benchmark的結(jié)果可以看出，Megalodon比Transformer在7B參數(shù)下有更好的數(shù)據(jù)學(xué)習(xí)效率。

計算效率

針對不同的4K和32K上下文長度，Megalodon這一架構(gòu)的預(yù)訓(xùn)練的計算效率也是非常強(qiáng)的。

學(xué)術(shù)基準(zhǔn)上短上下文評估

具體來說，研究人員在短上下文（4K token）的標(biāo)準(zhǔn)學(xué)術(shù)基準(zhǔn)上，對Megalodon與Llama 2，以及開源基礎(chǔ)模型進(jìn)行了比較。

在相同的2萬億token訓(xùn)練后，Megalodon-7B的表現(xiàn)明顯優(yōu)于Llama2-7B。

長上下文評估

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

圖5顯示了，驗(yàn)證數(shù)據(jù)集在4K到2M各種上下文長度下的困惑度（PPL）。

在Scroll數(shù)據(jù)集中的長上下文QA任務(wù)中，Megalodon在NaQA上獲得最佳F1，并與Llama 2 Long相競爭。

中等規(guī)�；鶞�(zhǔn)評估

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

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

一些感想

這里quote一下這項研究原作者的一些感悟和經(jīng)歷：

這個工作從有想法到最終完成，經(jīng)歷了近兩年的時間。期間經(jīng)歷數(shù)次失敗，也學(xué)習(xí)到了很多大規(guī)模預(yù)訓(xùn)練時代正確做科研的方法。

通過這個項目，研究者們也體會到了在大模型時代做新的模型架構(gòu)時要注意的問題�？偨Y(jié)來說：

-對于兩個不同模型架構(gòu)的比較必須要在數(shù)據(jù)完全相同的條件下才有說服力。當(dāng)數(shù)據(jù)不同的時候，哪怕不同的比例很小（<10%)，最后的結(jié)果也可能有明顯的差別。包括training loss和下游任務(wù)的結(jié)果，都受到訓(xùn)練數(shù)據(jù)的很大影響。

-對于不同的架構(gòu)，一定要在模型得到充分訓(xùn)練的條件下的比較才有意義。例如對于7B大小的模型，2T的訓(xùn)練數(shù)據(jù)幾乎是基本要求。有的模型可能在數(shù)據(jù)少的時候表現(xiàn)的很好，但是數(shù)據(jù)規(guī)模增大后反而落后其他模型。因此，對于大模型架構(gòu)的比較，結(jié)果有說服力的前提是充分的訓(xùn)練。

-對于架構(gòu)差別很大的模型，傳統(tǒng)的基于flops的scaling law的比較意義在降低。原因是兩個不同架構(gòu)的模型，即使有相同的flops，他們的實(shí)際速度可能差幾倍。這個和架構(gòu)算法本身是不是適合在最先進(jìn)的GPU上計算有很大的關(guān)系。因此，真正貼合實(shí)際的比較方法是像本文中那樣分成數(shù)據(jù)學(xué)習(xí)效率和計算效率兩個方面。但是這樣在實(shí)際中對于研究員的工程能力有很高的要求。在大模型時代，新算法的開發(fā)已經(jīng)和系統(tǒng)等方面高度結(jié)合在一起。

參考資料：

https://arxiv.org/abs/2404.08801https://zhuanlan.zhihu.com/p/692682649

文章來源：新智元

Megalodon Transformer LLM預(yù)訓(xùn)練

新智元

原創(chuàng)欄目