Post-Transformer Architectures

Background

Transformers quadratic attention complexity limits sequence length - driving research into sub-quadratic alternatives
Mamba (Gu & Dao, 2023) introduced selective state space models as a competitive RNN-like alternative
Mixture of Experts (MoE) routes tokens to specialized sub-networks, scaling capacity without proportional compute

The Concept

Post-transformer architectures explore alternatives to or extensions of the standard Transformer block. Key directions: state space models (linear-time sequence modeling), mixture of experts (conditional computation), and hybrid architectures that combine attention with convolutions or SSMs.

How it Links

SSMs (Mamba, RWKV) connect to control theory and the theory of linear time-invariant systems
MoE connects to ensemble methods and conditional computation in classical ML
Flash Attention and linear attention are efficiency improvements within the Transformer framework itself

Concrete Example

Processing a book (100,000 tokens): a standard Transformer requires 100,000^2 = 10 billion attention computations. Mamba processes the same sequence with linear complexity, maintaining a compressed state that captures the essential information from past tokens without attending to all of them explicitly.

Prerequisites

Transformers, RNNs, LSTMs, GRUs, Attention Mechanisms

Learning outcomes

Derive the O(n²) memory and compute cost of standard attention and explain why this limits long-context applications.
Explain the S4/Mamba selective state space model and the hardware-aware parallel scan that makes it efficient.
Describe how RWKV reconciles RNN-style recurrent inference with Transformer-style parallel training.
Compare RetNet, Hyena, and hybrid architectures on training throughput, long-context scaling, and hardware efficiency.
Explain FlashAttention's IO-aware tiling approach and why it achieves speedups without changing the mathematical output.

TL;DR

Standard self-attention [8] costs $O(n^2)$ in memory and compute for sequence length $n$ , making it expensive for contexts beyond 8k tokens. State space models — S4 [1] and Mamba [2] — achieve $O(n)$ in both by replacing attention with a selective recurrence that can be parallelised via an efficient scan. RWKV [3] reinterprets an RNN as a Transformer-equivalent that trains in parallel. FlashAttention [6] takes a different angle: same $O(n^2)$ math, but restructured to minimise slow HBM reads/writes, achieving 2—4x wall-clock speedup on GPUs.

The problem with attention

Standard scaled dot-product attention [8] computes:

\text{Attention}(Q, K, V) = \text{softmax}\!\left(\frac{QK^\top}{\sqrt{d_k}}\right) V.

For a sequence of length $n$ and head dimension $d_k$ :

The score matrix $QK^\top \in \mathbb{R}^{n \times n}$ requires $O(n^2)$ compute and $O(n^2)$ memory.
For $n = 128{,}000$ (128k context, as in Claude 3 or Gemini 1.5) and float16, the score matrix alone requires $128{,}000^2 \times 2 \approx 32\text{ GB}$ — more than the HBM on most GPUs.

Two separate bottlenecks therefore arise:

Compute: FLOPs scale as $O(n^2 d)$ per layer.
Memory bandwidth: writing and reading the $n \times n$ score matrix to/from GPU high-bandwidth memory (HBM) is slow, even when the compute could be done fast with local SRAM.

The post-Transformer architecture landscape attacks one or both of these bottlenecks.

State space models: S4

A state space model (SSM) maps an input sequence $u(t) \in \mathbb{R}$ to output $y(t) \in \mathbb{R}$ via a hidden state $x(t) \in \mathbb{R}^N$ :

\dot{x}(t) = A x(t) + B u(t), \qquad y(t) = C x(t) + D u(t).

Discretising with step size $\Delta$ gives a recurrence usable for sequences:

x_k = \bar{A} x_{k-1} + \bar{B} u_k, \qquad y_k = C x_k,

where $\bar{A} = e^{\Delta A}$ and $\bar{B} = (e^{\Delta A} - I) A^{-1} B$ . The recurrence runs in $O(n N)$ with state dimension $N$ — linear in sequence length.

S4 [1] makes this practical via three ideas:

Structured A matrix: parameterize $A$ as a diagonal plus low-rank (DPLR) matrix, enabling efficient matrix exponentials.
Convolutional view for training: the recurrence is equivalent to a causal 1D convolution with kernel $\bar{K}_k = C \bar{A}^k \bar{B}$ . Compute this kernel once, then train with FFT-based convolution in $O(n \log n)$ — faster than $O(n^2)$ attention for long sequences.
HiPPO initialization: initialize $A$ as a matrix designed to optimally memorise input history via orthogonal polynomial projections, giving gradients a good starting landscape.

Mamba: selective state spaces

S4’s key limitation: $A$ , $B$ , $C$ are fixed across all positions. The model cannot selectively attend to specific tokens — a critical capability for language.

Mamba [2] introduces input-dependent (selective) SSM parameters: $B_k = B(u_k)$ , $C_k = C(u_k)$ , and $\Delta_k = \Delta(u_k)$ are all functions of the current input token. This recovers the Transformer’s ability to selectively focus on relevant tokens, while retaining the $O(n)$ recurrent inference.

Mamba's selective state space and hardware-aware parallel scan

With input-dependent parameters, the recurrence becomes:

x_k = \bar{A}_k x_{k-1} + \bar{B}_k u_k, \qquad y_k = C_k x_k,

where $\bar{A}_k = \exp(\Delta_k A)$ and $\bar{B}_k = \Delta_k B_k$ (simplified zero-order-hold discretisation). The selectivity mechanism is $\Delta_k$ : when $\Delta_k$ is large, the new input $u_k$ has high influence and past state is discarded; when $\Delta_k$ is small, state persists and the current input is largely ignored.

The training challenge: because $\bar{A}_k$ now varies per token, the convolutional view no longer applies. Naively unrolling the recurrence is $O(n)$ sequential — no longer parallelisable.

Parallel prefix scan: the recurrence $x_k = \bar{A}_k x_{k-1} + \bar{B}_k u_k$ has the form of a linear recurrence, which can be evaluated in parallel using the associative scan algorithm. For a sequence of $n$ elements, the parallel scan computes all $x_k$ in $O(\log n)$ parallel steps with $O(n)$ total work. On a GPU with thousands of cores, this achieves near-linear wall-clock time.

Hardware-aware implementation: Mamba’s implementation avoids materialising the $n \times N$ state tensor in slow HBM. Instead, it keeps $x_k$ in fast SRAM, performs the scan there, and writes only the final outputs to HBM. This is the “hardware-aware” aspect: the algorithm is restructured to minimise HBM memory traffic, analogous to FlashAttention’s approach for standard attention.

Result: Mamba achieves $O(n)$ FLOPs and $O(n)$ memory, with $O(\log n)$ parallel depth. At sequence length 32k, Mamba runs 5x faster than a Transformer of similar parameter count.

RWKV: RNN that trains like a Transformer

RWKV [3] (Receptance Weighted Key Value) reframes an RNN as a Transformer-equivalent by observing that a specific form of attention can be computed either as a parallel matrix operation (for training) or as an $O(1)$ -per-token recurrence (for inference).

The core WKV mechanism at time $t$ for a single head is:

\text{wkv}_t = \frac{\sum_{i \le t} e^{-(t-1-i)w + k_i} v_i}{\sum_{i \le t} e^{-(t-1-i)w + k_i}},

where $w > 0$ is a learnable per-channel decay and $k_i, v_i$ are key and value vectors. This is a softmax attention with a fixed exponential decay prior on past tokens — no learned query-key interaction per pair, just positional decay. The numerator and denominator satisfy a linear recurrence:

\text{num}_t = e^{-w} \cdot \text{num}_{t-1} + e^{k_t} v_t, \qquad \text{den}_t = e^{-w} \cdot \text{den}_{t-1} + e^{k_t},

so inference runs in $O(1)$ per token with fixed-size state. Training parallelises this via a parallel prefix scan (similar to Mamba), achieving Transformer-like training throughput.

RWKV-6 adds a data-dependent mechanism to the channel-mixing and time-mixing blocks, bringing it closer to Mamba in expressiveness while retaining the clean RNN inference mode. RWKV has been trained at 14B parameter scale and shows competitive language model perplexity to similarly sized Transformers.

RetNet

Retentive Network (RetNet) [4] introduces a “retention” mechanism as a structural alternative to attention. Retention has three computation modes:

