Beyond Dense Transformers: MoE, SSMs, Hybrids

Chapter 8 described a causal decoder with two dense sublayers: attention mixes positions, and a feed-forward network (FFN) transforms each position. This chapter varies those two sublayers along independent axes.

A mixture-of-experts (MoE) FFN stores many expert networks but evaluates only a selected subset for each token. It changes parameter activation; it does not make attention sparse. A state-space or linear-recurrent layer changes how the sequence is mixed and what state is retained during decoding; it does not by itself add expert routing. A hybrid can use either change, or both, in one stack. Keeping the axes separate prevents three quantities from being confused: stored parameters, arithmetic executed for one token, and state retained for one sequence.

Two independent architecture axes

For a dense Transformer block, every token evaluates the same FFN weights, and every attention layer retains keys and values for every cached position. That is a specific design, not a definition of a language model. Alternatives ask two different questions:

These substitutions introduce new costs. MoE routing creates dispatch, imbalance, and communication. A recurrence compresses history into a state whose width is fixed by the architecture, so past positions are no longer individually represented in a growing key-value cache. The relevant comparison is therefore a workload-level one: quality, training arithmetic, memory, communication, batch shape, and serving latency all count.

MoE routes tokens through conditional FFNs

Sparse expert models predate Transformers. Shazeer et al. introduced a sparsely gated MoE layer at large scale in 2017 inside recurrent language models (Shazeer et al. 2017). GShard then combined conditional experts with Transformer models and automatic sharding (Lepikhin et al. 2020), while Switch Transformer studied top-1 routing and trillion-parameter models (Fedus et al. 2022). The durable idea is narrower than the surrounding scale claims: store more expert weights than one token evaluates.

Routing is a defined mathematical operation

In one common top-$k$ formulation, used by Mixtral, the first two terms score and select experts (Jiang et al. 2024):

$$z_t=W_rh_t, \qquad \mathcal{S}_t=\operatorname{TopK}(z_t,k).$$

The normalized gate for each selected expert is

$$g_{t,e}= \begin{cases} \dfrac{\exp(z_{t,e})} {\sum_{j\in\mathcal{S}_t}\exp(z_{t,j})}, & e\in\mathcal{S}_t,\\[6pt] 0, & e\notin\mathcal{S}_t, \end{cases}$$

Each selected expert then contributes to the output:

$$m_t=\sum_{e=1}^{E}g_{t,e}F_e(h_t).$$

Here $h_t\in\mathbb{R}^{d}$ is the hidden vector at token position $t$; $E$ is the number of routed experts; $W_r\in\mathbb{R}^{E\times d}$ is the router matrix; $z_{t,e}$ is expert $e$'s routing logit; $\mathcal{S}_t$ is the set of $k$ selected expert indices; $g_{t,e}$ is the selected expert's normalized gate weight; $F_e:\mathbb{R}^{d}\rightarrow\mathbb{R}^{d}$ is expert $e$; and $m_t\in\mathbb{R}^{d}$ is the MoE output. A residual connection and any shared expert are outside this equation.

This is one router contract, not a universal definition. Switch uses top-1 routing; some systems apply softmax before selection, some after it, and some use sigmoid affinities. The gate value may or may not be renormalized after top-$k$. These choices change gradients and checkpoint semantics, so a model description must state them.

Stored parameters, selected weights, and runtime are different

Suppose every expert is a bias-free SwiGLU FFN with model width $d$ and expert hidden width $d_f$. From Chapter 8, one expert contains approximately

$$P_e=3dd_f$$

parameters, where $P_e$ counts the gate, value, and down projections; $d$ is the input and output width; and $d_f$ is the expert's intermediate width. With $E$ routed experts and no shared expert, the layer stores and evaluates

$$P_{\mathrm{stored}}=EP_e+Ed, \qquad P_{\mathrm{evaluated}}(t)=kP_e+Ed.$$

Here $Ed$ is the router matrix, whose $E$ scores are all evaluated for token $t$; $EP_e$ counts every stored expert; and $kP_e$ counts only the selected expert weights. Biases and always-on block parameters are omitted. Reported "active parameter" counts often include additional embeddings, attention, normalization, and shared experts, so the reporting convention must be stated.

