Mixture-of-Mamba: Hybrid SSM-MoE Architectures

• Mixture-of-Mamba is an architecture combining selective state-space modeling with sparse expert routing to deliver efficient, scalable performance.
• It reduces training FLOPs and memory overhead by replacing quadratic self-attention with near-linear operations blended with MoE components.
• The design is applied in time series forecasting, language modeling, and computer vision, achieving significant empirical gains in accuracy and throughput.

A Mixture-of-Mamba architecture refers to the integration of selective state-space models (notably Mamba layers) with Mixture-of-Experts (MoE) and/or hybrid architectural elements such as self-attention, convolution, and feed-forward layers. This approach is designed to combine the near-linear efficiency of Mamba's state-space modeling with the expressive capacity and dynamic specialization of sparse expert selection. Mixture-of-Mamba architectures have been applied to time series forecasting, language modeling, multi-modal pretraining, computer vision, and generative modeling, demonstrating significant reductions in training FLOPs, improved scalability, and strong empirical performance on both standard and long-context benchmarks.

1. Core Principles of Mamba and Its Integration in MoE Frameworks

Mamba is a selective state-space model (SSM) module that operates on input token sequences $x\in\mathbb{R}^{T\times d}$, achieving a complexity of $O(Td^2)$ per forward pass by eschewing full $T\times T$ attention. The layer executes two principal operations: (a) gated projections using $\mathrm{SiLU}$-activated $\mathrm{Conv1D}$ and $\mathrm{Linear}$ projections, and (b) a time-variant SSM recurrence where parameters $(\bar{A}_t, \bar{B}_t, C_t)$ are computed online via small selector networks from token-wise features. Discretization employs a matrix exponential solution and provides per-token update dynamics.

In mixture-based frameworks, Mamba layers can serve two complementary functions:

• As global dependency extractors (linear in sequence length), augmenting or replacing quadratic-cost self-attention.
• As components whose projections or recurrences are “expertized,” i.e., replaced by sparse sets of modality- or task-specific parameters, with token- or segmentwise expert selection via routers.

Practically, mixture-of-Mamba models deploy Mamba either in sequence with, or interleaved among, other modules (feed-forward, convolution, self-attention) and MoE MLP blocks to maximize both capacity and efficiency (Peng et al., 2024, Pióro et al., 2024, Anthony et al., 2024, Lieber et al., 2024).

2. Mixture-of-Mamba in Time Series Forecasting: Mixture of Universals (MoU)

The "Mixture of Universals" (MoU) (Peng et al., 2024) architecture demonstrates an explicit, sequential Mixture-of-Architectures (MoA) block for time-series forecasting. MoA decomposes sequence modeling into:

1. Mamba Layer: Captures time-variant, partial, and periodic dependencies with dynamic SSM recurrence.
2. Feed-Forward Network: Injects nonlinearity to enhance representation capacity.
3. Convolution Layer: Expands each token's receptive field to neighboring patches, bridging short-term and medium-range context.
4. Self-Attention: Globally integrates all tokens, but with only a single T^2 cost per block.

The MoA block is formulated strictly as a sequential stack, without learnable soft mixing. Short-term dependencies are addressed by a Mixture-of-Feature-Extractors (MoF) submodule, which adaptively routes patches to specialized short-term extractors and ensures fine-grained context is preserved. The complexity analysis confirms that the quadratic cost associated with self-attention is invoked only once per block, sharply reducing the overall compute versus Transformer stacks.

Empirical results on seven multivariate forecasting benchmarks show MoU achieving mean squared error (MSE) reductions of ∼20% over PatchTST and ModernTCN, and mean absolute error (MAE) reductions of ∼15%, winning 35/56 evaluated tasks. Visualization reveals that Mamba layers in this context focus on near-diagonal, periodic structure, while the global self-attention stage aggregates across the entire input sequence.

3. Mixture-of-Mamba for Sparse Capacity Scaling in LLMs

Sparse Mixture-of-Experts mechanisms interleaved with Mamba SSM enable parameter-efficient scaling beyond transformer or SSM-only models. Key instantiations include:

