Transformer 王座即將被取而代之！Meta、USC、CMU 和 UCSD 聯(lián)合提出了革命性新架構(gòu) Megalodon，能夠處理無限上下文，在 2 萬億 token 訓(xùn)練任務(wù)中，性能超越 Llama2-7B 實(shí)現(xiàn)了非凡的效率。
繼 Mamba 之后，又一敢于挑戰(zhàn) Transformer 的架構(gòu)誕生了！
來自 Meta、南加州大學(xué)（USC）、CMU 和 UCSD 的研究人員提出了全新的神經(jīng)網(wǎng)絡(luò)架構(gòu) ——Megalodon（巨齒鯊）。
這是專為有效處理「無限上下文」長(zhǎng)度的 LLM 預(yù)訓(xùn)練，以及推理而設(shè)計(jì)的架構(gòu)。
我們都知道，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)，就是為了解決無限處理上下文的難題。
同時(shí)，它可以同時(shí)實(shí)現(xiàn)高效訓(xùn)練（減少通信和計(jì)算量），以及高效推理（保持恒定的 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 開啟了計(jì)算效率和性能的新時(shí)代。
GPT-3 發(fā)布以來最大里程碑
網(wǎng)友表示，先是谷歌，又是 Meta，無限上下文離我們更進(jìn)一步，LLM 將會(huì)釋放出無限潛力。
還有人認(rèn)為「無限上下文長(zhǎng)度，絕對(duì)是游戲規(guī)則的改變者」！
更有甚者，初創(chuàng)公司 CEO 稱，「這是自 GPT-3 發(fā)布以來最大的里程碑，但卻沒有任何動(dòng)靜？！
Megalodon 就相當(dāng)于是 AGI 的基礎(chǔ)」。
「Meta 的 Megalodon 是一項(xiàng)突破性進(jìn)展，對(duì) AGI 具有重要意義。它的無限上下文長(zhǎng)度模擬了人類的認(rèn)知，實(shí)現(xiàn)了無縫任務(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)，但它卻無法直接應(yīng)用于 Transformer 的自回歸序列建模，因未來信息會(huì)通過時(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í)過程，提高訓(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 萬億 token 的大規(guī)模 LLM 預(yù)訓(xùn)練中。
此外，作者還在中 / 小參數(shù)規(guī)模的序列建?；鶞?zhǔn)上進(jìn)行了實(shí)驗(yàn)，包括 Long Range Arena (LRA) 、Speech Commands 上的原始語(yǔ)音分類、ImageNet-1K 上的圖像分類，以及 WikiText-103 和 PG19 上的語(yǔ)言建模。
結(jié)果顯示，在這些任務(wù)中，Megalodon 在各種數(shù)據(jù)模式下的表現(xiàn)明顯優(yōu)于所有最先進(jìn)的基線模型。
數(shù)據(jù)學(xué)習(xí)效率
通過訓(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)估
具體來說，研究人員在短上下文（4K token）的標(biāo)準(zhǔn)學(xué)術(shù)基準(zhǔn)上，對(duì) Megalodon 與 Llama 2，以及開源基礎(chǔ)模型進(jìn)行了比較。
在相同的 2 萬億 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ǔ)音分類、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í)代正確做科研的方法。
通過這個(gè)項(xiàng)目，研究者們也體會(huì)到了在大模型時(shí)代做新的模型架構(gòu)時(shí)要注意的問題?？偨Y(jié)來說：
對(duì)于兩個(gè)不同模型架構(gòu)的比較必須要在數(shù)據(jù)完全相同的條件下才有說服力。當(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é)果有說服力的前提是充分的訓(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í)代，新算法的開發(fā)已經(jīng)和系統(tǒng)等方面高度結(jié)合在一起。
參考資料：
https://arxiv.org/abs/2404.08801
https://zhuanlan.zhihu.com/p/692682649
廣告聲明：文內(nèi)含有的對(duì)外跳轉(zhuǎn)鏈接（包括不限于超鏈接、二維碼、口令等形式），用于傳遞更多信息，節(jié)省甄選時(shí)間，結(jié)果僅供參考，IT之家所有文章均包含本聲明。