Efficient Architecture Design

Chapter 1: Sequence Modeling Primitives

Before we build new architectures, we need to understand the three fundamental ways to mix information across a sequence. Every sequence model is built from some combination of these primitives: convolutions, recurrences, and attention.

Convolutions

A convolution slides a kernel (a small weight vector) across the input and computes weighted sums at each position. The output at position n is a dot product of the kernel with a window of the input centered (or ending) at n.

Where h is the kernel and x is the input sequence. The kernel weights are fixed — the same weights are applied at every position. This means the mixing pattern is input-independent.

Properties: parallelizable during training (all positions computed at once via FFT), but the kernel has finite reach. A kernel of size K can only look K steps back.

Recurrences

A recurrence processes the sequence one token at a time, maintaining a hidden state that summarizes everything seen so far.

Properties: the hidden state has potentially infinite context (it can remember anything from the past), inference is O(1) per step (just update the state), but training is sequential — you can't compute h_k without h_k−1. Unless the function f is linear, in which case we can parallelize it.

Attention

Attention computes a weighted combination of all values, where the weights depend on pairwise similarity between queries and keys.

Properties: fully parallelizable, input-dependent mixing (the mixing matrix changes with every input), but O(N²) in both compute and memory.

*Linear recurrences are input-independent. Mamba makes them input-dependent — we'll get there in Chapter 6.

Chapter 2: State Space Models

A state space model (SSM) is a system borrowed from control theory. It describes how a hidden state evolves over time in response to an input signal, and how we observe the output. The continuous-time form has four matrices:

Think of it like this: you're in a dark room with a complex system of springs and masses (the state x). You push it with an input u. The springs respond according to their physics (matrix A). You can only observe the system through a small window (matrix C). The state captures everything happening inside, but you only see a projection.

Multi-Channel Processing

In practice, a sequence model has D input channels (e.g., D=768 for a typical model dimension). Each channel gets its own independent SSM with its own A, B, C matrices. So a single SSM layer is actually D parallel SSMs, each with hidden state size N. The total hidden state is [D, N].

Chapter 3: Discretization

The continuous SSM is beautiful in theory, but computers work with discrete sequences. We need to convert x'(t) = Ax(t) + Bu(t) into a discrete recurrence: x_k = Ax_k-1 + Bu_k. This process is called discretization.

Step 1: Zero-Order Hold (ZOH)

The simplest approach: assume the input u(t) is constant between sample points. Between time t_k and t_k+1 = t_k + Δ, the input is u_k. Under this assumption, the continuous ODE has an exact solution:

So A = e^AΔ and B = (e^AΔ − I) A⁻¹ B. The step size Δ controls the temporal resolution: small Δ = fine-grained sampling, large Δ = coarse sampling.

