---
title: 'MossNet: Dual MoE State-Space Models'
url: https://www.emergentmind.com/topics/mossnet
type: topic
---

# MossNet: Dual MoE State-Space Models

MossNet refers to a recent architectural paradigm in large language models (LLMs) that implements a dual Mixture-of-Experts (MoE) mechanism within a recurrent state-space modeling (SSM) framework. MossNet is designed to directly emulate linear multi-head attention (MHA) via parallel channel-wise and time-wise expert routing, offering a balance of expressiveness, efficiency, and hardware adaptability for language modeling. This approach addresses longstanding limitations of conventional SSM and gated-recurrent models, which typically approximate only a single attention head and thus restrict long-range contextual modeling. MossNet was introduced in "MossNet: Mixture of State-Space Experts is a Multi-Head Attention" [2510.26182].

## 1. Theoretical Motivation and SSM-MoE Equivalence to MHA

Traditional SSM-based models (such as Mamba) formalize sequence modeling using discretized ODEs of the form:
$$
x_{t} = \overline{A}_{t} x_{t-1} + \overline{B}_{t} u_{t}, \quad y_{t} = C_{t} x_{t} + D_{t} u_{t}
$$
where $x_{t}$ is the latent state, $u_{t}$ is the input, $y_{t}$ the output, and $A,B,C,D$ are parameter matrices (or functions). In prior art, SSM variants make these parameters input-dependent using lightweight gating networks, but they typically employ a single implicit attention head over the entire sequence.

MossNet augments this by integrating two MoE mechanisms:
1. Channel-mixing (MLP-MoE): Multiple MLP experts are routed per token to increase model capacity and representational diversity in the feed-forward (channel mixing) path.
2. Time-mixing (SSM-MoE): Multiple SSM parameter sets ("experts") are dynamically routed per token during state updates to instantiate multiple, independent "attention heads" in the temporal domain.

The crucial result demonstrated in [2510.26182] is that this Mixture-of-Experts parameterization of SSM kernels mathematically recovers linear multi-head attention, as the unrolled SSM with MoE gating yields a double sum analogous to a multi-query, multi-head attention operator:
$$
y_t = \sum_{m,n} \sum_{i=1}^t \langle q_t^m, k_i^n \rangle v_i
$$
with $q_t^m$ and $k_i^n$ forming query/key projections over routing-induced expert heads, and values $v_i$ corresponding to the source input. Thus, MossNet equips the SSM with a trainable, efficient surrogate for linear MHA.

## 2. Architectural Design

MossNet incorporates dual MoE paths within a Mamba-like SSM backbone. For each block, it consists of:
- **MLP-MoE sublayer:** The standard two-layer MLP is replaced by a top-$k$ MoE, where for each token, the softmax router activates $k$ experts, and their outputs are linearly combined according to expert-wise probabilities and a load balancing loss coefficient $\alpha$.
- **SSM-MoE sublayer:** The SSM parameter functions $\overline{B}_t$ and $C_t$ (and optionally $\overline{A}_t$) are implemented as weighted sums over $M$ and $N$ independent kernels. The gating is based again on the input and allows a small number of experts ($k$ active per-token) to be selected.

Key architectural hyperparameters across reported variants are:
- Number of SSM experts per block (typ. 8)
- Top-$k$ routing (typ. $k=2$ or $3$)
- Hidden width, number of heads, number of layers (e.g., 128–1024 hidden, 2–16 heads, 16–30 layers)
- Sparse activation: only the top-$k$ experts per token are active, preserving computational efficiency

In all designs, MossNet provides standard dense and top-$k$ MoE operation modes, allowing models to trade off expressiveness with computational footprint.

## 3. Training Regimes and Scaling

Empirical results are provided for MossNet models of various scales:
- **MossNet-8×8M:** 19.7M total/9.9M active params, 16 layers, 2 heads, Cosmopedia (22B tokens)
- **MossNet-8×20M:** 63.9M total/26.1M active params, 4 heads
- **MossNet-8×66M:** 325.9M total/102.9M active params, 8 heads
- **MossNet-8×200M⁺:** 1.5B total/0.5–0.7B active params, 16 heads (large variant), 2.8T tokens