Figure 9.2 evaluates these equations for $d=4096$, $d_f=14336$, and $k=2$. The selected-expert term stays fixed as $E$ grows, but the router term grows, and neither line measures wall-clock time. Runtime also depends on token dispatch, padding or ragged kernels, all-to-all communication, memory traffic, and the slowest expert shard.

The runnable cell prints the same accounting without requiring plotting libraries. Change $E$, $k$, or either width to inspect another layer.

model_width = 4096
expert_width = 14_336
selected_experts = 2

def moe_parameter_counts(experts):
    if not 1 <= selected_experts <= experts:
        raise ValueError("selected_experts must be between 1 and experts")
    per_expert = 3 * model_width * expert_width  # bias-free SwiGLU
    router = experts * model_width
    stored = experts * per_expert + router
    evaluated = selected_experts * per_expert + router
    return stored, evaluated

for experts in (8, 16, 32, 64):
    stored, evaluated = moe_parameter_counts(experts)
    print(
        f"E={experts:2d}, k={selected_experts}: "
        f"stored={stored / 1e9:.3f}B, "
        f"evaluated/token={evaluated / 1e9:.3f}B, "
        f"ratio={stored / evaluated:.2f}x"
    )

Routing creates a dispatch problem

Expert weights are commonly partitioned across an expert-parallel group. A forward pass then has four stages: compute router scores, group token copies by destination expert, exchange those copies to the expert owners, and return the expert outputs to their original token order. Implementations often use one all-to-all collective for dispatch and another for the return path, but an all-to-all is not intrinsic when all experts are local or the placement differs.

The arithmetic saved inside expert FFNs can therefore be replaced by network time or small, poorly utilized matrix multiplications. Useful measurements include bytes dispatched, tokens per expert, padded versus useful expert slots, collective time, overlap with computation, and tail latency. "Active parameters" alone proves none of these.

For a routing group with $T$ tokens, $E$ experts, and $k$ assignments per token, the even-share load is $kT/E$ assignments per expert. A fixed-capacity implementation may reserve

$$C_e=\left\lceil c\frac{kT}{E}\right\rceil$$

slots per expert, where $c>0$ is the capacity factor; $C_e$ is the assignment capacity of expert $e$; $T$ is the number of tokens in the routing group; $kT$ is the total number of routed assignments; and $E$ is the expert count. An overflow policy may skip, reroute, or defer assignments. Switch's published formula is the $k=1$ case (Fedus et al. 2022). Other systems use dropless ragged or block-sparse execution instead of a fixed cap (Gale et al. 2023).

Balance losses change optimization, not just utilization

Unregularized routing can concentrate assignments, under-train some experts, and make overloaded expert ranks into stragglers. Shazeer et al. used separate importance and load objectives (Shazeer et al. 2017). Switch introduced a simpler top-1 auxiliary loss

$$\mathcal{L}_{\mathrm{bal}} =\alpha E\sum_{e=1}^{E}f_eq_e, \qquad f_e=\frac{1}{T}\sum_{t=1}^{T} \mathbf{1}\!\left[\arg\max_j p_{t,j}=e\right], \qquad q_e=\frac{1}{T}\sum_{t=1}^{T}p_{t,e}.$$

Here $p_t=\operatorname{softmax}(z_t)$ is the full router distribution for token $t$; $f_e$ is the realized fraction of tokens sent to expert $e$; $q_e$ is its mean router probability; $T$ is the routing-group token count; $E$ is the expert count; $\alpha\ge0$ is the loss coefficient; and $\mathbf{1}[\cdot]$ is an indicator. The product penalizes jointly large traffic and probability mass. Top-$k$ systems use related but not identical definitions.

The coefficient is a real optimization choice: too little balance can create poor utilization, while too much pressure can override task-driven routing. DeepSeek-V3 primarily steers top-$k$ selection with per-expert biases that are adjusted from recent load, while retaining a small sequence-wise auxiliary loss to prevent extreme imbalance (DeepSeek-AI 2024; Wang et al. 2024). The bias affects selection, not the gate weight used to combine expert outputs. This is more precise than saying the model removed all balancing losses.

