Deep Dive — Mixture of Experts & State Space Models (Phase 02)

"MoE and SSMs are not incremental improvements — they are fundamental rethinks of how sequence models should scale."

Section 1: Mixture of Experts — From Dense to Sparse

1.1 The Scaling Problem with Dense Models

Every parameter in a standard transformer is activated for every input token. A 70B dense model runs 70B multiplications per token. This is expensive, and evidence from scaling laws (Hoffmann et al., Chinchilla) shows that most parameters could be replaced by conditional computation — only activate what you need.

Key question: Can we get the capacity of a large model while spending compute equivalent to a smaller one?

Yes: Mixture of Experts.

1.2 The MoE Layer Architecture

Replace the Feed-Forward Network (FFN) in each Transformer block with $N$ independent FFN "experts" and a router that chooses which experts process each token.

Standard Transformer Block:
  Token x → MHA → LayerNorm → FFN → LayerNorm → output

MoE Transformer Block:
  Token x → MHA → LayerNorm → [Router] → [Expert 0, Expert 1, ..., Expert N-1]
                                              ↑ only K experts activated per token

The architecture for a single transformer layer:

class TransformerBlock(nn.Module):
    def __init__(self, d_model, n_heads, n_experts=8, top_k=2):
        super().__init__()
        self.attn = MultiHeadAttention(d_model, n_heads)
        self.norm1 = nn.LayerNorm(d_model)
        self.norm2 = nn.LayerNorm(d_model)
        self.moe   = MoEBlock(d_model, d_ff=4*d_model, n_experts=n_experts, top_k=top_k)
    
    def forward(self, x):
        x = x + self.attn(self.norm1(x))
        moe_out, aux_loss = self.moe(self.norm2(x))
        x = x + moe_out
        return x, aux_loss

Parameter count:

• Dense FFN: $2 \times d_{model} \times d_{ff}$ (e.g., 2×4096×16384 = 134M params)
• MoE with 8 experts: $8 \times 2 \times d_{model} \times d_{ff/8}$ = same total parameters but only top-2 experts activate → 2/8 = 25% of compute

This is how Mixtral 8×7B has 47B parameters but computes like a ~13B model.

1.3 The Router: Soft Routing

The router is a learned linear projection from $d_{model}$ to $N$ expert logits, normalized with softmax:

Input x ∈ ℝ^{d_model}
Logits:  g = W_router @ x ∈ ℝ^N
Probs:   p = softmax(g) ∈ ℝ^N
Top-K:   (indices, weights) = topk(p, K)
Renorm:  weights = weights / sum(weights)

In code:

class TopKRouter(nn.Module):
    def __init__(self, d_model, n_experts, top_k=2):
        super().__init__()
        self.n_experts = n_experts
        self.top_k = top_k
        # NOTE: shape is (n_experts, d_model) so output is logit per expert
        self.router_weight = nn.Linear(d_model, n_experts, bias=False)

    def forward(self, x):
        # x: [batch*seq, d_model]
        logits = self.router_weight(x)           # [N, n_experts]
        probs  = F.softmax(logits, dim=-1)       # [N, n_experts]
        topk_w, topk_idx = probs.topk(self.top_k, dim=-1)  # [N, K] each
        
        # Renormalize weights to sum to 1
        topk_w = topk_w / topk_w.sum(dim=-1, keepdim=True)
        return topk_idx, topk_w

Walkthrough with numbers: Suppose N=6 tokens, 4 experts, top-K=2.

x = [6, d_model]
logits = router(x) → [6, 4]

Example for token 0:  logits = [2.1, 0.3, -0.5, 1.8]
                      probs  = [0.50, 0.09, 0.04, 0.37]
                      topk   = indices=[0,3], weights=[0.50, 0.37]
                      renorm = weights=[0.575, 0.425]
→ Token 0 goes to experts 0 and 3 with weights 0.575, 0.425

The MoE block output for each token: $$y_i = \sum_{k \in \text{TopK}(i)} w_{ik} \cdot E_k(x_i)$$

