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

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

• 論文標題：MixCon: A Hybrid Architecture for Efficient and Adaptive Sequence Modeling
• 論文地址：https://zhouchenlin.github.io/Publications/2024-ECAI-MixCon.pdf

一、現(xiàn)有序列建模模型的困境

線性注意力 Transformer

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

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

線性 RNN 模型

線性 RNN 模型如 Mamba 等通過將序列表示為狀態(tài)空間并利用掃描操作，以線性時間復雜度提供了序列建模的新解決方案。

然而，它們可能缺乏復雜序列建模任務所需的適應性和動態(tài)特性，并且像傳統(tǒng)序列模型一樣，缺少反饋機制和自適應控制。

MoE 模型

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

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

二、MixCon 的核心架構與技術

Conba 模型架構

1. 狀態(tài)空間方程

Conba 將序列建模任務表示為狀態(tài)空間系統(tǒng)，狀態(tài)空間定義為和，其中、和分別為時間步的狀態(tài)、輸入和輸出，和是非線性函數(shù)，可由神經(jīng)網(wǎng)絡近似。

，其中和是可學習參數(shù)矩陣。，是可學習參數(shù)矩陣。

為處理長序列，Conba 采用選擇性狀態(tài)空間機制，以及引入延遲狀態(tài)和動態(tài)狀態(tài)縮放機制。

最后狀態(tài)空間系統(tǒng)表示為來捕捉長程依賴和適應序列動態(tài)變化。

2. 自適應控制機制

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

控制增益矩陣通過更新，其中是跟蹤誤差向量的 2 范數(shù)，是學習率。

3. 實施細節(jié)

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

MixCon 模型架構

MixCon 是結合注意力機制的 Transformer 層、Conba 層和 MoE 組件的創(chuàng)新混合解碼器架構。

在內(nèi)存使用方面，通過平衡注意力和 Conba 層，相比 Mamba 可將 KV 緩存減少 32 倍。例如，在 256K 令牌上下文環(huán)境中，MixCon 仍能保持較小的 KV 緩存優(yōu)勢（如表 1 所示）。

在吞吐量方面，處理長序列時，Conba 層計算效率更高，增加其比例可提高整體吞吐量。

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

對于 Conba 層實施，采用 RMSNorm 等技術，模型詞匯量為 256K，使用 BPE 進行訓練，每個數(shù)字為單獨令牌。

模型架構圖如下所示：

三、MixCon 的實驗與評估

實施細節(jié)

選擇特定配置適應單塊 80GB A800 NVIDIA GPU 的計算能力，實現(xiàn)質量和吞吐量的優(yōu)化。

序列由 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 所示）。

數(shù)據(jù)集評估

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

MixCon 性能與類似或更大規(guī)模的先進公開模型相當或更優(yōu)，盡管總參數(shù)比 Llama - 2 少，但作為稀疏模型，其活躍參數(shù)僅 5B，處理長序列時 KV 緩存僅需 2GB，而 Mixtral 需 32GB（如表 2 所示）。

消融實驗

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

以 HellaSwag（10 - shot）、WinoGrande（5 - shot）、Natural Questions（NQ，5 - shot）為指標，MixCon 表現(xiàn)穩(wěn)健（如表 3 所示），MixCon（無 MoE）訓練過程損失更低（如圖 6 所示）。

長上下文評估

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

在 NarrativeQA、LongFQA、Natural Questions（NQ）、CUAD 等數(shù)據(jù)集上評估，MixCon 在多數(shù)數(shù)據(jù)集上優(yōu)于 Mixtral 和 Jamba，平均性能優(yōu)越，且在長上下文任務中具有更好的吞吐量（如表 4 所示）。

結合注意力和 Conba 的優(yōu)勢及混合專家的影響

1. 注意力和 Conba 比例研究

用 13 億參數(shù)模型在 2500 億令牌上訓練，MixCon 性能優(yōu)于純注意力或純 Mamba，注意力和 Conba 層比例為 2:6 或 1:7 時性能差異?。ㄈ绫?5 所示）。

2. 混合專家的影響

當在 MixCon 架構的大規(guī)模情境（5B 參數(shù)，在 50B 令牌上訓練）中應用 MoE 技術時，性能有顯著提升（如表 6 所示）。

四、MixCon 的優(yōu)勢與展望

MixCon 作為創(chuàng)新的混合序列建模架構，通過整合多種技術，在處理復雜動態(tài)序列時具有高效的計算效率，在各項任務中展現(xiàn)出顯著優(yōu)勢，能高效處理長序列、內(nèi)存使用低且吞吐量高，具有高可擴展性和實用性。然而，它仍有改進空間，如進一步優(yōu)化狀態(tài)空間表示、長序列的自適應控制、特定領域微調以及訓練算法等。

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