Modern model architectures: MoE and state-space models

Go beyond the dense Transformer block and understand the design axes of the models agents run on: encoder / decoder / encoder-decoder and their training objectives; Mixture-of-Experts (top-k routing, sparsity, load balancing, capacity); State Space Models and Mamba (discretization, selectivity, the parallel scan); hybrids; and the complexity trade-offs — above all the structural tension between inference memory and retrieval capacity.

Track

Research

Module

9 · Agent cognition & reasoning

Prerequisites

The Transformer and the residual stream

An agent with long contexts — tool traces, retrieved documents, multi-step histories — pays for every token in memory and latency. The architectures in this lesson all promise to make that cheaper. The scientific question is whether they do so without losing the capability that matters most in those contexts: recovering, with precision, a datum that appeared many tokens earlier.

The short answer, which the lesson justifies, is that there is a structural trade-off between inference memory and retrieval capacity. Attention stores the whole past (a key/value cache that grows with length) and can retrieve any element. A recurrent model compresses the past into a fixed-size state and can therefore only retrieve what fits in it. Mixture-of-Experts acts on a different axis — the per-token cost of the feed-forward network — and does not change that trade-off at all.

The three classic families share the same block (attention, MLP, residual, normalization) and differ in the attention mask and the training objective.

Encoder (bidirectional)

Every position sees every other. Trained by masked-language modeling: corrupt a fraction of positions and reconstruct them. Only the masked positions produce gradient, so it is less token-efficient, but each representation integrates context from both sides — good for classification, extraction, and retrieval encoders. It does not naturally generate text.

Decoder (causal)

Position t sees only positions up to t. Trained by next-token prediction, which factorizes the joint probability exactly and yields a loss term at every position in a single pass. Generation is natural: sample the next token and repeat. This is the dominant family for general-purpose generation and for agents.

Encoder-decoder

A bidirectional encoder reads the input; a causal decoder generates the output while cross-attending to the encoder. Good when the task is genuinely input-to-output (translation, summarization); awkward in multi-turn dialogue, where the encoder cache must be recomputed when the input changes.

The decoder prevailed for general-purpose use because every token carries training signal and any task can be cast as continuation. That is why "the model an agent runs on" is almost always a causal decoder with the modern refinements from the previous lesson (Pre-LN, RMSNorm, gated FFN variants, RoPE).

In a dense block the MLP holds most of the parameters and most of the FLOPs. A MoE layer replaces it with E independent experts and a router that, for each token, picks k of them (k much smaller than E). The output is the convex combination of the chosen experts' outputs, weighted by the router's affinities. The appeal is accounting: FLOPs per token scale with k, while total parameters scale with E — so you store far more knowledge than you compute per token.

synthetic — top-k router with capacity and load-balancing loss (pseudocode)
# x_t: token in the residual stream (dim d);  W_r: (d, E) router
s = softmax(x_t @ W_r)          # affinity of this token to each expert
chosen = top_k(s, k)            # keep the k best experts
w = s[chosen] / s[chosen].sum() # renormalize over the k chosen
y_t = sum(w[e] * expert[e](x_t) for e in chosen)

# Capacity: each expert accepts at most C = ceil(gamma * T * k / E)
# tokens per batch; OVERFLOW tokens are DROPPED for that slot and
# pass through the residual connection only (silent quality loss).

# Load-balancing auxiliary loss (encourages uniform expert use):
#   L_bal = alpha * E * sum_i( f_i * P_i )
#   f_i = fraction of tokens whose top expert is i   (not differentiable)
#   P_i = mean router probability for expert i        (differentiable)
# The gradient pushes down the probability of overloaded experts.

The hard part is load balancing. If a few experts start with slightly higher scores, they receive more tokens, train more, improve, and receive still more — a positive feedback loop that ends in collapse (a few active experts, the rest idle). Two remedies appear in the literature: an auxiliary balancing loss (plus a router "z-loss" that penalizes large logits and stabilizes low-precision training), and an auxiliary-loss-free scheme that adds a per-expert bias used only for selection and nudges it after each batch according to load. The source reports that the bias scheme can achieve better balance and perplexity than the auxiliary loss in the studied regimes — treat these as MEASURED in those conditions, not as universal laws.

