北大林宙辰團隊全新混合序列建模架構MixCon：效能遠超Mamba

机器之心發表於2024-10-15
架構

圖片

AIxiv專欄是機器之心釋出學術、技術內容的欄目。過去數年，機器之心AIxiv專欄接收報導了2000多篇內容，覆蓋全球各大高校與企業的頂級實驗室，有效促進了學術交流與傳播。如果您有優秀的工作想要分享，歡迎投稿或者聯絡報導。投稿郵箱：liyazhou@jiqizhixin.com；zhaoyunfeng@jiqizhixin.com


在自然語言處理、語音識別和時間序列分析等眾多領域中，序列建模是一項至關重要的任務。然而，現有的模型在捕捉長程依賴關係和高效建模序列方面仍面臨諸多挑戰。

因此，北京大學林宙辰、徐鑫提出了一種全新混合序列建模架構 ——MixCon，它為解決這些難題帶來了創新性的方案。經實驗驗證，其效能遠超 Mixtral、Mamba 和 Jamba。論文已在 European Conference on Artificial Intelligence (ECAI) 2024 上發表。
圖片
  • 論文標題：MixCon: A Hybrid Architecture for Efficient and Adaptive Sequence Modeling
  • 論文地址：https://zhouchenlin.github.io/Publications/2024-ECAI-MixCon.pdf

一、現有序列建模模型的困境

線性注意力 Transformer

線性注意力 Transformer 旨在透過近似注意力機制來提高原始 Transformer 模型的效率，將計算複雜度從圖片降低到圖片或圖片，但在處理長序列時可能會面臨效能下降和計算開銷增加的問題。

例如，早期利用區域性敏感雜湊方案雖降低複雜度，但引入大常數因子；近期透過改變計算順序等方法近似 Softmax 函式，但仍存在效能不如 Softmax 注意力且可能增加額外開銷的情況。

線性 RNN 模型

線性 RNN 模型如 Mamba 等透過將序列表示為狀態空間並利用掃描操作，以線性時間複雜度提供了序列建模的新解決方案。

然而，它們可能缺乏複雜序列建模任務所需的適應性和動態特性，並且像傳統序列模型一樣，缺少反饋機制和自適應控制。

MoE 模型

MoE 模型透過結合專家模組，能有效處理長序列並保持計算效率，根據輸入資料自適應選擇專家模組。

但 MoE 模型的專家模組稀疏啟用可能導致訓練穩定性問題，部分引數不常使用降低引數效率，在處理長序列時可能在計算效率和訓練穩定性方面面臨挑戰，且對動態變化適應性不足。

二、MixCon 的核心架構與技術

Conba 模型架構

1. 狀態空間方程

Conba 將序列建模任務表示為狀態空間系統，狀態空間定義為圖片和圖片，其中圖片、圖片和圖片分別為時間步圖片的狀態、輸入和輸出，圖片和圖片是非線性函式，可由神經網路近似。

圖片，其中圖片和圖片是可學習引數矩陣。圖片，圖片是可學習引數矩陣。

為處理長序列，Conba 採用選擇性狀態空間機制圖片，以及引入延遲狀態圖片和動態狀態縮放機制圖片。

最後狀態空間系統表示為圖片來捕捉長程依賴和適應序列動態變化。

2. 自適應控制機制

設計目標是最小化實際輸出圖片和期望輸出圖片之間的跟蹤誤差圖片。

控制增益矩陣圖片透過圖片更新，其中圖片是跟蹤誤差向量圖片的 2 範數，圖片是學習率。

3. 實施細節

神經網路近似：

利用 MLP 近似狀態轉移函式圖片的非線性部分和觀察函式圖片，透過端到端訓練確定最佳引數。

選擇性狀態空間函式圖片用 SwiGLU（基於 Swish 和 GLU 的混合啟用函式）近似，其公式為圖片。

進一步擴充套件和增強穩健性：引入過程噪聲圖片和觀察噪聲圖片， 圖片，增強模型對擾動的適應性。

4. 模型架構圖如下所示：
圖片
MixCon 模型架構

MixCon 是結合注意力機制的 Transformer 層、Conba 層和 MoE 元件的創新混合解碼器架構。

在記憶體使用方面，透過平衡注意力和 Conba 層，相比 Mamba 可將 KV 快取減少 32 倍。例如，在 256K 令牌上下文環境中，MixCon 仍能保持較小的 KV 快取優勢（如表 1 所示）。
圖片
在吞吐量方面，處理長序列時，Conba 層計算效率更高，增加其比例可提高整體吞吐量。