where $E_k$ is the $k$-th expert FFN.

1.4 Load Balancing: The Critical Problem

Expert collapse: without any penalty, the router quickly learns to always route to the same 1-2 experts (because they get more training signal and become better). Other experts starve and become useless. Your MoE degenerates to a dense FFN.

Measurement: define load imbalance = (tokens to most popular expert) / (ideal even split).

For 8 experts with 100 tokens: ideal = 100/8 = 12.5 tokens each. If expert 0 gets 80 tokens: imbalance = 80/12.5 = 6.4× — near-total collapse.

Solution: Switch Transformer auxiliary loss (Fedus et al., 2021):

For expert $i$:

• $f_i$ = fraction of tokens assigned to expert $i$ (counting top-K assignments)
• $P_i$ = mean router probability for expert $i$ across all tokens

$$\mathcal{L}{aux} = N{experts} \cdot \sum_{i=1}^{N_{experts}} f_i \cdot P_i$$

The minimum of $\mathcal{L}_{aux}$ is $1.0$ (achieved when routing is perfectly uniform). Multiplying $f_i \cdot P_i$ creates a smooth gradient: even though the top-K selection is non-differentiable, $P_i$ is differentiable and carries gradient back to the router weights.

In code:

def compute_aux_loss(probs, topk_indices, n_experts):
    """
    probs:       [N, n_experts]  — soft router probs (differentiable)
    topk_indices:[N, top_k]      — hard selections (non-differentiable)
    """
    N = probs.shape[0]
    
    # f_i: fraction of tokens dispatched to expert i (from hard selections)
    # Count each top-k assignment equally
    one_hot = torch.zeros(N, n_experts, device=probs.device)
    for k in range(topk_indices.shape[1]):
        one_hot.scatter_add_(1, topk_indices[:, k:k+1],
                             torch.full((N,1), 1.0/topk_indices.shape[1], device=probs.device))
    f = one_hot.mean(dim=0)   # [n_experts] — fraction of tokens to each expert
    
    # P_i: mean router probability (differentiable!)
    P = probs.mean(dim=0)     # [n_experts]
    
    aux_loss = n_experts * (f * P).sum()
    return aux_loss

Walkthrough: 4 experts, 4 tokens, top-1:

Routing: [E0, E1, E0, E3]
f = [0.5, 0.25, 0.0, 0.25]    (E0 gets 2/4=0.5)
P = [0.35, 0.20, 0.15, 0.30]  (average softmax probs)
aux = 4 * (0.5*0.35 + 0.25*0.20 + 0.0*0.15 + 0.25*0.30)
    = 4 * (0.175 + 0.05 + 0 + 0.075)
    = 4 * 0.30 = 1.20  (>1.0 → imbalanced, gradient pushes toward 1.0)

In training: add α * L_aux to the main loss. Typical α = 0.01.

1.5 Expert Execution: Dispatching Tokens

The naive implementation loops over every token × every K expert — O(N×K) sequential calls. Production implementations batch tokens per expert:

# Production MoE dispatch (simplified)
def moe_forward_batched(x, router, experts):
    B, T, d = x.shape
    x_flat = x.reshape(B*T, d)        # [N, d]
    N = x_flat.shape[0]
    K = router.top_k
    
    indices, weights = router(x_flat)  # [N, K], [N, K]
    output = torch.zeros_like(x_flat)
    
    # Group tokens by expert for batched execution
    for e in range(router.n_experts):
        # Find all (token, k) pairs that route to expert e
        token_mask = (indices == e)  # [N, K] bool
        if not token_mask.any():
            continue
        
        # Gather tokens going to this expert
        tok_ids, k_ids = token_mask.nonzero(as_tuple=True)
        expert_input = x_flat[tok_ids]    # [M, d] where M = tokens to this expert
        expert_out   = experts[e](expert_input)   # [M, d]
        
        # Scatter back with weights
        output.scatter_add_(0,
            tok_ids.unsqueeze(1).expand_as(expert_out),
            expert_out * weights[tok_ids, k_ids].unsqueeze(1))
    
    return output.reshape(B, T, d)

