Routing Mamba: Scaling State Space Models with Mixture-of-Experts Projection

We thank Reviewer VBDT for acknowledging the clarity of our problem statement and the strength of our experimental results.

Weakness and Q1: The motivation behind several design choices in RoM is unclear... why the naïve approach fails.

The naive MoE approach fails due to what we term functional fragmentation. In Mamba, the Conv, Gate, and Output projection layers are not independent modules; they are functionally interdependent and work in concert. A naive MoE with separate routers for each layer shatters this synergy. For any given token, its representation might be sent to Conv Expert A, Gate Expert B, and Output Expert C. Since these experts are trained independently, they cannot learn to coordinate, resulting in an incoherent and fragmented processing pathway. This creates an exceptionally difficult learning problem for the routers, which must somehow learn three independent, yet functionally-aligned, routing policies.

The empirical consequence of this fragmentation is not just suboptimal performance, but a near-complete failure to learn effectively, as demonstrated in Figure 2 and the training dynamics table in our response to Q2. The Independent MoE model's perplexity stagnates and remains worse than the dense baseline, providing clear evidence of this difficult, fragmented learning problem.

RoM's shared routing directly solves this. By using a single router to select a unified expert pathway (e.g., Conv Expert A, Gate Expert A, Output Expert A), we restore the functional coherence of the Mamba block. This drastically simplifies the router's learning task from coordinating three independent policies to learning a single, unified one. This allows the expert pathways to co-adapt and learn specialized functions synergistically, which is why RoM learns efficiently and achieves superior performance. We have further expanded our intuition and discussion in Section 5.3, Appendix A.7 and A.8 and will ensure it is highlighted more clearly in the main text.

Q2: The claim that reusing routing decisions helps reduce the difficulty of router learning could be made more convincing... A comparison between independent and shared routers... would help.

We thank the reviewer for pointing out how we can provide more direct evidence. Based on our training results, we have compiled the PPL at various training steps for the Samba 421M model. The comparison is between the MoE-Mamba (Conv, Gate, Out) strategy (Independent MoE) and our RoM (Conv, Gate, Out) strategy (Shared Router).

As the table clearly shows, while the Independent MoE's perplexity does improve over time, its learning is significantly slower and less effective than our RoM. At every stage of training, RoM achieves a substantially lower perplexity, demonstrating its superior learning efficiency. The final 9000-step evaluation results for both models are also reported in Table 1, Figure 2 and Table 4 in Appendix A.1. At the end of training, RoM reaches a PPL of 9.77, whereas the Independent MoE only manages 10.26, a result that is still worse than the final dense baseline (10.05 PPL).

This large and consistent performance gap provides direct empirical support for our claim that the shared routing strategy reduces the difficulty of router learning, allowing the model to effectively and efficiently leverage the expert capacity from the very beginning of training. We will add this table and analysis to the appendix and refer to it in the main text. We also commit to adding a PPL vs. training steps plot in the final camera-ready version to visualize this compelling trend.

Dear reviewer VBDT,

As we approach the end of the discussion period, we would greatly appreciate your input on the paper. We hope our responses and additional results address your concerns and welcome any further questions or suggestions.

Thank you for your time and effort.

We thank Reviewer KHjA for recognizing our thorough investigation and thoughtful design choices.

Weaknesses: The paper does not extensively analyze the trade-offs... as the number of experts increases or selecting at least 2 experts per-token.

This is an excellent point. To directly address the reviewer's question, we have conducted additional experiments to analyze these trade-offs. The following table presents the measured training throughput for various RoM configurations, extending the setting from Table 11 in Appendix A.6.

Analysis of Trade-offs:

1. Increasing Total Experts (e.g., 8 to 16): Doubling the total number of experts while keeping k=1 (8top1 vs. 16top1) results in a ~13% throughput reduction. This demonstrates that increasing the model's total parameter capacity incurs a manageable overhead.
2. Increasing Activated Experts (e.g., top-1 to top-2): Doubling the number of activated experts for a fixed set of 8 total experts (8top1 vs. 8top2) results in a ~26% throughput reduction.

