Linear Attention, Mamba & SSMs

The Quadratic Attention Wall

Standard Transformer self attention compares every token to every other token, scaling quadratically $O(N^2)$.

At sequence lengths of 1 Million tokens, self attention crashes due to compute and memory limits.

State Space Models (SSMs) like Mamba (Gu & Dao, 2023) achieve linear $O(N)$ training time and constant $O(1)$ generation memory per token!

  TRANSFORMER ATTENTION:    N x N Pairwise Matrix ──► Quadratic O(N^2) Compute & Memory Growth
  MAMBA STATE SPACE MODEL:  Selective State Dynamic ──► Linear O(N) Compute & O(1) Step Memory!

Continuous State Space Foundations

Derived from classical control theory, a continuous State Space Model maps an input signal $x(t) \in \mathbb{R}$ to an output signal $y(t) \in \mathbb{R}$ via a hidden state $h(t) \in \mathbb{R}^N$:

$$\text{State Equation: } h'(t) = A h(t) + B x(t)$$

$$\text{Output Equation: } y(t) = C h(t)$$

• $A$: State transition matrix (dictates how hidden state evolves).
• $B$: Input projection matrix.
• $C$: Output projection matrix.

Using discretization step $\Delta$, continuous equations are converted into discrete recurrence steps:

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

The Two Views of SSMs: Convolution vs Recurrence

SSMs possess a mathematical superpower:

  TRAINING MODE (Parallel Convolution):
  y = x * K   (Convolve entire sequence x with 1D long filter kernel K in O(N log N) time on GPUs!)

  INFERENCE MODE (Step-by-Step Recurrence):
  h_t = A_bar * h_{t-1} + B_bar * x_t   (Update tiny state vector in O(1) constant time per token!)

This combines the fast GPU parallel training of Transformers with the efficient $O(1)$ memory generation of RNNs!

Mamba: Selective State Space Models

Legacy SSMs (S4) used time invariant matrices ($A, B, C$ were constant for all tokens).

Because matrices never changed, S4 could not perform content selection (for example, it could not filter out filler words or copy exact facts from context).

Mamba (Gu & Dao, 2023) introduced Selection Mechanisms:

Make matrices $B(x_t)$, $C(x_t)$, and step size $\Delta(x_t)$ data dependent functions of input token $x_t$!

1. If token $x_t$ is an important keyword, $\Delta(x_t)$ is large $\to$ Store $x_t$ deeply in state $h_t$.
2. If token $x_t$ is a filler word, $\Delta(x_t)$ is near zero $\to$ Ignore $x_t$ and pass state $h_t$ through unchanged.

Say this out loud

State Space Models like Mamba replace quadratic self attention with linear time state dynamics. SSMs act as parallel 1D convolutions during GPU training and constant O(1) recurrent steps during token generation. Mamba introduces Selective SSMs by making state transition matrices data dependent, allowing models to process million token context windows efficiently.

Followups to expect

1. What is Hardware Aware Scan in Mamba? Mamba executes data dependent recurrent updates in GPU SRAM cache using a custom parallel scan kernel, avoiding slow HBM memory read write bottlenecks.
2. Can Mamba replace Transformers completely? Hybrid architectures (like Jamba) combining Mamba layers for long context efficiency with Transformer attention layers for exact retrieval achieve optimal overall performance.