ST-MoE also introduced a router z-loss (Zoph et al. 2022):

$$\mathcal{L}_{z} =\frac{\beta}{T}\sum_{t=1}^{T} \left(\log\sum_{e=1}^{E}\exp z_{t,e}\right)^2.$$

Here $z_{t,e}$ is the router logit defined earlier; the log-sum-exp term is the router log-partition for token $t$; $T$ and $E$ retain their earlier meanings; and $\beta\ge0$ controls the penalty. Switch also reported selectively using float32 for router computations inside a bfloat16 model (Fedus et al. 2022). Those are experiment-backed stabilizers, not proof that hard top-$k$ selection or low precision is the sole cause of every MoE failure.

Refinements solve different routing problems

Several named designs are alternatives, not a single progression toward one default:

• Top-1 and top-$k$. Switch evaluates one expert; Mixtral evaluates two. Increasing $k$ raises expert arithmetic and communication, approaching dense expert evaluation only as $k$ approaches $E$.
• Expert-choice and balanced assignment. Expert Choice gives each expert a fixed token bucket, eliminating expert-side overflow but giving tokens a variable number of selected experts, possibly zero (Zhou et al. 2022). BASE Layers instead formulate routing as a balanced assignment problem (Lewis et al. 2021).
• Fine-grained and shared experts. DeepSeekMoE divides expert capacity into smaller routed experts and adds always-evaluated shared experts (Dai et al. 2024). In that construction, shared outputs are added without a routed gate. The paper reports improved specialization in its experiments; the label "shared" does not guarantee that a network has learned only general knowledge.
• Sparse upcycling. A dense FFN can be copied into several experts, followed by continued training with a new router (Komatsuzaki et al. 2023). This changes initialization cost, not the eventual dispatch and serving contract.

GLaM and Mixtral demonstrate other points in the design space (Du et al. 2022; Jiang et al. 2024). Their results do not establish a universal best expert count, $k$, capacity factor, or routing loss. Those values depend on the training budget, batch shape, topology, memory, kernels, and target latency.

State-space layers replace a growing history with a recurrence

The second axis concerns sequence mixing. State-space models (SSMs) originate in dynamical-systems and control formulations, long before their use in neural sequence models. S4 showed in 2022 that a structured linear state-space layer could model long sequences efficiently (Gu et al. 2022). Mamba then made key parts of the recurrence input-dependent and supplied a hardware-aware selective scan (Gu and Dao 2024).

From a continuous system to a discrete layer

A linear time-invariant state-space system is

$$\frac{dh(t)}{dt}=Ah(t)+Bx(t), \qquad y(t)=Ch(t)+D_{\mathrm{skip}}x(t).$$

Here continuous time is $t$; $x(t)\in\mathbb{R}^{d_{in}}$ is the input; $h(t)\in\mathbb{R}^{n}$ is the state; $y(t)\in\mathbb{R}^{d_{out}}$ is the output; $A\in\mathbb{R}^{n\times n}$ is the state transition; $B\in\mathbb{R}^{n\times d_{in}}$ writes the input into state; $C\in\mathbb{R}^{d_{out}\times n}$ reads the state; and $D_{\mathrm{skip}}\in\mathbb{R}^{d_{out}\times d_{in}}$ is a direct input path. The dimensions $d_{in}$, $n$, and $d_{out}$ are fixed by the layer.

After discretization at a fixed step, the same layer can be written

$$h_t=\bar A h_{t-1}+\bar Bx_t, \qquad y_t=Ch_t+D_{\mathrm{skip}}x_t.$$

Here integer $t$ is a token position; $x_t$, $h_t$, and $y_t$ are the sampled input, state, and output; and $\bar A$ and $\bar B$ are the discretized transition and input matrices. They are not the continuous matrices $A$ and $B$, even when an implementation derives them from those parameters.

Fixed coefficients give recurrence-convolution equivalence