These results provide a clear quantitative picture of the trade-offs involved. Increasing the total number of experts is a relatively cheap way to scale model capacity, while activating more experts is more costly but may yield higher performance. We thank the reviewer for this valuable suggestion; we will add this table and analysis to the experimental results section in the final manuscript to provide a more complete picture of RoM's scaling properties.

Dear reviewer KHjA,

As we approach the end of the discussion period, we would greatly appreciate your input on the paper. We hope our responses and additional results address your concerns and welcome any further questions or suggestions.

Thank you for your time and effort.

We thank Reviewer bjjv for the positive feedback on the clarity of our paper and the comprehensive nature of our experiments.

Weakness: The contribution, while effective, could be viewed as incremental... an "engineering trick".

We respectfully argue that our contribution is more fundamental than an "engineering trick." The core challenge we address is that naive MoE integration in Mamba fails, a fact we demonstrate empirically in Figure 2 and which is consistent with prior work. Simply applying MoE to Mamba's projection layers independently, arguably the real straightforward "engineering" approach (Independent MoE), which not only fails to improve upon the dense baseline but actively degrades performance (10.26 PPL vs. 10.05 PPL).

Our primary contribution is identifying this critical failure mode and proposing a novel architectural solution: a shared routing strategy that creates a coherent expert pathway. This design is non-trivial and is the key that unlocks, for the first time, the effective scaling of Mamba-style SSMs with MoE. It is a targeted solution to a documented instability. The resulting 2.3 $\times$ gain in active parameter efficiency (Figure 3) underscores the significance as a novel and effective architectural approach, not merely an incremental tweak.

Lack of theoretical results to support the method... explanation is a bit heuristic... no formal analysis to prove... more stable training or better expert specialization..

We thank the reviewer for this insightful comment and agree that our paper's primary contributions are empirical. While we do not provide a formal theoretical proof, our explanation is grounded in a testable hypothesis that our experiments were designed to validate.

Our central hypothesis is that the naive MoE approach fails due to functional fragmentation. The Mamba block's projection layers (Conv, Gate, Output) are functionally interdependent. Using independent routers forces the model to learn three separate, uncoordinated routing policies, creating an exceptionally difficult optimization problem. Our results provide direct evidence for this failure. As shown in our response to Reviewer VBDT (Q2), the Independent MoE model exhibits an unstable learning curve and its final perplexity of 10.26 is worse than the dense baseline (10.05). This outcome is a direct validation of our hypothesis: the independent routers fail to converge to an effective policy.

In contrast, RoM's shared routing mechanism is designed specifically to enforce functional coherence. By using a single router, we simplify the learning problem to finding one unified policy, allowing the expert pathways to co-adapt. The empirical validation is clear: the RoM model learns stably, consistently outperforms the naive approach at every checkpoint, and achieves a superior final PPL of 9.77.

Therefore, the clearly different outcomes of these two controlled experimental conditions provide strong, causal evidence for our explanation, moving it beyond a mere heuristic to an empirically-verified principle for scaling SSMs.

We would like to highlight that our results provide direct empirical evidence that validates our approach leads to both more stable training and better expert specialization.

Evidence for More Stable Training: Our experiments provide two key pieces of evidence for training stability.

1. Consistent Learning Dynamics: As shown in the PPL vs. training steps table (in our response to Reviewer VBDT), the naive Independent MoE model exhibits unstable learning, with its final perplexity (10.26) being worse than the dense baseline (10.05). In contrast, our RoM model shows a smooth, consistent decrease in perplexity at every configuration, demonstrating a stable and effective learning process.

2. Inherent Load Balancing: As highlighted in Section 4.3 and Appendix A.3 (Table 6), RoM trains effectively without requiring an auxiliary load-balancing loss, a common mechanism used to stabilize MoE training. This inherent stability is a direct result of our shared routing design.

Evidence for Better Expert Specialization: The superior performance of RoM is a direct consequence of more effective expert specialization.

1. Superior Final Performance: The significant gap in final perplexity between RoM (9.77) and the naive approach (10.26) strongly indicates that RoM's experts have learned more specialized and effective functions. If the specialization were not better, the performance would not be superior.