A linear State Space Model (SSM) maps an input signal to an output through a hidden state that evolves by a linear recurrence: the state is updated from the previous state and the current input, and the output reads the state. In continuous time this is a differential equation; to use it in a network it is discretized with a step Δ, giving a discrete recurrence h_t = Ā·h_{t−1} + B̄·u_t and y_t = C·h_t. The step Δ has a physical reading: a small Δ makes Ā close to the identity and B̄ near zero, so the state is preserved and the input barely matters; a large Δ overwrites the state with the current input. That reading is the seed of selectivity.

The key practical fact is a duality. If the SSM parameters do not depend on t (a linear time-invariant system), unrolling the recurrence turns it into a convolution with a fixed kernel. The same operator then admits TWO algorithms: a recurrent one, constant state per step, ideal for generation; and a convolutional one, computed with an FFT, parallel over the sequence, ideal for training. This dual nature is what made deep SSMs practical, and a good choice of the state-transition structure (so the state approximates the history of the signal rather than forgetting exponentially) is what made them competitive on language.

synthetic — the two equivalent algorithms for a time-invariant SSM
# RECURRENT form (generation): O(N) per step, fixed-size state h
h = zeros(N)
for t in range(L):
    h = Abar @ h + Bbar * u[t]      # update fixed-size state
    y[t] = C @ h                    # read it out

# CONVOLUTIONAL form (training): precompute the kernel, then FFT
K = [C @ matrix_power(Abar, j) @ Bbar for j in range(L)]  # impulse resp.
y = fft_convolve(K, u)              # parallel over the whole sequence
# Both produce the SAME output (a PROVEN identity for the LTI case).

A time-invariant SSM treats every token identically. On a selective-copy task — memorize certain tokens scattered among filler — a time-invariant filter cannot decide "this token matters, that one does not", because its response depends on relative position, not on content. Mamba makes the step Δ and the input/output projections DEPEND on the input. If Δ_t goes to zero the token is ignored and the state is preserved; if Δ_t is large the state is reset with that token — a gate, like in an LSTM, but with linear state dynamics.

1. 1Make Δ, B and C functions of the current token, so the model can decide how much each token writes into and reads from the state.
2. 2Pay the price: with input-dependent parameters the system is no longer time-invariant, the fixed convolution kernel is gone, and only the sequential recurrence remains.
3. 3Recover parallelism: the update h_t = a_t·h_{t−1} + b_t is a linear recurrence, and its (a, b) pairs compose under an associative operator, so all prefixes can be computed with a parallel scan in logarithmic depth.
4. 4Make the scan hardware-aware: keep the expanded state in fast on-chip memory rather than materializing it in main GPU memory — the same memory-hierarchy idea as FlashAttention.

There is a clarifying duality with attention: restricting the state transition appropriately, a selective SSM is equivalent to a linear attention with a structured causal mask, whose state is a matrix that accumulates associations — an associative memory. That form explains both the efficiency (fixed-size state) and the limit (how many associations fit). With near-orthogonal keys of dimension N, interference grows roughly like sqrt(M)/N, so retrieval works while the number of associations M is at most about N and degrades afterward — the classic behavior of linear associative memories.

Per layer, for sequence length L, dimension d and state size N: exact attention costs on the order of L²·d in time and O(L²) memory if materialized (O(L·d) with FlashAttention), with an O(L·d) cache at inference. A time-invariant SSM trains in L·log L time (FFT) and generates in O(N·d) time with an O(N·d) state that does NOT grow with length. A selective SSM (Mamba) trades the FFT for a parallel scan but keeps the fixed-size state. MoE scales its cost with k experts and its parameter count with E, on a different axis entirely.

Two readings. In training, quadratic attention does not dominate at moderate context because the MLP also costs on the order of L·d²; the SSM advantage appears when L greatly exceeds d. At inference the difference is structural: the attention cache grows with every token and with batch size, limiting concurrent requests, while the SSM state does not grow. Formal results in the literature sharpen this: a shallow Transformer can copy strings of length exponential in its number of heads, while any fixed-size-state model is capped by its state size; and both Transformers and linear SSMs sit in a complexity class that cannot express certain state-tracking problems despite the recurrent form. These are expressivity results — PROVEN about what a mechanism CAN represent — not statements about what training finds or about the lengths that matter in practice.

