北大林宙辰團隊全新混合序列建模架構MixCon：性能遠超Mamba
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 及其他領域的應用開闢了新道路。未來，我們期待它在更多領域發揮更大的作用，為技術發展帶來更多的突破和創新。
