Research map / Architectures and efficiency / Mixture of experts
Routing Mamba: Scaling State Space Models with Mixture-of-Experts Projection
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes Routing Mamba (RoM), which turns Mamba's Conv, Gate and Out projections into sparse MoE experts that share one router, scaling SSMs and hybrid Samba models with fewer active parameters.
Problem
Naive MoE integration into Mamba's projection layers fails to improve and often degrades performance, so scaling SSM expressiveness with sparse experts remains unsolved.
Approach
Use Mamba's Conv, Gate and Out projections as sparse linear-projection experts governed by a single shared router whose top-K decisions are reused across these projections, with x Proj, dt Proj and 1D conv parameters shared across experts.
Key results
- RoM at 1.3B active (10B total) matches dense Mamba requiring up to 2.3x more active parameters in validation perplexity at 115M scale and 2.0x at 353M scale.
- Samba (expand=4) + RoM (Conv, Gate, Out) reaches 9.28 PPL at 4K context vs 9.78 for the dense Samba expand=4 model, with 23% FLOPS savings compared to dense Mamba scaling.
- RoM on Mamba with 353M active / 2.5B total improves PPL at 4K from 11.18 to 10.00.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Apply shared-routing RoM to Conv, Gate, and Output projections | Samba (expand=2) · SlimPajama validation, context length 4096 | perplexity: 10.05 -> 9.77 | Helped |
| Use independent MoE routing for Conv, Gate, and Output projections | Samba (expand=2) · SlimPajama validation, context length 4096 | perplexity: 10.05 -> 10.26 | Hurt |
| Apply RoM to Conv, Gate, and Output projections | Samba (expand=4) · SlimPajama validation, context length 4096 | perplexity: 9.78 -> 9.28 | Helped |
| Expertize the smaller dt and x projections rather than sharing their parameters (ablation) | RoM on Samba (expand=4) · SlimPajama validation, context length 16384 | perplexity: 8.85 -> 8.87 | Hurt |
| Add load-balance loss to RoM (Conv, Gate, Out) (ablation) | Samba (expand=4) · SlimPajama validation, context length 4096 | perplexity: 9.28 -> 9.30 | Hurt |
| Add load-balance loss to RoM (Conv, Gate, dt, x, Out) (ablation) | Samba (expand=4) · SlimPajama validation, context length 4096 | perplexity: 9.30 -> 9.26 | Helped |
| Add RoM to the Mamba projections | Mamba · SlimPajama validation, context length 4096 | perplexity: 11.18 -> 10.00 | Helped |
| Combine RoM with FFN-MoE (8top1) instead of FFN-MoE (16top1) | Samba 511M · Average downstream accuracy | accuracy: 50.2 -> 50.1 | No clear effect |
| Combine RoM with FFN-MoE (16top1) instead of FFN-MoE (32top1) | Samba 511M · Average downstream accuracy | accuracy: 49.5 -> 49.2 | Hurt |
| Apply RoM to Conv, Gate, and Output projections | Samba (expand=2) · Training throughput | tokens/s (× 10^5): 4.42 -> 3.57 | Hurt |
Surprising findings
- Naive independent MoE on Mamba projections (MoE-Mamba) generally degrades perplexity relative to the dense Samba baseline.
- Adding a load balance loss does not improve RoM and slightly hurts performance, suggesting routing balances naturally.
Limitations
- Optimal RoM configuration remains uncertain, per the authors.
- Applicability to broader SSM variants and self-attention or linear attention architectures is not established.
Experimental setup
Datasets: SlimPajama
Benchmarks: SlimPajama validation perplexity LAMBADA HellaSwag PIQA ARC-Easy ARC-Challenge WinoGrande
Metrics: perplexity accuracy FLOPS training throughput
Baselines: Mamba Samba MoE-Mamba SwitchHead Mixture of Attention Heads Llama-2 FFN-MoE Mamba2 Gated DeltaNet
Builds on: Mamba Samba Mixture of Experts MoE-Mamba Megablocks
Closest papers by meaning
- Block-Biased Mamba for Long-Range Sequence Processing (NeurIPS 2025 · similarity 0.91)
- Swimba: Switch Mamba Model Scales State Space Models (NeurIPS 2026 · similarity 0.89)
- The Mamba in the Llama: Distilling and Accelerating Hybrid Models (NeurIPS 2024 · similarity 0.89)
- LongMamba: Enhancing Mamba's Long-Context Capabilities via Training-Free Receptive Field Enlargement (ICLR 2025 · similarity 0.89)
- MambaLRP: Explaining Selective State Space Sequence Models (NeurIPS 2024 · similarity 0.88)
- Achilles' Heel of Mamba: Essential difficulties of the Mamba architecture demonstrated by synthetic data (NeurIPS 2025 · similarity 0.88)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.
