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

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

因此，北京大學林宙辰、徐鑫提出了一種全新混合序列建模架構 ——MixCon，它爲解決這些難題帶來了創新性的方案。經實驗驗證，其性能遠超 Mixtral、Mamba 和 Jamba。論文已在 European Conference on Artificial Intelligence (ECAI) 2024 上發表。

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

線性 RNN 模型

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

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

MoE 模型

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

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

二、MixCon 的核心架構與技術

Conba 模型架構

1. 狀態空間方程

3. 實施細節

神經網絡近似：

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 及其他領域的應用開闢了新道路。未來，我們期待它在更多領域發揮更大的作用，爲技術發展帶來更多的突破和創新。