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(L2) 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.
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., Bt(e),Ct(e),Xt(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±1.4
60.3±1.4
59.5±1.4
63.5±1.4
Arc-Easy
84.0±0.8
83.7±0.8
84.5±0.8
84.1±0.8
BoolQ
85.5±0.6
—
84.9±0.6
—
Hellaswag
61.9±0.5
81.2±0.4
64.2±0.5
84.9±0.4
MMLU
71.7±0.4
—
75.0±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×1010
Swimba-14B
1.51282×1010
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 (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.
Ramprasath Ganesaraja, Sahil Dilip Panse, Swathika N
State-space models (SSMs) have emerged as an efficient strategy for building powerful language models, avoiding the quadratic complexity of computing attention in transformers. Despite their promise, the interpretability and steerability of modern SSMs remain relatively underexplored. We take a major step in this direction by identifying activation subspace bottlenecks in the Mamba family of SSM models using tools from mechanistic interpretability. We then introduce a test-time steering intervention that simply multiplies the activations of the identified bottlenecks by a scalar. Across 7 SSMs and 6 diverse benchmarks, this intervention improves performance by an average of 8.27%, without requiring any task-specific tuning. Finally, we validate that the identified bottlenecks are indeed hindering performance by modifying them to yield an architecture we call Stable-Mamba, which achieves long-context performance gains when retrained from scratch.
Vamshi Sunku Mohan, Kaustubh Gupta, Aneesha Das +1
AppViewX, Bangalore · Independent Researcher · Microsoft Research, Redmond
Scaling inference-time compute has emerged as an important driver of LLM performance, making inference efficiency a central focus of model design alongside model quality. While the current Transformer-based models deliver strong model quality, their quadratic compute and linear memory make inference expensive. This has spurred the development of sub-quadratic models with reduced linear compute and constant memory requirements. However, many recent linear models trade off model quality and capability for algorithmic efficiency, failing on tasks such as state tracking. Moreover, their theoretically linear inference remains hardware-inefficient in practice. Guided by an inference-first perspective, we introduce three core methodological improvements inspired by the state space model (SSM) viewpoint of linear models. We combine: (1) a more expressive recurrence derived from SSM discretization, (2) a complex-valued state update rule that enables richer state tracking, and (3) a multi-input, multi-output (MIMO) formulation for better model performance without increasing decode latency. Together with architectural refinements, our Mamba-3 model achieves significant gains across retrieval, state-tracking, and downstream language modeling tasks. At the 1.5B scale, Mamba-3 improves average downstream accuracy by 0.6 percentage points compared to the next best model (Gated DeltaNet), with Mamba-3's MIMO variant further improving accuracy by another 1.2 points for a total 1.8 point gain. Across state-size experiments, Mamba-3 achieves comparable perplexity to Mamba-2 despite using half of its predecessor's state size. Our evaluations demonstrate Mamba-3's ability to advance the performance-efficiency Pareto frontier.
Aakash Lahoti, Kevin Y. Li, Berlin Chen +5
Carnegie Mellon University · Princeton University · Together AI +1