基本配置單位是 MixCon 塊，由 Conba 或注意力層組合而成，每個層包含注意力模組或 Conba 模組，後接 MLP 或 MoE 層。MixCon 中的 MLP 層被 MoE 層替換，以增加模型容量同時保持較低計算負載。

對於 Conba 層實施，採用 RMSNorm 等技術，模型詞彙量為 256K，使用 BPE 進行訓練，每個數字為單獨令牌。

模型架構圖如下所示：
圖片
三、MixCon 的實驗與評估

實施細節

選擇特定配置適應單塊 80GB A800 NVIDIA GPU 的計算能力，實現質量和吞吐量的最佳化。

序列由 4 個 MixCon 塊組成，每個 MixCon 塊含 8 層 L = 8，注意力層和 Conba 層比例為 2:6 (a:c = 2:6)，每隔一層 (e = 2) 用 MoE 替換 MLP 模組，模型有 16 個專家 (n = 16)，每個令牌使用 2 個頂級專家 (K = 2)。

上下文長度分析

MixCon 在單塊 80GB A800 GPU 上的最大上下文長度是 Jamba 的兩倍、Mixtral 的四倍、Llama - 2 - 70B 的十四倍（如圖 3 所示）。
圖片
吞吐量分析

1. 配置一：考慮不同批大小，在單塊 A800 80GB GPU（int8 量化）、8K 上下文長度下生成 512 個輸出令牌，MixCon 吞吐量是 Mixtral 的三倍、Jamba 的兩倍（如圖 4 所示）。
圖片
2. 配置二：單批次（批大小 = 1）、四塊 A800 GPUs（無量化）、不同上下文長度下生成 512 個輸出令牌，處理 128K 令牌時，MixCon 吞吐量是 Jamba 的 1.5 倍、Mixtral 的 4.5 倍（如圖 5 所示）。
圖片
資料集評估

本文在一系列標準學術基準測試中評估 Conba 效能，包括常識推理任務（如 HellaSwag、WinoGrande、ARC - E、ARC - Challenge）、閱讀理解任務（如 BoolQ、QuAC）、聚合基準測試（如 MMLU、BBH），採用不同的學習策略。

MixCon 效能與類似或更大規模的先進公開模型相當或更優，儘管總引數比 Llama - 2 少，但作為稀疏模型，其活躍引數僅 5B，處理長序列時 KV 快取僅需 2GB，而 Mixtral 需 32GB（如表 2 所示）。
圖片
消融實驗

展示注意力和 Conba 層結合的優勢及最佳比例和交織技術。純 Conba 模型在上下文學習有困難，Attention - Conba 混合模型有類似純 Transformer 模型的上下文學習能力。

以 HellaSwag（10 - shot）、WinoGrande（5 - shot）、Natural Questions（NQ，5 - shot）為指標，MixCon 表現穩健（如表 3 所示），MixCon（無 MoE）訓練過程損失更低（如圖 6 所示）。
圖片
長上下文評估

利用問答基準測試評估 MixCon 處理長上下文能力，使用 L - Eval 中最長上下文資料集的五個資料集，以少樣本格式（每個實驗用三個例子）進行實驗。

在 NarrativeQA、LongFQA、Natural Questions（NQ）、CUAD 等資料集上評估，MixCon 在多數資料集上優於 Mixtral 和 Jamba，平均效能優越，且在長上下文任務中具有更好的吞吐量（如表 4 所示）。
圖片
結合注意力和 Conba 的優勢及混合專家的影響

1. 注意力和 Conba 比例研究

用 13 億引數模型在 2500 億令牌上訓練，MixCon 效能優於純注意力或純 Mamba，注意力和 Conba 層比例為 2:6 或 1:7 時效能差異小（如表 5 所示）。
圖片
2. 混合專家的影響

當在 MixCon 架構的大規模情境（5B 引數，在 50B 令牌上訓練）中應用 MoE 技術時，效能有顯著提升（如表 6 所示）。
圖片
四、MixCon 的優勢與展望

MixCon 作為創新的混合序列建模架構，透過整合多種技術，在處理複雜動態序列時具有高效的計算效率，在各項任務中展現出顯著優勢，能高效處理長序列、記憶體使用低且吞吐量高，具有高可擴充套件性和實用性。然而，它仍有改進空間，如進一步最佳化狀態空間表示、長序列的自適應控制、特定領域微調以及訓練演算法等。

總體而言，MixCon 為序列建模提供了新的解決方案，在複雜序列處理方面表現出色，為 NLP 及其他領域的應用開闢了新道路。未來，我們期待它在更多領域發揮更大的作用，為技術發展帶來更多的突破和創新。

相關文章