1.6 MoE in Real Models

Interview question: "Why does Mixtral use top-2, not top-1?"
Answer: Top-1 is simpler but creates extreme routing variance — small perturbations change which expert is used entirely. Top-2 provides smoother gradients and more stable training while still maintaining sparsity.

Section 2: State Space Models — The Linear Alternative to Attention

2.1 Why Attention Has a Quadratic Problem

For a sequence of length $T$:

• Attention matrix: $T \times T$ — quadratic memory, quadratic compute
• At T=32768 and FP16: 32768² × 2 bytes = 2 GB just for attention weights
• This makes long-context training extremely expensive

Can we do sequence modeling in O(T) memory and O(T) compute? State Space Models say yes.

2.2 Classical SSMs — The Continuous Foundation

State Space Models from control theory describe a dynamical system:

$$\dot{h}(t) = Ah(t) + Bx(t)$$ $$y(t) = Ch(t) + Dx(t)$$

Where:

• $x(t) \in \mathbb{R}$ — scalar input signal at time $t$
• $h(t) \in \mathbb{R}^N$ — hidden state vector (the "memory" of the system)
• $y(t) \in \mathbb{R}$ — output
• $A \in \mathbb{R}^{N \times N}$ — state transition matrix
• $B \in \mathbb{R}^{N \times 1}$ — input projection
• $C \in \mathbb{R}^{1 \times N}$ — output projection
• $D$ — skip connection scalar (often D=1 or D=0)

Concrete example: $N=2$ state, like a damped oscillator:

A = [[-0.5, 1.0],    B = [[0.0],    C = [[1.0, 0.0]]
     [-1.0, -0.5]]        [1.0]]

This system: state h = [position, velocity]
             x = external force
             y = position

The negative eigenvalues of A (-0.5 ± i) ensure the system is stable (oscillations decay).

2.3 Discretization via Zero-Order Hold (ZOH)

Computers process discrete sequences, not continuous signals. We need to convert $A, B$ to their discrete counterparts $\bar{A}, \bar{B}$ for time step $\Delta$:

$$\bar{A} = e^{\Delta A}$$ $$\bar{B} = (e^{\Delta A} - I) A^{-1} B$$

For diagonal A (which is what S4/Mamba use for efficiency), this simplifies to element-wise:

$$\bar{A}_i = e^{\Delta \cdot A_i}$$ $$\bar{B}_i = \frac{e^{\Delta \cdot A_i} - 1}{A_i} B_i$$

Why diagonal A? Computing $e^{A}$ for a dense $N \times N$ matrix is $O(N^3)$. For diagonal A, it's $O(N)$.

Stability requirement: $A_i < 0$ ensures $|\bar{A}_i| = |e^{\Delta \cdot A_i}| < 1$ — the state decays over time (like a leaky memory). S4 initializes A using the HIPPO matrix theory (approximating Legendre polynomials for long-range memory), then parameterizes $A_i = -\text{softplus}(A_\log_i)$ to maintain negativity.

Code walkthrough:

def discretize_zoh(A, B, delta):
    """
    A:     [d_state]  diagonal of continuous-time state matrix (negative)
    B:     [d_state]  input projection
    delta: scalar     time step size (learned or fixed)
    
    Returns:
    A_bar: [d_state]  discrete state transition
    B_bar: [d_state]  discrete input projection
    """
    # A_bar = exp(delta * A)    — element-wise since A is diagonal
    A_bar = torch.exp(delta * A)   
    
    # B_bar = (exp(delta*A) - 1) / A * B
    # Note: A < 0, so (A_bar - 1) < 0, and A < 0, giving B_bar > 0
    B_bar = (A_bar - 1.0) / A * B
    
    return A_bar, B_bar

# Example:
A = torch.tensor([-1.0, -2.0, -0.5])  # 3-state SSM, all negative
B = torch.tensor([1.0, 1.0, 1.0])
delta = 0.1  # time step

