谷歌之後，Meta 也來卷無限長上下文。
Transformers 的二次複雜度和弱長度外推限制了它們擴充套件到長序列的能力，雖然存線上性注意力和狀態空間模型等次二次解決方案，但從以往的經驗來看，它們在預訓練效率和下游任務準確性方面表現不佳。
長文字是大語言模型一直在努力的方向。近日，谷歌提出的 Infini-Transformer 引入有效方法，可以將基於 Transformer 的大型語言模型 (LLM) 擴充套件到無限長輸入，而不增加記憶體和計算需求，吸引了人們的關注。
幾乎就在同時，Meta 也提出了一種無限長文字技術。
論文地址：https://arxiv.org/pdf/2404.08801.pdf
論文標題：MEGALODON: Efficient LLM Pretraining and Inference with Unlimited Context Length
程式碼：https://github.com/XuezheMax/megalodon
在 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 效能明顯領先。
更詳細內容請參考論文原文。