What Comes After Transformers?

The transformer architecture, built on scaled dot-product self-attention, has dominated sequence modeling since Vaswani et al. (2017), yet its O(n²) attention complexity, static computational graphs, and inability to perform true iterative reasoning impose fundamental limits that motivate a new generation of architectures. Post-transformer research spans state-space models, linear attention variants, neural Turing machines redux, diffusion-based reasoners, and hybrid systems that selectively abandon the all-pairs attention bottleneck while preserving or exceeding empirical performance. The counterintuitive fact: several leading candidates — particularly structured state-space models — were derived from classical control theory and signal processing, fields largely ignored by deep learning for decades.

The core insight

The transformer's power comes from its ability to route information between any two tokens in a single layer, but this universality is also its curse: the quadratic cost of computing all pairwise interactions becomes prohibitive at long contexts, and the architecture has no inductive bias toward the kind of sequential, recursive computation that many reasoning tasks demand. The central question driving post-transformer research is whether we can design sequence models that are simultaneously expressive (matching transformer quality), efficient (sub-quadratic in sequence length), and compositionally structured (capable of multi-step algorithmic reasoning). The emerging answer is that no single architecture wins on all three axes — instead, the frontier consists of hybrid systems that dynamically allocate different computational strategies to different parts of a problem. Understanding what comes after transformers therefore requires understanding why transformers fall short, not just what replaces them.

Common trapCommon Misconception

The most common mistake is treating "post-transformer" as a single architectural replacement — assuming that one model class (e.g., Mamba/SSMs) will simply supersede transformers the way transformers superseded LSTMs. This is wrong for two reasons. First, the failure modes of transformers are task-dependent: on long-document retrieval, quadratic attention costs dominate; on formal reasoning, the lack of iterative depth is the bottleneck; on streaming inference, the KV-cache memory is the constraint — and no single architecture addresses all three simultaneously. Second, the empirical record shows that hybrids (e.g., Jamba, Griffin, MambaFormer) consistently outperform pure SSM or pure attention baselines at matched parameter counts, meaning the field is moving toward mixture-of-computational-primitives rather than a clean succession. The correct mental model is an ensemble of inductive biases selected dynamically, not a linear architectural genealogy.

How It Works

Post-transformer architectures attack three core mathematical bottlenecks. We develop each with precision.

1. The Quadratic Attention Bottleneck

Standard scaled dot-product attention computes:

\quad \mathrm{Attn}(Q,K,V) = \mathrm{softmax}\!\left(\frac{QK^\top}{\sqrt{d_k}}\right)V

Explanation: For a sequence of length n, the matrix $QK^\top \in \mathbb{R}^{n \times n}$ requires O(n²d) time and O(n²) space, making attention quadratic in sequence length. Here $Q = XW_Q$ , $K = XW_K$ , $V = XW_V$ with $X \in \mathbb{R}^{n \times d}$ the input and $W_Q, W_K, W_V \in \mathbb{R}^{d \times d_k}$ learned projections. The softmax normalizes over the key dimension. For $n = 10^6$ (genomic sequences, hour-long audio), the $n^2$ term makes this computationally intractable.

2. Linear Attention and Kernel Approximation

The key insight (Katharopoulos et al., 2020) is that softmax can be approximated via a feature map $\phi : \mathbb{R}^{d_k} \to \mathbb{R}^r$ such that $\exp(q^\top k / \sqrt{d_k}) \approx \phi(q)^\top \phi(k)$ :

\quad \widetilde{\mathrm{Attn}}(Q,K,V)_i = \frac{\phi(q_i)^\top \sum_{j=1}^{n} \phi(k_j) v_j^\top}{\phi(q_i)^\top \sum_{j=1}^{n} \phi(k_j)}

Explanation: By associativity of matrix multiplication, the inner sums $\sum_j \phi(k_j)v_j^\top \in \mathbb{R}^{r \times d_v}$ and $\sum_j \phi(k_j) \in \mathbb{R}^r$ can be precomputed once, reducing total complexity to O(nrd) — linear in n. This reformulation is exact when $\phi$ perfectly factors the softmax kernel, and approximate (but empirically competitive) with random Fourier features or polynomial kernels. The challenge is that removing the softmax normalization destabilizes training; several works (RetNet, RWKV, GLA) propose different stabilization strategies.