Parallel mode (training): computes as a causal matrix product in $O(n^2)$ — same as attention, but with a fixed positional decay mask instead of a learned softmax. Fully parallelisable.
Recurrent mode (inference): reduces to an $O(1)$ -per-token state update.
Chunk-wise recurrent mode (long-context): processes chunks of tokens in parallel, then passes a summary state between chunks.

The decay mask in parallel mode is $D_{nm} = \gamma^{n-m}$ for $n \ge m$ , zero otherwise ( $\gamma \in (0, 1)$ ). This is a fixed exponential decay without the softmax normalization of attention, making the mechanism strictly linear and easily invertible to a recurrence.

RetNet achieves throughputs comparable to or better than Transformers in training mode and dramatically lower inference latency at batch size 1 (no KV cache needed), but its fixed decay may limit its ability to selectively attend to distant tokens — an expressiveness trade-off compared to Mamba’s selective gating.

Hyena

Hyena [5] generalises convolution-based sequence models by replacing the fixed convolutional filter with a long convolution parameterised by a small neural network. At each layer, Hyena applies $p$ gated convolutions in sequence, where each convolution kernel is generated by an MLP given position embeddings:

y = (h_p * (x_p \cdot \ldots \cdot (h_1 * (x_1 \cdot u))\ldots)),

where $*$ denotes causal convolution and $\cdot$ denotes element-wise multiplication (gating). This structure is strictly sub-quadratic ( $O(n \log n)$ via FFT) while being more expressive than a single linear convolution. Hyena has been scaled to 1B parameters and matches Transformer perplexity at sequence lengths above 8k tokens, where its sub-quadratic compute advantage becomes significant.

FlashAttention: IO-aware exact attention

FlashAttention [6] does not change the mathematical output of attention — it computes exactly the same result as standard attention. Instead, it restructures the computation to minimise memory bandwidth usage, which is the actual bottleneck on modern GPUs.

FlashAttention IO complexity and tiling

Modern GPUs have two memory levels:

HBM (high-bandwidth memory): large (40—80 GB), bandwidth 1.5—3 TB/s.
SRAM (on-chip): tiny (20—40 MB), bandwidth 19 TB/s.

Standard attention:

Read $Q, K$ from HBM. Write $S = QK^\top / \sqrt{d_k}$ to HBM: $O(n^2)$ write.
Read $S$ from HBM. Compute $P = \text{softmax}(S)$ . Write $P$ to HBM: $O(n^2)$ write.
Read $P, V$ from HBM. Write $O = PV$ to HBM: $O(nd)$ write.

Total HBM accesses: $O(n^2 + nd)$ .

FlashAttention tiling: process blocks of $Q, K, V$ that fit in SRAM simultaneously. For a tile of queries $Q_i$ (block size $B_r$ ) and keys/values $(K_j, V_j)$ (block size $B_c$ ):

Load $Q_i, K_j, V_j$ into SRAM.
Compute the tile of scores $S_{ij} = Q_i K_j^\top / \sqrt{d_k}$ .
Apply online softmax: maintain running max $m_i$ and running sum $\ell_i$ to normalise incrementally across tiles without materialising the full $n \times n$ score matrix.
Accumulate the partial output $O_i \mathrel{+}= P_{ij} V_j$ in SRAM.
Write $O_i$ to HBM.

The $n \times n$ score matrix never exists in HBM. Total HBM accesses: $O(n^2 d / M)$ where $M$ is the SRAM size, which is $O(nd)$ for practical values (much less than $O(n^2)$ ).

FlashAttention-2 [7] reduces non-matmul FLOPs, improves parallelism across the sequence dimension, and achieves 2—4x throughput improvement over FlashAttention on A100 GPUs, reaching 50—73% utilization of the theoretical 312 TFLOPS peak.

The crucial point: FlashAttention’s IO complexity is $O(n^2 d / M)$ vs $O(n^2 + nd)$ for standard attention. For large $M$ (SRAM size), this is a substantial win. But the compute complexity is still $O(n^2 d)$ , so FlashAttention does not help for extremely long sequences where compute (not memory bandwidth) becomes the bottleneck.

Hybrid architectures

No single architecture wins on all axes. Current practitioners often combine Mamba layers with attention layers in hybrid models:

Jamba (AI21 Labs, 2024): alternates Mamba blocks with Transformer attention blocks in a fixed ratio. Uses mixture-of-experts in the FFN layers. Achieves near-Transformer quality with substantially lower KV-cache memory at long context.
Zamba (Zyphra, 2024): a 7B hybrid model with a single global attention layer every 6 Mamba layers. The single attention layer provides global information mixing; Mamba layers handle efficient local state propagation.

The insight motivating hybrids: Mamba is excellent at compressing and propagating local patterns efficiently, but may miss long-range dependencies that require exact cross-position comparison. Sparse attention layers plug this gap at modest cost.

Architecture comparison

Architecture	Training complexity	Inference per token	Long context (100k+)	Hardware efficiency	Selective attention	Notes
Transformer	O(n² d)	O(n d) — reads full KV cache	Expensive (O(n²) memory)	Good with FlashAttention	Yes (full softmax)	Dominant for quality
S4	O(n log n) via FFT convolution	O(N) per step (fixed state size)	Excellent (O(n) memory)	Good (FFT-friendly)	No (fixed params)	Strong on audio/DNA
Mamba	O(n) via parallel scan	O(N) per step	Excellent (O(n) memory)	Excellent (hardware-aware scan)	Yes (input-dependent)	Best SSM to date
RWKV	O(n log n) via parallel scan	O(1) per step (fixed state)	Good (O(n) memory)	Good	Partial (position decay only)	Pure RNN at inference
RetNet	O(n²) parallel or O(n) chunk-recurrent	O(1) per step	Good (chunk mode)	Good	Partial (fixed decay)	Flexible compute modes
Hyena	O(n log n) via FFT	O(n log n) per token	Good (sub-quadratic)	Good	Via gating	Best for very long sequences
Jamba (hybrid)	Mixed (Mamba + sparse attn)	Mixed	Very good (sparse KV cache)	Very good	Yes (attention layers)	Current practical best

When to use which

Use standard Transformer + FlashAttention when:

Training with context length below 16k tokens.
Maximum model quality is the priority.
Hardware supports FlashAttention (CUDA GPU).

Use Mamba or Mamba hybrid when:

Training with context lengths above 32k.
Memory-constrained inference (edge, consumer hardware).
Streaming inference (autoregressive state update in $O(N)$ ).
Domains with strong local structure (audio, genomics).

Use RWKV when:

Pure RNN inference with constant memory per token is required.
Transformer-quality language modeling with RNN inference latency is the goal.

Use hybrid architectures (Jamba, Zamba) when:

You want near-Transformer quality with substantially reduced KV-cache at long context.
Mixture-of-experts sparsity is acceptable.

Mamba quality vs Transformer at scale

As of 2024, Mamba-pure models at 7B+ parameter scale still lag behind comparably sized Transformers on few-shot benchmarks (MMLU, HellaSwag). The gap is narrower for hybrid models. Whether this is a fundamental architecture issue or a training recipe gap is an open question. For mission-critical quality applications, Transformer-based models with FlashAttention remain the safe choice.

Check yourself

Score: 0 / 5(0 answered)

1. Why does standard self-attention have O(n²) memory cost?
2. What is "selective" about Mamba's state space model, compared to S4?
3. Why can RWKV be parallelised during training despite being formulated as an RNN?
4. FlashAttention achieves speedups over standard attention primarily by:
5. What is the key motivation for hybrid architectures like Jamba that mix Mamba and attention layers?

Derive the memory complexity of standard attention and explain how Mamba achieves O(n) memory instead.onsite

Standard attention: for a sequence of length $n$ , head dimension $d_k$ , the score matrix $S = QK^\top / \sqrt{d_k}$ has shape $n \times n$ and must be stored for the softmax and the subsequent multiplication with $V$ . Each entry is a float (2 bytes in BF16), so the score matrix costs $n^2 \times 2$ bytes. For $n = 32{,}000$ , that is 2 GB per layer per head — and a Transformer has $L \times H$ heads. The softmax output (same shape) doubles this. Total attention memory: $O(n^2)$ per layer.