All models are trained using the AdamW optimizer, with typical settings of $\text{lr}=2\times10^{-4}$, 1% warmup, cosine decay to 10% final, and $\alpha=1\times10^{-3}$ for MoE load-balance. Context lengths of up to 32,000 tokens are supported, with linear memory and runtime scaling guaranteed by the SSM formulation.

## 4. Experimental Outcomes and Benchmark Results

### Language Modeling
MossNet models consistently outperform both pure SSMs and Transformer/MoE architectures of similar size. On the Cosmopedia holdout set:
- MossNet-8×8M: Perplexity 13.1, versus 13.4 (Mixtral-8×8M), 13.5 (Mamba2)
- MossNet-8×66M: MMLU (5-shot avg.) 20.1%, versus 14.7 (MoE-Mamba), ≤16% (other ~60M baselines)

### Commonsense and Downstream Tasks
For downstream zero-shot performance on tasks such as ARC, BoolQ, HellaSwag, PIQA, WinoGrande:
- MossNet-8×8M achieves 37.1% average (vs 36.4% Mixtral-8×8M, 36.2% Mamba-8M)
- Large variant (MossNet-8×200M⁺) outperforms Qwen2.5-0.5B by 5.8 points average on a 7-task suite (53.5% vs 47.7%), and in top-3 MoE mode (∼700M) scores 55.4% vs. Mamba-790M’s 43.8%

### Device Profiling
Real-device tests on NVIDIA A100 (FP16, FlashAttention 2) and Samsung Galaxy S24 Ultra (CPU, Q8) reveal:
- MossNet-8×200M⁺ memory footprint: ∼8.4GB at 32K tokens (A100), ∼1.6GB at 32K tokens (mobile)
- Prefill/generation throughput competitive with or better than Llama3, Mamba
- Constant memory and high throughput are maintained as sequence length grows, affirming suitability for long-context and on-device inference

## 5. Strengths, Limitations, and Distinctive Properties

### Advantages
- **Expressiveness:** Multi-head routing enables modeling of richer temporal phenomena than single-head SSMs, overcoming bottlenecks in capturing long-range or parallel dependencies.
- **Hardware Adaptability:** Linear scaling in both runtime and memory; MoE amplitudes mean only top-$k$ experts are evaluated, maintaining practical efficiency for diverse devices.
- **Scalability:** Shown to scale from ∼10M to ∼1.5B total parameters (∼500–700M active), maintaining empirical gains across size regimes.
- **Resource Stability:** Fixed memory and graceful throughput degradation for long sequences, critical for both cloud and mobile/edge LLM deployment.

### Limitations
- Increased architectural and implementation complexity due to routing in both channel and temporal (SSM) domains
- MoE efficiency gains saturate as batch heterogeneity increases (i.e., diminishing returns for highly diverse server workloads)
- Evaluation limited to text-only tasks; no results reported for multimodal or RL extensions
- Hardware-specific profile; generality to other chipsets not established

## 6. Comparison to Prior Models and Future Directions

MossNet's theoretical advance is the explicit recovery of linear multi-head attention within the provably efficient SSM context, a property not possessed by prior SSM or gated recurrent memory models. Unlike conventional Transformer-based MoE, which typically applies gating only in the channel-mixing path and incurs dense compute in attention, MossNet distributes expert capacity across both feed-forward and temporal modeling layers.

Research directions highlighted include:
- MoE on additional SSM parameters ($A$) for finer temporal control
- Dense or adaptive multi-query/multi-key patterns
- Cross-token or grouped-query router mechanisms
- Specialist pruning and hardware-aware MoE routing
- Extension to non-text domains (multimodal, RL)

A plausible implication is that this architectural template generalizes beyond language to any sequential modeling domain where hardware efficiency and rich attention-style modeling are required [2510.26182].

Source: https://www.emergentmind.com/topics/mossnet