Swimba：开关Mamba模型扩展状态空间模型
机器学习
2026-03-10 v1
摘要
混合专家(MoE)是增加参数容量的常见方法，但将MoE应用于状态空间模型(SSM)token混合器会增加递归状态更新的成本。我们研究如何在保持计算效率的同时引入专家专业化到选择性SSM中。我们表明，MoE-SSM可以指两种设计：(1)在分离的SSM上进行MoE，维持多个状态轨迹，因此计算量随专家数量而扩展;和(2)MoE参数化SSM，将专家在参数空间中混合，维持单个状态轨迹，并仅评估一次递归。我们的方法，Switch Mamba(Swimba)，遵循第二种设计，通过路由到专家产生的SSM流来实现。理论上，我们为MoE参数化SSM建立了明确性和稳定性，并刻画了两种设计之间的关系。实验上，我们在标准基准任务上评估了Swimba，衡量了实际的吞吐量和延迟。在匹配的FLOPs下，Swimba的平均性能略高于基线，尽管在实际延迟和吞吐量上有轻微减慢。总体而言，这些结果表明，参数空间的MoE可以在保持主要递归成本固定的同时增加SSM的容量。
引用
@article{arxiv.2603.06938,
title = {Swimba: Switch Mamba Model Scales State Space Models},
author = {Zhixu Du and Krishna Teja Chitty-Venkata and Murali Emani and Venkatram Vishwanath and Hai Helen Li and Yiran Chen},
journal= {arXiv preprint arXiv:2603.06938},
year = {2026}
}