JOVANA
Explore Library Glossary Getting Started Three Levels Fields How it works Mission
Join the mission
All guides

Beyond softmax: linear attention, SSMs, and circuits

The frontier asks two questions: can we replace quadratic attention with something linear, and can we read the circuits inside a trained model? Linear attention, Mamba, and induction heads are the three best answers we have.

The quadratic wall and the recurrence dream

Softmax attention compares every token with every other, so cost grows with the square of sequence length, and its KV cache grows linearly forever. The old RNN family had the opposite profile: a fixed-size hidden state and linear-time processing, but no parallel training and a short effective memory. The dream of this whole research thread is to get an RNN's constant-memory inference and a transformer's parallel training and long reach.

The recurrence dream revived: like an RNN unrolled over time, linear attention and SSMs carry a compressed running state instead of an ever-growing cache.

An RNN unrolled across time steps, each step passing a hidden state to the next.

Since Attention Is All You Need this was treated as a dead idea — until two lines of work revived it. Both keep a compressed running state instead of an unbounded cache, and both are written so that training still parallelizes across the sequence. They differ in how that state is computed.

Linear attention: drop the softmax, keep a state

Linear attention starts from a simple algebraic move. Softmax sits between the queries and keys and forces you to build the full score matrix before multiplying by values. Replace the softmax with a kernel feature map applied separately to queries and keys, and the associativity of matrix multiplication lets you fold the keys and values into a single running summary first — turning quadratic cost into linear.

The same equation has two faces. During training you run it in parallel like attention; during generation you run it as a recurrence that updates a fixed-size state per token, giving constant memory and constant per-token cost no matter how long the context. The honest catch is that the fixed state is a lossy summary, so early linear-attention models traded away some of the precise recall that full softmax attention provides.

\text{out}_i=\frac{\phi(q_i)^{\top}\sum_{j\le i}\phi(k_j)\,v_j^{\top}}{\phi(q_i)^{\top}\sum_{j\le i}\phi(k_j)},\qquad S_i=S_{i-1}+\phi(k_i)\,v_i^{\top}

Dropping softmax for a kernel feature map φ lets you reorder the products and keep a fixed-size running state Sᵢ, turning generation into a constant-memory recurrence.

# softmax attention (quadratic): build S = softmax(Q @ K.T), then S @ V
# linear attention (linear): phi = feature_map
state = 0
for t in range(T):           # recurrent inference form
    state = state + outer(phi(k[t]), v[t])   # fixed-size running summary
    y[t] = phi(q[t]) @ state                 # O(1) memory per step
Linear attention exposes a recurrent inference form with a fixed-size state.

State-space models and Mamba: selective scanning

State-space models (SSMs) come at the same goal from control theory. A continuous linear system maps an input sequence to an output through a hidden state governed by a few learned matrices; discretized, it becomes a long convolution that trains in parallel and an efficient recurrence that runs in linear time. Early SSMs matched transformers on long-range synthetic tasks but lagged on language.

Mamba closed that gap with one key change: make the state-space parameters input-dependent so the model can selectively remember or forget based on content, instead of applying the same fixed dynamics to every token. That selectivity is what attention had and classic SSMs lacked. The price is that the convolution view no longer applies, so Mamba relies on a hardware-aware parallel scan to stay fast — the same IO-aware spirit as FlashAttention, applied to a recurrence.

h_t=\bar{A}\,h_{t-1}+\bar{B}\,x_t,\qquad y_t=C\,h_t,\qquad (\bar{A},\bar{B},C)=f(x_t)

An SSM maps the sequence through a hidden state hₜ; Mamba makes the parameters input-dependent so the model selectively remembers or forgets based on content.

Induction heads: a circuit you can read

The other frontier is not a new architecture but a new microscope. Mechanistic interpretability reverse-engineers the algorithms a trained transformer implements, and its landmark discovery is the induction head: a two-layer circuit that finds the previous place a token appeared and predicts whatever followed it last time — in short, 'I've seen this pattern; here's what came next.'

Induction heads are an attention circuit: click a token to see it look back to a previous occurrence and copy what followed it.

Interactive self-attention diagram showing a token attending back to a prior matching position.

Induction heads matter beyond cuteness. They emerge during a sharp, identifiable phase change in training, and that moment coincides with a sudden jump in the model's in-context learning ability. The leading hypothesis is that this simple copy-and-continue circuit is a foundational mechanism behind learning from the prompt — a rare case where we can name a concrete structure and link it to an emergent capability.

Tie the thread together. Every efficiency trick in this track — RoPE, GQA, FlashAttention-2, MoE, sliding windows, YaRN, sinks, ring attention, linear attention, Mamba — reshapes how a transformer spends compute and memory. Interpretability asks the orthogonal question of what it computes. The research that matters next lives where these meet: efficient architectures whose internal circuits we can still read, trust, and steer.