Mamba: the state space recurrence maintains a hidden state $x_k \in \mathbb{R}^{d \times N}$ where $N$ is the state dimension (typically 16). The recurrence $x_k = \bar{A}_k x_{k-1} + \bar{B}_k u_k$ processes one token at a time, never materialising a matrix of shape $n \times n$ . Total state memory: $O(n N d)$ for the full sequence, but in practice only $O(N d)$ is needed for autoregressive inference (you only need the current state). Training uses a parallel scan which computes all $x_k$ in $O(\log n)$ parallel steps with $O(n N d)$ intermediate storage — linear in $n$ , not quadratic.

Explain what 'hardware-aware' means in FlashAttention and why tiling makes attention faster even though the number of FLOPs is the same.staff

“Hardware-aware” refers to structuring the algorithm around the memory hierarchy of real GPUs, not around minimising FLOPs alone.

Modern GPUs have two memory levels: HBM (the main GPU memory, 1.5—3 TB/s bandwidth) and SRAM (on-chip registers and shared memory, ~19 TB/s bandwidth, 20—40 MB capacity). Standard attention writes the $n \times n$ score matrix to HBM and reads it back for the softmax and value aggregation. For $n = 4096$ , this is three separate HBM round-trips for a 32 MB tensor. HBM bandwidth is the bottleneck: even though the raw FLOP count for attention is small relative to the FFN, the memory transactions make it slower.

FlashAttention tiles $Q$ into blocks of rows (block size $B_r$ ) and $(K, V)$ into blocks of columns (block size $B_c$ ), chosen so that a tile of $Q_i, K_j, V_j$ fits in SRAM. For each tile pair, it:

Computes the score block $S_{ij} = Q_i K_j^\top / \sqrt{d_k}$ entirely in SRAM.
Updates the running softmax statistics (max and log-sum-exp) using the online softmax algorithm from Milakov and Gimelshein (2018), without materialising the full row.
Accumulates partial outputs $O_i$ in SRAM.
Writes the final $O_i$ once to HBM when all $K, V$ tiles for row block $i$ are processed.

The $n \times n$ score matrix never touches HBM. Total HBM reads/writes: $O(nd)$ for $Q, K, V$ inputs and $O(nd)$ for the output — linear in $n$ . The FLOPs are unchanged at $O(n^2 d)$ , but the GPU spends far less time waiting for memory. On an A100, attention is memory-bandwidth-bound at practical sequence lengths, so this IO reduction translates directly to wall-clock speedup (2—4x for FlashAttention, 2—4x additional for FlashAttention-2 via better parallelism and warp specialisation).

References

Gu, A., Goel, K., and Re, C. (2021). Efficiently Modeling Long Sequences with Structured State Spaces. ICLR 2022 (arXiv preprint 2021). arXiv:2111.00396
Gu, A. and Dao, T. (2023). Mamba: Linear-Time Sequence Modeling with Selective State Spaces. arXiv preprint. arXiv:2312.00752
Peng, B., Alcaide, E., Anthony, Q., Albalak, A., Arcadinho, S., Biderman, S., Cao, H., Cheng, X., Cox, M., Ding, N., and others (2023). RWKV: Reinventing RNNs for the Transformer Era. EMNLP 2023 (Findings). arXiv:2305.13048
Sun, Y., Dong, L., Huang, S., Ma, S., Xia, Y., Xue, H., Wang, L., and Wei, F. (2023). Retentive Network: A Successor to Transformer for Large Language Models. arXiv preprint. arXiv:2307.08621
Poli, M., Massaroli, S., Nguyen, E., Fu, D. Y., Dao, T., Baccus, S., Bengio, Y., Ermon, S., and Re, C. (2023). Hyena Hierarchy: Towards Larger Convolutional Language Models. ICML 2023. arXiv:2302.10866
Dao, T., Fu, D. Y., Ermon, S., Rudra, A., and Re, C. (2022). FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness. NeurIPS 2022. arXiv:2205.14135
Dao, T. (2023). FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning. ICLR 2024 (arXiv preprint 2023). arXiv:2307.08691
Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L., Gomez, A. N., Kaiser, L., and Polosukhin, I. (2017). Attention Is All You Need. NeurIPS 2017. arXiv:1706.03762

World Models Capstone: RAG System