让状态空间模型用专家路由提升容量,同时不增加核心计算负担。
Swimba: Switch Mamba Model Scales State Space Models
- 通过参数空间专家路由,仅维护单一状态轨迹,避免重复计算。
- 在相同算力下,性能略优于基线,延迟和吞吐量小幅下降。
- 适合追求高效扩展的序列建模任务,如长文本生成与推理。
混合专家(MoE)是提升模型参数容量的常用方法,但将其应用于状态空间模型(SSM)的令牌混合器会显著增加递归状态更新的成本。本文研究如何在保持计算效率的前提下,将专家专业化引入选择性SSM。提出两种设计:(1) 分离式SSM上的MoE,需维护多条状态轨迹,计算量随专家数增长;(2) 参数空间嵌入的MoE-SSM,仅维持单一状态轨迹,递归仅执行一次。本文提出的Switch Mamba(Swimba)采用第二种设计,通过路由机制在专家生成的SSM流间切换。理论上,建立了参数空间MoE-SSM的良定义性和稳定性,并厘清了两种设计的关系。实验表明,在标准基准任务上,匹配FLOPs条件下,Swimba性能略优,实时吞吐量和延迟仅有轻微下降。结果表明,参数空间MoE可在固定主循环成本下有效扩展SSM容量。
原文摘要 · Abstract (English)
Mixture-of-experts (MoE) is a common approach for increasing parameter capacity, but applying MoE to state space model (SSM) token mixers can multiply the cost of the recurrent state update. We study how to introduce expert specialization into selective SSMs while preserving computational efficiency. We show that MoE--SSM can refer to two designs: (1) MoE over separated SSMs, which maintains multiple state trajectories and thus scales compute with the number of experts; and (2) MoE-parameterized SSM, which mixes experts in parameter space, maintains a single state trajectory, and evaluates the recurrence once. Our method, Switch Mamba (Swimba), follows the second design by routing over expert-produced SSM streams. Theoretically, we establish well-definedness and stability for MoE-parameterized SSMs and characterize the relationship between the two designs. Empirically, we evaluate Swimba on standard benchmark tasks and measure real-time throughput and latency. Under matched FLOPs, Swimba achieves slightly better average performance than the baseline, with a small slowdown in real-time latency and throughput. Overall, these results suggest that parameter-space MoE can increase SSM capacity while keeping the dominant recurrence cost fixed.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。