Meta無限長文本大模型來了：參數僅7B，已開源

機器之心報道

編輯：澤南、陳萍

Transformers 的二次複雜度和弱長度外推限制了它們擴展到長序列的能力，雖然存在線性注意力和狀態空間模型等次二次解決方案，但從以往的經驗來看，它們在預訓練效率和下游任務準確性方面表現不佳。

長文本是大語言模型一直在努力的方向。近日，谷歌提出的 Infini-Transformer 引入有效方法，可以將基於 Transformer 的大型語言模型 (LLM) 擴展到無限長輸入，而不增加內存和計算需求，吸引了人們的關注。

幾乎就在同時，Meta 也提出了一種無限長文本技術。

在 4 月 12 日提交的一篇論文中，來自 Meta、南加州大學、CMU、UCSD 等公司、機構引入了 MEGALODON，一種用於高效序列建模的神經架構，上下文長度不受限制。

MEGALODON 繼承了 MEGA（帶有門控注意力的指數移動平均）的架構，並進一步引入了多種技術組件來提高其能力和穩定性，包括複數指數移動平均（CEMA）、時間步歸一化層、歸一化注意力機制和具有兩個特徵的預歸一化（pre-norm）殘差配置。

在與 LLAMA2 的直接比較中，MEGALODON 在 70 億參數和 2 萬億訓練 token 的規模上取得了比 Transformer 更好的效率。MEGALODON 的訓練損失達到 1.70，處於 LLAMA2-7B (1.75) 和 13B (1.67) 之間。MEGALODON 相對於 Transformers 的改進在不同任務和模式的一系列基準測試中表現強勁。

MEGALODON 本質上是一種改進的 MEGA 架構（Ma et al., 2023），它利用了門控注意力機制和經典指數移動平均（EMA）方法。爲了進一步提高 MEGALODON 在大規模長上下文預訓練方面的能力和效率，作者提出了多種技術組件。首先，MEGALODON 引入了複數指數移動平均 (CEMA) 組件，將 MEGA 中的多維阻尼 EMA 擴展到複數域。其次，MEGALODON 提出了時間步長歸一化層，它將組歸一化層推廣到自迴歸序列建模任務，以允許沿順序維度進行歸一化。

爲了提高大規模預訓練的穩定性，MEGALODON 進一步提出了歸一化注意力，以及通過修改廣泛採用的預歸一化和後歸一化方法而具有兩跳殘差配置的預歸一化。通過簡單地將輸入序列分塊爲固定塊，如 MEGA-chunk 中所做的那樣，MEGALODON 在模型訓練和推理中實現了線性計算和內存複雜性。

在與 LLAMA2 的直接比較上，控制了數據和計算的同時，MEGALODON-7B 在訓練困惑度方面顯著優於用於訓練 LLAMA2-7B 的最先進的 Transformer 變體。在對長上下文建模的評估上，包括高達 2M 的各種上下文長度中的困惑度以及 Scrolls 中的長上下文 QA 任務證明了 MEGALODON 對無限長度序列進行建模的能力。中小型基準的其他實驗結果，包括 LRA、ImageNet、Speech Commands、WikiText-103 和 PG19 證明了 MEGALODON 在體量和多模態上的能力。

方法介紹

首先，文章簡單回顧了 MEGA（ Moving Average Equipped Gated Attention ）架構中的關鍵組件，並討論了 MEGA 中存在的問題。

MEGA 將 EMA（ exponential moving average ） 組件嵌入到注意力矩陣的計算中，以納入跨時間步長維度的歸納偏差。具體而言，多維阻尼 EMA 首先通過擴展矩陣 將輸入序列 X 的每個維度單獨擴展爲 h 維，然後將阻尼 EMA 應用於 h 維隱藏空間。形式如下：

爲了降低全注意力機制中的二次複雜度，MEGA 簡單地將 (14-16) 中的查詢、鍵和值序列拆分爲長度爲 c 的塊。(17) 中的注意力單獨應用於每個塊，產生線性複雜度 O (kc^2 ) = O (nc)。

從技術上講，MEGA 中的 EMA 子層有助於捕獲每個 token 附近的本地上下文信息，從而緩解了在超出塊邊界的上下文中丟失信息的問題。儘管 MEGA 取得了令人深刻的印象，但面臨如下問題：

i）由於 MEGA 中 EMA 子層的表達能力有限，具有塊級注意力的 MEGA 性能仍然落後於全注意力 MEGA。

ii) 對於不同的任務、數據類型，最終的 MEGA 架構可能存在架構差異，例如，不同的歸一化層、歸一化模式和注意力函數 f (・) 。