3. State-Space Models: The Recurrence Formulation

Structured State-Space Models (S4, Mamba) parameterize sequence-to-sequence maps as continuous-time linear systems, discretized to yield a recurrence. The core SSM is:

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

Explanation: The hidden state $h_t \in \mathbb{R}^N$ summarizes all past inputs via the state transition matrix $\bar{A} \in \mathbb{R}^{N \times N}$ , enabling O(1) per-step inference (like an RNN) while the convolution view enables O(n \log n) parallel training. Here $\bar{A}, \bar{B}$ are obtained by discretizing continuous-time matrices $A, B$ with step size $\Delta$ : $\bar{A} = \exp(\Delta A)$ , $\bar{B} = (\Delta A)^{-1}(\bar{A} - I)\Delta B$ . The critical innovation in S4 (Gu et al., 2021) is constraining $A$ to the HiPPO matrix, which provably minimizes the approximation error of polynomial projections of the input history. Mamba (Gu & Dao, 2023) extends this by making $B$ , $C$ , and $\Delta$ input-dependent (selective), breaking the linear time-invariance constraint to enable content-based gating analogous to attention.

4. Expressivity Bounds and the Memorization Gap

A fundamental theoretical result characterizes the separation between SSMs and transformers:

\quad \mathrm{TC}^0 \subsetneq \mathrm{SSM\text{-}expressible} \subseteq \mathrm{L} \subseteq \mathrm{Transformer\text{-}expressible} \subseteq \mathrm{TC}^1

Explanation: Fixed-precision transformers with $\mathrm{poly}(n)$ layers and heads can simulate TC¹ circuits, while bounded-state SSMs are limited to regular languages (TC⁰), establishing a provable expressivity gap on formal reasoning tasks like parity and modular arithmetic. This result (Merrill & Sabharwal, 2023; Grazzi et al., 2024) implies SSMs cannot, in general, solve tasks requiring unbounded counting or explicit state tracking — precisely where transformers with chain-of-thought prompting succeed. Hybrid architectures that interleave attention layers with SSM layers (Jamba, Griffin) achieve TC¹ expressivity while retaining O(n) amortized cost for most tokens.

5. Diffusion and Flow-Based Sequence Models

An orthogonal post-transformer direction treats generation as iterative denoising in a continuous embedding space:

\quad p_\theta(x_0) = \int p(x_T) \prod_{t=1}^{T} p_\theta(x_{t-1} \mid x_t)\, dx_{1:T}

Explanation: Rather than autoregressive left-to-right factorization, diffusion language models (MDLM, PLAID) learn to reverse a noising process over discrete token embeddings, enabling non-autoregressive parallel decoding with iterative refinement. Recent masked diffusion language models (Shi et al., 2024; Sahoo et al., 2024) achieve competitive perplexity with GPT-2 scale transformers while enabling variable-length non-autoregressive generation, though scaling to GPT-4 class remains undemonstrated.

researcher · Albert Gu | b. 1995 | Structured State-Space Models and Mamba

Albert Gu completed his PhD at Stanford under Christopher Ré, where he developed the HiPPO framework for online function approximation and the S4 architecture that demonstrated SSMs could match transformers on long-range dependency benchmarks for the first time. His 2023 Mamba paper introduced selective state spaces with hardware-aware parallel scan algorithms, enabling training throughput competitive with FlashAttention while achieving linear inference cost — the most cited post-transformer architecture paper as of 2024. Gu's background in numerical analysis and control theory led him to derive the HiPPO matrix from Legendre polynomial approximation theory, a connection that surprised the deep learning community.

researcher · Ashish Vaswani | b. 1983 | Transformer Architecture