A_bar, B_bar = discretize_zoh(A, B, delta)
# A_bar ≈ [0.905, 0.819, 0.951]  — memory coefficients, all < 1
# B_bar ≈ [0.095, 0.181, 0.049]  — input gain

2.4 The Discrete Recurrence

Once discretized, the SSM runs as a linear recurrence:

$$h_t = \bar{A} \cdot h_{t-1} + \bar{B} \cdot x_t$$ $$y_t = C \cdot h_t + D \cdot x_t$$

For each time step $t$ we need $h_{t-1}$, so this must run sequentially in RNN mode.

But: for training, we can also run it as a convolution! Unrolling the recurrence:

$$h_1 = \bar{B} x_1$$ $$h_2 = \bar{A}\bar{B}x_1 + \bar{B}x_2$$ $$h_T = \sum_{k=1}^{T} \bar{A}^{T-k} \bar{B} x_k$$

The output $y_T = C h_T = \sum_k C \bar{A}^{T-k} \bar{B} x_k$.

Define the SSM kernel: $\bar{K} = (C\bar{B},\ C\bar{A}\bar{B},\ C\bar{A}^2\bar{B},\ \ldots,\ C\bar{A}^{T-1}\bar{B})$

Then $y = x * \bar{K}$ — a standard 1D convolution! This can be computed in $O(T \log T)$ with FFT.

Training (parallel):  y = conv(x, K)   → O(T log T)  ← convolve entire sequence at once
Inference (causal):   y_t = C*h_t      → O(1) per step ← just maintain state h

This dual-mode property is the key advantage: fast training via convolution, fast inference via recurrence.

2.5 The S4 Model

S4 (Structured State Space Sequence, Gu et al., 2021) makes SSMs practical:

1. Diagonal-plus-low-rank (DPLR) parameterization of A — can be efficiently computed
2. HIPPO initialization: A initialized to approximate memorizing the input using Legendre polynomial projections
3. Convolutional mode for training on GPUs

The S4 layer processes each feature channel independently with its own SSM:

Input:  x ∈ [B, T, d_model]
                    ↓ (process each feature channel independently)
For each feature f ∈ [0, d_model):
    SSM_f: x[:, :, f] → y[:, :, f]
    (each has its own A_f, B_f, C_f ∈ ℝ^{d_state})
Output: y ∈ [B, T, d_model]

2.6 Mamba: Selective State Spaces

The problem with S4: A, B, C are fixed for all inputs. The model can't selectively remember or forget based on content.

Mamba (Gu & Dao, 2023) makes B, C, and Δ input-dependent — "selective":

$$B_t = \text{Linear}_B(x_t)$$ $$C_t = \text{Linear}C(x_t)$$ $$\Delta_t = \text{softplus}(\text{Linear}\Delta(x_t))$$

Now A is still learned (but fixed at inference), while the way inputs are projected into state and read out changes per-token. This gives the model ability to:

• Focus: high $\Delta_t$ → strong input absorption (large $\bar{B}_t$)
• Forget: small $\Delta_t$ → nearly identity state update ($\bar{A}_t \approx I$)
• Select: B, C can attend to specific content in the input

Comparison to attention:

2.7 The Selective Scan: How Mamba is Actually Computed

At inference time: simple recurrence (O(1) per step)

At training time: naive sequential loop is O(T), but GPUs need parallelism. Mamba uses parallel prefix scan (a.k.a. "scan" or "cumulative operation"):

Parallel prefix scan for associative operation $\oplus$:

• Input: $[a_1, a_2, a_3, a_4]$
• Output: $[a_1, a_1 \oplus a_2, a_1 \oplus a_2 \oplus a_3, a_1 \oplus a_2 \oplus a_3 \oplus a_4]$

For SSM recurrence: $h_t = A_t h_{t-1} + B_t x_t$

Define pair operation: $(A_2, B_2) \oplus (A_1, B_1) = (A_2 A_1,\ A_2 B_1 + B_2)$