iii) 沒有經驗證據表明 MEGA 可擴展用於大規模預訓練。

CEMA：將多維阻尼 EMA 擴展到複數域

爲了解決 MEGA 面臨的問題，該研究提出了 MEGALODON。

具體而言，他們創造性地提出了復指數移動平均 CEMA（ complex exponential moving average ），將上式（1）改寫爲如下形式：

並將（2）中的 θ_j 參數化爲：

時間步（Timestep）歸一化

儘管層歸一化與 Transformer 相結合的性能令人印象深刻，但很明顯，層歸一化不能直接減少沿空間維度（也稱爲時間步長或序列維度）的內部協變量偏移。

在 MEGALODON 中，該研究通過計算累積均值和方差將組歸一化擴展到自迴歸情況。

圖 2 說明了層歸一化和時間步標準化。

MEGALODON 中的歸一化注意力

此外，該研究還提出了專門爲 MEGA 定製的歸一化注意力機制，以提高其穩定性。形式如下：

則上式 (17) 中的注意力操作改爲：

具有 Two-hop 殘差的預範數（Pre-Norm）

通過調查發現，擴大模型大小會造成預歸一化不穩定問題。基於 Transformer 塊的預歸一化可以表示爲（如圖 3 (b) 所示）：

在原始 MEGA 架構中， 將 φ (19) 用於門控殘差連接 (21) 以緩解此問題。然而，更新門 φ 引入了更多的模型參數，當模型規模擴大到 70 億時，不穩定問題仍然存在。MEGALODON 引入了一種名爲 pre-norm 的新配置，具有 two-hop 殘差，它只是簡單地重新排列每個塊中的殘差連接，如圖 3（c）所示：

實驗

爲了評估 MEGALODON 在長上下文序列建模上的可擴展性和效率，本文將 MEGALODON 擴展到 70 億規模大小。

LLM 預訓練

爲了提高數據效率，研究者在訓練過程中顯示了 MEGALODON-7B、LLAMA2-7B 和 LLAMA2-13B 的負對數似然 (NLL)，如圖 1 所示。

在相同數量的訓練 token 下，MEGALODON-7B 獲得了比 LLAMA2-7B 明顯更好（更低）的 NLL，表現出更好的數據效率。

圖 4 說明了分別使用 4K 和 32K 上下文長度的 LLAMA2-7B 和 MEGALODON-7B 在每個設備上的平均 WPS（ word/token per second ）。對於 LLAMA2 模型，該研究使用 Flash-Attention V2 加速全注意力的計算。在 4K 上下文長度下，由於引入了 CEMA 和時間步歸一化，MEGALODON-7B 比 LLAMA2-7B 稍慢（約 6%）。當將上下文長度擴展到 32K 時，MEGALODON-7B 明顯比 LLAMA2-7B 快（約 32%），這證明了 MEGALODON 對於長上下文預訓練的計算效率。

短上下文評估

表 1 總結了 MEGALODON 和 LLAMA2 在學術基準上的結果，以及其他開源基礎模型，包括 MPT、RWKV 、Mamba 、 Mistral 和 Gemma 的比較結果。在相同的 2T token 上進行預訓練後，MEGALODON-7B 在所有基準測試中均優於 LLAMA2-7B。在某些任務上，MEGALODON-7B 的性能與 LLAMA2-13B 相當甚至更好。

長上下文評估

圖 5 顯示了驗證數據集在 4K 到 2M 各種上下文長度下的困惑度 (PPL)。可以觀察到 PPL 隨着上下文長度單調下降，驗證了 MEGALODON 在建模極長序列方面的有效性和魯棒性。

指令微調

表 3 總結了 7B 模型在 MT-Bench 上的性能。與 Vicuna 相比，MEGALODON 在 MT-Bench 上表現出優越的性能，並且與 LLAMA2-Chat 相當，而後者利用 RLHF 進行了進一步的對齊微調。

中等規模基準評估

爲了評估 MEGALODON 在圖像分類任務上的性能，該研究在 Imagenet-1K 數據集上進行了實驗。表 4 報告了驗證集上的 Top-1 準確率。MEGALODON 的準確率比 DeiT-B 提高了 1.3%，比 MEGA 提高了 0.8%。

表 5 說明了 MEGALODON 在 PG-19 上的字級困惑度 (PPL)，以及與之前最先進的模型，包括 Compressive Transformer 、Perceiver AR、Perceiver AR、塊循環 Transformer 和 MEGABYTE 等的對比。MEGALODON 性能明顯領先。

更詳細內容請參考論文原文。