2. Efficient Scaling: The consistent scaling results in Figure 3, where RoM achieves performance equivalent to a 2.3$\times$ larger dense model, would not be possible without effective specialization. Poor specialization would lead to diminishing returns, but RoM demonstrates a clear and efficient scaling advantage, which is a hallmark of successful expert specialization.

Therefore, while the initial intuition may seem a bit heuristic, it is directly and rigorously tested. We have empirically identified a specific failure mode in naive MoE for SSMs and demonstrated that our proposed architectural principle, shared routing for functional coherence, is the solution. We believe this data-driven approach provides a solid, non-heuristic foundation for our claims, with formal analysis being a promising avenue for future work.

Dear reviewer bjjv,

As we approach the end of the discussion period, we would greatly appreciate your input on the paper. We hope our responses and additional results address your concerns and welcome any further questions or suggestions.

Thank you for your time and effort.

We thank Reviewer ceRt for the positive assessment and for recognizing the novelty and significance of our results.

W1: Training throughput still suffers.

We thank the reviewer for this crucial point on training efficiency. While RoM with our current FSDP-only setup has a lower raw throughput than a dense model of the same active size, we argue this represents a highly favorable trade-off for achieving superior model quality.

Our choice to use a pure FSDP setup, as detailed in Section 5.1, was a deliberate methodological decision. This approach, while not the fastest for MoE models, ensures a fair and stable comparison by eliminating confounding variables like expert parallelism (EP), capacity factors, and token dropping. This allows us to isolate and measure the architectural benefits of RoM itself.

More importantly, when comparing RoM to a dense model of comparable performance, RoM is actually more computationally efficient. Based on Table 11 in Appendix A.6:

1. Our RoM (Conv, Gate, Out) model (421M active params) has a training throughput of $3.57 \times 10^5$ tokens/s.
2. To achieve a comparable performance, the dense model needs to be scaled up to Samba (expand=4) (511M active params), which has a training throughput of only $3.34\times10^5$ tokens/s.

This comparison clearly shows that our RoM model not only achieves better performance with fewer active parameters but is also faster to train than a comparable, larger dense model. Therefore, RoM offers a more efficient path to higher performance than simply scaling up a dense model.

Furthermore, the throughput of RoM can be significantly accelerated by integrating standard optimizations like expert parallelism or better CUDA/Triton kernels designed for MoE, which we leave for future engineering work. Our current results already demonstrate a compelling efficiency frontier, and with further optimization, the practical benefits of RoM will be even greater.

W2: Shared routing introduces new design challenges and could affect training stability.

This is a valid consideration. We are pleased to report that our experiments demonstrate remarkable training stability.

1. Consistent Performance: As shown in Figure 3, RoM delivers consistent and significant perplexity improvements across all tested model scales (115M/710M to 1.3B/10B) and context lengths (4K to 16K), indicating that the shared routing mechanism is stable and generalizes well.
2. Inherent Load Balancing: A key finding, highlighted in Section 4.3 and Appendix A.3 (Table 6), is that RoM does not require an auxiliary load-balancing loss for small scale. The routing mechanism naturally distributes the load across experts, which points to the inherent stability and robustness of our shared routing design. This simplifies the training process compared to many traditional MoE models.

Formatting Concern 1: "constant inference-time computation" in the abstract.

Thank you for this precise correction. You are right that the complexity is dependent on the sequence length. We will revise this statement in the abstract and introduction to be more accurate, clarifying that SSMs feature linear-time computation complexity with respect to sequence length and constant memory complexity during the per-step recurrence, which is the key advantage over the quadratic complexity of standard Transformers.

Formatting Concern 2: The definition of the SSM (lines 124-127).

We appreciate you spotting this notational misuse. We will revise this part and ensure the definitions are rigorous and clear.

Dear reviewer ceRt,

As we approach the end of the discussion period, we would greatly appreciate your input on the paper. We hope our responses and additional results address your concerns and welcome any further questions or suggestions.

Thank you for your time and effort.