With $h_0=0$ and fixed $\bar A$, $\bar B$, and $C$, the terms in the unrolled recurrence are

$$K_i=C\bar A^{i}\bar B, \qquad y_t=\sum_{j=1}^{t}K_{t-j}x_j+D_{\mathrm{skip}}x_t.$$

Here $i\ge0$ is a lag; $K_i\in\mathbb{R}^{d_{out}\times d_{in}}$ is the convolution kernel at that lag; $\bar A^i$ is the $i$-th matrix power; $j$ is an earlier token position; and all other symbols retain their definitions above. The recurrent and convolutional views compute the same linear operator. S4 exploits structure in $A$ to generate and apply long kernels efficiently; its theoretical filter-generation and FFT costs should not be collapsed into a blanket claim that every SSM implementation is simply $O(S)$.

Mamba makes the recurrence selective

For one simplified Mamba channel, the token-dependent terms produced by the selection mechanism are

$$(\Delta_t,B_t,C_t)=s_\theta(x_t), \qquad \bar A_t=\exp(\Delta_tA), \qquad \bar B_t= \int_{0}^{\Delta_t}\exp\!\bigl((\Delta_t-\tau)A\bigr)B_t\,d\tau,$$

These terms define the discrete recurrence:

$$h_t=\bar A_th_{t-1}+\bar B_tx_t, \qquad y_t=C_th_t+D_{\mathrm{skip}}x_t.$$

Here $s_\theta$ is a learned projection with parameters $\theta$; $\Delta_t>0$ is an input-dependent discretization step; $B_t$ and $C_t$ are input-dependent write and read maps; $A$ remains a learned but input-independent continuous transition; $\exp$ is the matrix exponential; $\tau$ is the integration variable; and $\bar A_t$ and $\bar B_t$ are the token-dependent discrete maps. Production Mamba layers batch many channels and use structured parameterizations, but the distinction is the same: $\Delta_t$, $B_t$, and $C_t$ are selective; continuous $A$ is not (Gu and Dao 2024).

Input dependence removes the single stationary convolution kernel $K$ from the previous section. Mamba recovers parallel training through a fused associative scan over affine recurrence updates, then uses the recurrent form for autoregressive decoding. Mamba-2 relates a restricted class of SSM transitions to semiseparable mixing matrices and uses chunked matrix multiplication; its paper reports a 2 to 8 times faster core layer than Mamba's selective scan in the tested settings (Dao and Gu 2024). That is a measured result under those implementations, not an architecture-independent speed ratio.

Linear attention is another fixed-state route

Kernelized causal attention can also reorder its computation into a recurrent state. A simplified gated form is

$$S_t=\lambda_tS_{t-1}+\phi(k_t)v_t^{\top}, \qquad z_t=\lambda_tz_{t-1}+\phi(k_t), \qquad y_t=\frac{\phi(q_t)^{\top}S_t} {\phi(q_t)^{\top}z_t+\varepsilon}.$$

Here $q_t$, $k_t$, and $v_t$ are query, key, and value vectors; $\phi$ is a feature map that makes the kernel factorable; $S_t$ is a fixed-shape key-value summary matrix; $z_t$ is its normalizer state; $\lambda_t\in[0,1]$ is an optional learned forgetting gate; $\varepsilon>0$ prevents division by zero; and $y_t$ is the output. The associative factorization, not merely removing softmax, eliminates the need to form every query-key pair. Gated DeltaNet adds targeted delta-rule updates to this family (Yang et al. 2025).

Linear sequence scaling is not automatic speed

With model dimensions fixed, dense full-sequence attention forms a number of query-key relationships proportional to $S^2$, while a recurrent scan performs a number of state updates proportional to $S$. Figure 9.4 plots only that sequence-length growth, normalized at 1,000 tokens. It does not plot FLOPs or measured runtime; projection width, state size, kernel fusion, memory traffic, and hardware determine the constants.

During decode, one dense attention layer compares a new query with a retained key for each cached position. A recurrent layer instead updates a state whose shape is independent of context length. "Independent of context length" does not mean small or free. In the original Mamba implementation, a useful approximation for persistent decode state is