Empirically, the source reports that pure SSMs match Transformers on aggregate perplexity at billion-parameter scales but fall behind on tasks that require copying or precise retrieval, and that hybrids — a few attention layers interleaved among many recurrent ones — recover those capabilities while keeping most of the efficiency. The optimal fraction of attention and its dependence on the task remain open (UNKNOWN).

A retrieval-augmented system expects the model to cite and use specific data from the passages it inserts — associative in-context retrieval, precisely the weakest capability of fixed-state models. An agent that accumulates tool traces (identifiers, amounts, paths) and must reuse them verbatim ten steps later likewise depends on exact copying. The practical consequence: choosing a recurrent or hybrid model for an agent must be evaluated with retrieval and copy tasks over the agent's OWN trace format — multi-needle "needle in a haystack" with structured data — not with perplexity or general benchmarks.

Consistent with this lab's thesis, data that must survive literally (an invoice number, an account id) is better kept in explicit state OUTSIDE the model and re-injected when needed, rather than trusting the model's memory — and this is doubly true for fixed-state architectures, where the fixed memory size bounds how much can be recovered at all. For MoE specifically, token dropping and batch-dependent outputs are a reproducibility hazard that an auditable agent must account for.

• ▸MoE: router collapse (few active experts), silent token dropping that degrades quality without an error, batch-composition non-determinism, and router logit instability in low precision.
• ▸SSM: interference once the number of stored associations exceeds the state capacity; a recency bias when the state forgets; and subtle interactions between the discretization step and the state spectrum that can destroy capacity (aliasing).
• ▸Hybrids: the placement and number of attention layers are high-impact, under-studied hyperparameters.
• ▸Most large-scale comparisons do not hold data, compute and recipe fixed, so "architecture A beats B" often confounds architecture with training data; the few controlled studies carry the most weight.
• ▸Expressivity theorems cut both ways: that a mechanism CAN represent a function does not mean training finds it, and that it cannot in the limit does not mean it fails at the lengths you care about.

Expected direction: The risk is that a fixed-size recurrent state cannot losslessly retain many distinct IDs, so exact recall of an arbitrary earlier ID may degrade even though language quality looks fine. I would build a retrieval/copy benchmark over the agent's own trace format — many IDs inserted, then a request to reproduce a specific one verbatim at realistic distances and ledger sizes — and measure exact-match accuracy, not perplexity. Perplexity is an aggregate next-token metric dominated by easy tokens; the retrieval gap is concentrated in exactly the tokens that require recovering something seen earlier, so it is invisible in the aggregate. I would also keep the IDs in explicit external state and re-inject them, so correctness does not depend on the model's memory regardless of architecture.

MoE and state-space models both "make long context cheaper". Why is that description misleading?

They act on different axes. MoE reduces per-token FLOPs of the feed-forward network by activating only k of E experts, but keeps attention and its growing cache, so it does not change long-context memory or retrieval. State-space models change the sequence mixer — a fixed-size state instead of a growing cache — which does affect long-context memory, but at the cost of bounded retrieval capacity. Only the second touches the inference-memory/retrieval trade-off.

What does Mamba gain by making its parameters input-dependent, and what does it lose in exchange?

It gains selectivity: the model can decide, per token, how much to write into and read from the state (a content-dependent gate), which solves selective-copy tasks a time-invariant filter cannot. It loses time-invariance, so the fixed convolution kernel disappears and the efficient FFT training form is gone; parallelism is recovered instead through an associative parallel scan over the linear recurrence.

Why can a fixed-size-state model never match attention on arbitrary exact retrieval, and how does that show up as interference?

A fixed state of size N×d can store only a bounded number of associations; attention keeps every key–value pair explicitly. When the number of stored associations M approaches N, reading one back returns the target plus an interference term from the others that grows roughly like sqrt(M)/N, so retrieval degrades once M exceeds about N. Selectivity helps decide WHAT to store but does not increase HOW MUCH fits.

Why is token dropping in MoE a problem for an auditable agent specifically?

Capacity is enforced per batch, so whether a given token overflows its expert and is dropped (passing through only the residual) can depend on which other tokens share its batch. That makes the output batch-dependent and non-deterministic, degrading quality silently and without an error — which undermines reproducibility and audit trails that an agent in a regulated or high-stakes setting needs.