太平洋科技要聞

首頁(yè) > 科技要聞 > 科技> 正文

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

計(jì)算效率

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

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

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

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

長(zhǎng)上下文評(píng)估

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

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

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

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

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

其他評(píng)測(cè)集，如原始語(yǔ)音分類(lèi)、ImageNet-1K、WikiText-103和PG-19的結(jié)果如下：

一些感想

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

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

通過(guò)這個(gè)項(xiàng)目，研究者們也體會(huì)到了在大模型時(shí)代做新的模型架構(gòu)時(shí)要注意的問(wèn)題�？偨Y(jié)來(lái)說(shuō)：

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

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

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

參考資料：

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

文章來(lái)源：新智元

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

新智元

原創(chuàng)欄目