MoE-Mamba: Efficient Selective State Space Models with Mixture of Experts

🎯 Pitch

Mamba state space models with Mixture of Experts layers reach the same performance as vanilla Mamba in 2.35× fewer training steps—without sacrificing Mamba’s linear-time inference advantage. This SSM–MoE combo also consistently beats a Transformer–MoE baseline at matched active parameters, and the gains grow with more experts.

1. Executive Summary

This paper introduces MoE-Mamba, a model that combines the Mamba state space model with a Switch Transformer Mixture of Experts layer by interleaving unconditional Mamba blocks with conditional MoE feed-forward layers. Trained on the C4 dataset at two scales (27M and 121M active parameters), MoE-Mamba achieves the same performance as vanilla Mamba in 2.35× fewer training steps while preserving Mamba's linear-time inference benefits — for instance, MoE-Mamba with 117M active parameters reaches a log perplexity of 2.81 compared to 2.99 for the 121M-parameter Mamba baseline, and also consistently outperforms a Transformer-MoE counterpart with matched active parameters. The paper establishes that SSM-MoE integration yields robust gains that grow with the number of experts and training duration, though the advantage is smaller at the 25M scale and depends critically on allocating sufficient capacity to the MoE layers, with a 3:3 ratio of Mamba-to-MoE active parameters emerging as the sweet spot before routing overhead degrades practicality.

2. Context and Motivation

The Core Problem: SSMs Don't Yet Scale Like Transformers

The fundamental question this paper tackles is a practical one: can we make State Space Models scale as efficiently as Transformers for large language model training? This matters because the field faces an ongoing tension between architectural innovation and practical scalability. Transformers have become dominant not just because they perform well, but because there exists a mature ecosystem of techniques for scaling them — Mixture of Experts being chief among them — that enable training trillion-parameter models with manageable computational budgets. State Space Models, particularly the recent Mamba architecture, have shown remarkable promise as alternatives to attention-based models, but they lack this scaling infrastructure entirely.

The gap is concrete and measurable. Mamba achieves linear-time inference (unlike the quadratic self-attention mechanism) and strong performance on language modeling benchmarks, but its scaling characteristics are unknown beyond the moderate sizes explored in the original work. As the paper notes in Section 5, scaling is believed to be a critical factor in developing powerful AI systems, citing Sutton's "Bitter Lesson" (2019). If SSMs cannot be scaled efficiently — or if their scaling properties fundamentally differ from Transformers — then their architectural advantages (faster inference, better long-context processing) remain academic curiosities rather than practical benefits for production LLMs.

Why This Problem Is Important: The Dual Bottleneck of LLM Scaling

The paper's motivation rests on two independent concerns that happen to align around MoE as a solution:

First, the inference bottleneck of Transformers. Self-attention scales quadratically in sequence length. This means that as LLMs are deployed on longer contexts (making sense of entire documents, reasoning over extended conversations, processing code repositories), Transformers become increasingly expensive at inference time. Mamba's linear-time scaling with context length addresses this directly — but only if the model itself can be trained to competitive quality in the first place. The paper positions MoE as the mechanism that bridges this gap: MoE allows scaling model capacity (total parameters) without proportionally scaling computational cost (active FLOPs per token), making it feasible to train Mamba models large enough to exploit their inference advantage.

