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

Google之後，Meta 也來卷無限長上下文。

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

近日，Google提出的Infini-Transformer引入了有效方法，可以將基於Transformer的大型語言模型（LLM）擴展到無限長輸入，而不增加儲存和運算需求，吸引了人們的注意。

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

• 論文網址：https://arxiv.org/pdf/2404.08801.pdf

• 論文標題：MEGALODON: Efficient LLM Pretraining and Inference with Unlimited Context Length

##########################################################################在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 性能明顯領先。

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

以上是Meta無限長文本大模型來了：參數僅7B，已開源的詳細內容。更多資訊請關注PHP中文網其他相關文章！