This is associative — can be computed in $O(\log T)$ depth with $O(T)$ work in parallel!

Mamba's actual implementation (in Triton/CUDA) uses hardware-aware recomputation during backward pass to avoid storing O(T) states.

2.8 Full Mamba Block Architecture

Input x ∈ [B, T, d_model]
         │
         ├──── in_proj ──→ x_in ∈ [B, T, d_inner]
         │                 │
         │         depthwise conv1d (causal)
         │                 │
         │           silu activation
         │                 │
         │         SelectiveSSM (Δ, B, C input-dependent)
         │                 │
         └──── in_proj ──→ z ∈ [B, T, d_inner]
                           │
                   (SSM_out * silu(z))  ← gating
                           │
                     out_proj ──→ output ∈ [B, T, d_model]

The depthwise conv provides local context (a few tokens of short-range dependency) while the SSM provides global context (entire sequence history in the state).

class MambaBlock(nn.Module):
    def __init__(self, d_model, d_state=16, d_conv=4, expand=2):
        super().__init__()
        d_inner = int(expand * d_model)
        self.d_inner = d_inner
        
        # Expand input to two streams: x_in and z
        self.in_proj = nn.Linear(d_model, 2 * d_inner, bias=False)
        
        # Causal depthwise conv for local context
        # padding = d_conv - 1 ensures output length = input length (causal)
        self.conv1d = nn.Conv1d(d_inner, d_inner, d_conv,
                                padding=d_conv-1, groups=d_inner)
        
        # SSM projections
        self.x_proj = nn.Linear(d_inner, d_state*2 + d_inner, bias=False)  # delta, B, C
        self.dt_proj = nn.Linear(d_inner, d_inner)   # refine delta
        
        # A: diagonal, initialized log-spaced for multi-scale memory
        A_log = torch.log(torch.arange(1, d_state+1).float()).repeat(d_inner, 1)
        self.A_log = nn.Parameter(A_log)
        self.D = nn.Parameter(torch.ones(d_inner))
        
        self.out_proj = nn.Linear(d_inner, d_model, bias=False)
    
    def forward(self, x):
        B, T, _ = x.shape
        
        # Split into two streams
        xz = self.in_proj(x)                # [B, T, 2*d_inner]
        x_in, z = xz.chunk(2, dim=-1)       # each [B, T, d_inner]
        
        # Causal depthwise conv (truncate to T to ensure causality)
        x_conv = self.conv1d(x_in.transpose(1, 2))[:, :, :T]  # [B, d_inner, T]
        x_act  = F.silu(x_conv).transpose(1, 2)               # [B, T, d_inner]
        
        # Selective SSM
        y = self.selective_scan(x_act)       # [B, T, d_inner]
        
        # Gating: multiply by silu(z) branch
        output = self.out_proj(y * F.silu(z))
        return output
    
    def selective_scan(self, x):
        B, T, d = x.shape
        A = -F.softplus(self.A_log)    # [d_inner, d_state], negative
        
        # Project x → delta, B_input, C_input
        delta_BC = self.x_proj(x)      # [B, T, d_state*2 + d_inner]
        delta_raw, B_in, C_in = delta_BC.split([d, self.A_log.shape[1], self.A_log.shape[1]], dim=-1)
        delta = F.softplus(self.dt_proj(delta_raw))  # [B, T, d_inner]
        
        # For each batch item, run sequential scan
        # (actual Mamba uses optimized CUDA parallel scan)
        h = torch.zeros(B, d, self.A_log.shape[1], device=x.device)
        ys = []
        for t in range(T):
            dt = delta[:, t, :, None]          # [B, d, 1]
            A_bar = torch.exp(dt * A[None])    # [B, d, d_state]
            B_bar = (A_bar - 1) / A[None] * B_in[:, t, None, :]  # [B, d, d_state]
            h = A_bar * h + B_bar
            y_t = (C_in[:, t, None, :] * h).sum(-1) + self.D * x[:, t, :]
            ys.append(y_t)
        return torch.stack(ys, dim=1)

