Swimba: Switch Mamba Model Scales State Space Models

Zhixu Du; Krishna Teja Chitty-Venkata; Murali Emani; Venkatram Vishwanath; Hai Helen Li; Yiran Chen

cs.LGMar 6, 2026

Swimba: Switch Mamba Model Scales State Space Models

Authors: Zhixu Du, Krishna Teja Chitty-Venkata, Murali Emani, Venkatram Vishwanath, Hai Helen Li, Yiran Chen

Organizations: Duke University · Red Hat, Inc. · Argonne National Laboratory

Abstract

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.

Figures & tables

Figure 1: Compute scaling view for token mixing. The x-axis is sequence length

L

and the y-axis is per-token computational overhead. Attention scales as

O(L^{2})

while SSMs scale as

O(L)

. Bubble area indicates parameter count, illustrating that MoE can increase parameters with only a small increase in per-token compute, motivating Swimba (MoE-SSM) designs.

Panel of Figure 2 — Figure 2: Two MoE–SSM designs with different state and compute scaling. (a) MoE-parameterized SSM (parameter-space mixing) keeps a single hidden-state trajectory and mixes expert-produced streams before a single recurrence evaluation. (b) MoE over separated SSMs maintains one state trajectory per expert and combines expert outputs, which requires advancing multiple recurrences resulting computation scales with number of experts.

Figure 3: Swimba layer: an MoE-parameterized SSM token mixer built from a Mamba-2 block. Each expert applies its own input projections to produce candidate token-dependent streams (e.g.,

B_{t}^{(e)},C_{t}^{(e)},X_{t}^{(e)}

). At each layer, only a subset of experts is selected per token. The selected experts are combined by a weighted sum to form a single stream for the SSM computation.

	Nemotron-H-8B		Swimba-14B
Task	Accuracy	Normalized accuracy	Accuracy	Normalized accuracy
Arc-Challenge	$56.5_{\pm 1.4}$	$60.3_{\pm 1.4}$	$\mathbf{59.5}_{\pm 1.4}$	$\mathbf{63.5}_{\pm 1.4}$
Arc-Easy	$84.0_{\pm 0.8}$	$83.7_{\pm 0.8}$	$\mathbf{84.5}_{\pm 0.8}$	$\mathbf{84.1}_{\pm 0.8}$
BoolQ	$\mathbf{85.5}_{\pm 0.6}$	—	$84.9_{\pm 0.6}$	—
Hellaswag	$61.9_{\pm 0.5}$	$81.2_{\pm 0.4}$	$\mathbf{64.2}_{\pm 0.5}$	$\mathbf{84.9}_{\pm 0.4}$
MMLU	$71.7_{\pm 0.4}$	—	$\mathbf{75.0}_{\pm 0.4}$	—

Table 1: R sesults on standard benchmarks. Subscripts report standard deviation. Nemotron-H-8B and Swimba-14B are compared at similar FLOPs per token. The last row reports the unweighted average across tasks (and across normalized accuracies where provided). Swimba-14B improves on the baseline on most tasks and on the average score.

Model	FLOPs / token
Nemotron-H-8B	$1.51263\times 10^{10}$
Swimba-14B	$1.51282\times 10^{10}$

Table 2: Per-token inference FLOPs. Swimba-14B has essentially the same FLOPs per token as Nemotron-H-8B, indicating that using a single selected expert per Swimba layer does not materially change the dominant inference compute.

Figure 4: Decoding throughput on vLLM versus batch size for multiple input-output sequence lengths, with the inset reporting the throughput ratio (Swimba/Nemotron). Swimba-14B remains close to the Nemotron-H-8B baseline in throughput with a slight drop.

Figure 5: Latency on vLLM versus input-output sequence length for multiple batch sizes, with the inset reporting the ratio (Swimba/Nemotron). Latency trends are similar across sequence lengths, demonstrating slightly worse end-to-end cost between Swimba-14B and Nemotron-H-8B.

State Space Models

State space models represent sequences through evolving latent states, enabling efficient temporal modeling, system identification, and forecasting.

See all 295 papers

Ternary Mamba: Grouped Quantization-Aware Training of W1.58A16 State Space Models
Jun 16, 2026Ramprasath Ganesaraja, Sahil Dilip Panse, Swathika NLLM Quantization Ternary Quantization
EdgeVerve Systems Limited
Interpreting and Steering State-Space Models via Activation Subspace Bottlenecks
Feb 26, 2026Vamshi Sunku Mohan, Kaustubh Gupta, Aneesha Das +1State Space Language Models State Space Models
AppViewX, Bangalore · Independent Researcher · Microsoft Research, Redmond
Mamba-3: Improved Sequence Modeling using State Space Principles
Mar 16, 2026Aakash Lahoti, Kevin Y. Li, Berlin Chen +5State Space Language Models Sequence Modeling
Carnegie Mellon University · Princeton University · Together AI +1

Jun 16, 2026J/K move · Enter open · S save

Ternary Mamba: Grouped Quantization-Aware Training of W1.58A16 State Space Models

Ramprasath Ganesaraja, Sahil Dilip Panse, Swathika N

EdgeVerve Systems Limited

LLM Quantization Ternary Quantization LLM Compression Quantization-Aware Training Neural Network Quantization State Space Models

Open paper AskPDF arXiv

State Space Models (SSMs) such as Mamba-2 offer linear-time inference but their memory footprint limits edge deployment. Prior ternary SSM work (Slender-Mamba) trains from scratch on 150B tokens; we show a pretrained checkpoint suffices, reducing the marginal token budget by 1,000x. Using grouped quantization-aware training (QAT) with knowledge distillation from a frozen FP16 teacher, we compress Mamba-2 1.3B to 3.61x (2,687 to 744 MB) and achieve 48.1% zero-shot accuracy (7-task average) in just 102M tokens (4 GPU-hours, single H100) -- approaching Bi-Mamba's 48.4% (within +/-0.9pp CI). This QAT-from-pretrained setting reveals zero-ratio collapse, a novel instability caused by learnable quantization scales that does not arise in from-scratch training. We further show that post-hoc correction strategies effective for Transformers fail for SSMs due to error accumulation through the recurrence. These results demonstrate that ternary SSMs do not require expensive from-scratch training: QAT from pretrained checkpoints with KD is a data-efficient alternative.

Abstract

State Space Models

Explore similar work

Ternary Mamba: Grouped Quantization-Aware Training of W1.58A16 State Space Models

Interpreting and Steering State-Space Models via Activation Subspace Bottlenecks

Mamba-3: Improved Sequence Modeling using State Space Principles