From Transformers to Mamba: Is Attention All We Need?
Published:
Learning Sequential Inputs
Given a sequence $\mathbf{x}_1, \ldots, \mathbf{x}_T \in \mathbb{R}^d$, we want representations for prediction (language modeling, translation, time series, etc.). Unlike graphs, where the symmetry group is permutations (see my GNN post), sequences are ordered: the model must be order-sensitive, handle variable length, and capture long-range dependencies. The architectures below differ mainly in how computation scales with $T$ and how far information can travel:
- RNN/LSTM: $O(T)$ compute, $O(1)$ memory per step, but inherently sequential and information decays over time;
- Attention: $O(T^2)$ compute, fully parallel, direct interaction between any two positions;
- SSM: $O(T)$ compute as recurrence (or $O(T \log T)$ as convolution), parallel training, constant-size state at inference.
Recurrent Neural Network (RNN)
The vanilla RNN maintains a hidden state:
\[\mathbf{h}_t = \tanh (\mathbf{W}_h \mathbf{h}_{t-1} + \mathbf{W}_x \mathbf{x}_t + \mathbf{b}), \; \mathbf{y}_t = \mathbf{W}_y \mathbf{h}_t.\]Training uses backpropagation through time. The gradient with respect to a distant state contains the product
\[\frac{\partial \mathbf{h}_t}{\partial \mathbf{h}_s} = \prod_{k = s+1}^{t} \text{diag} (\tanh' (\cdot)) \, \mathbf{W}_h^T,\]whose norm decays or explodes exponentially in $t - s$, depending on the spectrum of $\mathbf{W}_h$ — the vanishing/exploding gradient problem [Bengio et al. 1994, Pascanu et al. 2013]. Exploding gradients are handled by clipping; vanishing gradients require architectural changes.
Long Short-Term Memory (LSTM)
LSTM [Hochreiter et al. 1997] introduces a cell state $\mathbf{c}_t$ with additive (rather than multiplicative) updates, gated by forget/input/output gates:
\[\begin{align*} \mathbf{f}_t, \mathbf{i}_t, \mathbf{o}_t & = \sigma (\mathbf{W}_{\lbrace f,i,o \rbrace} [\mathbf{h}_{t-1}, \mathbf{x}_t] + \mathbf{b}_{\lbrace f,i,o \rbrace}), \\ \tilde{\mathbf{c}}_t & = \tanh (\mathbf{W}_c [\mathbf{h}_{t-1}, \mathbf{x}_t] + \mathbf{b}_c), \\ \mathbf{c}_t & = \mathbf{f}_t \odot \mathbf{c}_{t-1} + \mathbf{i}_t \odot \tilde{\mathbf{c}}_t, \\ \mathbf{h}_t & = \mathbf{o}_t \odot \tanh (\mathbf{c}_t). \end{align*}\]The path $\mathbf{c}_{t-1} \to \mathbf{c}_t$ is modulated only by $\mathbf{f}_t$, so gradients survive over long horizons when the forget gate stays open — a learned, input-dependent memory. GRU [Cho et al. 2014] is a cheaper two-gate variant. Keep the phrase “input-dependent gating” in mind: it returns in Mamba.
Transformers: Attention is All You Need
The transformer [Vaswani et al. 2017] removes recurrence entirely. Scaled dot-product self-attention computes, from token embeddings $\mathbf{X} \in \mathbb{R}^{T \times d}$,
\[\mathbf{Q} = \mathbf{X} \mathbf{W}_Q, \; \mathbf{K} = \mathbf{X} \mathbf{W}_K, \; \mathbf{V} = \mathbf{X} \mathbf{W}_V, \; \text{Attn} (\mathbf{X}) = \text{softmax} \bigg(\frac{\mathbf{Q} \mathbf{K}^T}{\sqrt{d_k}}\bigg) \mathbf{V}.\]Attention is a differentiable soft dictionary lookup: each query retrieves a convex combination of values, weighted by query-key similarity. Multi-head attention runs $h$ such lookups in parallel subspaces and concatenates. Since attention is permutation equivariant, order must be injected via positional encodings (sinusoidal in the original paper; learned, relative, or rotary [Su et al. 2021] in modern variants). The price is $O(T^2)$ time and memory in the sequence length, plus a KV cache growing linearly with context at inference.
Interactive illustration: attention as similarity
Six tokens live in a 2D embedding space (drag them!). The highlighted token is the query; edge thickness and the bar chart show its attention weights $\text{softmax} (\langle \mathbf{q}, \mathbf{k}_i \rangle / \tau)$ over all tokens. Move tokens closer to the query to attract its attention; lower the temperature to sharpen the lookup toward hard retrieval.
Low-Rank Adaptation (LoRA)
Fine-tuning all weights of a large pretrained transformer is wasteful: the update lives in a low intrinsic dimension. LoRA [Hu et al. 2021] freezes $\mathbf{W}_0$ and learns a rank-$r$ update
\[\mathbf{W} = \mathbf{W}_0 + \frac{\alpha}{r} \mathbf{B} \mathbf{A}, \; \mathbf{B} \in \mathbb{R}^{d \times r}, \; \mathbf{A} \in \mathbb{R}^{r \times k}, \; r \ll \min (d, k),\]typically applied to $\mathbf{W}_Q, \mathbf{W}_V$. The update merges into $\mathbf{W}_0$ at inference, so there is zero latency overhead. This is the fine-tuning analogue of the low-rank factorizations in my model compression post — and the choice of the rank $r$ is again the empirically-tuned hyperparameter that one would like to predict automatically.
Vision Transformers (ViT)
ViT [Dosovitskiy et al. 2020] treats an image as a sequence: split into $16 \times 16$ patches, linearly embed, prepend a [CLS] token, add positional embeddings, and feed to a standard transformer encoder. With no convolutional inductive bias (locality, translation equivariance), ViT underperforms CNNs on small data but overtakes them with large-scale pretraining — the bias is learned from data instead of built in. DeiT [Touvron et al. 2021] closes the gap on ImageNet-scale data via distillation (from a CNN teacher, see KD).
State Space Model (SSM): Maybe Attention isn’t All You Need
A continuous linear state space model is
\[\mathbf{h}' (t) = \mathbf{A} \mathbf{h} (t) + \mathbf{B} x (t), \; y (t) = \mathbf{C} \mathbf{h} (t),\]discretized with step $\Delta$ (e.g., zero-order hold: $\bar{\mathbf{A}} = \exp (\Delta \mathbf{A})$, $\bar{\mathbf{B}} = (\Delta \mathbf{A})^{-1} (\exp (\Delta \mathbf{A}) - \mathbf{I}) \Delta \mathbf{B}$) into a linear recurrence
\[\mathbf{h}_t = \bar{\mathbf{A}} \mathbf{h}_{t-1} + \bar{\mathbf{B}} x_t, \; y_t = \mathbf{C} \mathbf{h}_t.\]Because the recurrence is linear and time-invariant, unrolling it gives a convolution $y = \bar{\mathbf{K}} * x$ with kernel $\bar{\mathbf{K}} = (\mathbf{C} \bar{\mathbf{B}}, \mathbf{C} \bar{\mathbf{A}} \bar{\mathbf{B}}, \mathbf{C} \bar{\mathbf{A}}^2 \bar{\mathbf{B}}, \ldots)$: train as a convolution (parallel, FFT), infer as a recurrence (constant state, no KV cache). S4 [Gu et al. 2022] makes this stable and expressive by initializing $\mathbf{A}$ with HiPPO theory [Gu et al. 2020], which optimally compresses the input history onto polynomial bases — solving the long-range dependency problem that plagued RNNs.
Mamba [Gu et al. 2023] makes the SSM selective: $\mathbf{B}, \mathbf{C}, \Delta$ become functions of the input $x_t$. This is exactly the LSTM lesson (input-dependent gating: $\Delta \to \infty$ resets the state, $\Delta \to 0$ ignores the token), but applied to a linear state space. Selectivity breaks time-invariance, hence the convolutional view — Mamba compensates with a hardware-aware parallel scan, keeping $O(T)$ training and $O(1)$-state inference while matching transformers of similar size on language modeling.
So, is attention all we need? Attention performs exact, content-based retrieval over an uncompressed context at $O(T^2)$ cost; selective SSMs maintain a fixed-size compressed summary at $O(T)$. The two are increasingly understood as points on a spectrum — linear attention is a degenerate SSM [Katharopoulos et al. 2020], and the duality is formalized in Mamba-2 [Dao et al. 2024]. Hybrid stacks (e.g., Jamba [Lieber et al. 2024]) suggest the practical answer: attention is what you need for a few layers of precise retrieval, and an SSM is what you need everywhere else.