Swimba: Switch Mamba Model Scales State Space Models¶
会议: NeurIPS 2026
arXiv: 2603.06938
OpenReview: https://openreview.net/forum?id=izlpFIy1uX
论文: https://arxiv.org/abs/2603.06938
领域: 其他
关键词: 状态空间模型、混合专家、Mamba-2、参数空间混合、线性复杂度
一句话总结¶
针对传统 MoE 作用于选择性状态空间模型（SSM）时会导致核心循环状态更新计算量随专家数量成倍暴涨的瓶颈，本文提出 Switch Mamba（Swimba），在参数空间对专家生成的输入注入与读出流进行动态加权聚合，维持单条隐藏状态轨迹与单次循环求解，在推理 FLOPs 几乎不变的前提下实现了模型容量的高效扩展与全任务精度提升。
研究背景与动机¶
状态空间模型（SSM）近年来已发展成为长序列建模领域中替代标准自注意力机制的强有力方案。从早期基于 HiPPO 理论与对角加低秩约束的 S4，到引入内容自适应输入依赖机制的 Mamba，再到利用状态空间对偶性（SSD）将结构化半可分矩阵计算与张量核心矩阵乘法统一的 Mamba-2，线性时间复杂度的序列混合器展现出了比肩乃至超越 Transformer 的潜力。与此并行，大语言模型向百亿与千亿规模扩张的另一核心支柱是稀疏混合专家（MoE）技术，它通过条件激活部分子网络，在不显著增加每步前向浮点运算量（FLOPs）的前提下大幅跃升参数容量。然而，现有的 MoE 研究绝大多数局限于前馈网络（FFN/MLP）模块，将 MoE 直接引入序列混合器（Token Mixer）层仍然面临严重的计算重载障碍。
将 MoE 拓展至选择性 SSM 面临的根本矛盾在于：SSM 的核心计算开销高度集中于沿时序推进的隐藏状态循环更新。学术界以往常将“MoE-SSM”模糊指代两类在计算拓扑和扩展性上完全不同的设计：其一是“独立分离式 SSM 的 MoE”（MoE over separated SSMs），该方案为每个专家各自维护一条独立的状态轨迹，每步必须并行推进多个专家的递归更新，导致状态内存占用和核心循环计算直接与专家数（或激活专家数）成正比膨胀；其二是“参数空间混合的 SSM”（MoE-parameterized SSM），各专家仅作为参数映射分支，在递归之前动态融合成统一的有效 SSM，维持单条系统状态轨迹并仅执行一次循环扫描。现有工作在将 MoE 引入 Mamba 模块时往往偏向工程试凑，缺乏对二者在状态轨迹演化、计算复杂度界以及系统稳定性上的形式化定性与量化区分。
本文切入的角度是：状态空间模型的本质在于高维隐状态对时序历史的记忆与过滤，只要转移矩阵保持稳定收缩，专家的特异化表达完全可以通过时变的输入注入项与输出读出向量进行刻画，无须让循环更新本身发生分裂与倍增。核心 idea：提出参数空间混合的 Switch Mamba（Swimba）架构，通过共享状态转移矩阵并路由聚合各专家前向投影流，维持全局单一隐藏状态轨迹并仅执行单次 SSM 循环求解，在保证状态有界与模型表征力严格增益的同时，以不变的循环复杂度实现 SSM 参数容量的有效扩展。
方法详解¶
整体框架¶
Swimba 的核心设计建立在 Mamba-2 的状态空间对偶性（SSD）模块之上，将原架构中用于生成选择性 SSM 参数的单一大矩阵输入投影层替换为由路由调控的 MoE 模块。对于输入序列中的每一个 token 特征 \(X_t \in \mathbb{R}^P\)，Router 计算激活专家及其归一化权重；各个被激活的专家作为独立的线性投影分支，分别产生各自候选的输入注入流与输出读出流；随后系统在参数空间执行加权求和，聚合成当前时间步唯一的有效输入项 \(\tilde{U}_t\) 与有效读出向量 \(\tilde{C}_t\)；由于状态转移矩阵 \(A\) 在各专家及全时步间保持共享，底层的隐状态演化严格遵循单条轨迹推进，并直接复用 Mamba-2 硬件友好的块级（chunked）SSD 并行算法完成高效求解。
%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
A["输入 Token 特征<br/>X_t ∈ ℝ^P"] --> B["Token 级路由器 R(X_t)<br/>计算专家权重 π_t"]
A --> C["多专家线性输入投影<br/>生成候选流 (B_t^(e), C_t^(e), X_t^(e))"]
B --> D["参数空间混合<br/>加权聚合得到 U_t 与 C_t"]
C --> D
D --> E["单轨迹状态更新<br/>共享转移矩阵 A: h_t = A h_t-1 + U_t"]
E --> F["SSD 结构化读出与输出<br/>Y_t = (C_t)^T h_t"]关键设计¶
1. 参数空间混合与单轨迹循环：消除专家数倍增的循环开销
针对传统独立分离式 MoE-SSM 必须为各个专家维护独立隐状态 \(h_t^{(e)}\)、导致状态存储与循环扫描算力随专家数量 \(E\) 线性缩放（复杂度为 \(\mathcal{O}(T E \cdot \mathcal{C}_{\text{step}}(N, P))\)）的核心痛点，Swimba 严格贯彻“隐状态单一性”。设总专家数为 \(E\)，路由概率为 \(\pi_t = \text{softmax}(g_t)\)，激活集合为 \(\mathcal{K}_t\)（通常为 Top-\(k\) 稀疏激活）。每个专家 \(e\) 基于相同的输入特征 \(X_t\) 通过专家特异的线性投影生成候选流 \(B_t^{(e)} \in \mathbb{R}^{N \times P}\)、\(C_t^{(e)} \in \mathbb{R}^{N \times P}\) 与 \(X_t^{(e)} \in \mathbb{R}^P\)。Swimba 并非独立运行 \(E\) 次状态方程后再拼接输出，而是在参数空间对流进行显式加权聚合： $\(\tilde{U}_t = \sum_{e \in \mathcal{K}_t} \pi_{t,e} B_t^{(e)} X_t^{(e)}, \qquad \tilde{C}_t = \sum_{e \in \mathcal{K}_t} \pi_{t,e} C_t^{(e)}\)$ 整个系统的状态空间演化被精简为标准的单通道选择性 SSM： $\(h_t = A h_{t-1} + \tilde{U}_t, \qquad Y_t = \tilde{C}_t^\top h_t\)$ 该设计在理论上由定理 3.1（单 SSM 结构保持性）和定理 3.2（循环复杂度与专家数解耦）所保证：前向传播的总计算量被收敛至 \(\mathcal{O}(T \cdot \mathcal{C}_{\text{step}}(N, P) + T \cdot \mathcal{C}_{\text{mix}}(k, P))\)。当保持激活专家数 \(k\) 固定时，无论静态专家库 \(E\) 如何扩大，占主导地位的核心状态递归演化始终只执行一次，彻底消除了 MoE-SSM 的循环计算屏障。
2. 共享状态转移与解耦注入读出：保证系统收缩性与稳定性
频繁的时序动态路由切换容易引起隐状态轨迹的数值发散或梯度震荡。为解决这一潜在动力学不稳定性，Swimba 在参数空间解耦了“时序衰减转移”与“内容特征注入/读出”：令所有专家在所有时步共享同一对角化/标量恒等转移矩阵 \(A_t^{(e)} \equiv A\)。根据定理 3.3，若转移矩阵满足收缩性条件（在诱导矩阵范数下 \(\|A\| \le \rho < 1\)），且各专家投影层保证聚合后的注入流与读出流范数有界（\(\|\tilde{U}_t\| \le U\)、\(\|\tilde{C}_t\| \le C\)），则无论路由权重 \(\pi_t\) 如何随时间剧烈跳变，系统隐藏状态 \(h_t\) 及其输出 \(Y_t\) 均严格满足有界输入有界输出（BIBO）稳定性： $\(\|h_t\| \le \rho^t \|h_0\| + \frac{1 - \rho^t}{1 - \rho} U, \qquad \|Y_t\| \le C \left( \rho^t \|h_0\| + \frac{1 - \rho^t}{1 - \rho} U \right)\)$ 历史记忆以几何速率平滑衰减，路由机制对系统状态的调制完全局限在驱动项 \(\tilde{U}_t\) 的幅度之内。此外，定理 3.4 阐明了 Swimba 与独立分离式 MoE 的本质联系与偏差边界：当路由策略时间不变且读出权重一致时，二者在数学上完全等价；而在动态路由场景下，两者的输出误差 \(\|Y_t^{\text{sep}} - Y_t^{\text{mix}}\| \le C \sum_{e} \pi_{t,e} \|h_t^{(e)} - h_t\|\) 满足由单专家注入与混合注入差值驱动的稳定线性递归，证明了单轨迹近似的严密性与受控性。
3. 严格表征力增益与非多项式映射：以单次循环突破单专家表达极限
为了论证参数空间混合不仅是一项算力节省的工程技巧，更带来了实质性的模型容量跃升，定理 3.5 在单步特化场景（\(T=1, A=0, N=P=1\)）下对模型函数类进行了形式化对比。在仅采用单专家线性投影的层中，输出必为输入特征的多项式组合（阶数至多为 3 次），无法表达具有饱和特性的非线性边界；而在 Swimba 中，Softmax 路由器生成的权重与各专家线性流耦合后，能够精确构建 Logistic Sigmoid 等非多项式激活映射： $\(\pi_{1,2} = \frac{e^x}{1 + e^x} = \sigma(x)\)$ 这表明即使在只进行单次状态更新的强约束下，动态门控参数混合所能表征的函数集合依然严格真包含于任何单专家线性投影模型，从数学上解释了为什么在 FLOPs 几乎相同的情况下 Swimba 能显著提升语言建模的知识召回与推理能力。
损失函数 / 训练策略¶
Swimba 采用继续预训练/监督微调（SFT）两阶段策略融入主流大语言模型骨干： 1. 骨干初始化与参数冻结：基于 52 层的 Nemotron-H-8B 混合架构（全模型仅包含 4 层分散的自注意力层，其余层由 Mamba-2 序列混合器与 FFN 交替组成），直接复制原预训练权重，将每个 Mamba-2 层的输入投影替换为 4 专家 Swimba 模块（Top-1 激活）。在初始训练阶段冻结已复制的全部骨干参数，仅优化新引入的专家投影权重与 Router 参数。 2. 训练数据与目标：在 Llama Nemotron Post-Training Dataset 的 SFT 子集上进行教师强制（teacher-forced）自回归交叉熵训练，覆盖数学、代码、科学、指令遵循与安全等领域，采用标准语言建模负对数似然损失： $\(\mathcal{L}_{\text{SFT}} = - \sum_{t=1}^T \log P(x_t \mid x_{<t})\)$
实验关键数据¶
主实验¶
在评估基准上，Swimba-14B 与基线 Nemotron-H-8B 在 LM-Evaluation-Harness 评测套件下的 9 项主流基准上进行了端到端对比。Swimba-14B 采用 4 专家且每 token 激活 1 个专家的稀疏配置，其实际前向激活参数量与 Nemotron-H-8B 完全对齐。
| 数据集 | 任务类型 | 指标 | Nemotron-H-8B (8B) | Swimba-14B (Top-1) | 提升幅度 |
|---|---|---|---|---|---|
| Arc-Challenge | 小学科学推理 | Acc / Norm-Acc | 56.5±1.4 / 60.3±1.4 | 59.5±1.4 / 63.5±1.4 | +3.0 / +3.2 |
| Arc-Easy | 基础常识问答 | Acc / Norm-Acc | 84.0±0.8 / 83.7±0.8 | 84.5±0.8 / 84.1±0.8 | +0.5 / +0.4 |
| BoolQ | 自然是非问答 | Acc | 85.5±0.6 | 84.9±0.6 | -0.6 |
| HellaSwag | 常识推理延续 | Acc / Norm-Acc | 61.9±0.5 / 81.2±0.4 | 64.2±0.5 / 84.9±0.4 | +2.3 / +3.7 |
| MMLU | 综合学科知识 (57科) | Acc | 71.7±0.4 | 75.0±0.4 | +3.3 |
| OpenBookQA | 开卷科学问答 | Acc / Norm-Acc | 35.4±2.1 / 47.4±2.2 | 34.9±2.1 / 46.7±2.2 | -0.5 / -0.7 |
| PIQA | 物理常识推理 | Acc / Norm-Acc | 81.7±0.9 / 82.3±0.9 | 82.0±0.9 / 82.5±0.9 | +0.3 / +0.2 |
| RTE | 文本蕴含识别 | Acc | 71.8±2.7 | 73.2±2.7 | +1.4 |
| WinoGrande | 指代消歧推理 | Acc | 76.7±1.2 | 79.1±1.2 | +2.4 |
| 平均成绩 (Average) | 全任务无加权平均 | Acc / Norm-Acc | 69.47 / 70.98 | 70.81 / 72.34 | +1.34 / +1.36 |
消融与计算效率分析¶
为了验证 Swimba 在理论计算量与真实推理引擎中的性能表现，论文从算子级理论 FLOPs 和 vLLM 端到端实测两个维度给出了量化数据。
| 模型架构 | 激活专家机制 | 隐状态轨迹数 | 每 Token 推理 FLOPs | FLOPs 相对增量 | vLLM 解码延迟开销 (A100) |
|---|---|---|---|---|---|
| Nemotron-H-8B (Dense) | N/A (单专家基线) | 1 | \(1.51263 \times 10^{10}\) | 基准 (0.00%) | 基准 (1.00×) |
| Swimba-14B (Ours) | MoE 参数空间混合 (Top-1) | 1 | \(1.51282 \times 10^{10}\) | +0.012% (<0.2%) | 约 1.10× - 1.15× (仅慢约 10%) |
| MoE 分离式 SSM (理论推导) | 独立状态分支 (Top-1) | 1 / 每激活专家 | \(\approx 1.513 \times 10^{10}\) | 类似 | 状态分配与同步延迟较高 |
| MoE 分离式 SSM (理论推导) | 独立状态分支 (Dense/E=4) | 4 | \(\approx 2.4 \times 10^{10}\) | > +50% | 显存与循环算力爆炸性增长 |
关键发现¶
- 显著提升学科知识与复杂推理能力：在 MMLU（+3.3%）、Arc-Challenge（+3.0%）和 HellaSwag（+2.3%）等高度依赖知识储备和逻辑辨别能力的基准上，Swimba 实现了稳健的性能跨越，证明了参数空间专家特异化成功放大了模型容量。
- FLOPs 理论几乎完全对齐：PyTorch Profiler 实测表明 Swimba-14B 的单 token 前向运算量为 \(1.51282 \times 10^{10}\) FLOPs，相比稠密基线仅增加 0.012%，理论上完美兑现了“等算力扩展容量”的设计承诺。
- 真实系统延迟存在小幅折损：在单张 40GB A100 GPU 上利用 vLLM 进行端到端解码基准测试时，Swimba-14B 展现出了约 10% 的吞吐下降与延迟增加。该现象主要源自 GPU 上细粒度 Token 路由算子及内核启动（Kernel Launch）的固有开销，而非序列混合器底层的循环扫描瓶颈。
亮点与洞察¶
- 解构 MoE-SSM 的二元设计空间：首次在理论层面厘清了“独立分离式 SSM 的 MoE”与“参数空间混合的 SSM”的本质分野，指出隐藏状态轨迹的分裂是算力失控的根源，为后续 SSM 架构规模化指明了正确的演进范式。
- 巧用参数空间求和保留单次 SSD 求解：通过将专家权重加权操作前置到 \((B_t, C_t, X_t)\) 参数流，使 Mamba-2 高度优化的高性能半可分矩阵变换（SSD Chunked Scan）得以完全原样复用，实现零框架改造。
- 完备的理论稳定性与表征能力背书：不仅给出了收缩矩阵下的 BIBO 稳定性严格上限，还通过单步构造证明了 Softmax 门控参数空间混合能够突破多项式表征边界，从数学理论上筑牢了方法根基。
局限与展望¶
- 作者承认的局限：实测服务延迟受制于通用推理框架中的动态路由开销，在当前 vLLM 实现下未能完全达到 FLOPs 所预测的零损失吞吐；此外，骨干模型采用预训练后参数冻结的轻量微调方案，尚未从零启动超大规模全量预训练。
- 实验维度的潜在不足：主要评测基于 Top-1 激活策略与 4 专家配置，对于更激进的 Top-2 组合混合（Dense/Sparse 多专家融合）以及 8/16 专家的极限扩展性尚未给出细致消融；长上下文极限测试（如 64k/128k 针检索）尚待补全。
- 未来改进思路：开发针对 SSM 输入投影的定制化 Fused MoE 算子，消除路由器与参数切片的多次显存往返；探索将参数空间混合扩展至转移矩阵 \(A\) 的对角分量，在保障李雅普诺夫稳定性的前提下进一步解放动力学表达能力。
相关工作与启发¶
- vs MoE-Mamba / Jamba / Nemotron-H: 此类架构采用宏观层级（Macro-level）混合，仅在外部将标准 Mamba 序列混合器与 MoE-FFN 模块进行交替堆叠，SSM 内部依然是稠密且单容量的；Swimba 则是微观层级（Micro-level）的序列混合器内部 MoE，直接对状态空间的生成方程实施专家特异化。
- vs Routing Mamba (RoM): RoM 同样尝试在 Mamba 周围引入 MoE 投影，但偏向启发式工程堆叠，未能形式化区分独立轨迹与参数混合两种流派；Swimba 则建立了完整的数学框架（稳定性有界性、结构保持性及误差线性递归界），提供了坚实的理论指导。
评分¶
- 新颖性: ⭐⭐⭐⭐⭐ [清晰界定两类 MoE-SSM 范式，首次提出并理论论证了参数空间混合的单轨迹 Mamba-2 扩展架构]
- 实验充分度: ⭐⭐⭐⭐☆ [覆盖 9 大权威基准、算子 FLOPs 分析与 vLLM 端到端实测，但缺乏从零预训练的大规模消融]
- 写作质量: ⭐⭐⭐⭐⭐ [动机叙述逻辑严密，数学符号规范，定理证明完备且直击核心矛盾]
- 价值: ⭐⭐⭐⭐☆ [为 Mamba/SSM 家族的大模型化与容量缩放提供了极具前景的实用技术路线]