Step 2: The Bilinear Transform (Tustin's Method)

Computing matrix exponentials is expensive. The bilinear transform (also called Tustin's method or the trapezoidal rule) gives us an algebraic approximation that is both cheap and numerically stable.

The idea: approximate the derivative x'(t) using the trapezoidal rule — the average of the slope at the start and end of each interval:

Substituting x' = Ax + Bu and solving for x_k:

Worked Example: Scalar Case

Let A = −1, B = 1, Δ = 0.5. Then:

derivation
# Bilinear transform (scalar case)
A = -1, B = 1, delta = 0.5

A_bar = (1 - delta*A/2)**(-1) * (1 + delta*A/2)
      = (1 + 0.25)**(-1) * (1 - 0.25)
      = (0.8) * (0.75)
      = 0.6

B_bar = (1 - delta*A/2)**(-1) * delta * B
      = 0.8 * 0.5 * 1
      = 0.4

# Recurrence: x_k = 0.6 * x_{k-1} + 0.4 * u_k
# |A_bar| = 0.6 < 1, so the system is stable

Chapter 4: The HiPPO Framework

Here's a deceptively hard question: how should you initialize the matrix A? Random initialization won't work — the hidden state would either explode or forget everything instantly. We need A to do something principled: compress the input history into the state in an optimal way.

This is the problem that the HiPPO framework (High-Order Polynomial Projection Operators) solves. The insight is beautiful: choose A so that the hidden state x(t) stores the coefficients of an optimal polynomial approximation of the input history.

The HiPPO-LegS Matrix

The most important instance is HiPPO-LegS (Legendre with Scaled measure), which approximates the history over a sliding window. The resulting matrix has a specific structure:

This produces a lower-triangular matrix with entries that grow in a structured way. The key properties:

• The state captures both recent and ancient information (different state components operate at different timescales)
• Earlier components (small n) capture fast, recent features; later components capture slow, long-range features
• The polynomial basis gives a mathematically optimal compression of the history under the L2 norm

Worked Example: 4-Dimensional HiPPO

For N=4, the HiPPO-LegS matrix is:

python
import numpy as np

def hippo_legs(N):
    A = np.zeros((N, N))
    for n in range(N):
        for k in range(N):
            if n > k:
                A[n, k] = -np.sqrt(2*n+1) * np.sqrt(2*k+1)
            elif n == k:
                A[n, k] = -(n + 1)
    return A

A = hippo_legs(4)
# A = [[-1.0,  0.0,  0.0,  0.0],
#      [-1.73, -2.0,  0.0,  0.0],
#      [-2.24, -2.45, -3.0,  0.0],
#      [-2.65, -3.0, -3.32, -4.0]]

Chapter 5: S4 — Structured State Spaces

We now have all the ingredients: a continuous SSM (Ch 2), discretization (Ch 3), and HiPPO initialization (Ch 4). The S4 architecture (Gu et al., 2022) combines them with one crucial trick: it shows that the same SSM can be computed as either a recurrence or a convolution, and you can switch between them.

The Duality: Recurrence ↔ Convolution

After discretization, our SSM is a linear recurrence:

Let's unroll this recurrence starting from x₀ = 0:

See the pattern? Each output y_k is a weighted sum of all past inputs, with weights CA^jB. That's a convolution! The kernel is:

And the output is simply y = K * u (convolution).

FFT Convolution: Why It's Fast

Naive convolution of two length-N sequences is O(N²). But the FFT Convolution Theorem states that convolution in the time domain equals pointwise multiplication in the frequency domain:

The Fast Fourier Transform computes FFT in O(N log N). So the full pipeline is: FFT the kernel (N log N), FFT the input (N log N), pointwise multiply (N), inverse FFT (N log N). Total: O(N log N) instead of O(N²).

The Diagonal-Plus-Low-Rank Trick

There's one more obstacle: computing the kernel K = (CB, CAB, ..., CA^L-1B) requires L matrix-vector products, which is O(LN²) for a dense N×N matrix A. S4's key structural insight: decompose A into a diagonal matrix plus a low-rank correction (DPLR). This allows computing the kernel in O(N log² N) using a clever recursive doubling scheme. Later work (S4D) showed that simply using a diagonal A works almost as well, making the kernel computation trivially O(NL).

Chapter 6: Mamba — Selective SSMs

S4 was a breakthrough for long-range tasks like audio and signal processing. But when researchers tried it on language modeling, something frustrating happened: it consistently underperformed Transformers. The perplexity gap was significant — about 1-3 points worse than attention-based models of the same size.

Why? The answer lies in a task called recall.

The Recall Problem

Consider this sequence: "A 3 B 7 C 2 ... What is A?" To answer, the model needs to look back, find "A 3", and return "3". This requires the mixing matrix (how tokens influence each other) to be input-dependent. The model needs to route "A?" back to wherever "A 3" appeared.

In a Transformer, the attention matrix is input-dependent — queries and keys are computed from the input, so the model can learn to attend to the right location. In an SSM, the matrices A, B, C are fixed — the same mixing pattern applies regardless of input content. The SSM literally cannot vary its routing based on what it sees.

Mamba's Solution: Make B, C, and Δ Input-Dependent

Mamba (Gu & Dao, 2023) keeps the linear recurrence structure but makes three critical matrices functions of the input:

The Δ Knob: Selective Forgetting

The step size Δ is particularly powerful. Remember: Δ controls how much of the continuous dynamics to apply. A large Δ means the state changes a lot — the model is "paying attention" to this token, absorbing it into memory. A small Δ means the state barely changes — the model is "ignoring" this token, letting it pass through.

This is exactly the selection mechanism: when the model sees an important token (like "A 3" in a recall task), it can increase Δ to write it into the state. When it sees filler tokens, it can decrease Δ to preserve the existing state.

Hardware-Aware Parallel Scan

Without the FFT trick, how does Mamba train efficiently? It uses a parallel scan (also called prefix sum). The key insight: a linear recurrence x_k = a_kx_k-1 + b_k can be computed in O(log N) parallel steps using a tree reduction, even though it looks sequential.

Mamba's implementation is hardware-aware: it fuses the discretization, scan, and multiplication into a single CUDA kernel that minimizes HBM reads/writes — the same philosophy as FlashAttention but applied to recurrences. The result: Mamba trains at 3-5x the throughput of a Transformer of similar size.

Chapter 7: Architecture Comparisons

We now have a rich landscape of sequence architectures. Each makes different tradeoffs between quality, training efficiency, inference efficiency, and the ability to handle long-range dependencies. Let's put them side by side.

The Quality-Efficiency Frontier

When to Use What

The Scaling Picture

At 360M parameters trained on 10B tokens of the Pile, the perplexity comparison looks like this:

The gap is narrowing, but attention-based models still lead on language. The story is different on audio (S4 beats Transformers) and very long sequences (where Transformers simply can't run).

Chapter 8: Connections

We've traveled from the quadratic bottleneck of attention, through the primitives of sequence modeling, into the elegant mathematics of state space models, and emerged with architectures that process sequences in linear time. Here's the cheat sheet.

Key Takeaways from CS 229s

1. Efficiency is not just about FLOPs. Convolutions are O(N log N) vs attention's O(N²), but attention solves recall more efficiently. Think about efficiency per task, not just per sequence length.
2. Error analysis drives architecture design. The discovery that SSMs fail at recall came from carefully comparing prediction errors between attention and SSM models. This led directly to Mamba.
3. New primitives need new implementations. Just as FlashAttention made attention hardware-efficient, Mamba needed a custom CUDA kernel for its parallel scan. The algorithm-hardware co-design is essential.

What's Next

The field is moving fast. Mamba-2 introduces the State Space Duality framework (connecting SSMs to structured attention). Griffin and Hawk explore gated linear recurrences. Based and Gated Linear Attention close the gap from the attention side. The boundary between "SSM" and "attention" is blurring.

Related Lessons

• microMamba — A from-zero introduction to SSMs and Mamba (complementary to this systems-focused lecture)
• microTransformer — The attention mechanism in full detail