Second, the capacity-efficiency tradeoff of dense models. A dense 121M-parameter Mamba (the paper's □100M scale) achieves a log perplexity of 2.99. Making it larger would improve performance, but at linearly increasing training cost. The MoE approach — as demonstrated in Transformer-MoE systems like GLaM (Du et al., 2022), Switch Transformers (Fedus et al., 2022), and Mixtral (Jiang et al., 2024) — decouples total model capacity from per-token computation by activating only a subset of parameters for each input. A model can have billions of total parameters while using only a fraction per token, dramatically improving the FLOPs-to-performance ratio. The paper asks whether this same decoupling works for Mamba, or whether the SSM architecture interacts with MoE in ways that break the analogy to Transformers.

Where Prior Approaches Fall Short

The paper identifies three specific categories of prior work, each with a distinct limitation:

1. Dense SSMs (Mamba and predecessors) have no scaling mechanism beyond model width and depth. Gu & Dao (2023) demonstrated that Mamba can reach competitive performance with Transformers at matched parameter counts, but their largest model was 2.8B parameters — far below the frontier models that Transformers support via MoE scaling to hundreds of billions of parameters. The paper's implicit argument is that all successful LLM architectures will eventually need some form of conditional computation to reach frontier scale, and SSMs have not yet developed this capability. As Section 3.2 notes, vanilla Mamba does not even use feed-forward layers, which are the standard target for MoE insertion in Transformers. This absence creates a non-obvious design question: where does MoE go in an architecture that lacks the component it typically replaces?

2. Transformer-MoE systems assume an attention backbone. The extensive literature on MoE for language models — Shazeer et al. (2017), Lepikhin et al. (2020), Fedus et al. (2022), Du et al. (2022), Jiang et al. (2024) — integrates MoE exclusively into Transformer architectures, typically by converting some fraction of feed-forward layers into sparsely-activated expert layers. These methods rely on the Transformer's decomposition of each layer into attention (unconditional, sequence-mixing) and feed-forward (per-token, position-wise) components. The paper explicitly references this decomposition in Section 3.2 when describing their interleaving strategy as analogous to alternating vanilla and MoE feed-forward layers in Transformers. However, transferring this approach to Mamba is not straightforward — Mamba's layer structure combines sequence mixing (the SSM) and per-token projections differently than a Transformer's clean separation, raising questions about whether the same conditional processing benefits will materialize.

3. Prior SSM-MoE integration was contemporaneous and limited. The paper acknowledges (Section 2, final paragraph) that Anthony et al. (2024) presented a similar architecture (BlackMamba) concurrently with the first version of this work, which validates the general direction. However, the paper distinguishes its contribution by running "extensive ablations on the model architecture, number of experts, and other design choices" that BlackMamba did not explore. BlackMamba focused on downstream task performance; this paper focuses on understanding why and how SSM-MoE integration works through systematic variation of the ratio of Mamba-to-MoE parameters, parallel versus sequential placement, the number of experts, and alternative MoE insertion points within the Mamba block itself. This analytical depth is what the paper positions as its unique contribution relative to both the SSM and MoE literatures.

Conflicting Signals in Prior SSM Design

An important subtext in the paper's motivation comes from a counterintuitive finding in the original Mamba work: Gu & Dao (2023) reported that interleaving Mamba layers with standard feed-forward layers — which would be the dense equivalent of the MoE layers this paper proposes — actually degraded performance relative to pure Mamba without FF layers. This is cited in Section 4.5:

"This is consistent with Gu & Dao (2023) reporting that Mamba interleaved with feed-forward layers (which corresponds to a single-expert MoE layer) is worse than vanilla Mamba."

This creates a puzzle: if adding a single fixed feed-forward layer hurts performance, why would adding a bunch of them (as multiple experts) help? The resolution the paper implicitly relies on is that MoE introduces conditional computation — different tokens get processed by different experts, allowing the model to specialize its processing in ways a single FF layer cannot. A single FF layer adds parameters without adding representational flexibility; multiple experts add both parameters and the capacity to route different inputs to different sub-networks. The paper validates this interpretation in Figure 4: MoE-Mamba with a single expert underperforms vanilla Mamba, but performance monotonically improves as the number of experts increases from 1 to 32. The benefit comes from the sparsity and specialization of experts, not merely from adding FF computation.

This is a subtle but crucial point: the paper is not arguing that Mamba "needs" feed-forward layers. It is arguing that Mamba benefits from conditional, specialized feed-forward computation that a dense FF layer — which applies the same transformation to every token — cannot provide. This reframes MoE integration as introducing a new capability axis (adaptive per-token processing) rather than merely increasing capacity.

How This Paper Positions Itself

The paper positions itself as filling the intersection of two previously disconnected research threads:

From the SSM side: It extends Mamba's architectural innovations (selective mechanism, hardware-aware scan) by adding the scaling technique (MoE) that Transformers have used to reach extreme scales, without which SSMs risk hitting a practical ceiling well below Transformer-based systems.
From the MoE side: It demonstrates that conditional computation benefits are not specific to attention-based architectures — they transfer to SSMs with comparable or greater gains, validating MoE as a general scaling technique rather than a Transformer-specific optimization.

The paper's experimental design reflects this positioning. The key comparison is not just MoE-Mamba versus Mamba (which would show that MoE helps SSMs), but also MoE-Mamba versus Transformer-MoE (to show that SSM + MoE is competitive with or better than Transformer + MoE). Table 1 makes this explicit: MoE-Mamba100M achieves 2.81 log perplexity versus 2.88 for Transformer-MoE100M, with matched active parameters (117M vs. 114M). This is the paper's strongest claim: not just that MoE works with Mamba, but that Mamba is the better backbone for MoE scaling.

The paper also explicitly admits its limitations in scope. The experiments top out at 2.4B total parameters and 30B training tokens on C4, which is modest by LLM standards. Section 5 frames this as a necessary first step — establishing proof of concept and exploring the design space — while advocating for future work on scaling laws, distillation, and integration with attention mechanisms. The tone is not "we have built the next Llama" but rather "we have discovered a promising architecture and thoroughly characterized its behavior at small scale, creating a roadmap for much larger experiments."

The Unexplored Design Space

A final motivation, evident from the paper's extensive ablation studies, is that the design space of SSM-MoE integration is essentially terra incognita. Where should MoE layers be placed? What proportion of parameters should they control? Can MoE be embedded within the Mamba block rather than added alongside it? The paper systematically explores these questions through:

Sequential vs. parallel placement (Section 3.3 and 4.4): placing the MoE layer after Mamba (standard MoE-Mamba) versus running both in parallel and summing their outputs, finding the sequential variant consistently superior.
Ratio of Mamba-to-MoE active parameters (Section 4.3): sweeping from 1:5 to 5:1, discovering that 3:3 is the practical sweet spot where Mamba's sequence processing gets sufficient capacity while MoE's specialized processing yields diminishing returns beyond this point.
Inside-block MoE (Section 4.4): replacing the Conv, Gate, and Output projections within the Mamba block itself with MoE layers, finding that none of the seven possible combinations outperform the simpler interleaved design, but several approach or slightly exceed vanilla Mamba.

This systematic exploration of a previously empty design space is a defining feature of the paper. It treats SSM-MoE integration not as a single novel architecture to propose, but as a territory to survey. The 2.35× training speedup is the headline, but the architectural insights (where MoE does and doesn't help, how capacity should be allocated, why one-expert MoE fails) are arguably the deeper contribution.

3. Technical Approach

3.1 Reader Orientation (Approachable Technical Breakdown)

The system is a decoder-only language model for next-token prediction that replaces the standard Transformer block with an alternating stack of Mamba state space model layers and Switch Transformer Mixture of Experts (MoE) layers. It solves the problem of scaling state space models to higher capacity without proportionally increasing per-token computation cost: the Mamba layers handle sequence mixing with linear-time complexity, while the MoE layers provide massive conditional parameter capacity by routing each token to a specialized expert sub-network, yielding the same training loss as a dense Mamba model in roughly half the number of training steps while preserving Mamba's inference efficiency advantages over Transformers.

3.2 Big-Picture Architecture (Diagram in Words)

The architecture is a stack of identical blocks, each containing two sequential sub-layers:

Input: a sequence of token embeddings (vectors in $\mathbb{R}^{d_{\text{model}}}$ ) produced by a standard embedding lookup. The sequence length is 1024 tokens in all experiments, and $d_{\text{model}}$ is 512 for $\square$ 25M models and 768 for $\square$ 100M models.
Mamba layer (unconditional, per-token): Every token passes through a Mamba layer independently, which applies a selective state space mechanism to integrate information from the token's prefix while maintaining a compressed hidden state. This layer is "unconditional" in the sense that all tokens receive the same transformation — there is no routing or token-dependent parameter selection. It outputs a vector of the same dimensionality $d_{\text{model}}$ that is added to the residual stream.
MoE layer (conditional, per-token): Each token embedding is fed into a linear router that produces scores over $N_{\text{experts}}$ expert networks. A softmax converts these scores to probabilities, and the highest-scoring expert is selected to process the token (Switch routing with $k=1$ ). The selected expert — a feed-forward network with size $d_{\text{expert}} = 3d_{\text{model}}$ — applies a token-specific transformation. Tokens assigned to the same expert in a batch that exceed the expert's capacity are dropped (capacity factor = 1). A load-balancing auxiliary loss with weight $\alpha = 0.01$ is added to the training objective to encourage uniform expert utilization. The output is scaled by the router's softmax probability and added to the residual stream.
Repeated stacking: These two-layer blocks (Mamba → MoE) are repeated 8 times for $\square$ 25M models and 16 times for $\square$ 100M models. Vanilla Mamba uses 16 and 32 Mamba blocks respectively (double the number, consistent with Mamba blocks having roughly $6d_{\text{model}}^2$ parameters versus $8d_{\text{model}}^2$ for a standard Transformer FF layer). MoE-Mamba uses fewer blocks because each block contains both a Mamba layer and the MoE layer with its associated expert parameters.
Output: After the final block, a standard unembedding projection maps the residual stream back to vocabulary logits for next-token prediction.

The residual stream carries token representations sequentially through unconditional sequence processing (Mamba) followed by conditional, specialized processing (MoE), with the Mamba layer's prefix-integration enabling the router to make context-informed expert assignment decisions.

3.3 Roadmap for the Deep Dive

First, the Mamba layer mechanics: explain the selective state space model, its three internal projections (Gate, Conv, Output), and how it achieves linear-time sequence mixing, since this is the backbone that replaces self-attention. Understanding what the Mamba layer produces is essential for understanding what the MoE layer operates on.
Second, the Switch MoE layer: detail the router, expert networks, token dropping, load balancing loss, and how conditional computation decouples total parameters from active parameters, since this is the core scaling innovation being combined with Mamba. Understanding the capacity factor and auxiliary loss is key to why MoE-Mamba trains stably and why more experts improve performance.
Third, the interleaving strategy (MoE-Mamba versus alternatives): explain why Mamba and MoE are placed sequentially rather than in parallel, why the Mamba layer comes first, and how this design relates to the standard Transformer practice of alternating attention and feed-forward layers. Introduce the parallel MoE-Mamba variant and the inside-block MoE experiments to clarify the design space the paper explores.
Fourth, the parameter allocation system: explain the ratio of Mamba-to-MoE active parameters, the expansion factor $E$ , and how $d_{\text{expert}}$ is chosen to match total active parameters across comparison models, since this is the key architectural hyperparameter that controls the tradeoff between sequence processing capacity and expert specialization.
Fifth, the training pipeline: detail the optimization setup, learning rates, batch sizes, token budgets, and the EMA smoothing for loss reporting, since the paper's core metric (training step speedup) depends on precise comparison of training dynamics.

3.4 Detailed, Sentence-Based Technical Breakdown

This is primarily an architecture design and empirical ablation paper whose core idea is that interleaving unconditional Mamba layers with conditional Switch MoE layers produces a model that achieves the same training loss as dense Mamba in fewer steps while adding the capacity-scaling benefits that MoE provides to Transformers. The approach solves a concrete architectural design problem (where to place MoE in an SSM that lacks feed-forward layers) and systematically characterizes the resulting design space through controlled experiments varying placement, parameter ratios, and expert count.

Vanilla Mamba Layer Mechanics

The Mamba layer (Gu & Dao, 2023) replaces self-attention with a selective state space model that processes sequences in linear time while maintaining the ability to selectively attend to or ignore inputs based on content. Understanding its internal structure is necessary because MoE-Mamba's design — where MoE layers are added, how parameters are counted, and what expansion factor means — depends on the Mamba layer's architecture.

The Mamba layer receives a residual stream vector $x \in \mathbb{R}^{d_{\text{model}}}$ for each token position and produces an output of the same dimensionality that is added back to the residual stream (Figure 3, right). The layer has three trainable linear projections:

Gate projection: maps $x$ to a vector of dimension $E \cdot d_{\text{model}}$ , where $E$ is the expansion factor (default $E = 2$ in vanilla Mamba). This expanded representation is passed through a SiLU/Swish activation and then gated element-wise with the SSM output.
Conv projection: maps $x$ to a second vector of dimension $E \cdot d_{\text{model}}$ . This representation undergoes a 1D depthwise convolution (with kernel size 4) across the sequence dimension, followed by the selective SSM scan.
Output projection: maps the $E \cdot d_{\text{model}}$ -dimensional gated result back to $d_{\text{model}}$ .

The forward pass proceeds in four steps:

Expand: Compute $z = \text{SiLU}(\text{Gate}(x))$ and $h_{\text{in}} = \text{Conv}(x)$ , both in $\mathbb{R}^{E \cdot d_{\text{model}}}$ .
Convolve: Apply a causal 1D convolution (kernel width 4) to $h_{\text{in}}$ along the sequence dimension to produce $h_{\text{conv}}$ . This provides local context mixing analogous to the local attention patterns that emerge in Transformers.
Selective SSM scan: Feed $h_{\text{conv}}$ through the selective state space mechanism. For each position $t$ , the SSM maintains a hidden state that evolves according to input-dependent transition matrices. Crucially, the A, B, and C matrices of the state space model are computed from the input via learned projections — this is the "selective" mechanism that allows Mamba to attend to some inputs and ignore others, unlike fixed SSMs where the system matrices are input-independent. The scan is implemented via a work-efficient parallel scan (not sequential recurrence), making training efficient on GPUs; intermediate states needed for backpropagation are recomputed during the backward pass rather than stored, which the paper notes "reduces memory requirements" (Section 3.1).
Gate and project: Element-wise multiply the SSM output with the gating signal $z$ , then linearly project back to $d_{\text{model}}$ via the Output projection.

Parameter count: A single Mamba layer contains "slightly over $6d_{\text{model}}^2$ " parameters (Section 4.1). This falls between a Transformer attention layer ( $4d_{\text{model}}^2$ for the Q, K, V, and output projections) and a Transformer feed-forward layer ( $8d_{\text{model}}^2$ for the two FF weight matrices with standard expansion factor 4). The expansion factor $E$ controls the layer's capacity: doubling $E$ roughly doubles both FLOPs and parameters in the Mamba layer, since all three projections operate at dimension $E \cdot d_{\text{model}}$ .

Why Mamba lacks a separate FF layer: Unlike a Transformer block, which has two distinct sub-layers (self-attention for sequence mixing and a feed-forward network for per-position processing), the Mamba block integrates both functions. The SSM handles sequence mixing, and the gating mechanism (Gate projection + SiLU activation + element-wise multiplication) provides per-position non-linear processing that partially substitutes for a feed-forward layer. This is why vanilla Mamba has no FF layers — and why adding a single dense FF layer degrades performance, as the paper notes from prior work (Section 4.5).

Switch MoE Layer Mechanics

The Switch Transformer MoE layer (Fedus et al., 2022) is the conditional computation component that enables scaling total model parameters without proportionally increasing per-token FLOPs. It replaces a standard feed-forward layer with a set of $N_{\text{experts}}$ parallel FF networks, a router that selects which expert processes each token, and a load-balancing mechanism that ensures all experts are used during training.

The experts: Each expert $E_i$ for $i \in \{1, \dots, N_{\text{experts}}\}$ is a standard two-layer feed-forward network with hidden dimension $d_{\text{expert}}$ . All experts have identical architecture but separately trained parameters. In MoE-Mamba, $d_{\text{expert}} = 3d_{\text{model}}$ (Section 4.1), making each expert smaller than a standard Transformer FF layer (which typically uses expansion factor 4). This choice is deliberate: by using more smaller experts rather than fewer larger ones, the model can achieve finer-grained specialization while keeping the number of active parameters per token matched to baselines. For $\square$ 25M models with $d_{\text{model}} = 512$ , each expert has hidden dimension 1536; for $\square$ 100M models with $d_{\text{model}} = 768$ , each expert has hidden dimension 2304.

The router: For an input token embedding $x \in \mathbb{R}^{d_{\text{model}}}$ , the router computes logits via a linear layer $W \in \mathbb{R}^{d_{\text{model}} \times N_{\text{experts}}}$ and then applies softmax to produce expert selection probabilities:

$h(x) = Wx$

$p_i(x) = \frac{\exp(h(x)_i)}{\sum_{j=1}^{N_{\text{experts}}} \exp(h(x)_j)}$

where $h(x) \in \mathbb{R}^{N_{\text{experts}}}$ is the vector of raw router logits for token $x$ , $W$ is the learned router weight matrix, $p_i(x)$ is the softmax-normalized probability assigned to expert $i$ , and the denominator sums over all experts to ensure the probabilities form a valid distribution.

What this computes: For each token, the router outputs a probability distribution over the available experts. Very simply: "how likely is expert $i$ the right choice for this token, given the token's representation after the Mamba layer?" The softmax ensures these scores sum to 1, which prevents the model from assigning zero probability to all experts (which would cause token dropping to cascade) and makes the load-balancing loss well-defined.

Why softmax rather than sigmoid: A sigmoid (independent per-expert) would allow the model to route a token to zero experts or multiple experts simultaneously. Softmax forces exactly-one-expert semantics (in expectation, though inference takes argmax), which is the Switch Transformer's key simplification: setting $k=1$ forces hard routing with exactly one expert per token, reducing communication overhead and making the capacity analysis tractable.

Token routing and expert selection: For each token, the expert with the highest probability is selected:

$I = \arg\max_i p_i(x)$

The output of the MoE layer for that token is calculated as:

$y = p_I(x) \cdot E_I(x)$

where $I$ is the index of the selected expert, $p_I(x)$ is that expert's softmax probability (used as a scaling factor), and $E_I(x)$ is the selected expert's FF output for input $x$ .

Why multiply by the router probability: This scaling serves as a soft gating mechanism. If the router is uncertain (probabilities spread uniformly), the selected expert's output is down-weighted, preventing low-confidence expert decisions from dominating the residual stream. If the router is highly confident (probability near 1), the expert's full output passes through. This is the standard Switch Transformer formulation and aligns with the gradient flow: tokens routed to a particular expert receive gradients scaled by the router probability, encouraging the router to assign high probability to experts that produce useful outputs.

Capacity and token dropping: During batched execution with batch size $N$ , the fair distribution would assign $N / N_{\text{experts}}$ tokens to each expert. In practice, the router's assignments are not perfectly uniform. The Switch Transformer uses a capacity factor (set to 1.0 in all experiments) that limits each expert to processing at most $\text{capacity\_factor} \times N / N_{\text{experts}}$ tokens. Tokens routed to an expert whose capacity is already full are dropped: they pass through the MoE layer unchanged (identity mapping), contributing only the residual stream from previous layers. Dropped tokens are "not updated" for that layer (Section 3.1), meaning they skip the expert computation entirely and receive no gradient signal from the MoE layer.

This dropping mechanism is a training-time feature that prevents computational imbalance during batched execution and implicitly discourages the router from collapsing to a single expert — if an expert receives too many tokens, some are dropped, reducing the effective batch size that expert sees and thus its gradient quality.

Load balancing loss: To explicitly encourage uniform expert utilization, the paper adds an auxiliary load balancing loss with weight $\alpha = 0.01$ :

$\mathcal{L}_{\text{balance}} = \alpha \cdot N_{\text{experts}} \cdot \sum_{i=1}^{N_{\text{experts}}} f_i \cdot P_i$

where $f_i$ is the fraction of tokens in the batch routed to expert $i$ (the empirical assignment frequency), $P_i$ is the average softmax probability assigned to expert $i$ across all tokens in the batch (the router's average preference for expert $i$ ), and the factor $N_{\text{experts}}$ scales the loss to be independent of the number of experts (since $\sum_i f_i = 1$ and $\sum_i P_i = 1$ ). The $\alpha = 0.01$ weight makes this auxiliary loss contribute at roughly 1% of the magnitude of the primary cross-entropy loss, enough to guide the router toward uniform assignment without overwhelming the language modeling objective.

What this loss computes: It measures the correlation between actual expert usage ( $f_i$ ) and router preference ( $P_i$ ). When routing is uniform, $f_i \approx P_i \approx 1/N_{\text{experts}}$ for all $i$ , and the loss is low. When one expert is heavily preferred ( $P_i$ large) and actually used more often ( $f_i$ large), the product $f_i \cdot P_i$ is large for that expert, and the loss penalizes this imbalance.

Why this form: The product structure penalizes cases where the router both prefers an expert and routes tokens to it — if the router prefers expert $i$ but tokens can't reach it because of capacity constraints (so $P_i$ is high but $f_i$ is low), the penalty is smaller than when both are high. This targets genuine routing imbalance rather than penalizing capacity-constrained assignments.

Interleaving Strategy: Why Sequential Mamba → MoE

The core architectural decision in MoE-Mamba is interleaving: the model alternates between Mamba layers and MoE layers, with each block consisting of one Mamba layer followed by one MoE layer (Figure 2, rightmost panel). This arrangement separates unconditional token processing (handled by Mamba) from conditional, expert-specialized processing (handled by MoE).

Why sequential and not parallel: The paper experiments with a parallel variant ("Parallel MoE-Mamba," Section 3.3, Figure 3 left) where the Mamba layer and MoE layer operate on the same input independently and their outputs are summed. This is inspired by the parallel attention-FF arrangement in Wang (2021) and Chowdhery et al. (2023). The results (Figure 6, Table 7) show that sequential MoE-Mamba consistently outperforms parallel across all expert counts (e.g., at 16 experts, sequential achieves 3.67 log perplexity versus 3.70 for parallel after 1B tokens). The paper suggests (Section 4.4) that parallel might become attractive at larger scales due to hardware utilization benefits (reduced communication via fused input projections), but at the scales tested, the representational benefit of having the MoE layer condition on the Mamba layer's sequence-processed output dominates.

Why Mamba first then MoE: The ordering places the sequence-mixing operation before the conditional per-token computation. This means the router sees token representations that have already incorporated context from the token's prefix (via the SSM's hidden state), potentially enabling more informed expert assignments. If MoE came first, the router would operate on context-free token embeddings (or the previous block's residual stream), potentially missing signals about whether a token represents, say, the start of a mathematical expression versus a narrative sentence. The paper does not ablate this ordering, so the causal claim about why this order works better is speculative based on what each component contributes, but the principle aligns with standard practice in Transformer-MoE architectures where attention (sequence mixing) precedes FF/MoE in each block.

Why interleaving rather than one MoE after all Mamba blocks: The alternating pattern ensures that conditional, expert-specialized processing occurs at multiple levels of abstraction. Early blocks process low-level features (word identity, local syntax) after minimal context integration; late blocks process high-level features (semantic roles, discourse structure) after substantial context integration. Placing MoE layers only at the end would mean expert specialization can only operate on fully context-integrated representations, missing the opportunity to specialize at intermediate abstraction levels.

Relation to Transformer-MoE: This design mirrors the standard Transformer-MoE approach (Lepikhin et al., 2020; Fedus et al., 2022) where some feed-forward layers are replaced with MoE layers. However, the paper explicitly notes (Section 3.2) that "vanilla Mamba does not use feed-forward layers," making the MoE layer an addition rather than a replacement. This is a key distinction: in Transformers, MoE substitutes for an existing component (making the model sparser without changing the architectural skeleton); in Mamba, MoE introduces an entirely new processing stage, fundamentally altering the information flow.

Parallel MoE-Mamba: An Alternative Block Design

As an ablation, the paper explores placing the Mamba and MoE layers in parallel rather than sequentially (Section 3.3, Figure 3 left). In this design, both layers receive the same input from the residual stream simultaneously, and their outputs are summed:

$y = \text{Mamba}(x) + \text{MoE}(x)$

where $\text{Mamba}(x)$ is the Mamba layer's output (including its internal gating and SSM processing) and $\text{MoE}(x)$ is the routed expert output scaled by the router probability.

What this computes: The parallel variant treats the Mamba and MoE layers as independent transformations of the same input that are combined additively. Conceptually, Mamba provides sequence-contextualized representations, MoE provides token-specialized processing, and the model gets both simultaneously rather than sequentially.

Why it underperforms sequential: The paper reports (Figure 6) that parallel MoE-Mamba with 8 experts matches vanilla Mamba performance, but requires "between 2 and 4 times as many experts and total parameters to match the performance of the sequential variant." At 16 experts, sequential achieves 3.67 log perplexity after 1B tokens, while parallel achieves 3.70 (Table 7). The paper speculates (Section 4.4) that parallel might benefit from hardware-level optimizations at larger scales due to the ability to fuse input projections into a single matrix multiplication, but at the tested scales, the representational benefit of sequential conditioning dominates. This is an important corrective to the common assumption that parallel computation is always preferable — in this architecture, the information flow (sequential processing enabling the MoE router to see Mamba-processed representations) matters more than the computational independence that parallel placement provides.

Inner MoE: Conditional Computation Inside the Mamba Block

Pursuing a uniform layer design, the paper explores placing MoE directly inside the Mamba block by converting its internal linear projections (Gate, Conv, Output — see Section 3.4 and Figure 3 right) into MoE layers (Section 4.4). The Gate and Conv projections expand the input from $d_{\text{model}}$ to $E \cdot d_{\text{model}}$ , while the Output projection compresses it back. Replacing a projection with MoE means that instead of one weight matrix, the block has $N_{\text{experts}}$ weight matrices, and a router selects which expert processes each token for that projection.

Experimental setup: There are 7 possible designs (replacing one, two, or all three projections with MoE, with the constraint that at least one is MoE). The paper constrains total parameters and FLOPs across designs by ensuring the total number of expert FF layers across all MoE-equipped projections sums to 24. Specifically, if one projection is MoE-ified, that projection gets 24 experts; if two are MoE-ified, they get 12 experts each; if all three, they get 8 each. Additionally, inspired by Fedus et al. (2022), the paper tests placing MoE in every other Mamba block (half the blocks) with 48 experts to maintain total parameters.

Results (Table 3): None of the 7 designs outperforms the interleaved MoE-Mamba. Three designs achieve results "marginally better than vanilla Mamba": MoE in every other layer's Conv projection (3.71 log perplexity vs. 3.72 for vanilla Mamba), every other layer's Gate projection (3.70), and every layer's Output projection (3.70 — interestingly, this is tied with every-other-layer Gate MoE). The worst performers are designs that MoE-ify the Output projection (3.70–4.05, with full-layer Output MoE at 4.05 being substantially worse than vanilla Mamba's 3.72) and combinations involving Conv + Output projections, which perform significantly worse than vanilla Mamba.

Why inside-block MoE underperforms interleaving: The paper offers limited mechanistic explanation but notes that these results "suggest the most promising research directions in future work." A plausible interpretation is that the Mamba block's internal projections are tightly coupled — the Gate and Conv projections operate in parallel (both expanding the input to dimension $E \cdot d_{\text{model}}$ ), and their outputs interact via gating and the SSM scan. Introducing routing variability into one projection but not the others may break this balance, causing some tokens to have mismatched Gate and Conv representations (if, say, Gate is MoE but Conv is dense, different tokens within a sequence would have different Gate transformations but identical Conv transformations). The interleaved design avoids this by keeping the Mamba block internally consistent (all projections are dense) and adding MoE as a separate, self-contained processing stage.

Parameter Allocation: The Mamba-to-MoE Ratio and Expansion Factor

A central architectural hyperparameter in MoE-Mamba is how total active parameters are allocated between the Mamba layers and the MoE layers. This is controlled by two coupled variables: the Mamba expansion factor $E$ (which scales all internal Mamba projections) and the number and size of experts.

The scaling constraint: The paper matches active parameters per token across all comparison models (Table 1). For $\square$ 25M models, the target is approximately 26M active parameters per token. Mamba25M achieves this with 27M parameters (all active). MoE-Mamba25M uses 542M total parameters but only 26M active parameters per token — the other 516M are inactive (not used for that specific token). This is the core MoE benefit: total model capacity increases by ~20× while per-token compute increases only marginally.

The ratio sweep (Section 4.3): The paper varies the ratio of active parameters in Mamba to active parameters in MoE from 1:5 to 5:1 while holding total active parameters per token fixed (Table 2). This ratio determines:

Mamba expansion factor $E$ : A higher ratio allocates more active parameters to Mamba, requiring a larger expansion factor. At ratio 1:5, $E = 2$ (the vanilla default); at ratio 5:1, $E = 3\frac{1}{3}$ . The expansion factor scales the dimensionality of the Gate projection output, Conv projection output, and Output projection input, directly controlling Mamba's representational capacity.
Expert size and count: Since total MoE active parameters = expert size (since only one expert is active per token, with capacity factor 1), fixing the ratio and total active parameters determines $d_{\text{expert}}$ . With $d_{\text{expert}}$ set and the convention $d_{\text{expert}} = 3d_{\text{model}}$ , the number of experts $N_{\text{experts}}$ is determined by fitting the total MoE parameters (active + inactive) into experts of size $d_{\text{expert}}$ . At ratio 1:5, there are 19 experts of size 3 $d_{\text{model}}$ (2560 hidden dimension); at ratio 5:1, there are 96 experts of much smaller size (512 hidden dimension, which is only $\frac{1}{3}d_{\text{model}}$ of the FF hidden size, a problematic configuration).

Results (Figure 5): Increasing the ratio from 1:5 to 3:3 improves performance, but further increases to 5:1 yield "marginal" gains and are "impractical due to inefficient hardware utilization and high routing costs caused by a large number of experts." The 3:3 ratio emerges as the practical sweet spot: Mamba gets enough capacity for effective sequence processing, MoE gets enough capacity for meaningful expert specialization, and the number of experts (32 for the ratio-3:3 model in Table 2) is large enough to provide routing diversity without excessive routing overhead. This ratio is used as the default in all other experiments (Section 4.3).

The "6:0" hypothetical: The paper explicitly addresses (Appendix F) the question of whether putting all active parameters into Mamba (ratio 6:0, i.e., pure Mamba with no MoE) would be best. To match the total active parameter budget of ~26M without MoE, they would need a Mamba model with expansion factor $E=4$ and 8 layers (instead of 16). This model achieves "marginally worse final log perplexity than Mamba25M" (3.73), which is worse than MoE-Mamba with any ratio ≥ 2:4. This confirms that MoE provides benefits beyond what increasing Mamba's internal capacity alone can achieve — conditional expert specialization is genuinely additive, not merely capacity-competitive with making Mamba bigger.

Training Pipeline and Metric Conventions

Dataset and tokenization: All models are trained on the C4 dataset (Raffel et al., 2020) using the GPT2 tokenizer (Radford et al., 2019). The training objective is standard next-token prediction with cross-entropy loss. Context length is fixed at 1024 tokens for all experiments. The $\square$ 25M models train for 150K steps at batch size 64, processing approximately 10B tokens total; the $\square$ 100M models train for 30K steps at batch size 1024, processing approximately 30B tokens total (Table 5 and Table 6 in Appendix A).

Base model configurations: The $\square$ 25M models use $d_{\text{model}} = 512$ and are structurally equivalent to BERT-MEDIUM (Devlin et al., 2019; Turc et al., 2019). The $\square$ 100M models use $d_{\text{model}} = 768$ and copy BERT-BASE configuration while increasing the number of blocks from 12 to 16 (Appendix A). After determining $d_{\text{model}}$ , the number of blocks, number of experts, and expert sizes are set to achieve the target active parameter count.

Optimization: All models use AdamW (Loshchilov & Hutter, 2019) with weight decay 0.1, gradient clipping at 0.5, and a cosine learning rate schedule with 1% warmup steps and final LR ratio of 0.1. The maximum learning rate is tuned separately for each $\square$ 25M model:

Mamba25M: 1e-3
Transformer25M: 5e-4
Transformer-MoE25M: 5e-4
MoE-Mamba25M: 5e-4

For $\square$ 100M models, the $\square$ 25M learning rates are reused but divided by 2 for Transformer-MoE100M "due to repeated instabilities" (Section 4.1), resulting in 2.5e-4 for Transformer-MoE100M versus 1e-3 for Mamba100M and 5e-4 for MoE-Mamba100M.

Training infrastructure: Models are implemented in PyTorch (Paszke et al., 2019) with FSDP (Zhao et al., 2023b) for multi-GPU parallelism. The codebase is publicly available at the provided GitHub repository (Appendix I).

Loss reporting and speedup metric: The paper uses EMA-smoothed (exponential moving average with $\alpha = 0.001$ ) training log perplexity as the primary comparison metric rather than test perplexity, arguing (Section 4.1) that it is "a more fine-grained comparison metric than test log perplexity" because it provides a smooth, low-variance signal at every training step. The test log perplexity comparison for $\square$ 100M models is provided in Appendix G (Figure 9) and is consistent with training perplexity trends.

Speedup calculation: The headline "2.35× fewer training steps" is computed operationally as finding the number of training tokens at which MoE-Mamba100M first achieves the same smoothed log perplexity as vanilla Mamba100M at its final checkpoint (30K steps × 1024 batch size × 1024 context length tokens). From Figure 1, Mamba100M reaches approximately 2.99 log perplexity after 30K steps; MoE-Mamba100M reaches 2.99 log perplexity after approximately 12.8K steps, yielding a ratio of 30K/12.8K ≈ 2.35. The paper notes (Appendix D) that this speedup "increases as training continues" — early in training the advantage is smaller, growing as the MoE benefits compound (Figure 8). For $\square$ 25M models, the speedup "oscillates between 1.6 and 1.9" (Section 4.2), indicating that the advantage is present but less dramatic at smaller scale, possibly because the 10B token budget is insufficient for the MoE routing dynamics to fully mature.

Parameter counting convention: The paper follows Kaplan et al. (2020) in reporting non-embedding parameters only — that is, trainable parameters in the input embedding layer and output unembedding projection are excluded from all parameter counts in Table 1 and throughout the paper. This convention is used because embedding parameters scale with vocabulary size rather than model capacity per se, and their "relatively low importance for final performance has been noted by Lan et al. (2020)" (Appendix E). This means the true total parameters including embeddings are higher than reported, but the relative comparisons are unaffected since all models use the same GPT2 vocabulary.

Why C4 rather than a curated LLM dataset: The paper does not discuss this choice explicitly, but C4 (the Colossal Clean Crawled Corpus) is a standard, publicly available web-text dataset that provides a consistent training signal without the curation decisions that complicate comparisons on proprietary or mixed-source datasets. At the 10–30B token scale, data quality differences between C4 and alternatives like The Pile would likely be second-order relative to architectural differences, and C4's consistent formatting eliminates a confound. The brevity of training (well below one epoch through C4's ~156B tokens for the $\square$ 25M models) means the models see only a fraction of the available data, making data repetition not a concern.

4. Key Insights and Innovations

Innovation 1: SSMs and MoE Are Genuinely Complementary — Not Just Compatible

The paper's most fundamental conceptual contribution is demonstrating that state space models and mixture of experts are not merely interoperable components that can be bolted together, but are genuinely complementary mechanisms whose combination yields gains beyond what either achieves alone. This is a stronger claim than "MoE works with Mamba" — it's the assertion that MoE provides something Mamba fundamentally lacks, and vice versa.

Prior framing in the field. Before this work, the dominant assumption — implicit in Transformer-MoE architectures — was that MoE serves primarily as a capacity multiplier: it lets you pack more parameters into a fixed FLOPs budget by activating only a subset per token. Under this view, MoE is architecture-agnostic; you could slap it onto any model with a feed-forward layer and get roughly proportional benefits. The paper's experiments falsify this view in two directions.

First, Mamba lacks the very component MoE typically replaces. In Transformers, MoE swaps out existing dense FF layers for sparse ones — it's a drop-in substitution that preserves the architectural skeleton. Mamba has no FF layers at all (Section 3.2). This means MoE-Mamba introduces an entirely new computational stage rather than parallelizing an existing one. The paper's demonstration that this addition helps — and helps more than in the Transformer case where MoE is merely a replacement — is evidence that the SSM alone cannot capture certain representational capacities that even sparse expert FF layers can provide.

Second, the single-expert failure mode is the smoking gun. The paper reports (Section 4.5) that MoE-Mamba with a single expert underperforms vanilla Mamba (3.75 vs. 3.72 log perplexity after 1B tokens), consistent with Gu & Dao's (2023) finding that adding dense FF layers to Mamba degrades performance. But as the number of experts increases from 1 to 4 to 8 to 16 to 32, performance improves monotonically (Figure 4, Table 4). If MoE were merely adding FF capacity, one expert should be better than zero — it's more parameters doing the same computation. The fact that one expert hurts while many help reveals that the benefit comes from conditional, token-dependent specialization, not from additional FF computation per se. A single expert applies identical transformation to every token, which disrupts Mamba's carefully tuned internal dynamics. Multiple experts allow the model to route different tokens to different specialized transformations, creating a new capability axis (adaptive per-token processing) that Mamba's uniform SSM mechanism cannot provide.

Third, the evidence gap with Transformer-MoE. Table 1 shows that MoE-Mamba100M achieves 2.81 log perplexity versus 2.88 for Transformer-MoE100M with matched active parameters (117M vs. 114M). This is counterintuitive: Transformer-MoE has the benefit of both attention-based sequence mixing and conditional expert processing, yet MoE-Mamba outperforms it. This suggests that Mamba's SSM representations are more compatible with expert routing than Transformer attention representations. A plausible mechanism: Mamba's compressed state-space representation may produce token embeddings that are more linearly separable by expert function, making the router's task (a simple linear projection followed by softmax) more effective. The paper doesn't mechanistically prove this, but the consistent performance gap across scales points to a genuine architectural synergy rather than incidental hyperparameter differences.

This innovation is fundamental, not incremental. It reframes MoE from a capacity-scaling technique (which is how the Transformer literature treats it) to a capability-complementing technique for architectures that lack per-token conditional processing. It implies that SSMs and MoE may be even more naturally paired than Transformers and MoE — not despite Mamba's simpler structure, but because of it.

Innovation 2: The Parameter Allocation Between Unconditional and Conditional Processing Is a First-Class Architectural Hyperparameter with a Sharp Optimum

The paper introduces and systematically characterizes a design dimension that has no analog in standard Transformer-MoE architectures: the ratio of active parameters devoted to unconditional sequence processing (Mamba) versus conditional expert processing (MoE). This is a genuinely new architectural degree of freedom created by the Mamba-MoE combination.

Why this matters in a Transformer-free model. In Transformers, the split between attention and FF layers is largely fixed by convention: attention layers have $4d_{\text{model}}^2$ parameters, FF layers have $8d_{\text{model}}^2$ , and MoE replaces only FF layers, meaning the ratio of sequence-mixing to per-token-processing parameters is essentially fixed. You can vary the number of experts, but not the proportion of total capacity devoted to each role. In MoE-Mamba, because the Mamba layer's internal capacity is controlled by the expansion factor $E$ (which scales the Gate, Conv, and Output projections) and the MoE layer's capacity is controlled by expert size and count, the designer has explicit control over how much of the model's "thinking" is done by sequence integration versus token specialization.

The 3:3 sweet spot is genuinely informative. The ratio sweep (Section 4.3, Figure 5) finds that performance increases as Mamba receives a larger share of active parameters — from 1:5 to 3:3 — but gains become "marginal" beyond 3:3. At 5:1, the practical disadvantages (96 tiny experts, high routing overhead, inefficient hardware utilization) make the architecture impractical despite marginal perplexity improvements. The 3:3 ratio is thus not just an empirical optimum but a sweet spot where both components have sufficient representational capacity to be individually useful while their combination is greater than the sum of their parts.

This has a deeper implication: Mamba and MoE are not in a zero-sum capacity competition. If they were, we would expect either the ratio-1:5 or ratio-5:1 extremes to be best — either MoE is so powerful that all capacity should go to it, or Mamba's sequence processing is so important that MoE is a distraction. The fact that a balanced ratio is optimal means both components contribute independently valuable computations. The Mamba layer needs enough capacity to produce representations that make expert routing informative; the MoE layer needs enough capacity (and enough experts) to exploit that routing.

Comparison to prior work. This ratio-sweep analysis has no direct precedent in the MoE or SSM literature. Transformer-MoE scaling studies (Clark et al., 2022; Krajewski et al., 2024) explore expert count and granularity, but not the allocation of capacity between attention and FF/MoE components, because the Transformer architecture's parameter distribution is largely fixed by convention and the scaling laws literature (Kaplan et al., 2020) establishes standard aspect ratios. The paper's ratio analysis is conceptually analogous to the Chinchilla scaling laws (Hoffmann et al., 2022) — not in method, but in spirit: it identifies the optimal allocation of a fixed resource across architectural components, where allocation matters because the components have different scaling properties and the optimum is not at either extreme.

This innovation is fundamental in establishing a new design axis but incremental in the specific optimum found — the 3:3 ratio is almost certainly scale-dependent and would differ for larger models, longer training, or different datasets. The contribution is the framing (ratio as a hyperparameter to tune) and the methodology (sweeping it under fixed active parameter budgets), not the specific number.

Innovation 3: The Difficulty of Integrating Conditional Computation Inside the Mamba Block Is a Negative Result That Channels Future Research

Section 4.4's experiments on placing MoE inside the Mamba block — replacing its internal projections (Gate, Conv, Output) with expert layers — are the paper's most diagnostically informative results despite (or because of) being largely negative. Seven design variants are tested; none outperforms the simpler interleaved MoE-Mamba, and only three marginally beat vanilla Mamba (Table 3).

What this negative result teaches. The Mamba block's internal projections are tightly coupled: the Gate and Conv projections operate in parallel on the same expanded representation, and their outputs are combined through gating and the SSM state update. Introducing routing variability into one projection but not others breaks this coupling by creating mismatched representations — a token routed to expert 3 for the Gate projection and expert 7 for the Conv projection receives transformations that were never trained to be compatible, since the dense baseline always applies the same (paired) Gate and Conv transformations to every token. The interleaved design avoids this by keeping the Mamba block internally consistent and adding conditional computation as a separate, self-contained stage whose output is added to the residual stream.

Why this matters beyond Mamba. This result generalizes to a broader design principle: conditional computation should be added at module boundaries where internal coupling is low, not injected into tightly-integrated internal components. The Mamba block's internal projections are each other's coupling partners — they co-evolve during training to produce compatible representations. MoE-ifying them independently disrupts this co-evolution. This principle is not obvious a priori; many architecture papers attempt to make every component conditional (cf. mixture-of-attention-heads approaches) under the assumption that more sparsity is better. The paper's systematic failure of all seven inside-block variants provides concrete evidence that this assumption does not hold for SSM architectures, and likely does not generalize.

The partial successes are also informative. Among the three designs that marginally beat vanilla Mamba, two involve MoE in the Conv or Gate projections in every other layer (not all layers), and one involves MoE in the Output projection in every layer. The "every other layer" pattern suggests that occasional conditional processing within the block is tolerable, but every-layer conditional processing is disruptive — the model needs some layers with consistent, dense internal representations to maintain the representational scaffolding that the SSM relies on. The Output projection's relative tolerance for MoE-ification (3.70 log perplexity, tied with every-other-layer Gate MoE) may reflect that this projection operates after the SSM and gating, meaning it processes already-integrated representations rather than representations that still need to be combined.

Comparison to prior work. The inside-block MoE experiments have no direct precedent. The closest analog is the "Mixture of Attention Heads" approach occasionally explored in Transformers, but those modify the attention mechanism itself (which is already sequence-mixing) rather than its internal projections. The paper's systematic exploration of all 7 possible MoE insertion points, with controlled total parameters and FLOPs, sets a methodological standard for how to diagnose architectural compatibility that future SSM and non-Transformer architecture work can adopt.

This innovation is a negative result with strong prescriptive force. It doesn't advance the state-of-the-art numbers, but it answers a question that every subsequent SSM-MoE researcher would ask ("should we MoE-ify the Mamba block internals?") with empirical evidence rather than intuition, saving future work from dead ends.

Innovation 4: The Growth of MoE Benefits With Training Duration Reveals a Dynamic That Static Comparisons Miss

The paper reports a finding that, while not highlighted in the abstract, has significant methodological implications: the speedup of MoE-Mamba over vanilla Mamba increases as training progresses. For $\square$ 25M models, the speedup "oscillates between 1.6 and 1.9" (Section 4.2); for $\square$ 100M models, it "rises steadily" (Appendix D, Figure 8).

Why this is conceptually important. If MoE-Mamba's advantage were constant throughout training, we could attribute it simply to having more total parameters (even if fewer are active per token) — i.e., the MoE model has a better initialization-scale or a favorable parameter-to-FLOP ratio. The fact that the advantage grows with training implies something different: expert specialization is learned gradually, and its benefits compound as the router develops more refined token-to-expert assignments. Early in training, the router is essentially random (encouraged toward uniformity by the load-balancing loss), so MoE-Mamba behaves similarly to a model with a single averaged expert. As training proceeds, experts differentiate, the router learns which tokens benefit from which expert, and the conditional computation becomes genuinely more valuable than the equivalent dense computation.

The methodological implication. This dynamic means that short-training comparisons underestimate MoE benefits. If the paper had stopped at 1B tokens (as the ablation experiments do), the observed speedup would be substantially smaller — Table 4 shows a 1.21× speedup at 1B tokens for 16 experts versus the 2.35× at 30B tokens for the $\square$ 100M model. Extrapolating from early-training results would incorrectly suggest that MoE provides only modest benefits for SSMs. This is important for the research community: MoE studies that train for only a few billion tokens (common at academic scale) may systematically underreport the technique's potential, creating a pessimistic bias in the literature relative to what industrial-scale training would reveal.

Comparison to prior work. The Transformer-MoE literature has documented that MoE benefits increase with model scale (Fedus et al., 2022; Clark et al., 2022), but the training-duration dependence at fixed scale is less studied. Most MoE papers report final performance after fixed training budgets, and the scaling laws literature (Kaplan et al., 2020; Hoffmann et al., 2022) focuses on model size and data quantity, not the interaction of training duration with architectural sparsity. The paper's explicit tracking of speedup over training steps (Figure 8) provides a more dynamic picture: MoE is not just a capacity multiplier but a training accelerator whose benefits compound.

This innovation is incremental in the raw finding (it's a trend observed in experiments, not a novel mechanism) but fundamental in its methodological implications for how MoE-Mamba and MoE-SSM systems should be evaluated. It argues that fair comparisons require training to sufficient duration for expert specialization to mature, and that short-horizon results should be interpreted as lower bounds.

5. Experimental Analysis

Evaluation Methodology

Dataset. All models are trained on the C4 dataset (Raffel et al., 2020), a publicly available web-text corpus. The paper does not specify train/validation/test split sizes explicitly but notes that the models process less than 20% of the available tokens, making train-test discrepancy negligible (Appendix G). All models use the GPT2 tokenizer (Radford et al., 2019). Context length is fixed at 1024 tokens for all experiments.
Base model(s). The paper uses Mamba (Gu & Dao, 2023) as the primary SSM backbone, a decoder-only Transformer with Rotary Position Embeddings (Su et al., 2023) as the attention-based baseline, and Switch Transformer MoE (Fedus et al., 2022) as the conditional computation mechanism. Two scale configurations are tested: □25M models (~25M active non-embedding parameters, $d_{\text{model}} = 512$ ) equivalent to BERT-MEDIUM, and □100M models (~100M active parameters, $d_{\text{model}} = 768$ ) copying BERT-BASE while increasing layers from 12 to 16 (Appendix A). Mamba was chosen because it represents the state-of-the-art in SSM-based language modeling with demonstrated competitive performance against Transformers at matched parameter counts.
Metrics. The primary metric is EMA-smoothed training log perplexity (exponential moving average with $\alpha = 0.001$ ), which the paper argues is "a more fine-grained comparison metric than test log perplexity" (Section 4.1) because it provides a smooth, low-variance signal at every training step. Log perplexity is computed from cross-entropy loss on the next-token prediction objective. Test log perplexity is reported in Appendix G (Figure 9) and aligns with training trends. Speedup is measured as the ratio of training steps (or processed tokens) required for MoE-Mamba to reach the same smoothed log perplexity as a baseline model at its final checkpoint. For instance, if Mamba100M reaches log perplexity 2.99 at step 30K and MoE-Mamba100M reaches 2.99 at step ~12.8K, the speedup is 30K/12.8K ≈ 2.35×.
Baselines. The paper compares against three baselines:
- Vanilla Mamba (Gu & Dao, 2023): dense SSM-based model with no feed-forward layers, using the standard expansion factor $E = 2$ and 16 layers for □25M / 32 layers for □100M.
- Vanilla Transformer: standard decoder-only Transformer with Rotary Position Embeddings, 8 attention heads (□25M), and $d_{\text{ff}} = 2048$ .
- Transformer-MoE: Transformer with feed-forward layers replaced by Switch MoE layers, following the established approach of Lepikhin et al. (2020) and Fedus et al. (2022).
Generation budget / compute accounting. The paper measures compute in terms of training tokens processed (batch size × context length × training steps). This enables direct comparison across architectures with different per-token FLOPs. For parameter counting, the paper reports non-embedding parameters (excluding input embedding and output unembedding layers, following Kaplan et al., 2020) and active parameters per token (the subset of total parameters used to compute the output for a given token — essentially those in the Mamba layers plus one expert per MoE layer). The paper argues (Appendix B) that active parameters are roughly proportional to FLOPs and more relevant for hardware-optimized architectures like Mamba. The key constraint across all comparisons is that models are matched on active parameters per token to ensure fair per-step computation budgets.
Cross-validation / statistical protocol. The paper does not use cross-validation or statistical significance testing. Instead, it relies on the smoothness of EMA-smoothed training curves (which reduce noise from individual batch losses) and reports single training runs for each configuration. The primary comparisons are visual (training curves in figures) and operational (speedup computed as the ratio of steps to reach matched loss values). This is standard for language modeling scaling studies where the computational cost of multiple runs is prohibitive, but it means variance estimates are unavailable. The paper implicitly addresses this by reporting trends across multiple scales (□25M and □100M) and multiple design variants, where consistent patterns across configurations serve as informal confirmation.

Main Quantitative Results

Training Efficiency: MoE-Mamba vs. Vanilla Mamba

The paper's headline result is that MoE-Mamba100M reaches the same performance as vanilla Mamba100M in 2.35× fewer training steps (Table 1, Figure 1). Translated to log perplexity: MoE-Mamba100M achieves 2.81 final log perplexity compared to 2.99 for Mamba100M after 30K steps, and matches Mamba's 2.99 at approximately 12.8K steps. The speedup is calculated operationally as the ratio of training tokens required to reach matched loss values.

At the □25M scale, the speedup is smaller: between 1.6× and 1.9× (Section 4.2), with MoE-Mamba25M achieving 3.19 log perplexity versus 3.34 for Mamba25M after 10B tokens (Table 1). The paper attributes the reduced benefit at smaller scale to "a lower number of training tokens" (Section 4.2), consistent with the finding that MoE advantages grow with training duration (Appendix D, Figure 8).

For the Transformer-MoE baselines: MoE-Mamba100M outperforms Transformer-MoE100M (2.81 vs. 2.88 log perplexity, Table 1) with matched active parameters (117M vs. 114M). At □25M, the gap is smaller: 3.19 vs. 3.23. Transformer-MoE achieves a 1.56× speedup over Mamba25M at □25M and 1.79× at □100M — consistently less than MoE-Mamba's advantages of 1.76× and 2.35× respectively.

Key takeaway from Table 1: MoE-Mamba achieves the best log perplexity at both scales while using fewer active parameters for Mamba than the vanilla counterpart (26M vs. 27M at □25M; 117M vs. 121M at □100M) and substantially more total parameters (542M vs. 27M; 2439M vs. 121M). The total-to-active parameter ratio is approximately 20×, matching the standard MoE scaling factor.

Number of Experts and Performance Scaling

Figure 4 and Table 4 show the effect of varying the number of experts while keeping active parameters per token constant at ~26M. The finding is that performance improves monotonically with expert count from 4 to 32, with 32 experts achieving the best log perplexity (3.67 after 1B tokens). However, models with 1 expert underperform vanilla Mamba (3.75 vs. 3.72), confirming Gu & Dao's (2023) observation that adding a single dense FF layer degrades Mamba. Four experts roughly match vanilla Mamba (3.72 vs. 3.72), representing the threshold where conditional computation becomes beneficial. The corresponding speedups over vanilla Mamba at 1B tokens (Table 4) are: 1.03× at 4 experts, 1.10× at 8 experts, 1.21× at 16 experts, and 1.23× at 32 experts. The paper notes (Section 4.5) that it "expect[s] further gains with even more experts," extrapolating from the monotonic trend without testing beyond 32.

Key takeaway from Figure 4: The benefit of MoE is not simply about adding parameters — it requires having enough experts for routing-based specialization to matter. A single expert (equivalent to a dense FF layer) is actively harmful; two to four experts break even; eight or more provide clear gains that increase with count.

Optimal Mamba-to-MoE Active Parameter Ratio

Section 4.3 sweeps the ratio of active parameters devoted to Mamba layers versus MoE layers from 1:5 to 5:1, keeping total active parameters and total parameters fixed (Table 2). The results in Figure 5 show that increasing the proportion of active parameters in Mamba improves performance up to a 3:3 ratio, beyond which gains become "marginal" and the design becomes "impractical due to inefficient hardware utilization and high routing costs caused by a large number of experts" (Section 4.3). At 3:3, the model has expansion factor $E = 2$ , expert size $3d_{\text{model}}$ (1536 for □25M), and 32 experts (Table 2). At 5:1, the model requires 96 tiny experts (size 512), creating routing overhead that outweighs the modest perplexity improvement.

The paper explicitly tests the extreme case of putting all active parameters into Mamba (ratio "6:0", Appendix F), which corresponds to Mamba25M with $E = 4$ and 8 layers instead of 16. This model achieves 3.73 log perplexity — "marginally worse" than standard Mamba25M (3.72) and worse than any MoE-Mamba variant with ratio ≥ 2:4. This confirms that MoE provides benefits beyond what increased Mamba capacity alone can achieve.

Key takeaway from Figure 5: The 3:3 ratio is the practical optimum. It allocates sufficient capacity to both the unconditional sequence processing (Mamba) and the conditional expert processing (MoE) such that neither is a bottleneck. The fact that the optimum is interior (not at either extreme) demonstrates that both components contribute independently valuable computations.

Sequential vs. Parallel Placement of MoE

Section 4.4 and Figure 6 compare two placement strategies for the MoE layer: sequential (Mamba block output feeds into MoE layer) versus parallel (both operate on the same input, outputs summed). The sequential MoE-Mamba consistently outperforms the parallel variant across all tested expert counts (1 to 32). At 16 experts after 1B tokens, sequential achieves 3.67 log perplexity versus 3.70 for parallel (Table 7). The paper notes (Section 4.4) that "Parallel MoE-Mamba matches vanilla Mamba when $N_{\text{experts}} \geq 8$ while requiring between 2 and 4 times as many experts and total parameters to match the performance of the sequential variant." This is a substantial efficiency gap in favor of sequential placement.

Key takeaway from Figure 6: The representational benefit of having the MoE router condition on Mamba-processed representations (via sequential placement) dominates any potential hardware advantages of parallel execution at the tested scales. The paper speculates (Section 4.4) that parallel MoE-Mamba "may be an attractive alternative at larger scales due to potentially enabling more efficient use of hardware" through fused input projections (Wang, 2021; Chowdhery et al., 2023), but provides no empirical evidence at such scales.

Inside-Block MoE: Modifying Mamba's Internal Projections

Table 3 reports the results of converting Mamba's internal linear projections (Gate, Conv, Output — see Figure 3, right) into MoE layers. Seven possible designs are tested (one, two, or all three projections MoE-ified), with expert counts adjusted to maintain total parameters and FLOPs across variants (24 total expert FF layers per block, distributed across MoE-ified projections). All experiments are at 1B tokens.

The main finding is that none of the seven inside-block MoE designs outperforms the interleaved MoE-Mamba (which achieves 3.67 log perplexity with 16 experts, Table 4). Three designs achieve results marginally better than vanilla Mamba (3.72):

MoE in Conv projection, every other layer: 3.71
MoE in Gate projection, every other layer: 3.70
MoE in Output projection, every layer: 3.70

The worst performers involve MoE on the Output projection, especially when combined with other projections: Conv+Output MoE achieves 3.76 (every other layer) to 4.17 (all layers); Gate+Output MoE achieves 3.88 to 4.16; and full inside-block MoE (Conv+Gate+Output) achieves 3.88 to 4.39 — substantially worse than vanilla Mamba.

The "every other layer" pattern appears beneficial: for Conv MoE, every-other-layer achieves 3.71 while all-layer achieves only 3.79. For Gate MoE, the pattern reverses (every-other-layer at 3.70, all-layer at 3.89, though the all-layer number is worse). The paper does not mechanistic explain these asymmetries.

Key takeaway from Table 3: The Mamba block's internal projections are tightly coupled, and introducing routing variability into them independently degrades the representational coherence that the SSM mechanism relies on. The interleaved design, which keeps Mamba internally dense and adds conditional computation as a separate stage, is architecturally safer and empirically superior.

Accuracy vs. Perplexity Inconsistency

Section 4.6 and Appendix C report a curious discrepancy: MoE-Mamba25M with 32 experts maintains lower perplexity than Transformer-MoE25M throughout training, but Transformer-MoE "consistently achieves higher accuracy" (next-token prediction accuracy). The paper hypothesizes (Section 4.6) that this reflects a limitation of SSM-based models in verbatim token copying: "attention-based models are able to copy tokens verbatim, unlike SSM-based models, whose similar abilities might be hindered by the compression of the history into a finite hidden state." The paper notes (Appendix C) that this capability has been mechanistically linked to Transformer induction heads (Elhage et al., 2021; Olsson et al., 2022) and suggests Mamba may compensate for this weakness on other task types, potentially "outperform[ing] Transformers in 0-shot tasks in contrast to tasks allowing few-shot demonstrations or requiring in-context learning."

Key takeaway from Appendix C: Log perplexity and token-level accuracy can rank architectures differently. Mamba achieves better aggregate probability mass across the vocabulary (lower perplexity) but makes more discrete errors on tokens that require exact copying from context. This is a potential failure mode for SSM-based language models that requires further investigation.

Training Dynamics: Speedup Growth Over Time

Appendix D (Figure 8) tracks the speedup of MoE-Mamba relative to vanilla Mamba as training progresses, measured as the ratio of processed tokens to reach a given loss value. For □25M models, the speedup "oscillates between 1.6 and 1.9" after the initial training period. For □100M models, the speedup "rises steadily" (Section 4.2, Appendix D). The paper interprets this (Section 5, Future Work discussion) as evidence that expert specialization is learned gradually: early in training, the router is largely uniform, so MoE behaves similarly to a model with a single averaged expert. As training proceeds, experts differentiate, and the conditional computation becomes genuinely more valuable than equivalent dense computation.

Key takeaway from Figure 8: MoE benefits compound with training. Short-training experiments (like the 1B-token ablations) systematically underestimate the ultimate advantage. This has methodological implications: fair comparisons require training to sufficient duration for routing dynamics to mature.

Ablation Studies and Robustness Checks

Number of experts (Figure 4, Table 4): Performance improves monotonically from 1 to 32 experts. A single expert underperforms vanilla Mamba (3.75 vs. 3.72), confirming that conditional computation — not merely added FF capacity — drives the benefit. Four experts roughly break even; eight or more provide clear gains. The trend suggests further improvement beyond 32 experts is likely but untested.
Mamba-to-MoE active parameter ratio (Figure 5, Table 2): Sweeping from 1:5 to 5:1 shows that increasing Mamba's share of active parameters helps up to 3:3, with diminishing returns thereafter. The extreme "6:0" (no MoE) fails to match any ratio ≥ 2:4 (Appendix F), proving MoE provides benefits beyond increased Mamba capacity. The 3:3 ratio is selected as the default for all other experiments.
Sequential vs. parallel MoE placement (Figure 6, Table 7): Sequential placement consistently outperforms parallel by a widening margin as expert count increases. At 32 experts, sequential achieves 3.66 vs. 3.69 for parallel (Table 7). Parallel requires 2–4× more experts and total parameters to match sequential performance, though the paper notes potential hardware advantages at larger scales.
Alternative inside-block MoE designs (Table 3): Seven variants replacing Mamba's internal projections (Gate, Conv, Output) with MoE layers are tested. None outperforms interleaved MoE-Mamba. Three variants marginally beat vanilla Mamba (Conv-only every other layer at 3.71, Gate-only every other layer at 3.70, Output-only every layer at 3.70). Full inside-block MoE (all three projections) performs substantially worse than vanilla Mamba (3.88–4.39). The "every other layer" pattern appears beneficial for some projections, suggesting that occasional dense processing is important for maintaining representational coherence.
Scale robustness (Figure 1, Table 1): The speedup advantage holds at both □25M (1.76×) and □100M (2.35×), with larger gains at larger scale. The paper attributes the smaller □25M benefit partly to insufficient training tokens for routing dynamics to mature, consistent with the finding that speedup grows with training duration.
Model size matching (Table 1): All comparisons match active parameters per token (~25M or ~114–121M) while allowing total parameters to vary substantially (25M–545M for □25M; 121M–2454M for □100M). This controls for per-step computation cost, making the speedup metric valid. The paper acknowledges that parameter counts "are not matched exactly ... due to, among other reasons, the MoE models including routers and Mamba layer not containing precisely 6d² model parameters" but considers these differences "too small to be significant for our results" (Table 1 note).
Metric consistency check (Appendix G, Figure 9): Test log perplexity trends for □100M models align with training log perplexity, confirming that the EMA-smoothed training metric is not misleading due to overfitting or data repetition.
Parameter counting convention (Appendix E): All reported parameter counts exclude embedding and unembedding layers, following Kaplan et al. (2020). This convention does not affect relative comparisons since all models use the same vocabulary and tokenizer.
Learning rate sensitivity (Section 4.1, Appendix A): Learning rates are tuned separately for each □25M model and reused for □100M counterparts, with one exception: Transformer-MoE100M uses half the □25M learning rate "due to repeated instabilities" (2.5e-4 vs. 5e-4). This instability is specific to Transformer-MoE and not observed for MoE-Mamba, which may indicate that Mamba-based MoE training is more stable than Transformer-based MoE training at comparable scales.
Negative result: ReSTᵉᵐ attempt not applicable. The paper does not report reinforcement learning or iterative self-improvement experiments (unlike the reference example). The revision model and ReST experiments described in the example paper are from a different research context. This paper is purely an architectural study; the only "self-improvement" dynamic is the endogenous evolution of expert routing during standard supervised training.

Critical Assessment

Does MoE-Mamba genuinely achieve 2.35× training speedup over Mamba?

What the experiments demonstrate: At the □100M scale, with 30B training tokens, MoE-Mamba100M achieves a log perplexity of 2.81 versus Mamba100M's 2.99 (Table 1), and reaches Mamba's 2.99 at approximately 12.8K steps versus Mamba's 30K steps — a 2.35× ratio (Figure 1). This is a straightforward operational comparison on a matched active-parameter basis.

What is not demonstrated: The speedup is measured against a Mamba baseline trained with the same token budget (30B). We do not know whether Mamba100M would eventually catch up to MoE-Mamba's 2.81 if trained longer — the speedup might be larger (if Mamba plateaus) or smaller (if Mamba eventually achieves 2.81 with 4× training). The paper's Figure 1 shows Mamba's loss curve still declining at 30K steps, suggesting continued training would yield further improvements, but the relative slopes of the two curves beyond 30K steps are unknown. The 2.35× figure is therefore best interpreted as a lower bound on speedup at this loss threshold, not necessarily the asymptotic efficiency ratio.

Additionally, the speedup is measured in training steps, which is not the same as wall-clock time or total FLOPs. MoE-Mamba has approximately 20× more total parameters than Mamba (2439M vs. 121M, Table 1), which increases memory requirements and communication overhead during distributed training. The paper does not report actual training throughput (tokens per second) or total FLOPs for any model, making it impossible to assess whether the step-count speedup translates to end-to-end training time or energy savings. This is a significant omission for a paper whose central claim is about training efficiency.

Does MoE-Mamba consistently outperform Transformer-MoE?

What the experiments demonstrate: At matched active parameters (~114–117M), MoE-Mamba100M achieves 2.81 log perplexity versus 2.88 for Transformer-MoE100M (Table 1). This gap is consistent in direction across both scales (□25M: 3.19 vs. 3.23) and across the training curve (Figure 1 shows MoE-Mamba below Transformer-MoE throughout).

Caveats: The Transformer-MoE baseline uses Switch routing with the same capacity factor (1.0) and load-balancing weight (α=0.01) as MoE-Mamba, which is a reasonable but not exhaustive baseline. The paper does not compare against other MoE routing strategies (e.g., top-2 gating, expert choice, soft MoE), any of which might improve Transformer-MoE performance without changing the SSM vs. attention comparison. Additionally, Transformer-MoE100M required learning rate reduction (from 5e-4 to 2.5e-4) "due to repeated instabilities" (Section 4.1), while MoE-Mamba100M trained stably at 5e-4. This suggests Transformer-MoE training is more brittle in this setting, which is a meaningful practical advantage of MoE-Mamba, but also means the hyperparameter tuning effort was not symmetric — Transformer-MoE received more careful treatment to avoid divergence.

Are the design space ablations conclusive?

Ratio sweep (Section 4.3, Figure 5): The sweep is systematic and well-controlled, varying the ratio from 1:5 to 5:1 while holding total active parameters fixed. However, the number of experts changes dramatically across ratios (from 19 at 1:5 to 96 at 5:1, Table 2), which means ratio and expert count are confounded. The paper acknowledges this confound implicitly when noting that high ratios are "impractical due to ... a large number of experts," but does not disentangle the effect of ratio from the effect of expert count. A more controlled experiment would fix expert count and vary Mamba expansion factor independently, or fix ratio and vary expert count.

Inside-block MoE (Table 3): The seven design variants are tested at 1B tokens, which is much shorter than the main experiments (10–30B tokens). Given the paper's own finding that MoE benefits grow with training duration (Appendix D, Figure 8), it is possible that some inside-block designs would eventually outperform the interleaved variant if trained longer. The paper acknowledges this implicitly by describing the results as "suggest[ing] the most promising research directions in future work" (Section 4.4) rather than as definitive rejections, but the framing in the abstract ("via comprehensive studies, we confirm that the improvement ... is robust to ... design choices") overstates the conclusiveness of these negative results.

Parallel vs. sequential (Figure 6): The comparison is clean and consistent across expert counts, but only tests one parallel configuration (Mamba and MoE operating on the same input, outputs summed). Alternative parallel designs — such as gated combination of parallel outputs, or MoE operating on Mamba's intermediate representations rather than the same input — are not explored. The paper's claim that sequential placement is superior should be understood as applying to the specific parallel design tested, not to all possible parallel configurations.

Does the growing-speedup claim hold at scale?

What the experiments demonstrate: Figure 8 (Appendix D) shows that the speedup of MoE-Mamba100M over Mamba100M increases with training steps, from approximately 1.5× early in training to 2.35× at 30K steps. The curve is still rising at 30K steps, suggesting the advantage might grow further.

Limitations: This is observed at a single scale (□100M) with a single training budget (30B tokens). We do not know whether the trend continues, plateaus, or reverses at larger scales or longer training. The □25M speedup oscillates between 1.6 and 1.9 rather than rising steadily, which could indicate scale-dependent dynamics or simply higher variance from shorter training. The mechanistic interpretation — that experts are learning to specialize gradually — is plausible but untested; no analysis of expert specialization (e.g., measuring expert diversity or routing entropy over training) is provided.

Does the paper demonstrate that MoE-Mamba preserves Mamba's inference efficiency?

What the experiments demonstrate: The paper makes strong claims about inference benefits ("preserving the inference performance gains of Mamba against Transformer," abstract; "shares the inference benefits of Mamba," Section 6). However, no inference experiments are conducted. There are no measurements of inference latency, throughput, memory usage, or scaling with context length for any model. The claim rests entirely on the architectural argument that MoE-Mamba replaces attention with an SSM (providing linear rather than quadratic scaling with sequence length) and uses sparse activation (reducing per-token FLOPs relative to total parameters). While these architectural properties are inherited from Mamba and MoE respectively, their interaction at inference time — particularly the overhead of expert routing, the memory footprint of storing all expert parameters, and the impact on batch size and hardware utilization — is never empirically characterized.

This is a significant gap between claim and evidence. The paper's title is "Efficient Selective State Space Models with Mixture of Experts," and efficiency is the central selling point, but the efficiency evidence is exclusively about training step count. Inference efficiency is asserted based on architectural properties, not demonstrated through measurement.

Missing experiments that would strengthen the paper

Longer training runs for □25M models: The 10B-token budget may be insufficient for MoE routing dynamics to fully mature, as the paper itself acknowledges. Extending to 30B or more would clarify whether the smaller-scale speedup (1.6–1.9×) is a real scale effect or an artifact of undertraining.
Inference benchmarks: Latency, throughput, and memory measurements on standard hardware for MoE-Mamba vs. Mamba vs. Transformer-MoE at matched active parameters, ideally at multiple sequence lengths to demonstrate the linear-vs-quadratic scaling advantage.
Scaling beyond □100M: The largest model has 117M active parameters and 2.4B total parameters — far below the scale where MoE is typically deployed in production (Switch Transformer and Mixtral demonstrate MoE at 7B+ active parameters). Running at least one larger configuration (e.g., □300M active) would substantially strengthen claims about scaling prospects.
Downstream task evaluation: All results are on language modeling perplexity. The paper mentions Mamba's potential weakness in few-shot tasks requiring in-context learning (Appendix C) but provides no downstream evaluations. Assessing MoE-Mamba on standard LLM benchmarks (e.g., HellaSwag, PIQA, MMLU) would address whether the perplexity improvements translate to task performance and whether the suspected copying limitation affects practical utility.
Expert specialization analysis: The paper interprets MoE benefits as arising from expert specialization but provides no analysis of whether experts actually specialize (e.g., by measuring inter-expert similarity, routing entropy over training, or expert activation patterns by token type). Such analysis would transform the "growing speedup" observation from an empirical trend into a mechanistically understood phenomenon.

Summary assessment

The paper's central claim — that combining Mamba with MoE yields faster training convergence than either component alone — is supported by the training curve comparisons at two scales (Table 1, Figure 1). The design space ablations (ratio sweep, parallel vs. sequential, inside-block MoE) are systematic and informative, though confounded in places. The claim that MoE-Mamba outperforms Transformer-MoE is supported but rests on a single MoE configuration (Switch routing) and unequal hyperparameter treatment (Transformer-MoE needed learning rate reduction to avoid instability). The claim about inference efficiency is not empirically supported — it rests on architectural arguments without measurements. The claim that improvements are "robust to model sizes, design choices, and the number of experts" (abstract) is supported for expert count and sequential placement at the tested scales, but "robust to model sizes" is based on only two scale points (□25M and □100M), and "robust to design choices" is contradicted by the inside-block MoE results (most designs underperform). The paper's primary value is as a design space exploration and proof of concept rather than a definitive demonstration of scalable SSM-MoE efficiency.

6. Limitations and Trade-offs

The Difficulty of Integrating Conditional Computation Inside the Mamba Block Is a Negative Result That Channels Future Research

The assumption or constraint. The paper assumes, implicitly through its experimental design, that the optimal placement of MoE relative to Mamba can be determined by testing the interleaved, parallel, and inside-block configurations exhaustively at small scale. The inside-block experiments (Section 4.4) test all seven possible combinations of MoE-ifying Mamba's internal projections (Gate, Conv, Output) and find that none outperform the simpler interleaved design. The paper acknowledges this openly:

"none outperforming MoE-Mamba. These results suggest the most promising research directions in future work." (Section 4.4)

The consequence. The paper cannot distinguish between two very different interpretations of this result. Interpretation A: inside-block MoE is genuinely worse than interleaved MoE, and the tight coupling of Mamba's internal projections makes conditional computation within the block fundamentally incompatible with the SSM mechanism. Interpretation B: the tested inside-block configurations happen to be suboptimal, and a different expert allocation, routing strategy, or training duration would close the gap. The paper's own finding that MoE benefits grow with training duration (Appendix D, Figure 8) cuts against Interpretation A for the inside-block experiments: these were run for only 1B tokens (Table 3), versus 10–30B for the main experiments. If routing dynamics mature slowly, the 1B-token budget may be insufficient for inside-block experts to develop useful specializations. The paper cannot rule out that a 30B-token inside-block MoE run would outperform the interleaved variant, which means the design space exploration — while systematic in coverage — is not conclusive in its conclusions.

What evidence exists in the paper. Table 3 reports final log perplexity after 1B tokens for all seven inside-block designs, with the best (3.70 for Gate MoE every other layer and Output MoE every layer) only marginally better than vanilla Mamba (3.72) and worse than interleaved MoE-Mamba with 16 experts (3.67, Table 4). The speedup-over-training analysis (Appendix D, Figure 8) shows that MoE-Mamba's advantage over vanilla Mamba grows substantially between 1B tokens (1.21× speedup at 16 experts, Table 4) and 30B tokens (2.35×, Table 1). The paper does not apply this analysis to the inside-block variants, so we cannot tell whether their small advantage at 1B tokens would grow or shrink with longer training.

Mitigation status. The paper does not mitigate this — it acknowledges the limitation and delegates resolution to future work. A follow-up study could select the most promising inside-block design (e.g., Gate MoE every other layer at 3.70) and train it to 10–30B tokens to determine whether the interleaved design's current advantage is robust to training duration or a transient artifact of undertraining.

No Inference Efficiency Measurements Despite Central Claims

The assumption or constraint. The paper's title, abstract, and conclusions make strong claims about inference efficiency that rest entirely on architectural reasoning rather than empirical measurement. The abstract states MoE-Mamba "preserv[es] the inference performance gains of Mamba against Transformer," and Section 3.1 argues that "the advantages of Mamba over the attention mechanism are especially prominent during inference, as not only is the computational complexity lowered, but also the memory usage is not dependent on the context length." Section 5 (Future Work) speculates about "efficiency gains growing with the context length due to better hardware utilization."

The consequence. None of these claims are tested. The paper reports no inference latency, throughput, memory usage, or scaling-with-context-length measurements for any model — Mamba, MoE-Mamba, or Transformer-MoE. This matters because MoE introduces inference-time overhead that could partially or fully negate Mamba's architectural advantages:

Expert parameter storage: MoE-Mamba100M has 2.44B total parameters versus 121M for vanilla Mamba (Table 1). At inference time, all expert parameters must reside in memory (or be streamed from disk, adding latency), even though only a fraction are used per token. For a 20× total-to-active parameter ratio, this means MoE-Mamba requires ~20× more memory than a dense model with matched active parameters, which could eliminate Mamba's memory advantage over Transformers for deployment.
Routing overhead: For each token, the router must compute softmax probabilities over all experts and select the argmax. With 42 experts (the default for □100M models, Table 6), this adds a modest but non-zero per-token cost that does not exist in dense Mamba.
Load imbalance at inference: The capacity factor mechanism that drops tokens during training (Section 3.1) is not applied at inference, but expert utilization may still be non-uniform depending on the input distribution, potentially causing some experts to sit idle while others become bottlenecks — a hardware utilization problem that does not affect dense Mamba.

What evidence exists in the paper. None. Section 5 (Future Work) calls for exploring "synergies of Mamba and MoE" with respect to inference, but no inference experiments are presented. Appendix B notes that "the number of FLOPs is both harder to calculate and less relevant for hardware-aware architecture like Mamba with its optimizations, especially during inference," explicitly side-stepping the FLOPs comparison. The paper's claims about inference efficiency are therefore architectural hypotheses, not demonstrated facts.

Mitigation status. Not addressed. The paper acknowledges inference as a topic for future work (Section 5) but makes unconditional claims about inference benefits in the abstract and conclusions without caveats. A practitioner evaluating MoE-Mamba for deployment would need to run their own inference benchmarks, as the paper provides no guidance on expected latency, throughput, or memory requirements for any hardware configuration or sequence length.

The Headline 2.35× Speedup Excludes the Cost of Storing and Communicating 20× More Parameters

The assumption or constraint. The paper measures training efficiency in terms of training steps (or equivalently, tokens processed) until a target log perplexity is reached. This metric treats each training step as having equal cost regardless of the model's total parameter count. However, MoE-Mamba100M has 2,439M total parameters versus 121M for Mamba100M (Table 1) — a ~20× increase — while using only 117M active parameters per token (similar to Mamba's 121M). Training a 2.4B-parameter model is not 2.35× cheaper than training a 121M-parameter model just because it takes 2.35× fewer steps; the per-step cost in terms of memory, communication, and wall-clock time is substantially higher.

The consequence. The 2.35× training step reduction may not translate to a 2.35× reduction in actual training time or cost:

Memory: FSDP (Fully Sharded Data Parallelism, Zhao et al., 2023b) is used for multi-GPU training (Appendix A). With FSDP, model parameters are sharded across devices and all-gathered before each forward pass. MoE-Mamba's 20× larger parameter count means 20× more data must be communicated and stored in aggregate, even if each device only holds a shard. The communication volume for expert parameters during training (all-gathering the expert weights for whichever expert each token in the batch needs) is additional overhead not present in dense Mamba.
Optimizer states: AdamW maintains two momentum buffers per parameter. This means MoE-Mamba requires ~20× more optimizer state memory than Mamba — for the □100M model, that is roughly 4.9B additional float32 values (for the 2.44B non-embedding parameters) just for optimizer states, which must be distributed across GPUs.
Wall-clock time: The paper does not report training throughput (tokens per second) for any model, making it impossible to assess whether MoE-Mamba100M's 2.35× step advantage translates to faster end-to-end training. If MoE-Mamba100M trains at, say, 0.4× the tokens-per-second of Mamba100M due to communication overhead, the effective training time advantage would be 2.35 × 0.4 = 0.94× — essentially no speedup.

What evidence exists in the paper. None directly. The paper acknowledges (Appendix B) that it reports active parameters rather than FLOPs because "the number of FLOPs is both harder to calculate and less relevant for hardware-aware architecture like Mamba with its optimizations," but does not address the memory or communication cost of the 20× total parameter increase. The load-balancing auxiliary loss (Section 3.1, α = 0.01) and capacity factor mechanism (1.0) are described at the algorithmic level, but their impact on distributed training throughput (all-to-all communication for token routing, expert weight gathering) is never measured or discussed.

Mitigation status. Not addressed. The paper heavily emphasizes "efficiency" (the word appears in the title, abstract, and throughout) without distinguishing between step efficiency (which is measured) and hardware efficiency (which is not). Future work on scaling MoE-Mamba should include throughput measurements and a breakdown of training time by component (computation vs. communication vs. memory movement).

Single Benchmark, Single Model Family, Short Training

The assumption or constraint. All experiments use the C4 dataset for language modeling, the PaLM-derived Mamba architecture (with BERT-inspired sizing), training budgets of 10B tokens (□25M) or 30B tokens (□100M), and context lengths of 1024 tokens. The paper acknowledges this scope limitation implicitly through its future work section, which calls for "scaling" studies and exploration of "other modalities" (Section 5). There is no explicit acknowledgment that the results may not transfer to other datasets, domains, or training regimes.

The consequence. The paper's findings — 2.35× speedup, 3:3 optimal Mamba-to-MoE ratio, monotonic improvement with expert count, inferiority of inside-block MoE — are demonstrated under a very narrow set of conditions that differ substantially from modern LLM training:

Training tokens (10–30B) are orders of magnitude below Chinchilla-optimal for models of this size. At 121M parameters, Chinchilla scaling laws (Hoffmann et al., 2022) would prescribe roughly 2.4B tokens for compute-optimal training. The □100M models receive 30B tokens — well above this threshold — but at 2.4B total parameters (MoE-Mamba100M), Chinchilla-optimal training would require ~48B tokens. The 30B-token budget is 62% of that target. If MoE benefits grow with training duration (as the paper itself demonstrates, Figure 8), the observed speedup may be a lower bound — or, alternatively, the relative advantage may shift as models approach convergence, which the current data cannot resolve.
C4 is a web-text corpus. Web text has characteristic statistical properties (repetition, templated content, uniform formatting) that may influence expert specialization patterns. A dataset with more diverse domains (code, mathematics, multilingual text) could yield different routing dynamics — experts might specialize along domain boundaries, or the router might learn different allocation strategies altogether.
Context length of 1024 tokens is short by modern standards (8K–128K contexts are routine). Mamba's claimed advantage over Transformers is most pronounced at long context lengths, and the paper speculates (Section 5) about "synergies" where "Mamba alleviates computation and memory throughput issues stemming from larger context sizes, while MoE alleviates those same issues stemming from increasing number of parameters." This synergy cannot be evaluated at 1024-token contexts.

What evidence exists in the paper. The paper demonstrates robustness across two scale points (□25M and □100M), across multiple expert counts (1–32), and across several architectural variants (sequential, parallel, inside-block). These within-study replications increase confidence that the observed patterns are not artifacts of a single configuration, but they do not address generalization to substantially different training regimes, data distributions, or context lengths.

Mitigation status. The paper acknowledges these gaps indirectly through its future work section, which calls for "scaling laws," "exploration of different types of MoE," "long context utilization," and "other modalities" (Section 5). However, the abstract and conclusions present the findings as general ("we confirm that the effect is robust to the design choices") without qualifying the narrow scope of the experimental conditions under which robustness was demonstrated.

Training Instability in Transformer-MoE Complicates the Core Comparison

The assumption or constraint. A central claim of the paper is that MoE-Mamba outperforms Transformer-MoE at matched active parameters: 2.81 versus 2.88 log perplexity for □100M models (Table 1). This comparison assumes that both models are trained under their best-known hyperparameter configurations. However, Section 4.1 reveals an asymmetry:

"When training Transformer-MoE100M, we divide the learning rate by two due to repeated instabilities."

For □25M models, the learning rates are tuned separately for each architecture (Appendix A, Table 5). For □100M models, the □25M learning rates are reused but Transformer-MoE100M requires halving (from 5e-4 to 2.5e-4) to avoid training instability, while MoE-Mamba100M trains stably at 5e-4 (Table 6).

The consequence. This asymmetric hyperparameter treatment creates a confound in what should be the paper's cleanest head-to-head comparison. It is impossible to determine whether:

MoE-Mamba genuinely outperforms Transformer-MoE because Mamba is a better backbone for MoE, or
Transformer-MoE underperforms because it was trained with a suboptimal learning rate to avoid instability, and a more sophisticated stabilization technique (gradient clipping adjustment, different optimizer, different initialization) would close or reverse the gap.

The fact that Transformer-MoE exhibits training instability at all is itself practically informative — it suggests that Transformer-based MoE training is more brittle than Mamba-based MoE training in this regime — but it undermines the clean architectural comparison the paper aims to present. The instability could be specific to the interaction of Switch routing with the Transformer architecture at this particular scale and batch size, and might not generalize to larger models where MoE Transformers are routinely trained successfully (e.g., Mixtral, GLaM, Switch Transformer).

What evidence exists in the paper. Table 1 reports the final log perplexity numbers without mentioning the learning rate discrepancy. The discrepancy is disclosed only in Section 4.1 and Appendix A. The □25M comparison does not have this confound — both Transformer-MoE25M and MoE-Mamba25M use learning rate 5e-4 (Table 5) — and the gap there is 3.23 vs. 3.19 log perplexity, a smaller margin (0.04) than at □100M (0.07). This pattern is consistent with the learning rate confound partly explaining the larger gap at □100M. However, the paper does not report whether additional hyperparameter tuning was attempted for Transformer-MoE100M beyond learning rate reduction, or whether training runs at 5e-4 consistently diverged or only sporadically.

Mitigation status. The paper partially mitigates this by reporting the □25M comparison (where learning rates are matched) and by being transparent about the learning rate reduction in Section 4.1. However, the abstract, Figure 1, and Table 1 present the MoE-Mamba vs. Transformer-MoE comparison without qualifying the learning rate asymmetry. A more robust comparison would test Transformer-MoE with gradient clipping adjustments, warmup tuning, or different optimizers to find a configuration that trains stably at 5e-4, or would apply the same 2.5e-4 learning rate to MoE-Mamba to create a symmetric comparison (even if suboptimal for MoE-Mamba). Neither is attempted.

No Characterization of Expert Specialization or Routing Dynamics

The assumption or constraint. The paper attributes MoE-Mamba's benefits to conditional computation — the idea that different tokens are routed to different specialized experts, providing representational flexibility that a dense Mamba layer cannot achieve. This is stated explicitly in Section 4.5's interpretation of why a single expert fails while many succeed:

"This is consistent with Gu & Dao (2023) reporting that Mamba interleaved with feed-forward layers (which corresponds to a single-expert MoE layer) is worse than vanilla Mamba."

The implication is that the sparsity and specialization of multiple experts drives the benefit. However, the paper provides no analysis of whether experts actually specialize, how routing patterns evolve over training, or what (if anything) different experts learn to handle differently.

The consequence. Without analyzing expert behavior, the paper cannot validate its own mechanistic explanation, nor can it provide practical guidance for practitioners. Key questions go unanswered:

Do experts specialize by token type? For instance, do some experts handle punctuation, others handle nouns, others handle mathematical expressions? Or is routing essentially random/load-balanced with experts learning redundant representations?
Does routing entropy decrease over training? If experts genuinely specialize, the router should become more confident (lower entropy) in its assignments as experts differentiate. If routing remains near-uniform, the benefit may come from having multiple independent representations averaged (an ensemble effect) rather than from conditional computation.
Does the load-balancing loss (α = 0.01) prevent specialization? The auxiliary loss explicitly penalizes non-uniform expert utilization. If the router genuinely wants to send all mathematical tokens to expert 3, but expert 3 is capacity-limited (capacity factor = 1), tokens get dropped or redistributed, potentially preventing the expert from specializing in mathematics. The paper provides no data on token drop rates over training or effective expert utilization after load balancing.
Are more experts better because they enable finer specialization, or because they increase total parameters? The monotonic improvement with expert count (Figure 4) could reflect either mechanism. If experts are redundant (each learning a similar function), the benefit is purely from ensembling over more independent representations. If experts are specialized, the benefit is from more finely partitioned conditional computation. These have very different implications for optimal expert count and scaling behavior.

What evidence exists in the paper. None. The paper includes no analysis of expert activation patterns, inter-expert similarity, routing entropy, token drop rates, or expert specialization by token type. The "growing speedup with training" observation (Appendix D, Figure 8) is interpreted as evidence of gradual specialization, but no direct evidence of specialization is provided — the speedup growth could equally reflect other dynamics (e.g., the router learning to handle load balancing more efficiently, or the Mamba layers learning to produce more separable representations for routing).

Mitigation status. Not addressed. The paper positions expert specialization as a topic for future work only indirectly, through its call for "more in-depth studies of synergies of Mamba and MoE" (Section 5). This is a significant gap for a paper whose central architectural claim is that conditional computation via multiple experts provides benefits that a single dense FF layer cannot. Without evidence of specialization, an alternative explanation — that any sufficiently large set of independent FF representations, even if used uniformly rather than conditionally, would improve Mamba — remains unfalsified by the paper's experiments.

7. Implications and Future Directions

How This Work Changes the Landscape

This paper reshapes the conversation around state space model scaling from an empirical demonstration to a design-space cartography with a concrete research roadmap. Prior to this work, SSM scaling was an open question with no established mechanism — Mamba had demonstrated competitive performance at modest scale (2.8B parameters in the original work), but how to push SSMs to the 10B–100B+ parameter regime was unknown. The ML community had two disconnected toolkits: SSM architectures for efficient sequence processing, and MoE for conditional computation scaling. This paper is the first to systematically connect them, producing not just a single novel architecture but a characterized design space with identified sweet spots, dead ends, and open questions.

The magnitude of the shift is best understood as infrastructure-building rather than paradigm-breaking. The paper does not propose a fundamentally new learning algorithm, objective, or theoretical framework. Instead, it establishes that the scaling technique that made Transformers dominant (MoE) transfers to SSMs — and that the transfer requires careful architectural decisions (where to place MoE, how to allocate parameters, how many experts to use) that differ meaningfully from the Transformer case. This is analogous to early work showing that batch normalization works for ResNets after being developed for VGG-style architectures: the technique transfers, but the optimal configuration differs, and characterizing that configuration is what makes the transfer practical.

The paper resolves one specific contradiction in the prior literature that had no clear explanation before. Gu & Dao (2023) reported that adding standard feed-forward layers to Mamba degraded performance, which created a puzzle: if adding FF layers hurts, how could an architecture with MoE — which is essentially many FF layers with routing — help? MoE-Mamba resolves this empirically: one FF layer (one expert) hurts (3.75 vs. 3.72 log perplexity after 1B tokens, Table 4), but performance improves monotonically with expert count, with 4 experts breaking even and 32 experts providing a 1.23× speedup. The resolution is not that Mamba "needs" FF computation per se, but that it benefits from conditional, specialized FF computation that a single dense FF layer — which applies identical transformation to every token — cannot provide. This reframes the failure mode from "FF layers are harmful to SSMs" to "uniform per-token processing is harmful; conditional per-token processing is beneficial." This distinction was not recognized in prior work and has implications beyond Mamba: other SSM architectures (S4, H3, Hyena) that lack per-token conditional processing may similarly benefit from MoE integration, even if adding dense FF layers does not help them.

Several research directions become substantially more attractive after this work:

Scaling SSMs to 10B+ parameters is now on a credible path. Before MoE-Mamba, there was no demonstrated mechanism for SSMs to reach the parameter counts where Transformers dominate. The paper's demonstration that a 2.4B-parameter MoE-Mamba with 117M active parameters outperforms a 121M-parameter dense Mamba by 2.35× in training steps (Table 1) provides the first concrete evidence that MoE-based SSM scaling works and yields efficiency gains comparable to or exceeding those in the Transformer-MoE literature. This lowers the risk for industrial labs considering SSM-based LLM development.
The design space inside the Mamba block is now partially mapped. The systematic failure of all seven inside-block MoE variants (Section 4.4, Table 3) channels future research away from naively MoE-ifying Mamba's internal projections and toward either interleaved designs (which work) or more sophisticated inside-block approaches that preserve the coupling between Gate, Conv, and Output projections (e.g., routing all three projections jointly through the same expert, or using a single router for all internal projections).
The ratio of unconditional to conditional processing capacity emerges as a first-class architectural hyperparameter. In Transformer-MoE, this ratio is essentially fixed by convention (attention takes ~ $4d^2$ parameters, FF/MoE takes ~ $8d^2$ ). MoE-Mamba makes it tunable via the expansion factor $E$ and expert sizing. The paper's finding of a 3:3 sweet spot (Figure 5) establishes that this ratio matters and that the optimum is interior — neither all-Mamba nor all-MoE — which means future SSM-MoE architectures should treat this ratio as a hyperparameter to tune per-scale rather than inheriting it from Transformer conventions.

Conversely, some directions become less attractive:

Pure dense SSM scaling without conditional computation now faces a higher burden of proof. If MoE provides 2.35× training speedup at 100M active parameters and the advantage grows with training duration (Appendix D, Figure 8), then at larger scales where training budgets are measured in trillions of tokens, the efficiency gap between dense and MoE-based SSMs likely widens. A dense Mamba at 7B parameters competing against a MoE-Mamba at 7B active / 140B total parameters may face an insurmountable training efficiency disadvantage, similar to how dense Transformers at frontier scales have been largely superseded by MoE variants.
Naively adding FF layers to SSM architectures is empirically demonstrated to be counterproductive at one expert (Table 4) and only weakly beneficial at few experts, which means future SSM designs should either commit to MoE with sufficient expert count (≥8 based on Figure 4) or not add FF layers at all. The intermediate regime (2–4 experts) provides marginal benefits at best and adds architectural complexity without clear justification.
Exploring arbitrary MoE insertion points in SSM blocks without controlling for coupling is unlikely to be fruitful. The inside-block experiments (Table 3) show that random MoE-ification of internal projections produces results ranging from marginally-better-than-baseline (3.70) to substantially-worse (4.39), with no clear pattern beyond the vague observation that the Output projection is more tolerant of MoE-ification than the Conv or Gate projections. Future work on inside-block MoE should start from a mechanistic hypothesis about which projections are coupled and how to preserve that coupling under routing, rather than from an exhaustive combinatorial sweep.

The paper also provides a methodological template for how SSM-MoE design space exploration should be conducted: match active parameters across variants, sweep one architectural dimension at a time under fixed total parameter budgets, report both training dynamics (speedup curves, not just final loss) and static comparisons, and verify trends at multiple scales. This template was absent from prior SSM work, and its adoption could accelerate convergence on optimal SSM-MoE architectures much as the Chinchilla scaling laws provided a shared methodology for Transformer scaling studies.

Follow-Up Research This Work Enables

Characterize expert specialization directly through routing analysis. The paper attributes MoE-Mamba's gains to conditional computation but provides zero evidence that experts actually specialize — no measurements of inter-expert weight similarity, routing entropy over training, per-token-type activation patterns, or expert dropout effects. A priority follow-up would instrument a trained MoE-Mamba model (the publicly available checkpoints or a reproduction) to measure: (a) the distribution of expert assignments across token types (e.g., do punctuation tokens consistently route to expert 7 while mathematical expressions route to expert 3?), (b) whether routing entropy decreases over training (indicating genuine specialization rather than load-balanced uniformity), (c) the effective capacity utilization (what fraction of tokens get dropped at each MoE layer under capacity factor 1.0), and (d) whether experts can be pruned without loss (if experts are redundant, the benefit is ensemble-based, not specialization-based). The experiment would compare these metrics across expert counts (1, 4, 8, 16, 32) to determine whether the monotonic perplexity improvement in Figure 4 corresponds to finer specialization, larger ensemble, or both. The result would transform the speedup observation from an empirical trend into a mechanistically understood phenomenon and would directly inform whether expert count scaling should prioritize more experts (for finer specialization) or larger experts (for more capacity per specialization).

Extend the inside-block MoE experiments to the full 30B-token training budget. Table 3's inside-block MoE results are run for only 1B tokens, which the paper's own Appendix D (Figure 8) suggests is insufficient for MoE routing dynamics to mature. The most promising inside-block design — Gate MoE every other layer at 3.70 log perplexity after 1B tokens (vs. 3.67 for interleaved MoE-Mamba with 16 experts) — is only 0.03 away from the interleaved variant at a training horizon where the speedup advantage of MoE-Mamba over vanilla Mamba is only 1.21× (Table 4) rather than the 2.35× observed at 30B tokens. A 30B-token run of this variant would resolve whether inside-block MoE is genuinely inferior to interleaved MoE, or merely slower to converge because internal coupling requires more training for expert specialization to overcome the initial disruption. If the gap closes or reverses, it would reopen the inside-block design space that the paper currently treats as a dead end; if the gap widens, it would strengthen the paper's negative conclusion. The experiment should control for total parameters and active parameters to match the interleaved MoE-Mamba configuration, and should track routing entropy and expert diversity metrics alongside perplexity to diagnose whether any convergence difference stems from routing dynamics or representational capacity.

Scale MoE-Mamba to 1B+ active parameters with throughput benchmarks. The paper's largest model has 117M active parameters and 2.4B total parameters — well below the regime where MoE is typically deployed. A scaling run to ~1B active parameters with a total parameter budget of ~20B (maintaining the ~20× total-to-active ratio) would test several unverified claims: (a) does the 2.35× training speedup hold or grow at larger scale (the paper shows it grows from □25M to □100M, Figure 8, but extrapolation is speculative), (b) does the 3:3 optimal Mamba-to-MoE ratio shift with scale (larger models may benefit from different allocations), (c) do the training instability issues that plagued Transformer-MoE100M (requiring learning rate halving, Section 4.1) appear in MoE-Mamba at larger scale or is Mamba-based MoE genuinely more stable, and (d) crucially, what is the actual training throughput in tokens-per-second on standard hardware (e.g., 8×A100 or 8×H100 nodes), and how does the step-count speedup translate to wall-clock training time when accounting for the 20× parameter communication and memory overhead? This experiment would transform MoE-Mamba from a proof-of-concept at academic scale to a credible candidate for production LLM training. The benchmark should report: tokens-per-second per GPU, peak GPU memory usage, all-to-all communication volume for expert routing, and total training FLOPs.

Test MoE-Mamba on long-context tasks where Mamba's inference advantage should compound with MoE's capacity. The paper speculates extensively about the synergy of Mamba (linear-time inference, context-independent memory) and MoE (parameter scaling without proportional FLOP increase) for long-context processing (Section 5), but trains all models at context length 1024. A follow-up should evaluate MoE-Mamba at context lengths of 8K, 32K, and 128K tokens (matching modern deployment scenarios) on tasks that require long-range reasoning: document-level question answering (e.g., NarrativeQA, Qasper), long-context retrieval (Needle-in-a-Haystack), and repository-level code understanding. The key measurements are: (a) inference latency as a function of context length for MoE-Mamba vs. Transformer-MoE at matched active parameters (testing the claim that Mamba's linear scaling provides a growing advantage), (b) memory usage during inference for both model parameters and KV cache (or SSM hidden state), including the cost of storing all expert parameters even when only a subset are active, and (c) task performance (accuracy/ROUGE/F1, not just perplexity) to verify that the language modeling improvements translate to downstream utility. This experiment would convert the paper's architectural claims about inference efficiency into empirical evidence, filling the largest gap between the paper's assertions and its data.

Explore alternative MoE routing strategies beyond Switch top-1. The paper uses Switch routing (top-1 expert selection with capacity factor 1.0 and load-balancing loss α=0.01) exclusively, citing it as "well-established and easy-to-implement" (Section 3.1). But this choice confounds two mechanisms that could drive MoE-Mamba's benefits: sparsity (only one expert is active per token) and load balancing (the capacity factor and auxiliary loss force uniform expert utilization, possibly preventing specialization). A controlled comparison of routing strategies — top-2 gating (Shazeer et al., 2017), expert-choice routing (Zhou et al., 2022) where experts select tokens rather than vice versa, and soft MoE (Puigcerver et al., 2023) where all experts process token-weighted combinations — would isolate whether sparsity, load balancing, or the conditioning itself drives the improvement. The experiment should use the 16-expert MoE-Mamba configuration (where the benefit over vanilla Mamba is cleanly established, Table 4) and report final log perplexity, training stability (number of divergent runs), and expert utilization entropy for each routing strategy. A finding that top-2 routing outperforms Switch top-1 would suggest that the capacity for multiple experts per token is beneficial; a finding that expert-choice routing matches or exceeds Switch would shift the optimal routing paradigm for SSM-MoE architectures.

Stress-test the perplexity-accuracy discrepancy on tasks requiring verbatim recall. Appendix C reports that MoE-Mamba achieves lower perplexity but lower next-token accuracy than Transformer-MoE, with the paper hypothesizing that SSMs "struggle with tasks requiring detailed knowledge of the past input (e.g., some instances of copying)" due to the compression of history into a finite hidden state. This hypothesis is testable: construct synthetic sequences that require exact copying of a token presented k positions earlier (e.g., "[A] [distractor tokens] ... [A]" with varying k), and measure both models' accuracy at predicting the second occurrence of [A] as a function of distance k. If the SSM compression hypothesis is correct, MoE-Mamba's accuracy should degrade with k while Transformer-MoE's attention mechanism should maintain high accuracy (by attending directly to the earlier occurrence). The experiment should also test whether MoE helps or hurts this capability — does routing tokens through different experts disrupt the SSM's ability to maintain a consistent hidden state representation of earlier tokens? This experiment would quantify a potential fundamental limitation of SSM-based language models and inform whether MoE-Mamba should be preferred for tasks requiring in-context learning and few-shot demonstrations (which rely on verbatim recall of prompt tokens) versus zero-shot tasks (where the SSM's representational advantages may dominate).

Practical Applications and Downstream Use Cases

Cost-efficient training of mid-scale language models for domain-specific applications. For organizations training language models in the 100M–1B active parameter range on domain-specific corpora (legal documents, medical literature, scientific papers, code repositories), MoE-Mamba's 2.35× training step reduction (Table 1) directly translates to reduced compute budgets. A model that would require 30B training tokens with a dense Mamba could, in principle, reach the same loss with ~12.8B tokens using MoE-Mamba — a savings of ~17B tokens of training compute. If the domain corpus is limited in size (e.g., a company's internal documentation), this efficiency gain means the model can be trained to lower loss on the available data, potentially closing the quality gap with larger general-purpose models. The practical caveat is that MoE-Mamba's 20× parameter overhead (2.4B total parameters for 117M active) requires more GPU memory for training, so the cost savings depend on whether the organization's hardware is memory-constrained or compute-constrained. For cloud GPU instances where memory scales with the instance type, the tradeoff may be favorable; for on-premise clusters with fixed per-GPU memory, it may not be. The paper's lack of throughput benchmarks means practitioners must measure this themselves for their specific hardware configuration.

On-device or edge deployment of small models with expanded capacity. MoE-Mamba's architecture decouples total model capacity from per-token inference FLOPs, which is particularly relevant for deployment scenarios where memory is relatively abundant but compute is constrained — a common pattern in edge devices (modern smartphones have 8–16GB RAM but limited sustained GPU throughput) and in batch processing on CPUs. A MoE-Mamba model with 25M active parameters and 500M+ total parameters could, in principle, run on-device with the per-token latency of a 25M-parameter model while benefiting from the representational capacity of a much larger model, because only the active expert parameters need to be loaded into compute memory per token. The inactive expert parameters can remain in main memory or be streamed from storage with a prefetching scheme. This application is speculative — the paper provides no inference measurements — but the architectural properties (sparse activation, expert parameters that can be stored separately and loaded on-demand) make it a natural fit. The practical implementation would need to solve the expert parameter loading latency problem: if experts are stored on flash, the time to load an expert's weights on a cache miss could exceed the time to process a token, negating the compute savings. Expert prefetching based on the router's predicted next-token assignments (since the router produces a probability distribution, not just an argmax) could mitigate this.

Scaling SSM-based models to parameter counts where dense training is prohibitive. The paper's core contribution — demonstrating that MoE works with Mamba — has direct implications for research labs and companies that want to train SSM-based LLMs at scales where dense training budgets would be infeasible. A dense Mamba at 70B parameters (matching Llama 2 70B) would require training hardware budgets comparable to Meta's; a MoE-Mamba at 70B active parameters with 20× total parameters (1.4T total) could potentially achieve comparable or better performance with substantially less training compute, following the pattern established by Mixtral 8×7B (which matches Llama 2 70B performance with ~6× fewer active FLOPs per token). The paper provides the first evidence that this pattern holds for SSMs, not just Transformers, which means SSM-based architectures are now credible candidates for the next generation of open-source LLMs. The immediate practical step for a lab considering this direction would be to reproduce the paper's □100M results, then scale to ~1B active parameters to validate that the 2.35× speedup persists or grows, before committing to a 10B+ active parameter training run. The key risk is the unknown scaling behavior of the Mamba-to-MoE ratio optimum — if the 3:3 sweet spot shifts with scale, the architecture would need to be re-tuned at each scale point, adding engineering overhead.

When to Prefer This Method

The paper does not articulate explicit decision rules for choosing MoE-Mamba over alternatives (Mamba, Transformer-MoE, or dense Transformer). It presents MoE-Mamba as a new architecture that outperforms its baselines under the tested conditions, but does not characterize the conditions under which those advantages hold or reverse. The experiments are limited to language modeling perplexity on C4 at 10–30B token budgets with 1024-token contexts. The paper does compare against named alternatives (Table 1 shows MoE-Mamba > Transformer-MoE > Mamba > Transformer in log perplexity), but these comparisons are architecture rankings at fixed scale, not tradeoff analyses with conditions. The paper's Section 5 speculates about conditions where MoE-Mamba might excel (long contexts, inference efficiency) or face limitations (tasks requiring verbatim recall, few-shot learning), but these speculations are hypotheses for future work rather than empirically grounded decision rules. A "prefer X when Y" matrix based on the paper's actual data would be premature given the narrow experimental scope, and generating one would fabricate guidance the paper does not support.

🎯 Pitch

1. Executive Summary

2. Context and Motivation

The Core Problem: SSMs Don't Yet Scale Like Transformers

Why This Problem Is Important: The Dual Bottleneck of LLM Scaling

Where Prior Approaches Fall Short

Conflicting Signals in Prior SSM Design

How This Paper Positions Itself

The Unexplored Design Space

3. Technical Approach

3.1 Reader Orientation (Approachable Technical Breakdown)

3.2 Big-Picture Architecture (Diagram in Words)

3.3 Roadmap for the Deep Dive

3.4 Detailed, Sentence-Based Technical Breakdown

Vanilla Mamba Layer Mechanics

Switch MoE Layer Mechanics

Interleaving Strategy: Why Sequential Mamba → MoE

Parallel MoE-Mamba: An Alternative Block Design

Inner MoE: Conditional Computation Inside the Mamba Block

Parameter Allocation: The Mamba-to-MoE Ratio and Expansion Factor

Training Pipeline and Metric Conventions

4. Key Insights and Innovations

Innovation 1: SSMs and MoE Are Genuinely Complementary — Not Just Compatible

Innovation 2: The Parameter Allocation Between Unconditional and Conditional Processing Is a First-Class Architectural Hyperparameter with a Sharp Optimum

Innovation 3: The Difficulty of Integrating Conditional Computation Inside the Mamba Block Is a Negative Result That Channels Future Research

Innovation 4: The Growth of MoE Benefits With Training Duration Reveals a Dynamic That Static Comparisons Miss

5. Experimental Analysis

Evaluation Methodology

Main Quantitative Results

Training Efficiency: MoE-Mamba vs. Vanilla Mamba

Number of Experts and Performance Scaling

Optimal Mamba-to-MoE Active Parameter Ratio

Sequential vs. Parallel Placement of MoE

Inside-Block MoE: Modifying Mamba's Internal Projections

Accuracy vs. Perplexity Inconsistency

Training Dynamics: Speedup Growth Over Time

Ablation Studies and Robustness Checks

Critical Assessment

Does MoE-Mamba genuinely achieve 2.35× training speedup over Mamba?

Does MoE-Mamba consistently outperform Transformer-MoE?

Are the design space ablations conclusive?

Does the growing-speedup claim hold at scale?

Does the paper demonstrate that MoE-Mamba preserves Mamba's inference efficiency?

Missing experiments that would strengthen the paper

Summary assessment

6. Limitations and Trade-offs

The Difficulty of Integrating Conditional Computation Inside the Mamba Block Is a Negative Result That Channels Future Research

No Inference Efficiency Measurements Despite Central Claims

The Headline 2.35× Speedup Excludes the Cost of Storing and Communicating 20× More Parameters

Single Benchmark, Single Model Family, Short Training

Training Instability in Transformer-MoE Complicates the Core Comparison

No Characterization of Expert Specialization or Routing Dynamics

7. Implications and Future Directions

How This Work Changes the Landscape

Follow-Up Research This Work Enables

Practical Applications and Downstream Use Cases

When to Prefer This Method

Adjacent in domain

Mini-Gemini: Mining the Potential of Multi-modality Vision Language Models

Mixtral of Experts