• MoE-Mamba (Pióro et al., 2024): Alternates Mamba SSM layers with Switch-style sparse MoE MLPs, using per-token routing. Each token activates only its assigned expert, yielding high total capacity with constant per-token compute. MoE-Mamba reduces the steps required to reach a target loss by up to 2.35× over dense Mamba, preserving SSM inference benefits.
• BlackMamba (Anthony et al., 2024): Alternates input-dependent Mamba SSMs with MoE MLPs using Sinkhorn top-1 routing. BlackMamba achieves up to 3–5× faster inference and ∼2× lower training FLOPs versus dense Transformer or Transformer-MoE baselines for comparable accuracy, maintaining load-balanced expert activation and constant memory irrespective of sequence length.
• Routing Mamba (RoM) (Zhan et al., 22 Jun 2025): "Expertizes" not only the dense MLPs, but also the major Mamba linear projections (input, output, gating) via per-token, top-K sparse routing. A single router and mask are shared across all expertized projections, yielding overhead-proportional (K/N) parameter and FLOP reductions. RoM matches perplexity of dense Mamba requiring 2.3× more active parameters and yields 23% FL0PS savings at the 1.3B parameter scale.
• Hybrid Mamba–Transformer Models: Jamba (Lieber et al., 2024) and Nemotron 3 Nano (NVIDIA et al., 23 Dec 2025) alternate or interleave SSM (Mamba) sublayers, grouped-query or vanilla attention layers, and sparse MoE MLPs. MoE blocks use top-K sparse gating (e.g., 2/16 in Jamba, 6/128 in Nemotron), with router balancing and tokenwise dispatch. These architectures achieve memory savings (KV cache up to 8× smaller), linear-scaling inference for very long context (up to 1M tokens (NVIDIA et al., 23 Dec 2025)), and throughput improvements of up to 3.3× compared to open baseline LLMs of similar size.

4. Mixture-of-Mamba for Multi-Modal and Specialized Domains

Mixture-of-Mamba extends to modality-aware, domain-specialized, or generative systems:

• Modality-aware Sparsity (Liang et al., 27 Jan 2025): All projection layers in Mamba (e.g., $W_{\mathrm{in}}, W_x, W_{dt}, W_{\mathrm{out}}$) are replaced with small, modality-specific blocks. This injects MoE-style sparsity without requiring a learned router: each token directly selects its modality-specific parameterization via a one-hot modality mask. Across three multi-modal pretraining settings, this design realizes comparable loss at 34–42% of reference FLOPs.
• Spatio-Temporal Mixture-of-Mamba (Chen et al., 17 Aug 2025): For complex spatio-temporal time-series, STM3 uses a Mixture-of-Experts (MMM) where each expert is a Multiscale Mamba unit specialized to capture distinct temporal dynamics. Routing is performed by node embeddings, yielding smooth and stable expert assignments, with additional causal contrastive loss to enforce pattern disentanglement across both experts and temporal scales. STM3 surpasses leading baselines in multi-step forecasting accuracy and efficiency.
• Vision and Restoration Applications: In image restoration, Mamba is fused with MoE and CNN experts via multi-stage routers (CLIP-guided (Wang et al., 9 Jun 2025), DA-CLIP (Wang et al., 16 Mar 2026)) to dynamically assign tasks or pixels to specialized restoration experts. These models achieve state-of-the-art results on composite and real-world image degradation tasks, balancing global (Mamba, SSM) and local (CNN) recovery.

5. Theoretical and Empirical Properties

Mixture-of-Mamba models demonstrate several notable architectural properties and empirical phenomena:

• Expressivity/Capacity: Injection of (sparse or modality-gated) MoE components allows much larger overall parameter counts ("total parameters") without increasing the per-token compute or "active parameters" at inference (Pióro et al., 2024). When MoE is applied to SSM internal projections (Zhan et al., 22 Jun 2025), capacity and trade-offs can be further optimized across tasks.
• Efficiency: When quadratic-cost attention is replaced or minimized (e.g., one per block in MoU (Peng et al., 2024), 1:7 ratio in Jamba (Lieber et al., 2024)), memory and FLOP scaling shifts toward $O(Td^2)$ (linear in sequence length), speeding up training/inference for tasks with long input contexts and shrinking KV cache footprint.
• Expert Assignment and Routing: Various strategies are used, including softmax-based top-K (Lieber et al., 2024, NVIDIA et al., 23 Dec 2025), Sinkhorn-normalized (Anthony et al., 2024), CLIP/DA-CLIP-guided (Wang et al., 9 Jun 2025, Wang et al., 16 Mar 2026), modality-determined (Liang et al., 27 Jan 2025), or node-embedding-based (Chen et al., 17 Aug 2025). Where gating is fixed (e.g., by modality), no additional balancing loss is necessary. For learned routers, load-balancing regularization or bias updates stabilize allocations.
• Pattern Disentanglement and Specialization: The inclusion of causal contrastive loss (Chen et al., 17 Aug 2025) or motif-level supervision sharpens expert specialization (e.g., scale or region), as visualized by t-SNE cluster separation in expert outputs.
• Empirical Performance: Mixture-of-Mamba systems routinely surpass dense SSM or Transformer baselines in convergence speed, memory efficiency, and downstream task score (e.g., ∼20%–30% lower error in forecasting (Peng et al., 2024), improved perplexity at $\frac{1}{2}$–$O(Td^2)$0 the compute in LMs (Pióro et al., 2024, Anthony et al., 2024, NVIDIA et al., 23 Dec 2025), +4–5% accuracy/AUC gain in high-resolution vision (Bayatmakou et al., 23 Jul 2025)).

6. Computation, Memory, and Scaling Law Observations

A recurring theme is that Mixture-of-Mamba architectures support scaling up total network capacity (number of experts, per-expert hidden size, etc.) while maintaining constant or modestly increasing active compute. For example, routing K out of N experts for each token yields a parameter cost reduction by a factor of $O(Td^2)$1, enabling practical models with 10–50B total parameters and $O(Td^2)$23.2B active per token on commodity hardware (NVIDIA et al., 23 Dec 2025, Lieber et al., 2024). Memory use for sequence modeling grows only with sequence length in Mamba, unlike $O(Td^2)$3 or $O(Td^2)$4-cache growth in Transformer-style attention.

Tables from benchmark studies consistently show ranking improvements, inference time reductions, and more robust error growth on long-term or long-context tasks.

Model/Domain	MoE type	Mamba Role	Routing Method	Main Efficiency Gain
MoU (Peng et al., 2024)	Sequential	SSM block	Fixed per-token	1 global attention; rest SSM
MoE-Mamba (Pióro et al., 2024)	Switch (MLP)	SSM layer, dense	Softmax top-1, balanced	2.35× speedup vs Mamba
Routing Mamba (Zhan et al., 22 Jun 2025)	Projected MoE	All major SSM projections	Shared tokenwise softmax	23% FLOP savings, 1/8 params
Jamba (Lieber et al., 2024)	Interleaved	SSM block	Top-2 over 16	3× throughput @ 256K ctx
STM3 (Chen et al., 17 Aug 2025)	Node-embedding	Multiscale SSM	Stable node routing	3–7% error reduction
M2Restore (Wang et al., 9 Jun 2025)	Pixel-wise	SSM+CNN expert	Prompt+CLIP+softmax	+0.5 dB avg PSNR, local adapt.
Nemotron 3 Nano (NVIDIA et al., 23 Dec 2025)	Top-6 of 128	SSM head, GQA attn	2-layer MLP softmax	3.3× throughput, 1M context

7. Outlook and Open Directions

Mixture-of-Mamba establishes a flexible family of architectures that unite efficient linear state-space modeling with sparse expert specialization. Open questions remain concerning:

• Scaling laws for SSM-MoE systems at very large data/model sizes (Anthony et al., 2024, NVIDIA et al., 23 Dec 2025);
• Theoretical foundations of expressivity versus sparsity under fixed parameter budgets (Zhan et al., 22 Jun 2025);
• Optimum ratios and placements of Mamba, attention, and MoE blocks in hybrid networks (Lieber et al., 2024, NVIDIA et al., 23 Dec 2025);
• Task-specific expert specialization and transfer across domains (e.g., self-supervised vision, audio, agentic reasoning) (Bayatmakou et al., 23 Jul 2025, Liang et al., 27 Jan 2025, NVIDIA et al., 23 Dec 2025).

A plausible implication is that Mixture-of-Mamba-style models will continue to supersede dense Transformers for domains with long-range dependencies and dynamic contextual structure under strict memory or latency budgets. Continued empirical and analytic investigation will be required to fully elucidate the inductive biases and practical limits of this design paradigm.