Meta無限長文本大模型來了：參數(shù)僅7B，已開源
Transformers 的二次復(fù)雜度和弱長度外推限制了它們擴(kuò)展到長序列的能力，雖然存在線性注意力和狀態(tài)空間模型等次二次解決方案，但從以往的經(jīng)驗(yàn)來看，它們在預(yù)訓(xùn)練效率和下游任務(wù)準(zhǔn)確性方面表現(xiàn)不佳。
長文本是大語言模型一直在努力的方向。近日，谷歌提出的 Infini-Transformer 引入有效方法，可以將基于 Transformer 的大型語言模型 (LLM) 擴(kuò)展到無限長輸入，而不增加內(nèi)存和計算需求，吸引了人們的關(guān)注。
幾乎就在同時，Meta 也提出了一種無限長文本技術(shù)。
- 論文地址：https://arxiv.org/pdf/2404.08801.pdf
- 論文標(biāo)題：MEGALODON: Efficient LLM Pretraining and Inference with Unlimited Context Length
- 代碼：https://github.com/XuezheMax/megalodon
在 4 月 12 日提交的一篇論文中，來自 Meta、南加州大學(xué)、CMU、UCSD 等公司、機(jī)構(gòu)引入了 MEGALODON，一種用于高效序列建模的神經(jīng)架構(gòu)，上下文長度不受限制。
MEGALODON 繼承了 MEGA（帶有門控注意力的指數(shù)移動平均）的架構(gòu)，并進(jìn)一步引入了多種技術(shù)組件來提高其能力和穩(wěn)定性，包括復(fù)數(shù)指數(shù)移動平均（CEMA）、時間步歸一化層、歸一化注意力機(jī)制和具有兩個特征的預(yù)歸一化（pre-norm）殘差配置。
在與 LLAMA2 的直接比較中，MEGALODON 在 70 億參數(shù)和 2 萬億訓(xùn)練 token 的規(guī)模上取得了比 Transformer 更好的效率。MEGALODON 的訓(xùn)練損失達(dá)到 1.70，處于 LLAMA2-7B (1.75) 和 13B (1.67) 之間。MEGALODON 相對于 Transformers 的改進(jìn)在不同任務(wù)和模式的一系列基準(zhǔn)測試中表現(xiàn)強(qiáng)勁。
MEGALODON 本質(zhì)上是一種改進(jìn)的 MEGA 架構(gòu)（Ma et al., 2023），它利用了門控注意力機(jī)制和經(jīng)典指數(shù)移動平均（EMA）方法。為了進(jìn)一步提高 MEGALODON 在大規(guī)模長上下文預(yù)訓(xùn)練方面的能力和效率，作者提出了多種技術(shù)組件。首先，MEGALODON 引入了復(fù)數(shù)指數(shù)移動平均 (CEMA) 組件，將 MEGA 中的多維阻尼 EMA 擴(kuò)展到復(fù)數(shù)域。其次，MEGALODON 提出了時間步長歸一化層，它將組歸一化層推廣到自回歸序列建模任務(wù)，以允許沿順序維度進(jìn)行歸一化。
為了提高大規(guī)模預(yù)訓(xùn)練的穩(wěn)定性，MEGALODON 進(jìn)一步提出了歸一化注意力，以及通過修改廣泛采用的預(yù)歸一化和后歸一化方法而具有兩跳殘差配置的預(yù)歸一化。通過簡單地將輸入序列分塊為固定塊，如 MEGA-chunk 中所做的那樣，MEGALODON 在模型訓(xùn)練和推理中實(shí)現(xiàn)了線性計算和內(nèi)存復(fù)雜性。
在與 LLAMA2 的直接比較上，控制了數(shù)據(jù)和計算的同時，MEGALODON-7B 在訓(xùn)練困惑度方面顯著優(yōu)于用于訓(xùn)練 LLAMA2-7B 的最先進(jìn)的 Transformer 變體。在對長上下文建模的評估上，包括高達(dá) 2M 的各種上下文長度中的困惑度以及 Scrolls 中的長上下文 QA 任務(wù)證明了 MEGALODON 對無限長度序列進(jìn)行建模的能力。中小型基準(zhǔn)的其他實(shí)驗(yàn)結(jié)果，包括 LRA、ImageNet、Speech Commands、WikiText-103 和 PG19 證明了 MEGALODON 在體量和多模態(tài)上的能力。
方法介紹
首先，文章簡單回顧了 MEGA（ Moving Average Equipped Gated Attention ）架構(gòu)中的關(guān)鍵組件，并討論了 MEGA 中存在的問題。
MEGA 將 EMA（ exponential moving average ） 組件嵌入到注意力矩陣的計算中，以納入跨時間步長維度的歸納偏差。具體而言，多維阻尼 EMA 首先通過擴(kuò)展矩陣
將輸入序列 X 的每個維度單獨(dú)擴(kuò)展為 h 維，然后將阻尼 EMA 應(yīng)用于 h 維隱藏空間。形式如下：
為了降低全注意力機(jī)制中的二次復(fù)雜度，MEGA 簡單地將 (14-16) 中的查詢、鍵和值序列拆分為長度為 c 的塊。(17) 中的注意力單獨(dú)應(yīng)用于每個塊，產(chǎn)生線性復(fù)雜度 O (kc^2 ) = O (nc)。
從技術(shù)上講，MEGA 中的 EMA 子層有助于捕獲每個 token 附近的本地上下文信息，從而緩解了在超出塊邊界的上下文中丟失信息的問題。盡管 MEGA 取得了令人深刻的印象，但面臨如下問題：
i）由于 MEGA 中 EMA 子層的表達(dá)能力有限，具有塊級注意力的 MEGA 性能仍然落后于全注意力 MEGA。
ii) 對于不同的任務(wù)、數(shù)據(jù)類型，最終的 MEGA 架構(gòu)可能存在架構(gòu)差異，例如，不同的歸一化層、歸一化模式和注意力函數(shù) f (?) 。
iii) 沒有經(jīng)驗(yàn)證據(jù)表明 MEGA 可擴(kuò)展用于大規(guī)模預(yù)訓(xùn)練。
CEMA：將多維阻尼 EMA 擴(kuò)展到復(fù)數(shù)域
為了解決 MEGA 面臨的問題，該研究提出了 MEGALODON。
具體而言，他們創(chuàng)造性地提出了復(fù)指數(shù)移動平均 CEMA（ complex exponential moving average ），將上式（1）改寫為如下形式：
并將（2）中的 θ_j 參數(shù)化為：
時間步（Timestep）歸一化
盡管層歸一化與 Transformer 相結(jié)合的性能令人印象深刻，但很明顯，層歸一化不能直接減少沿空間維度（也稱為時間步長或序列維度）的內(nèi)部協(xié)變量偏移。
在 MEGALODON 中，該研究通過計算累積均值和方差將組歸一化擴(kuò)展到自回歸情況。
圖 2 說明了層歸一化和時間步標(biāo)準(zhǔn)化。
MEGALODON 中的歸一化注意力
此外，該研究還提出了專門為 MEGA 定制的歸一化注意力機(jī)制，以提高其穩(wěn)定性。形式如下：
則上式 (17) 中的注意力操作改為：
具有 Two-hop 殘差的預(yù)范數(shù)（Pre-Norm）
通過調(diào)查發(fā)現(xiàn)，擴(kuò)大模型大小會造成預(yù)歸一化不穩(wěn)定問題?；?Transformer 塊的預(yù)歸一化可以表示為（如圖 3 (b) 所示）：
在原始 MEGA 架構(gòu)中， 將 φ (19) 用于門控殘差連接 (21) 以緩解此問題。然而，更新門 φ 引入了更多的模型參數(shù)，當(dāng)模型規(guī)模擴(kuò)大到 70 億時，不穩(wěn)定問題仍然存在。MEGALODON 引入了一種名為 pre-norm 的新配置，具有 two-hop 殘差，它只是簡單地重新排列每個塊中的殘差連接，如圖 3（c）所示：
實(shí)驗(yàn)
為了評估 MEGALODON 在長上下文序列建模上的可擴(kuò)展性和效率，本文將 MEGALODON 擴(kuò)展到 70 億規(guī)模大小。
LLM 預(yù)訓(xùn)練
為了提高數(shù)據(jù)效率，研究者在訓(xùn)練過程中顯示了 MEGALODON-7B、LLAMA2-7B 和 LLAMA2-13B 的負(fù)對數(shù)似然 (NLL)，如圖 1 所示。
在相同數(shù)量的訓(xùn)練 token 下，MEGALODON-7B 獲得了比 LLAMA2-7B 明顯更好（更低）的 NLL，表現(xiàn)出更好的數(shù)據(jù)效率。
圖 4 說明了分別使用 4K 和 32K 上下文長度的 LLAMA2-7B 和 MEGALODON-7B 在每個設(shè)備上的平均 WPS（ word/token per second ）。對于 LLAMA2 模型，該研究使用 Flash-Attention V2 加速全注意力的計算。在 4K 上下文長度下，由于引入了 CEMA 和時間步歸一化，MEGALODON-7B 比 LLAMA2-7B 稍慢（約 6%）。當(dāng)將上下文長度擴(kuò)展到 32K 時，MEGALODON-7B 明顯比 LLAMA2-7B 快（約 32%），這證明了 MEGALODON 對于長上下文預(yù)訓(xùn)練的計算效率。
短上下文評估
表 1 總結(jié)了 MEGALODON 和 LLAMA2 在學(xué)術(shù)基準(zhǔn)上的結(jié)果，以及其他開源基礎(chǔ)模型，包括 MPT、RWKV 、Mamba 、 Mistral 和 Gemma 的比較結(jié)果。在相同的 2T token 上進(jìn)行預(yù)訓(xùn)練后，MEGALODON-7B 在所有基準(zhǔn)測試中均優(yōu)于 LLAMA2-7B。在某些任務(wù)上，MEGALODON-7B 的性能與 LLAMA2-13B 相當(dāng)甚至更好。
長上下文評估
圖 5 顯示了驗(yàn)證數(shù)據(jù)集在 4K 到 2M 各種上下文長度下的困惑度 (PPL)?？梢杂^察到 PPL 隨著上下文長度單調(diào)下降，驗(yàn)證了 MEGALODON 在建模極長序列方面的有效性和魯棒性。
指令微調(diào)
表 3 總結(jié)了 7B 模型在 MT-Bench 上的性能。與 Vicuna 相比，MEGALODON 在 MT-Bench 上表現(xiàn)出優(yōu)越的性能，并且與 LLAMA2-Chat 相當(dāng)，而后者利用 RLHF 進(jìn)行了進(jìn)一步的對齊微調(diào)。
中等規(guī)?；鶞?zhǔn)評估
為了評估 MEGALODON 在圖像分類任務(wù)上的性能，該研究在 Imagenet-1K 數(shù)據(jù)集上進(jìn)行了實(shí)驗(yàn)。表 4 報告了驗(yàn)證集上的 Top-1 準(zhǔn)確率。MEGALODON 的準(zhǔn)確率比 DeiT-B 提高了 1.3%，比 MEGA 提高了 0.8%。
表 5 說明了 MEGALODON 在 PG-19 上的字級困惑度 (PPL)，以及與之前最先進(jìn)的模型，包括 Compressive Transformer 、Perceiver AR、Perceiver AR、塊循環(huán) Transformer 和 MEGABYTE 等的對比。MEGALODON 性能明顯領(lǐng)先。
本文轉(zhuǎn)自 機(jī)器之心 ，作者：機(jī)器之心