$$M_{\mathrm{rec}} =B_{\mathrm{batch}}L_{\mathrm{rec}}d_{\mathrm{inner}} \bigl(n_{\mathrm{state}}+d_{\mathrm{conv}}\bigr)b \quad\text{bytes}.$$

Here $B_{\mathrm{batch}}$ is the number of resident sequences; $L_{\mathrm{rec}}$ is the number of recurrent layers; $d_{\mathrm{inner}}$ is the number of recurrent channels; $n_{\mathrm{state}}$ is the SSM state width per channel; $d_{\mathrm{conv}}$ is the short-convolution buffer length; and $b$ is bytes per stored element. Exact layouts vary by implementation. The state scales with batch, layers, channel width, state width, convolution width, and precision, but not with cached context length.

The trade is representational. Attention retains a separate cached representation for each past position and permits direct content-based access. A recurrent layer repeatedly compresses prior information into a fixed-shape state, so information can interfere or be lost. Neither mechanism guarantees recall. Retrieval quality must be measured on the target distribution, at the target length, after matched training.

Hybrids choose a layer schedule

MoE and recurrent sequence mixing compose because they occupy different block slots. A hybrid may use attention in some layers, a recurrence in others, and either dense or expert FFNs after each mixer. There is no required ratio.

Released models illustrate different schedules and reporting conventions:

These examples prove that the components can be trained and implemented at large scale. They do not prove that a hybrid dominates full attention on every quality measure or serving workload. MiniMax later chose full attention for M2 and reported that hybrid weaknesses appeared on complex multi-hop tasks at larger scale, while low-precision state, prefix caching, speculative decoding, kernel maturity, and evaluation remained practical concerns (Sun 2025). That first-party retrospective is one team's evidence, not a universal verdict against recurrent or hybrid models.

Record sparse and recurrent architecture as a contract

A reproducible architecture description should name at least:

• which layers use attention, an SSM, linear attention, or another mixer;
• every attention mask, head layout, position rule, and KV-cache data type;
• recurrent state shapes, discretization or gate parameterization, convolution width, reset semantics, and state precision;
• expert count, shared-expert count, selected routed experts, expert FFN shapes, and which layers are sparse;
• router scoring, normalization, top-$k$ tie-breaking, gate renormalization, and output-combination rule;
• balance and z-loss equations with coefficients, plus any bias-update rule;
• routing-group size, capacity factor, overflow or dropless policy, and token ordering after combine;
• expert placement, expert-parallel group, dispatch collective, router precision, and reported total/activated-parameter convention.

Names such as "MoE," "Mamba hybrid," or "80B-A3B" leave most of this contract undefined. Checkpoint tensors, training code, inference kernels, and model cards must agree on it.

Validate each axis before scaling

1. Check parameter accounting. Reconcile stored, selected, shared, router, and always-on parameters with the serialized checkpoint.
2. Test routing exactly. On a small batch, compare sharded dispatch and combine with a single-device reference, including ties, overflow, padding, and zero-token experts.
3. Measure the load distribution. Report assignments per expert, imbalance, dropped or rerouted assignments, padding, entropy, and the slowest expert rank over training and evaluation data.
4. Measure communication. Record dispatch and combine bytes, collective time, compute overlap, and end-to-end tokens per second at realistic batch sizes.
5. Compare scan and recurrence. Confirm the parallel training path and token-by-token recurrent path agree within declared numerical tolerances.
6. Test state boundaries. Reset state at document and sequence boundaries; test packed samples, prefix reuse, state copying, and speculative rollback.
7. Stress retrieval and state tracking. Evaluate direct lookup, order, repeated keys, distractors, and multi-hop use across the deployed length, rather than relying on one needle test.
8. Benchmark the complete workload. Compare quality, training throughput, prefill, decode, memory, batch capacity, and failure recovery against dense and attention-only baselines under matched budgets.

MoE and recurrent layers each remove one dense assumption, but neither removes resource accounting. The useful abstraction is a contract: which weights a token evaluates, where those weights live, how sequence state evolves, and which costs move into communication or serving.

Further reading