2.9 Memory Comparison: O(T²) vs O(T)

Let's make this concrete:

At 32k tokens, attention needs 2 GB for a single head, while the SSM state is 1 MB. This is why Mamba-like architectures are the leading approach for truly long-context models (1M+ tokens).

2.10 Interview Pitfalls

Q: "What's the fundamental trade-off between MoE and SSMs?"

MoE addresses width scaling — more parameters, same compute. SSMs address temporal scaling — longer sequences, same memory. GPT-4 reportedly uses both: MoE blocks for width efficiency, plus the base transformer architecture. Future models will likely combine MoE (for capacity) with SSM or linear attention (for long context).

Q: "Why can't you just use a very large d_state in SSMs to approximate attention?"

Theoretically yes, but: (1) SSMs have a finite-impulse response limitation — with $N$ state dimensions, you can represent at most $N$ distinct "memories". Attention can attend to any of $T$ previous tokens. (2) The selectivity of attention (content-based retrieval) is more powerful than SSM's "smooth decay + input projection" model. In practice, Mamba-2 and similar hybrid models alternate between SSM layers and full attention layers to get the best of both.

Q: "How does Mixtral achieve competitive quality with a dense model 4× larger?"

The key insight is that language modeling doesn't require every parameter to be active on every token. Different "topics" or "styles" of computation naturally cluster to different experts. A math-reasoning token activates different experts than a poetry token. The model effectively partitions its capacity by input domain.

Q: "Explain load balancing loss without looking at the code."

"The naive router collapses to using one expert because that expert gets more gradient signal and becomes better, attracting more traffic — a positive feedback loop. The aux loss breaks this by penalizing the product $f_i \times P_i$: $f_i$ is the fraction of tokens going to expert $i$ (from the hard top-K selection, non-differentiable) and $P_i$ is the mean softmax probability for expert $i$ (differentiable). Their product peaks when load is uneven, so minimizing $\sum f_i P_i$ pushes the router toward uniform load. The neat trick is that even though the routing decision itself has no gradient, $P_i$ does — so we can still backpropagate."

Section 3: Quick Setup Guide

Installing dependencies

pip install torch>=2.1.0
# For actual Mamba (hardware-optimized):
pip install mamba-ssm  # requires CUDA
# Or the pure-PyTorch reference:
pip install mamba-minimal

Running your first MoE forward pass

import torch
from solution import MoEBlock

torch.manual_seed(42)
B, T, d = 2, 16, 128
moe = MoEBlock(d_model=d, d_ff=512, n_experts=4, top_k=2)

x = torch.randn(B, T, d)
out, aux_loss = moe(x)

print(f"Input:    {x.shape}")           # [2, 16, 128]
print(f"Output:   {out.shape}")         # [2, 16, 128]
print(f"Aux loss: {aux_loss.item():.4f}")  # > 1.0 if imbalanced

# Check expert utilization
from solution import compute_expert_utilization
util = compute_expert_utilization(moe.router, x)
print(f"Expert loads: {[f'{l:.2f}' for l in util['expert_load']]}")
print(f"Imbalance:    {util['load_imbalance']:.2f}x")

Running your first SSM forward pass

import torch
from solution import MambaBlock, discretize_zoh, s4_recurrence

# Test S4 recurrence
d = 32
A = -torch.ones(d)              # stable (negative diagonal)
B = torch.randn(d)
C = torch.randn(d)
A_bar, B_bar = discretize_zoh(A, B, delta=0.1)

x_seq = torch.randn(64, d)     # sequence of 64 timesteps
y_seq = s4_recurrence(x_seq, A_bar, B_bar, C)
print(f"S4 output: {y_seq.shape}")  # [64, 32]

# Test MambaBlock
mamba = MambaBlock(d_model=32, d_state=8)
x = torch.randn(1, 64, 32)     # [B, T, d_model]
y = mamba(x)
print(f"Mamba output: {y.shape}")   # [1, 64, 32]