Ashish Vaswani led the team at Google Brain that published "Attention Is All You Need" (2017), establishing the transformer as the dominant sequence modeling paradigm and directly motivating the question of what might succeed it. His subsequent work at Adept AI on action transformers for computer use demonstrated that transformer limitations in long-horizon planning were fundamental, not merely engineering challenges. Vaswani left Google shortly after the landmark paper — one of the most-cited works in all of computer science — to pursue startup research, and has spoken publicly about his view that the architectural search space beyond attention remains vastly underexplored.

researcher · Tri Dao | b. 1994 | FlashAttention, Mamba Hardware Co-design

Tri Dao (Stanford PhD, advised by Christopher Ré and Atri Rudra) developed FlashAttention (2022), which showed that the practical bottleneck in transformer training is not FLOPs but memory-bandwidth — a reframing that simultaneously extended transformers' practical reach and clarified why hardware co-design is essential for any successor architecture. His collaboration with Albert Gu on Mamba explicitly designed the selective scan algorithm around A100/H100 memory hierarchy, achieving 5× throughput over attention at sequence length 2K. Dao's insight that IO-awareness is as important as algorithmic complexity has become a foundational principle of post-transformer hardware-software co-design.

Figure

Post-Transformer Architecture Landscape

Figure

SSM Dual-Mode: Recurrence vs Convolution

Intuition

The Mental Model

Imagine you're a researcher trying to understand a 10,000-page document. The transformer approach is like hiring a team of analysts who read every page simultaneously and can instantly compare any sentence to any other — incredibly powerful, but requires a conference room proportional to the square of the document length. State-space models are like hiring a single expert who maintains a running mental summary (the state vector $h_t$ ) and updates it as each new page arrives — she can't recall verbatim details from 5,000 pages ago, but she processes each new page in constant time and uses constant memory. Linear attention is like a compromise: the analyst keeps a compressed "sketch matrix" of everything seen so far, which she updates in O(1) per token and queries in O(1) — but the sketch loses information about fine-grained token ordering that softmax attention preserves. Hybrid architectures are like having the expert maintain her running summary (SSM layers) for most of the document, but periodically stopping to conduct a detailed cross-reference meeting (attention layers) at chapter breaks — you get near-constant memory cost with occasional exact recall. The diffusion approach is entirely different: rather than reading left to right, you start with a blurry, noisy version of the whole document simultaneously and iteratively sharpen your understanding across all positions at once — more like an editor refining a draft than a reader processing a stream.

Historical Development

Origins · 1 minhow this was discovered

History

1990

Elman introduces the simple RNN, establishing the recurrence paradigm for sequence modeling that SSMs will later rigorously ground in control theory.

1997

Hochreiter & Schmidhuber publish LSTM, which adds gating to solve the vanishing gradient problem — the same problem that motivates HiPPO's polynomial projection approach forty years later.

2017

Vaswani et al., "Attention Is All You Need," replaces recurrence entirely with scaled dot-product self-attention, achieving state-of-the-art on machine translation and launching the transformer era.

2020

Katharopoulos et al. introduce linear attention via kernel feature maps; Choromanski et al. propose Performer with random Fourier features — first serious theoretical attack on the O(n²) bottleneck.

2021

Gu et al. publish S4, deriving the HiPPO structured state matrix and achieving transformer-competitive performance on Long Range Arena benchmarks, establishing SSMs as a credible alternative.

2022

Dao et al. publish FlashAttention, demonstrating that IO-awareness enables 3× faster exact attention — temporarily extending transformers' practical range and clarifying that hardware co-design is essential for any successor; Sun et al. introduce RetNet with recurrent/parallel duality.

2023

Gu & Dao publish Mamba with selective state spaces and hardware-aware scan; Lieber et al. (AI21) release Jamba, the first production hybrid SSM-Transformer model; RWKV-4 scales to 14B parameters as a pure linear-attention RNN.

2024

Sahoo et al. (MDLM) and Shi et al. (masked diffusion) demonstrate competitive language model perplexity with diffusion-based sequence models; De et al. publish Griffin (Google DeepMind) showing that linear recurrences + local attention match Mamba at matched compute; Behrouz et al. introduce Titans with neural long-term