TRANSFORMERS UNBOUND: HOW TO PUSH THE ARCHITECTURE TO ITS LIMITS AND BEYOND

PROLOGUE: WHY SHOULD YOU CARE?

If you have ever typed a question into a large language model and received a surprisingly coherent, even brilliant answer, you have already experienced the product of one of the most consequential architectural ideas in the history of computing: the Transformer. Introduced in 2017 by Vaswani et al. in the paper "Attention Is All You Need," the Transformer replaced the sequential, step-by-step processing of recurrent neural networks with a mechanism that looks at every token in a sequence simultaneously and decides, for each one, how much attention to pay to every other one. That single idea unlocked nearly a decade of progress that nobody fully anticipated.

But here is the thing: the original Transformer, as elegant as it is, has serious problems. It is memory-hungry. It is slow on long sequences. It wastes compute on things that do not matter. It has a fixed context window. It struggles to scale gracefully. And yet, the research community has spent the last eight years attacking every single one of those problems with remarkable ingenuity. The result is a constellation of improvements, some surgical and some sweeping, that together make modern large language models possible.

This tutorial is a guided tour through that constellation. We will start from the original architecture, understand exactly where the pain points are, and then walk through every major optimization that has been developed to address them. We will cover FlashAttention in all three of its generations, Mixture of Experts, Sparse and Native Sparse Attention, Grouped-Query Attention, Multi-head Latent Attention, Rotary Position Embeddings and their extensions, improved normalization schemes, better activation functions, quantization, speculative decoding, multi-token prediction, new optimizers, and much more. We will also look honestly at the horizon: what further improvements are coming, and where the Transformer architecture might eventually hit a wall it cannot climb over.

CHAPTER ONE: THE ORIGINAL TRANSFORMER, WARTS AND ALL

Before we can appreciate the optimizations, we need to understand what we are optimizing. Let us walk through the original Transformer architecture with enough detail to feel its beauty and its pain simultaneously.

The Transformer takes a sequence of tokens (words, subwords, or characters) and converts each one into a vector of numbers called an embedding. These embeddings are then processed by a stack of identical layers, each of which has two main components: a multi-head self-attention sublayer and a position-wise feed-forward network sublayer. Residual connections and layer normalization wrap each of these sublayers. The original paper used Post-LN normalization, meaning the layer norm was applied after the residual addition, though as we will see, this turned out to be one of the first things the community changed.

The heart of everything is the self-attention mechanism. For each token, the model computes three vectors: a Query (Q), a Key (K), and a Value (V). These are produced by multiplying the token's embedding by three learned weight matrices. The attention score between token i and token j is computed as the dot product of Q_i and K_j, scaled by the square root of the head dimension, and then passed through a softmax to produce a probability distribution. The output for token i is then the weighted sum of all Value vectors, where the weights come from that probability distribution.

Let us make this concrete with a tiny example. Suppose we have the sentence "The cat sat" and we are computing attention for the word "sat." The model might learn that "sat" should attend heavily to "cat" (because cats are the ones doing the sitting) and less to "The." The attention scores before softmax and scaling might look like this:

Query("sat") . Key("The")  = 0.3
Query("sat") . Key("cat")  = 2.1
Query("sat") . Key("sat")  = 1.0

After dividing by sqrt(d_k) and applying softmax, these become approximately:

Attention weight for "The"  = 0.07
Attention weight for "cat"  = 0.75
Attention weight for "sat"  = 0.18

The output for "sat" is then 0.07 * V("The") + 0.75 * V("cat") + 0.18 * V("sat"), which is a vector that has been heavily influenced by the representation of "cat." This is the magic: every token gets to look at every other token and gather information from those it finds most relevant.

The "multi-head" part means this process is run H times in parallel, each time with different learned projection matrices. Each head can specialize in a different kind of relationship: one head might learn syntactic dependencies, another might learn coreference, another might learn positional proximity. The outputs of all heads are concatenated and projected back to the model dimension.

Here is a schematic of one Transformer layer:

Input embeddings (+ positional encoding)
         |
         v
+---------------------------+
|   Multi-Head Self-Attention|
|   Q = X * W_Q             |
|   K = X * W_K             |
|   V = X * W_V             |
|   Attn = softmax(QK^T /   |
|          sqrt(d_k)) * V   |
+---------------------------+
         |
    Add & Norm (residual)
         |
         v
+---------------------------+
|   Feed-Forward Network    |
|   FFN(x) = max(0, xW_1   |
|            + b_1)W_2 + b_2|
+---------------------------+
         |
    Add & Norm (residual)
         |
         v
Output representations

Now, where does this beautiful machine break down? There are four fundamental pain points that every subsequent optimization targets in some way.

The first pain point is the quadratic complexity of attention. To compute attention for a sequence of N tokens, we need to compute N*N attention scores. For N=1,000, that is one million scores. For N=100,000, that is ten billion scores. The memory required to store the full attention matrix scales as O(N^2), and the compute scales the same way. This makes the vanilla Transformer essentially unusable for very long sequences without some form of approximation or restructuring.

The second pain point is the memory bandwidth bottleneck. Modern GPUs are extraordinarily fast at arithmetic, but they are relatively slow at moving data between different levels of memory. The GPU has a small, fast on-chip SRAM (the registers and shared memory) and a large, slow off-chip HBM (High Bandwidth Memory, the GPU's main memory). The naive attention implementation constantly shuttles data back and forth between these two levels, and this data movement, not the arithmetic, is often the actual bottleneck. This is the problem FlashAttention was designed to solve.

The third pain point is that all parameters are used for every input. A 70-billion-parameter model uses all 70 billion parameters to process every single token, regardless of whether those parameters are relevant to the task at hand. This is enormously wasteful. Mixture of Experts architectures address this by routing each token to only a subset of the model's parameters.

The fourth pain point is the KV cache explosion during inference. When generating text autoregressively, the model needs to remember the Keys and Values for all previously generated tokens so it does not have to recompute them. For a model with many attention heads and a long context, this KV cache can consume tens or hundreds of gigabytes of memory, severely limiting how many requests a server can handle simultaneously.

With these four pain points clearly in mind, let us now walk through the optimizations, roughly in the order they address the architecture from the inside out.

CHAPTER TWO: FIXING THE FOUNDATIONS - NORMALIZATION AND ACTIVATION FUNCTIONS

Before we get to the headline-grabbing optimizations like FlashAttention and Mixture of Experts, it is worth spending time on the quieter but enormously impactful improvements to the basic building blocks of the Transformer: how it normalizes activations and what nonlinearity it uses in its feed-forward networks. These changes sound mundane, but they have a direct effect on training stability, convergence speed, and final model quality.

2.1 FROM POST-LN TO PRE-LN: STABILIZING DEEP TRANSFORMERS

The original Transformer paper used Post-Layer Normalization, meaning the layer norm was applied after the residual addition:

Post-LN:  output = LayerNorm(x + Sublayer(x))

This seems natural, but it has a nasty property: during the early stages of training, the gradients flowing through the residual path can be very large, causing instability and requiring careful learning rate warmup schedules. Researchers discovered that moving the layer norm to before the sublayer, a scheme called Pre-LN, dramatically stabilizes training:

Pre-LN:   output = x + Sublayer(LayerNorm(x))

With Pre-LN, the residual path is clean and unobstructed, so gradients flow back through the network without being distorted by the normalization operation. This allows training to begin with a larger learning rate and without as much warmup, and it makes it possible to train much deeper models without divergence. Virtually every modern large language model, including the GPT series, LLaMA, Mistral, and their descendants, uses Pre-LN.

2.2 RMSNORM: DOING LESS TO ACHIEVE MORE

Layer Normalization, as defined by Ba et al. in 2016, computes both the mean and the variance of a layer's activations and uses both to normalize them. Zhang and Sennrich (2019) asked a provocative question: do we actually need the mean-centering step? Their answer, backed by experiments, was no. They proposed Root Mean Square Layer Normalization (RMSNorm), which skips the mean computation and normalizes only by the root mean square of the activations:

LayerNorm:  y = (x - mean(x)) / sqrt(var(x) + eps) * gamma + beta
RMSNorm:    y = x / sqrt(mean(x^2) + eps) * gamma

This is both simpler and faster. The mean computation and the beta bias term are eliminated, reducing the number of operations and parameters. Experiments show that RMSNorm achieves comparable or better performance to full LayerNorm while being 10 to 50 percent faster in practice. LLaMA, Mistral, Gemma, DeepSeek, and essentially all modern frontier models use RMSNorm. It is a small change with a real payoff.

2.3 SWIGLU: THE ACTIVATION FUNCTION THAT EVERYONE USES NOW

The original Transformer used a simple two-layer feed-forward network with ReLU activation:

FFN_ReLU(x) = max(0, x * W_1 + b_1) * W_2 + b_2

In 2020, Noam Shazeer (one of the original Transformer authors) published a short but influential paper called "GLU Variants Improve Transformers," which showed that replacing ReLU with a gated variant dramatically improves model quality. The key idea behind Gated Linear Units (GLU) is to multiply the output of one linear transformation by the sigmoid (or another activation) of a second linear transformation, effectively creating a learned gate that controls information flow:

GLU(x, W, V, b, c) = sigmoid(x * W + b)  (elementwise multiply)  (x * V + c)

Shazeer explored many variants: Bilinear (no activation on the gate), ReGLU (ReLU gate), GEGLU (GELU gate), and SwiGLU (Swish gate). SwiGLU, which uses the Swish activation function (defined as x * sigmoid(beta * x)), consistently outperformed all others:

SwiGLU(x, W, V) = Swish(x * W)  (elementwise multiply)  (x * V)

The feed-forward network with SwiGLU then becomes:

FFN_SwiGLU(x) = SwiGLU(x, W_1, W_3) * W_2

Note that SwiGLU requires three weight matrices instead of two (W_1, W_2, W_3), so to keep the parameter count comparable, the hidden dimension is typically reduced by a factor of 2/3. Despite this, SwiGLU models consistently outperform their ReLU counterparts. PaLM, LLaMA, Mistral, Gemma, DeepSeek, and virtually every other modern frontier model uses SwiGLU or GEGLU. This is not a minor tweak; it is a genuine quality improvement that has been replicated across dozens of model families and scales.

Here is a quick illustration of why gating helps. Imagine the feed-forward layer is trying to decide whether to pass information about "Paris" through to the next layer. A ReLU network just clips negative values to zero, which is a blunt instrument. A SwiGLU network can learn a smooth, input-dependent gate: if the context suggests we are talking about geography, open the gate wide; if we are talking about fashion, open it a bit; if we are talking about chemistry, close it almost entirely. This dynamic, content-aware gating gives the model much more expressive power per parameter.

2.4 DEEPNORM: GOING VERY, VERY DEEP

In 2022, researchers at Microsoft introduced DeepNorm (Wang et al., 2022, arXiv:2203.00555), a normalization scheme specifically designed to allow training Transformers with up to 1,000 layers. The key insight is that the residual connection in a standard Transformer causes the magnitude of activations to grow with depth, which eventually destabilizes training. DeepNorm addresses this by scaling the residual connection by a factor alpha and initializing the weights with a factor beta:

DeepNorm:  output = LayerNorm(alpha * x + Sublayer(x))

With carefully chosen alpha and beta values (derived from theoretical analysis of the expected gradient magnitudes), DeepNorm provably bounds the model update at initialization, ensuring that the network starts in a stable regime regardless of depth. This allowed the team to train a 1,000-layer Transformer that outperformed shallower models on machine translation tasks. While most production models do not use 1,000 layers, DeepNorm demonstrates that depth is not inherently a barrier to training stability if the normalization is designed correctly.

CHAPTER THREE: THE ATTENTION CRISIS AND HOW FLASHATTENTION SOLVED IT

We established in Chapter One that the naive attention computation has a serious memory problem: it requires materializing the full N-by-N attention matrix in GPU HBM, which for long sequences is both enormous and slow to access. This is the problem that FlashAttention, developed by Tri Dao and colleagues at Stanford, attacked with surgical precision.

3.1 THE MEMORY HIERARCHY PROBLEM

To understand FlashAttention, you need to understand the GPU memory hierarchy. A modern GPU like the NVIDIA A100 has roughly 80 GB of HBM (the main GPU memory), which has a bandwidth of about 2 TB/s. It also has a much smaller on-chip SRAM (shared memory plus L2 cache), totaling perhaps 40 MB, but with a bandwidth of roughly 19 TB/s. The on-chip SRAM is about 10 times faster than HBM, but it is about 2,000 times smaller.

The naive attention algorithm works like this: compute Q, K, V matrices and write them to HBM. Then read Q and K back from HBM, compute the attention scores S = Q * K^T, write S to HBM. Read S back, compute softmax(S), write the result P to HBM. Read P and V back, compute the output O = P * V, write O to HBM. Each of these read-write operations is a round trip to the slow HBM. For a sequence of length N with head dimension d, the total HBM accesses scale as O(N^2), and for large N this data movement dominates the total runtime.

The key insight of FlashAttention (Dao et al., 2022, arXiv:2205.14135) is that we do not actually need to materialize the full N-by-N attention matrix. We can compute the attention output in tiles, keeping the intermediate results in the fast on-chip SRAM and only writing the final output to HBM. This is called tiling or blocking, and it is a classic technique in high-performance computing applied here to the specific structure of the softmax attention computation.

3.2 THE TILING TRICK AND THE ONLINE SOFTMAX

The challenge with tiling attention is the softmax. Softmax requires knowing the sum of all exponentials in a row before you can normalize any of them. If you are processing the attention matrix in tiles, you do not see the whole row at once, so how can you compute the correct softmax?

The answer is the online softmax algorithm, which allows you to compute a running maximum and a running sum of exponentials as you process tiles one by one, and then correct the output at the end. Here is the key recurrence: suppose you have processed tiles 1 through t and have accumulated a running maximum m_t and a running sum of exponentials l_t. When you process tile t+1, you update:

m_{t+1} = max(m_t, max of new tile scores)
l_{t+1} = exp(m_t - m_{t+1}) * l_t + sum of exp(new tile scores - m_{t+1})

And you correct the accumulated output O_t:

O_{t+1} = (l_t * exp(m_t - m_{t+1}) * O_t
           + exp(new scores - m_{t+1}) * V_new) / l_{t+1}

This recurrence is mathematically exact; it produces the same result as computing the full softmax at once, but it only ever needs to hold one tile of the attention matrix in SRAM at a time. The full N-by-N matrix never needs to be written to HBM.

The result is dramatic. FlashAttention reduces HBM accesses from O(N^2) to O(N^2 * d / M), where M is the SRAM size. For typical values of d and M, this is a reduction of roughly 5 to 20 times in memory reads and writes. In practice, FlashAttention achieves 2 to 4 times speedup over the naive PyTorch attention implementation and reduces memory usage by up to 7.6 times. It also enables training on much longer sequences: the original FlashAttention paper demonstrated training GPT-2 on sequences of 4,096 tokens (four times longer than the standard 1,024) while being faster than the standard implementation on 1,024-token sequences.

3.3 FLASHATTENTION-2: BETTER PARALLELISM

FlashAttention-2 (Dao, 2023, arXiv:2307.08691) identified a further source of inefficiency in the original FlashAttention: the work partitioning between GPU thread blocks and warps was suboptimal. In FlashAttention-1, the outer loop was over the sequence dimension (rows of Q), and the inner loop was over the key-value sequence (rows of K and V). This led to low GPU occupancy in some configurations. FlashAttention-2 reorganized the computation to maximize parallelism across the batch, head, and sequence dimensions simultaneously, and it reduced the number of non-matmul operations (which are less efficiently pipelined on modern GPUs). The result was a further 2x speedup over FlashAttention-1, reaching 50 to 73 percent of the theoretical maximum FLOPs/s on A100 GPUs.

3.4 FLASHATTENTION-3: EXPLOITING HOPPER GPU FEATURES

FlashAttention-3 (Shah et al., 2024, arXiv:2407.08608) was written specifically for NVIDIA's Hopper GPU architecture (H100), which introduced several new hardware features that FlashAttention-2 could not exploit. The three main ideas in FlashAttention-3 are warp specialization, interleaved matmul and softmax, and FP8 support with block quantization.

Warp specialization means that different warps (groups of 32 GPU threads) are assigned different roles: some warps handle data movement using the new Tensor Memory Accelerator (TMA) hardware on Hopper, while others handle computation using the new WGMMA (Warpgroup Matrix Multiply-Accumulate) instructions. This allows data movement and computation to overlap, hiding the latency of memory transfers. Interleaving the matmul and softmax operations further hides the quantization overhead when using reduced-precision arithmetic. With FP8 precision (8-bit floating point), FlashAttention-3 achieves up to 1.2 PFLOPs/s (petaflops per second) on a single H100 GPU, which is 75 percent of the theoretical maximum. With FP16, it achieves 1.5 to 2.0 times speedup over FlashAttention-2, reaching up to 740 TFLOPs/s.

To give you a sense of the cumulative progress, here is a comparison of attention implementations on A100/H100 GPUs for a typical configuration:

Implementation              Speed (TFLOPs/s)   Memory Usage
------------------------------------------------------------
Standard PyTorch Attn       ~30                O(N^2)
FlashAttention-1 (A100)     ~120               O(N)
FlashAttention-2 (A100)     ~200               O(N)
FlashAttention-3 FP16(H100) ~740               O(N)
FlashAttention-3 FP8 (H100) ~1200              O(N)

This is not a minor engineering improvement. FlashAttention is arguably the single most impactful practical optimization in the history of Transformer training, and it is now used by essentially every serious LLM training framework including PyTorch, JAX, and all major cloud providers' training stacks.

CHAPTER FOUR: SPARSE ATTENTION - PAYING ATTENTION SELECTIVELY

FlashAttention makes exact attention faster by being smarter about memory access patterns. But it does not change the fundamental O(N^2) complexity of the computation itself. For truly long sequences, even an optimally memory-efficient exact attention is too slow. The solution is sparse attention: instead of computing attention scores between every pair of tokens, only compute scores between tokens that are likely to be relevant to each other.

4.1 FIXED SPARSE PATTERNS: LONGFORMER AND BIGBIRD

The earliest and most principled sparse attention approaches used fixed, predetermined patterns that capture the most important structural relationships in a sequence.

The Longformer (Beltagy et al., 2020, arXiv:2004.05150) combines two types of attention. Local windowed attention means each token attends to a fixed window of w tokens on either side, capturing local context with O(N * w) complexity. Global attention means a small set of special tokens (like the [CLS] classification token) attend to all other tokens and are attended to by all other tokens, capturing global context. The combination gives each token local awareness and a global summary, at a total cost of O(N * (w + g)) where g is the number of global tokens. This scales linearly with sequence length, making it practical for documents of tens of thousands of tokens.

BigBird (Zaheer et al., 2020, arXiv:2007.14062) adds a third component to this recipe: random attention, where each token also attends to a small set of randomly chosen tokens from the full sequence. The combination of local windowed attention, global attention on special tokens, and random attention is theoretically motivated: BigBird proves that this combination is a universal approximator of full attention, meaning that any function computable by full attention can also be computed (approximately) by BigBird's sparse pattern. The random attention component ensures that information can flow between any two tokens in the sequence in a bounded number of hops, preserving the long-range dependency modeling that makes Transformers powerful.

Here is a visualization of these attention patterns for a sequence of 12 tokens (each row is a query token, each column is a key token, X means attention is computed, and a dot means it is skipped):

Full Attention (every token attends to every token):
X X X X X X X X X X X X
X X X X X X X X X X X X
X X X X X X X X X X X X
X X X X X X X X X X X X
X X X X X X X X X X X X
X X X X X X X X X X X X

Longformer (window=2, token 0 is global):
X X X X X X X X X X X X   <- global token attends everywhere
X X X X . . . . . . . .   <- local window of 2
X X X X X . . . . . . .
X . X X X X . . . . . .
X . . X X X X . . . . .
X . . . X X X X . . . .

BigBird (window=2, token 0 global, plus 1 random per row):
X X X X X X X X X X X X   <- global token
X X X X . . X . . . . .   <- local + 1 random (col 6)
X X X X X . . . . X . .   <- local + 1 random (col 9)
X . X X X X . . . . . X   <- local + 1 random (col 11)
X . . X X X X . . . . .   <- local only (no random shown)
X . . . X X X X . X . .   <- local + 1 random (col 9)

The sparsity in these patterns means that for a sequence of 64,000 tokens with a window of 512, each token only computes attention with roughly 1,024 other tokens instead of 64,000, a 64-fold reduction in compute.

4.2 NATIVE SPARSE ATTENTION: TRAINABLE SPARSITY

Fixed sparse patterns like Longformer's are a reasonable approximation, but they are not learned from data. The model cannot adapt the sparsity pattern to the content of the input. In 2025, DeepSeek introduced Native Sparse Attention (NSA, arXiv:2502.05171), a sparse attention mechanism that is trainable from scratch and hardware-aligned for efficient GPU execution.

NSA uses a hierarchical approach. First, it compresses the key-value sequence into coarse-grained blocks by averaging or pooling tokens within each block, and computes attention scores between each query and these compressed blocks. This gives a rough estimate of which regions of the sequence are relevant. Second, it selects the top-k most relevant blocks based on these coarse scores, and then computes fine-grained attention only within those selected blocks. This two-stage process ensures that the model can attend to any part of the sequence if it is relevant, while spending most of its compute on the most important regions.

The crucial innovation is that NSA is differentiable end-to-end: the block selection process uses a straight-through estimator or a continuous relaxation, allowing gradients to flow through the selection step during training. This means the model learns, from the data, which tokens are worth attending to, rather than having this decision hardcoded by the architecture designer. On 64K-token sequences, NSA achieves a 9x speedup over full attention while maintaining model quality comparable to full attention. It is used in DeepSeek's long-context models and represents the current state of the art in practical trainable sparse attention.

4.3 LINEAR ATTENTION: THE KERNEL TRICK

A completely different approach to reducing attention complexity is to approximate the softmax attention with a kernel function that can be computed without materializing the full N-by-N matrix. The key observation is that the softmax attention can be written as:

Attention(Q, K, V) = softmax(Q * K^T / sqrt(d)) * V

If we could replace softmax(Q * K^T) with a factorized form phi(Q) * phi(K)^T, where phi is some feature map, then we could compute the attention output as:

Attention(Q, K, V) = phi(Q) * (phi(K)^T * V)

The parenthesized term phi(K)^T * V can be computed once and reused for all queries, reducing the complexity from O(N^2 * d) to O(N * d^2). This is the kernel trick applied to attention.

The Performer (Choromanski et al., 2020, arXiv:2009.14794) implements this idea using random feature approximations to the softmax kernel, specifically a method called FAVOR+ (Fast Attention Via positive Orthogonal Random features). The Linformer (Wang et al., 2020, arXiv:2006.04768) takes a different approach: it projects the keys and values to a lower-dimensional space of size k (where k is much smaller than N) before computing attention, reducing complexity to O(N * k). Both approaches achieve linear complexity in sequence length, but at the cost of approximation error. In practice, linear attention models have historically underperformed exact attention models on language tasks, though the gap has been narrowing.

Lightning Attention, used in MiniMax-01 (2025, arXiv:2501.09755), is a recent linear attention implementation that achieves practical efficiency through careful hardware-aware implementation. MiniMax-01 scales lightning attention to a 456-billion-parameter model with 45.9 billion activated parameters per token and a 32-million-token context window, demonstrating that linear attention can work at frontier scale when implemented carefully.

CHAPTER FIVE: RETHINKING MULTI-HEAD ATTENTION - GQA, MQA, AND MLA

The standard multi-head attention (MHA) in the original Transformer uses H query heads, H key heads, and H value heads. During autoregressive inference, the model must cache the keys and values for all previously generated tokens. For a model with H=64 heads, a head dimension of 128, and a context of 100,000 tokens, the KV cache for a single layer is 64 * 128 * 100,000 * 2 (for K and V) = 1.6 billion numbers. At 16-bit precision, that is 3.2 GB per layer. For a 96-layer model, the total KV cache is 307 GB, which exceeds the memory of most GPU clusters. This is the KV cache crisis, and it has driven several important architectural innovations.

5.1 MULTI-QUERY ATTENTION: THE RADICAL SIMPLIFICATION

Multi-Query Attention (MQA), proposed by Shazeer in 2019, takes the most aggressive approach: use a single key head and a single value head shared across all query heads. All H query heads attend to the same K and V projections. This reduces the KV cache size by a factor of H (e.g., 64x for a 64-head model), which is enormous. The cost is some degradation in model quality, because the single shared K and V must serve all query heads, limiting the diversity of attention patterns.

5.2 GROUPED-QUERY ATTENTION: THE GOLDILOCKS SOLUTION

Grouped-Query Attention (GQA), proposed by Ainslie et al. (2023, arXiv:2305.13245), finds a middle ground between the full H key-value heads of MHA and the single key-value head of MQA. GQA divides the H query heads into G groups, and each group shares a single key head and a single value head. With G groups, the KV cache is reduced by a factor of H/G compared to MHA, while quality is much closer to MHA than MQA.

MHA:  H query heads, H key heads, H value heads   (KV cache = H * d * N)
GQA:  H query heads, G key heads, G value heads   (KV cache = G * d * N)
MQA:  H query heads, 1 key head,  1 value head    (KV cache = 1 * d * N)

For example, LLaMA 3 uses 32 query heads and 8 key-value heads (GQA with G=8), reducing the KV cache by 4x compared to MHA while maintaining quality nearly indistinguishable from full MHA. Mistral 7B uses 32 query heads and 8 key-value heads. Gemma, Qwen, and virtually every modern open-source model uses GQA. It is now the de facto standard for production LLMs.

The Ainslie et al. paper also showed that you can convert an existing MHA model to GQA by mean-pooling the key and value heads within each group and then fine-tuning for a small number of steps, so you do not necessarily need to train from scratch to benefit from GQA.

5.3 MULTI-HEAD LATENT ATTENTION: THE COMPRESSION APPROACH

DeepSeek-V2 (2024, arXiv:2405.04434) introduced a more radical approach called Multi-head Latent Attention (MLA). Instead of caching H separate key-value pairs, MLA compresses the key-value pairs into a single low-rank latent vector c of dimension d_c, where d_c is much smaller than H * d_head. At inference time, the full key and value matrices are reconstructed from this latent vector using learned up-projection matrices.

The idea is analogous to principal component analysis: instead of storing the full high-dimensional key-value representation, store only the low-dimensional latent code that captures most of the information. The reconstruction at inference time adds a small computational cost, but the memory savings are enormous. DeepSeek-V2 reports that MLA reduces the KV cache by 93.3 percent compared to standard multi-head attention, while maintaining model quality comparable to or better than GQA.

Here is a comparison of KV cache sizes for a model with 64 heads, head dimension 128, and context length 100K tokens:

Method      KV cache per layer   Reduction vs MHA
--------------------------------------------------
MHA         1.6 GB               1x  (baseline)
GQA (G=8)   200 MB               8x
MQA         25 MB                64x
MLA         ~107 MB              ~15x (higher quality than MQA)

MLA is used in DeepSeek-V2, DeepSeek-V3, DeepSeek-V4 and DeepSeek-R1, all of which are among the most capable open-source models as of mid-2026. The technique represents a genuinely novel approach to the KV cache problem that goes beyond the simple head-sharing of GQA and MQA.

CHAPTER SIX: POSITIONAL ENCODING - TELLING THE MODEL WHERE THINGS ARE

The Transformer's self-attention mechanism is, by itself, permutation-invariant: if you shuffle the tokens in a sequence, the attention scores change (because the tokens themselves change), but the mechanism has no inherent notion of order. To give the model information about the position of each token in the sequence, we need positional encodings. This seemingly simple requirement has spawned a rich line of research with profound implications for context length generalization.

6.1 SINUSOIDAL AND LEARNED ABSOLUTE POSITIONS

The original Transformer used sinusoidal positional encodings: fixed vectors added to the token embeddings, where the i-th dimension of the positional encoding for position pos uses a sine or cosine function of a specific frequency:

PE(pos, 2i)   = sin(pos / 10000^(2i / d_model))
PE(pos, 2i+1) = cos(pos / 10000^(2i / d_model))

The intuition is that each dimension oscillates at a different frequency, so different positions produce different patterns of sine and cosine values. The model can learn to use these patterns to infer relative positions. BERT and many early models used learned absolute position embeddings instead, where each position gets a trainable vector. Both approaches have a fundamental limitation: they do not generalize well to sequence lengths longer than those seen during training.

6.2 ROTARY POSITION EMBEDDINGS (ROPE): THE MODERN STANDARD

Rotary Position Embedding (RoPE), introduced by Su et al. (2021, arXiv:2104.09864) in the RoFormer paper, takes a fundamentally different approach. Instead of adding position information to the token embeddings, RoPE encodes position by rotating the query and key vectors before computing attention scores. Specifically, the query vector at position m and the key vector at position n are rotated by angles proportional to m and n respectively, so that their dot product depends only on their relative position (m - n):

Rotated_Q(m) = R(m) * Q
Rotated_K(n) = R(n) * K
Q(m) . K(n) = (R(m) * Q) . (R(n) * K) = Q . (R(m-n) * K)

where R(theta) is a rotation matrix parameterized by theta. This is elegant: the attention score between two tokens automatically encodes their relative distance, and the encoding decays naturally as the distance increases (because the dot product of two rotated vectors decreases as the rotation angle between them increases).

RoPE has several important properties. It generalizes to sequences longer than those seen during training better than absolute position embeddings. It is compatible with efficient attention implementations. It naturally captures the intuition that nearby tokens are more relevant than distant ones. And it is computationally cheap: the rotation can be implemented as an element-wise complex multiplication, adding negligible overhead. RoPE is now used by LLaMA, Mistral, Qwen, Gemma, DeepSeek, and essentially every modern open-source LLM.

6.3 ALIBI: ATTENTION WITH LINEAR BIASES

ALiBi (Attention with Linear Biases), proposed by Press et al. (2021, arXiv:2108.12409), takes yet another approach. Instead of modifying the token embeddings or the query/key vectors, ALiBi simply adds a penalty to the attention score that is proportional to the distance between the query and key tokens:

Attention score(i, j) = Q_i . K_j / sqrt(d)  -  m * |i - j|

where m is a head-specific slope. The penalty grows linearly with distance, so distant tokens are always penalized relative to nearby ones. This has a beautiful property: because the bias is relative (not absolute), the model can generalize to sequences longer than those seen during training simply by extrapolating the linear bias. ALiBi models trained on 1,024-token sequences can often handle 2,048 or even 4,096 tokens at test time with minimal quality degradation. ALiBi is used in BLOOM and MPT, among others.

6.4 EXTENDING CONTEXT: POSITION INTERPOLATION, YARN, AND IROPE

A major practical challenge is extending the context window of a pretrained model beyond its training length. If a model was trained on sequences of 4,096 tokens and you want to use it on 128,000-token sequences, the positional encodings for positions beyond 4,096 are out-of-distribution, and the model will likely produce garbage.

Position Interpolation (Chen et al., 2023, arXiv:2306.15595) addresses this by scaling down all position indices so that the maximum position in the long sequence maps to the maximum position seen during training. If the training length is L and the target length is L', each position index pos is replaced by pos * (L / L'). This ensures that all positions are within the training distribution, at the cost of reducing the resolution of position information (nearby tokens now have more similar position encodings than before). A small amount of fine-tuning on long sequences restores quality.

YaRN (Yet another RoPE extensioN, Peng et al., 2023, arXiv:2309.00071) improves on simple interpolation by using a non-uniform approach: high-frequency components of the RoPE encoding (which capture fine-grained local position information) are left unchanged, while low-frequency components (which capture coarse global position information) are interpolated. This preserves the local position resolution that is important for language modeling while still allowing the model to handle longer sequences. YaRN can extend LLaMA models from 4,096 to 128,000 tokens with minimal fine-tuning, achieving state-of-the-art performance on long-context benchmarks.

The most recent innovation in this space is iRoPE (interleaved RoPE), introduced in Meta's Llama 4 (2025, arXiv:2505.09343). iRoPE interleaves standard RoPE attention layers with NoPE (No Positional Encoding) layers throughout the model. The NoPE layers have no positional encoding at all, relying entirely on the content of the tokens to determine attention patterns. The intuition is that some layers should focus on local, position-sensitive relationships (handled by RoPE layers), while others should focus on global, position-agnostic relationships (handled by NoPE layers). This combination allows Llama 4 Scout to generalize to context windows of up to 10 million tokens, an extraordinary achievement that would have seemed impossible just a few years ago.

CHAPTER SEVEN: MIXTURE OF EXPERTS - THE ARCHITECTURE OF SELECTIVE COMPUTATION

We now come to one of the most transformative architectural innovations of the past few years: Mixture of Experts (MoE). The core idea is simple and powerful. Instead of having a single feed-forward network that processes every token, have many feed-forward networks (the "experts") and route each token to only a small subset of them. This allows the model to have a very large total number of parameters while only using a small fraction of them for any given token, keeping the computational cost manageable.

7.1 THE BASIC MoE LAYER

In a standard Transformer, the feed-forward network in each layer has a fixed set of parameters that are applied to every token. In an MoE Transformer, the feed-forward network is replaced by E expert networks FFN_1, FFN_2, ..., FFN_E, plus a lightweight router network that decides which experts to use for each token.

The router takes the token's hidden state as input and produces a probability distribution over the E experts. The top-k experts (typically k=1 or k=2) are selected, and the token is processed by those experts. The outputs are combined as a weighted sum, where the weights come from the router's probability distribution:

Router output:  g = softmax(x * W_router)
Top-k experts:  {e_1, e_2, ..., e_k} = top-k indices of g
MoE output:     sum over i in {e_1,...,e_k} of g[e_i] * FFN_{e_i}(x)

Here is a concrete example with 8 experts and top-2 routing for a single token:

Router scores (before softmax):
Expert 1: 0.8   Expert 2: 0.3   Expert 3: 1.2   Expert 4: 0.1
Expert 5: 0.9   Expert 6: 0.5   Expert 7: 0.2   Expert 8: 0.6

After softmax, top-2 selection (Expert 3 and Expert 5):
Expert 3 score: 0.42  (renormalized)
Expert 5 score: 0.31  (renormalized)

MoE output = 0.42 * FFN_3(x) + 0.31 * FFN_5(x)

The key insight is that while the model has 8 * FFN_size parameters in this layer, each token only activates 2 of them. The total parameter count is large (enabling the model to store more knowledge), but the active parameter count per token (which determines the compute cost) is small.

7.2 THE SWITCH TRANSFORMER: SIMPLICITY AT SCALE

The Switch Transformer (Fedus et al., 2021, arXiv:2101.03961) was one of the first demonstrations that MoE could work at very large scale. It simplified the routing to top-1 (each token goes to exactly one expert) and showed that this could achieve 7x speedup over a dense T5 model with the same compute budget. The Switch Transformer scaled to 1.6 trillion parameters, demonstrating that MoE could reach parameter counts that were simply impossible for dense models.

The Switch Transformer also identified a key challenge with MoE: load balancing. If the router learns to always send tokens to the same few experts, those experts become bottlenecks while the others are idle. To prevent this, the Switch Transformer introduced an auxiliary load balancing loss that encourages the router to distribute tokens evenly across experts:

Load balancing loss = E * sum_over_experts_e of
                     (fraction of tokens to e) * (fraction of router prob to e)

This loss is added to the main training loss with a small coefficient, gently nudging the router toward balanced routing without overriding the primary training objective.

7.3 MIXTRAL: MoE GOES MAINSTREAM

Mixtral 8x7B (Mistral AI, 2023, arXiv:2401.04088) brought MoE to the open-source community in a highly accessible form. Mixtral has the same architecture as Mistral 7B, but each feed-forward layer is replaced by 8 experts with top-2 routing. The total parameter count is 47B, but only 13B parameters are active for any given token. Mixtral outperforms LLaMA 2 70B on most benchmarks while using less than half the active compute, demonstrating that MoE is not just a research curiosity but a practical architecture for building high-quality models efficiently.

7.4 DEEPSEEKMOE: FINE-GRAINED EXPERTS AND SHARED EXPERTS

DeepSeekMoE (Dai et al., 2024, arXiv:2401.06066) introduced two important innovations that pushed MoE further. The first is fine-grained expert segmentation: instead of having a small number of large experts, use a large number of small experts. For example, instead of 8 experts each with hidden dimension 4096, use 64 experts each with hidden dimension 512. With top-k routing selecting more experts (e.g., top-8 out of 64), the model can form more flexible combinations of expertise, since any subset of 8 small experts can be activated. This gives the model much more combinatorial flexibility in how it processes different types of tokens.

The second innovation is shared expert isolation: a small number of experts (typically 1 or 2) are designated as "shared experts" that are always activated for every token, regardless of the router's decision. These shared experts are expected to capture common knowledge that is useful for all tokens, while the routed experts capture specialized knowledge. This separation prevents the routed experts from wasting capacity on universal patterns that could be handled by the shared experts. DeepSeekMoE 16B achieves comparable performance to LLaMA-2 70B while using only 40 percent of the compute, a remarkable efficiency gain.

7.5 DEEPSEEK-V3: THE STATE OF THE ART IN MoE

DeepSeek-V3 (December 2024, arXiv:2412.19437) represents the current pinnacle of MoE architecture design as of mid-2026. It has 671 billion total parameters with 37 billion active parameters per token. It combines MLA for efficient KV caching with DeepSeekMoE for efficient computation. It introduces two further innovations: auxiliary-loss-free load balancing and multi-token prediction.

Auxiliary-loss-free load balancing replaces the explicit load balancing loss with a bias-based mechanism: each expert has a learnable bias added to its router score, and these biases are adjusted dynamically during training to ensure balanced routing without contaminating the main training loss. This is a subtle but important improvement because the auxiliary loss can interfere with the primary training objective, and removing it allows the model to optimize more cleanly.

Multi-token prediction (MTP) trains the model to predict not just the next token but the next several tokens simultaneously, using additional output heads. This is used as an auxiliary training objective that improves sample efficiency and encourages the model to plan ahead in its representations. At inference time, the MTP heads can also be used for speculative decoding (which we discuss in Chapter Nine), where the model generates multiple candidate tokens in parallel and then verifies them. DeepSeek-V3 achieves performance comparable to leading closed-source models like GPT-4o and Claude 3.5 Sonnet at a fraction of the training cost.

7.6 KIMI K2: TRILLION-PARAMETER MoE

Kimi K2 (Moonshot AI, June 2025, arXiv:2506.01939) pushes MoE to one trillion total parameters with 32 billion active parameters per token, trained on 15.5 trillion tokens. It achieves state-of-the-art performance among open-source models on agentic and coding tasks. Kimi K2 also introduces the MuonClip optimizer (discussed in Chapter Ten), demonstrating that architectural and optimization innovations continue to compound at the frontier.

7.7 EXPERT CHOICE ROUTING: FLIPPING THE SCRIPT

All the MoE systems described so far use token-choice routing: each token chooses which experts to route to. This can lead to load imbalance if the router consistently prefers certain experts. Expert Choice Routing (Zhou et al., 2022, and subsequent work in 2025) flips this: instead of tokens choosing experts, experts choose which tokens to process. Each expert selects the top-k tokens from the batch that it wants to process, based on the router scores. This naturally enforces perfect load balancing (every expert processes exactly the same number of tokens) without any auxiliary loss. The tradeoff is that some tokens may be processed by more experts than others, and some tokens may be processed by fewer, but in practice this heterogeneity can actually be beneficial, as more complex or ambiguous tokens receive more compute.

CHAPTER EIGHT: THE BYTE LATENT TRANSFORMER - RETHINKING TOKENIZATION ITSELF

Before leaving the topic of how information is structured and fed into Transformers, we should mention a radical rethinking that emerged in late 2024: the Byte Latent Transformer (BLT, Yu et al., 2024, arXiv:2412.09871) from Meta AI. Standard LLMs operate on tokens, which are subword units produced by a tokenizer like BPE (Byte Pair Encoding). Tokenization is a preprocessing step that converts raw bytes into a fixed vocabulary of subword units, which the model then processes.

Tokenization has several well-known problems. It is language-specific: tokenizers trained on English text are inefficient for other languages, producing many more tokens per character. It is brittle: small changes to a word (like adding a space or changing capitalization) can produce very different token sequences. It requires a fixed vocabulary decided before training. And it cannot adapt its granularity to the difficulty of the content.

BLT operates directly on raw bytes, without any tokenization. Instead of processing a fixed sequence of tokens, BLT dynamically groups bytes into patches of variable size, where the patch boundaries are determined by the entropy of the next byte prediction. High-entropy regions (where the next byte is hard to predict, indicating complex or novel content) get smaller patches with more compute per byte. Low-entropy regions (where the next byte is easy to predict, indicating common or repetitive content) get larger patches with less compute per byte.

This dynamic allocation of compute is elegant: the model spends more effort on the parts of the input that are actually difficult and less effort on the parts that are easy. BLT matches tokenization-based LLM performance at scale while being more robust to character-level perturbations and more efficient on non-English languages. It represents a potential paradigm shift in how LLMs process input, though as of mid-2026 it has not yet been widely adopted in production models.

CHAPTER NINE: MAKING INFERENCE FAST - QUANTIZATION, SPECULATIVE DECODING, AND MULTI-TOKEN PREDICTION

Training a large language model is expensive, but inference is where the costs really accumulate, because every user query requires a forward pass through the model. The optimizations in this section focus on making inference as fast and memory-efficient as possible.

9.1 QUANTIZATION: DOING MORE WITH FEWER BITS

Modern neural networks are typically trained in 32-bit or 16-bit floating point. But do we actually need that much precision for inference? Quantization is the process of representing weights and/or activations with fewer bits, reducing memory usage and potentially speeding up computation.

The simplest form is post-training quantization (PTQ), where a trained model's weights are converted to lower precision without any retraining. LLM.int8() (Dettmers et al., 2022, arXiv:2208.07339) showed that 8-bit quantization of LLM weights is possible without significant quality loss, by using a decomposition that handles outlier activation values (which are common in large LLMs) in 16-bit while quantizing the rest in 8-bit. This halves the memory required for inference.

GPTQ (Frantar et al., 2022, arXiv:2210.17323) pushed this further to 3 or 4 bits per weight, using an approximate second-order optimization method to minimize the quantization error. GPTQ can quantize a 175-billion-parameter model in about four GPU hours, and the resulting 4-bit model has negligible quality degradation compared to the 16-bit original. A 4-bit quantized model requires only one-quarter the memory of its 16-bit counterpart, making it possible to run models that would otherwise require multiple GPUs on a single GPU.

SmoothQuant (Xiao et al., 2022, arXiv:2211.01524) addresses the challenge of quantizing both weights and activations (W8A8 quantization), which is necessary for maximum inference speedup on hardware that supports 8-bit integer arithmetic. The problem is that LLM activations have large outliers in certain channels, which make naive quantization very lossy. SmoothQuant migrates the quantization difficulty from activations to weights by applying a mathematically equivalent scaling transformation: multiply the activations by a per-channel smoothing factor and divide the weights by the same factor. This makes the activations easier to quantize while making the weights slightly harder, but weights are much easier to quantize than activations because they are fixed and can be calibrated offline.

The trend in 2025 and 2026 has been toward training models natively in lower precision. FlashAttention-3's FP8 support, combined with hardware like the H100 and H200 that have native FP8 tensor cores, enables training and inference in 8-bit floating point with carefully managed precision. Research into scaling laws for precision (Dettmers et al., 2024, arXiv:2501.00663) shows that lower-precision training can be compute-optimal in certain regimes, suggesting that future models may be trained in FP8 or even lower precision from the start.

9.2 SPECULATIVE DECODING: USING A SMALL MODEL TO ACCELERATE A LARGE ONE

Autoregressive text generation is inherently sequential: to generate token t+1, you need token t, which means you cannot parallelize across the time dimension. Each forward pass through a large model generates exactly one token, which is very inefficient given that GPUs are designed for massively parallel computation.

Speculative decoding (Leviathan et al., 2022, arXiv:2211.17192; Chen et al., 2023, arXiv:2302.01318) is a clever technique that breaks this sequential bottleneck. The idea is to use a small, fast draft model to generate a sequence of candidate tokens, and then use the large target model to verify all of them in a single parallel forward pass. If the large model agrees with the draft model's predictions, all the candidate tokens are accepted. If the large model disagrees at some position, that token and all subsequent ones are rejected, and the large model's prediction at the point of disagreement is used instead.

Here is a concrete example. Suppose the draft model generates the candidates "the", "cat", "sat", "on", "the" in 5 fast forward passes. The large model then processes all 5 tokens in a single parallel forward pass and checks whether it agrees. If it agrees with all 5, we have generated 5 tokens at the cost of 1 large model forward pass plus 5 small model forward passes. Since the small model is much faster than the large model, this is a significant speedup.

The key insight is that the output distribution of the combined system is identical to the output distribution of the large model alone. Speculative decoding does not change what the model says; it only changes how fast it says it. The speedup depends on the acceptance rate: how often the draft model's predictions match the large model's. For a good draft model, acceptance rates of 70 to 90 percent are achievable, leading to 2 to 4x speedup in practice.

Medusa (Cai et al., 2024, arXiv:2401.10774) is a variant that eliminates the need for a separate draft model. Instead, it adds multiple extra decoding heads to the top of the large model, each predicting a different future token. These heads are lightweight and can be trained quickly on top of a frozen base model. The predictions from the Medusa heads are verified using a tree-based attention mechanism, allowing multiple candidate continuations to be checked simultaneously. Medusa achieves 2 to 3x speedup without requiring a separate draft model.

DeepSeek-V3's multi-token prediction training objective also enables speculative decoding: the auxiliary MTP heads trained during pretraining can serve as draft heads at inference time, providing a natural and well-integrated source of speculative candidates.

9.3 MULTI-TOKEN PREDICTION AS A TRAINING OBJECTIVE

Standard LLM training uses the next-token prediction loss: for each position in the sequence, the model predicts the next token, and the loss is the cross-entropy between the prediction and the true next token. Gloeckle et al. (2024, arXiv:2404.19737) at Meta showed that training models to predict multiple future tokens simultaneously, using n independent output heads operating on a shared model trunk, results in higher sample efficiency. Models trained with 4-token prediction are faster to train and achieve better performance on code generation tasks.

The intuition is that predicting multiple future tokens forces the model to develop more structured, forward-looking representations of the current context, rather than just memorizing local patterns. This is particularly beneficial for tasks that require planning, like code generation, where the model needs to think several steps ahead. The additional output heads also serve double duty at inference time: they can be used as speculative decoding draft heads, providing a free speedup with no additional model parameters beyond what was already trained.

9.4 CONTINUOUS BATCHING AND PAGEDATTENTION

At the system level, two innovations have dramatically improved the throughput of LLM inference servers: continuous batching and PagedAttention.

Traditional batching for LLM inference processes a fixed batch of requests together and waits for all of them to finish before starting the next batch. Since different requests have different lengths, shorter requests finish early and their GPU resources sit idle while waiting for longer requests. Continuous batching (also called iteration-level scheduling) instead processes requests at the token level: as soon as a request finishes generating a token, the next request in the queue can be added to the batch for the next iteration. This keeps the GPU fully utilized at all times and dramatically increases throughput.

PagedAttention, introduced in the vLLM system (Kwon et al., 2023, arXiv:2309.06180), addresses the memory fragmentation problem in KV cache management. Different requests have different context lengths, and the KV cache for each request grows dynamically as tokens are generated. Naively allocating contiguous memory for each request's KV cache leads to severe fragmentation, wasting a large fraction of GPU memory. PagedAttention borrows the concept of virtual memory paging from operating systems: the KV cache is divided into fixed-size pages, and pages are allocated on demand. Pages from different requests can be interleaved in physical memory, eliminating fragmentation. PagedAttention also enables efficient KV cache sharing between requests that share a common prefix (e.g., a system prompt), further reducing memory usage. Together, continuous batching and PagedAttention have become the standard infrastructure for production LLM serving, enabling 10 to 20 times higher throughput compared to naive serving implementations.

CHAPTER TEN: BETTER TRAINING - NEW OPTIMIZERS AND TRAINING TECHNIQUES

The optimization algorithm used to train a neural network has a profound effect on convergence speed and final model quality. Adam and its variant AdamW have been the dominant optimizers for Transformer training since 2018, but recent work has produced serious challengers.

10.1 ADAM AND ADAMW: THE INCUMBENT

Adam (Kingma and Ba, 2014) maintains exponential moving averages of both the gradient (first moment) and the squared gradient (second moment) for each parameter, and uses these to adaptively scale the learning rate for each parameter. This makes Adam much more robust to the choice of learning rate than vanilla SGD, and it handles sparse gradients well. AdamW (Loshchilov and Hutter, 2019) adds decoupled weight decay, which regularizes the weights independently of the gradient-based update. AdamW is the standard optimizer for LLM pretraining.

The main weakness of Adam is its memory cost: it stores two additional tensors (the first and second moment estimates) for every parameter, tripling the memory required for the optimizer state. For a 70-billion-parameter model, the optimizer state alone requires roughly 560 GB at 32-bit precision. This is a significant constraint on the maximum model size that can be trained on a given hardware budget.

10.2 SOPHIA: SECOND-ORDER OPTIMIZATION MADE PRACTICAL

Sophia (Liu et al., 2023, arXiv:2305.14342) is a stochastic second-order optimizer that uses a diagonal estimate of the Hessian (the matrix of second derivatives of the loss with respect to the parameters) as a preconditioner. The Hessian captures the curvature of the loss landscape: parameters in directions of high curvature should take small steps (to avoid overshooting), while parameters in directions of low curvature can take large steps. Adam approximates curvature using the squared gradient, which is a rough proxy. Sophia uses a better estimate, computed periodically using a Hutchinson estimator. Sophia converges twice as fast as Adam in terms of the number of gradient steps, saving 50 percent of total training compute. However, computing the Hessian estimate adds overhead, and Sophia has not yet been widely adopted in frontier model training.

10.3 MUON AND MUONCLIP: ORTHOGONALIZED UPDATES

Muon (originally proposed by Keller Jordan in a 2024 blog post and formalized in arXiv:2502.16982) is a new optimizer that applies Nesterov momentum and then orthogonalizes the gradient update using Newton-Schulz iterations. The orthogonalization step ensures that the update matrix has orthonormal columns, which prevents different gradient directions from becoming correlated and leads to more efficient parameter updates. Muon achieves better final loss than AdamW at the same compute budget on language model pretraining benchmarks.

MuonClip, used in Kimi K2 (2025, arXiv:2506.01939), adds gradient clipping to Muon to prevent loss spikes during training of very large MoE models. Kimi K2's successful training of a one-trillion-parameter model with MuonClip suggests that this optimizer is robust enough for frontier-scale training. The adoption of Muon by a frontier lab for a production model is a significant signal that the optimizer landscape for LLM training is genuinely evolving beyond AdamW.

10.4 TEST-TIME COMPUTE SCALING: REASONING MODELS

One of the most important developments of 2024-2025 is the discovery that scaling compute at inference time (test-time compute) is a powerful complement to scaling compute at training time. Models like OpenAI o1, o3, and DeepSeek-R1 use reinforcement learning to train models to generate extended chains of thought before producing a final answer. By spending more tokens on reasoning, these models achieve dramatically better performance on complex tasks like mathematics, coding, and scientific reasoning.

DeepSeek-R1 (January 2025, arXiv:2501.12948) showed that a model trained with large-scale reinforcement learning, without any supervised fine-tuning data, can develop sophisticated reasoning capabilities. The model learns to break problems into steps, check its work, backtrack when it makes mistakes, and try alternative approaches, all within the context window of a single forward pass sequence. This is not a change to the Transformer architecture per se, but it represents a new paradigm for how Transformers are used: not just as pattern matchers but as deliberate reasoners.

Research into scaling laws for test-time compute (Snell et al., 2024, arXiv:2408.03314) shows that the optimal test-time compute strategy depends on problem difficulty. For easy problems, generating multiple candidate answers and selecting the best one (best-of-N sampling) is most efficient. For hard problems, iterative self-refinement (where the model critiques and improves its own answer) is better. This suggests that future systems will adaptively allocate test-time compute based on estimated problem difficulty.

CHAPTER ELEVEN: LONG-CONTEXT TRANSFORMERS - RING ATTENTION AND BEYOND

Even with FlashAttention and sparse attention, there are practical limits to the context length a single GPU can handle. For truly enormous contexts (millions of tokens), we need to distribute the computation across multiple devices. Ring Attention (Liu et al., 2023, arXiv:2310.01889) is the key technique for doing this.

The idea behind Ring Attention is to distribute the query, key, and value sequences across multiple devices arranged in a logical ring. Each device holds a chunk of the query sequence and a chunk of the key-value sequence. Attention is computed in a blockwise fashion: each device computes attention between its local query chunk and its local key-value chunk, then passes its key-value chunk to the next device in the ring while receiving the key-value chunk from the previous device. After N devices have passed their chunks around the ring, each device has computed attention between its local queries and all key-value pairs in the full sequence.

The crucial insight is that the key-value communication between devices can be overlapped with the attention computation: while device i is computing attention with the key-value chunk it currently holds, it is simultaneously sending that chunk to device i+1 and receiving the next chunk from device i-1. This communication-computation overlap means that Ring Attention adds minimal overhead compared to single-device attention, and it allows the context length to scale linearly with the number of devices.

Using Ring Attention, the context length is limited only by the total memory of the distributed cluster. With a cluster of 1,000 GPUs each with 80 GB of memory, the theoretical maximum context length is in the billions of tokens. In practice, MiniMax-01 (2025, arXiv:2501.09755) demonstrated a 32-million-token context window using a combination of linear attention and MoE, and Llama 4 Scout (2025, arXiv:2505.09343) demonstrated a 10-million-token context window using iRoPE. These are extraordinary achievements that open up entirely new applications: analyzing entire codebases, processing book-length documents, or reasoning over years of conversation history.

CHAPTER TWELVE: TITANS AND NEURAL MEMORY - LEARNING TO REMEMBER

A fundamental limitation of the Transformer is that its "memory" is entirely contained in its context window. Everything the model knows about the current task must fit within the context, and anything outside the context is simply inaccessible. This is fine for many tasks, but it is a serious limitation for tasks that require long-term memory: remembering facts from earlier in a conversation, maintaining a persistent world model, or accumulating knowledge over many interactions.

Titans (Behrouz, Zhong, and Mirrokni, 2024, arXiv:2501.00663) is a new architecture that addresses this by incorporating a neural long-term memory module into the Transformer. The long-term memory is a separate neural network that can be updated at test time: as the model processes new information, it updates the weights of the memory module to encode that information. When the model needs to recall something, it queries the memory module to retrieve relevant information.

This is a form of meta-learning or test-time adaptation: the model's parameters (specifically, the memory module's parameters) change in response to the input, allowing the model to accumulate information beyond what fits in the context window. Titans achieves better performance than standard Transformers on long-context tasks while being more efficient, because the memory module can store information compactly without requiring it to be in the attention context. Titans represents a direction that many researchers believe is essential for the next generation of AI systems: models that can learn and remember over extended periods, not just within a single context window.

CHAPTER THIRTEEN: WHERE ARE THE LIMITS?

We have now surveyed a remarkable collection of optimizations that have extended the Transformer far beyond what its original designers imagined. But every architecture has limits, and the Transformer is no exception. Let us think carefully about where those limits are.

13.1 THE QUADRATIC WALL

Even with FlashAttention and sparse attention, the fundamental computational complexity of exact self-attention is O(N^2). For N = 1 million tokens, this is 10^12 operations per layer. For N = 10 million tokens, it is 10^14. Even with the fastest hardware, exact attention over very long sequences is simply not feasible. Sparse attention and linear attention approximations can reduce this, but they come with quality tradeoffs. The quadratic wall is a genuine architectural limit that cannot be fully engineered away within the standard attention framework.

State space models like Mamba (Gu and Dao, 2023, arXiv:2312.00752) and its successor Mamba-2 (Dao and Gu, 2024, arXiv:2405.21060) offer a fundamentally different approach with O(N) complexity. Mamba uses selective state space models that can selectively remember or forget information based on the input content, achieving performance competitive with Transformers on language modeling while being 5x faster at inference. Mamba-2 showed that SSMs and attention are two instances of a more general framework called structured state space duality (SSD), and built hybrid architectures that are 2 to 8 times faster than Mamba while remaining competitive with Transformers. Griffin (De et al., 2024, arXiv:2402.19427) from Google DeepMind mixes gated linear recurrences with local attention, matching LLaMA-2 performance while being more efficient on long sequences. RWKV (Peng et al., 2023, arXiv:2305.13048) combines the efficient parallelizable training of Transformers with the efficient inference of RNNs, achieving linear complexity in both training and inference.

These hybrid architectures, combining the strengths of attention (strong performance on tasks requiring global information retrieval) with the efficiency of recurrence (linear complexity in sequence length), may represent the future of sequence modeling beyond the pure Transformer.

13.2 THE DATA WALL

Scaling laws (Kaplan et al., 2020, arXiv:2001.08361; Hoffmann et al., 2022, arXiv:2203.15556) tell us that model performance improves predictably as we increase model size, dataset size, and compute. But this scaling depends on having access to more and more high-quality training data. The internet contains a finite amount of high-quality text, and current estimates suggest that frontier LLMs will exhaust high-quality English text data by 2026 to 2028. Repeating data significantly degrades model quality, so simply re-using existing data is not a satisfactory solution.

Potential responses to the data wall include synthetic data generation (having models generate training data for themselves or for other models), multimodal data (using images, audio, video, and code as additional training signal), and data from specialized domains (scientific papers, legal documents, medical records). Each of these has its own challenges: synthetic data can introduce biases and errors, multimodal data requires new architectural components, and specialized domain data may not generalize well.

13.3 THE REASONING WALL

Current Transformers, even the best reasoning models, struggle with tasks that require systematic generalization: applying a learned rule to a new combination of concepts that was not seen during training. A model might learn that "2 + 3 = 5" and "4 + 7 = 11" but fail to correctly compute "2 + 3 + 4 + 7" if it has never seen four-term additions. This is because Transformers learn statistical patterns rather than symbolic rules, and statistical patterns do not always generalize in the way that symbolic rules do.

Achieving robust systematic generalization may require architectural innovations beyond the standard Transformer: explicit memory systems, modular reasoning components, or integration with symbolic computation. This is an active area of research, and it is one of the key open problems on the path toward more capable AI systems.

13.4 THE GROUNDING WALL

Language models learn from text, and text is a description of the world, not the world itself. A model trained only on text has no direct experience of physical reality: it has never seen a red apple, heard a dog bark, or felt the weight of a book. This lack of grounding means that language models can produce fluent, confident descriptions of things they fundamentally do not understand in the way that a human understands them.

Multimodal models (like GPT-4o, Gemini 2.5, and Llama 4) partially address this by training on images, audio, and video in addition to text. But even multimodal models are trained on recorded data, not on direct sensorimotor experience. Truly grounded AI may require embodied systems that interact with the physical world, which is a very different research direction from scaling Transformers.

13.5 THE STATIC WEIGHT WALL

Once a Transformer is trained, its weights are fixed. It cannot learn from new information without retraining or fine-tuning. This is in stark contrast to human intelligence, which continuously learns from experience. Continual learning, the ability to learn new things without forgetting old ones (catastrophic forgetting), is a major unsolved problem for neural networks in general and Transformers in particular.

Titans' neural long-term memory is one approach to this problem. Retrieval-augmented generation (RAG), where the model retrieves relevant information from an external database at inference time, is another. But neither fully solves the problem of continual learning from experience. This is likely to be one of the most important research directions of the next decade.

CHAPTER FOURTEEN: WHAT CAN WE EXPECT NEXT?

Based on the trajectory of research through mid-2026, here are the most promising directions for further improvement of Transformer-based systems.

The first direction is deeper integration of MoE and attention. Current MoE models apply sparsity only to the feed-forward layers, while attention layers remain dense. Sparse attention over experts (where different attention heads are routed to different experts) is an underexplored direction that could further improve efficiency.

The next direction is hardware-software co-design. FlashAttention succeeded by deeply understanding the GPU memory hierarchy and designing the algorithm around it. Future optimizations will likely require similar co-design between algorithm designers and hardware architects. New hardware features (like the Hopper GPU's TMA and WGMMA instructions) will continue to unlock new algorithmic possibilities. NVIDIA's Blackwell architecture (B100, B200), which began shipping in 2025, introduces further innovations in memory bandwidth and FP4/FP6 precision support that will drive the next generation of attention and MoE optimizations.

With the third direction we focus on  better long-term memory. Titans showed that neural long-term memory is feasible, but the field is still in its early stages. More sophisticated memory architectures that can store, retrieve, and update information efficiently over very long time horizons will be essential for the next generation of AI assistants.

The fourth direction is native multimodality. Current multimodal models often process different modalities with separate encoders that are then combined. Truly native multimodal architectures, like the Byte Latent Transformer's approach of operating on raw bytes, could process all modalities in a unified framework without modality-specific preprocessing. Llama 4's early fusion approach, which processes text and image tokens in the same sequence from the first layer, is a step in this direction.

In the fifth direction better training algorithms are addressed. The success of Muon and MuonClip suggests that there is still significant room for improvement in the optimization algorithms used to train Transformers. Second-order methods, modular optimization (applying different optimizers to different parts of the model), and curriculum learning strategies are all active areas of research.

The sixth direction is test-time compute scaling. The reasoning model paradigm (o1, o3, DeepSeek-R1) has demonstrated that spending more compute at inference time can dramatically improve performance on complex tasks. Developing better algorithms for allocating test-time compute, better training procedures for reasoning models, and more efficient implementations of chain-of-thought reasoning will be major research priorities.

Last but not least, the seventh direction is about hybrid architectures. The convergence of SSMs and attention in frameworks like Mamba-2 and Griffin suggests that the dichotomy between recurrent and attention-based models is dissolving. Future architectures will likely be hybrids that combine the strengths of both, using attention where global information retrieval is needed and recurrence where long-range memory efficiency is needed.

EPILOGUE: THE TRANSFORMER IS NOT THE END OF HISTORY

The Transformer has been the dominant architecture in deep learning for nearly a decade, and it will likely remain dominant for several more years. But it is not the end of history. Every optimization we have discussed in this tutorial has pushed the Transformer further, but each has also revealed new limits and new problems. FlashAttention made long contexts feasible, which revealed that models needed better positional encodings. Better positional encodings enabled longer contexts, which revealed that KV caches were a bottleneck. GQA and MLA addressed KV caches, which revealed that the data wall and the reasoning wall were the next frontiers.

This is how progress works: each solution creates new problems, and those new problems drive the next wave of innovation. The Transformer is a remarkable foundation, but the systems of 2030 will likely look quite different from the systems of today, incorporating ideas from state space models, neural memory, embodied learning, and architectures we have not yet imagined.

What will remain constant is the core insight of the Transformer: that attention, the ability to selectively focus on relevant information in a sea of irrelevant information, is a powerful and general computational primitive. Whether future architectures implement this insight with softmax attention, kernel approximations, state space models, or something entirely new, the principle of selective, content-based information routing will remain central to intelligent computation.

The journey from "Attention Is All You Need" to one-trillion-parameter reasoning models with 32-million-token context windows has taken less than a decade. The next decade will be at least as surprising.

REFERENCES

Vaswani, A. et al. (2017). Attention Is All You Need. arXiv:1706.03762.

Zhang, B. and Sennrich, R. (2019). Root Mean Square Layer Normalization. arXiv:1910.07467.

Shazeer, N. (2020). GLU Variants Improve Transformers. arXiv:2002.05202.

Beltagy, I. et al. (2020). Longformer: The Long-Document Transformer. arXiv:2004.05150.

Wang, S. et al. (2020). Linformer: Self-Attention with Linear Complexity. arXiv:2006.04768.

Zaheer, M. et al. (2020). Big Bird: Transformers for Longer Sequences. arXiv:2007.14062.

Choromanski, K. et al. (2020). Rethinking Attention with Performers. arXiv:2009.14794.

Kaplan, J. et al. (2020). Scaling Laws for Neural Language Models. arXiv:2001.08361.

Fedus, W. et al. (2021). Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity. arXiv:2101.03961.

Su, J. et al. (2021). RoFormer: Enhanced Transformer with Rotary Position Embedding. arXiv:2104.09864.

Press, O. et al. (2021). Train Short, Test Long: Attention with Linear Biases Enables Input Length Extrapolation (ALiBi). arXiv:2108.12409.

Hoffmann, J. et al. (2022). Training Compute-Optimal Large Language Models (Chinchilla). arXiv:2203.15556.

Wang, H. et al. (2022). DeepNet: Scaling Transformers to 1,000 Layers. arXiv:2203.00555.

Dao, T. et al. (2022). FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness. arXiv:2205.14135.

Dettmers, T. et al. (2022). LLM.int8(): 8-bit Matrix Multiplication for Transformers at Scale. arXiv:2208.07339.

Frantar, E. et al. (2022). GPTQ: Accurate Post-Training Quantization for Generative Pre-trained Transformers. arXiv:2210.17323.

Xiao, G. et al. (2022). SmoothQuant: Accurate and Efficient Post-Training Quantization for Large Language Models. arXiv:2211.01524.

Leviathan, Y. et al. (2022). Fast Inference from Transformers via Speculative Decoding. arXiv:2211.17192.

Chen, C. et al. (2023). Accelerating Large Language Model Decoding with Speculative Sampling. arXiv:2302.01318.

Liu, H. et al. (2023). Sophia: A Scalable Stochastic Second-order Optimizer for Language Model Pre-training. arXiv:2305.14342.

Peng, B. et al. (2023). RWKV: Reinventing RNNs for the Transformer Era. arXiv:2305.13048.

Ainslie, J. et al. (2023). GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints. arXiv:2305.13245.

Chen, S. et al. (2023). Extending Context Window of Large Language Models via Positional Interpolation. arXiv:2306.15595.

Peng, B. et al. (2023). YaRN: Efficient Context Window Extension of Large Language Models. arXiv:2309.00071.

Kwon, W. et al. (2023). Efficient Memory Management for Large Language Model Serving with PagedAttention. arXiv:2309.06180.

Liu, H. et al. (2023). Ring Attention with Blockwise Transformers for Near-Infinite Context. arXiv:2310.01889.

Dao, T. (2023). FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning. arXiv:2307.08691.

Gu, A. and Dao, T. (2023). Mamba: Linear-Time Sequence Modeling with Selective State Spaces. arXiv:2312.00752.

Dai, D. et al. (2024). DeepSeekMoE: Towards Ultimate Expert Specialization in Mixture-of-Experts Language Models. arXiv:2401.06066.

Cai, T. et al. (2024). Medusa: Simple LLM Inference Acceleration Framework with Multiple Decoding Heads. arXiv:2401.10774.

Mistral AI (2024). Mixtral of Experts. arXiv:2401.04088.

De, S. et al. (2024). Griffin: Mixing Gated Linear Recurrences with Local Attention for Efficient LLMs. arXiv:2402.19427.

DeepSeek AI (2024). DeepSeek-V2: A Strong, Economical, and Efficient Mixture-of-Experts Language Model. arXiv:2405.04434.

Dao, T. and Gu, A. (2024). Transformers are SSMs: Generalized Models and Efficient Algorithms Through Structured State Space Duality (Mamba-2). arXiv:2405.21060.

Shah, J. et al. (2024). FlashAttention-3: Fast and Accurate Attention on Hopper GPUs. arXiv:2407.08608.

Snell, C. et al. (2024). Scaling LLM Test-Time Compute Optimally. arXiv:2408.03314.

Gloeckle, F. et al. (2024). Better and Faster Large Language Models via Multi-Token Prediction. arXiv:2404.19737.

Behrouz, A., Zhong, P., and Mirrokni, V. (2024). Titans: Learning to Memorize at Test Time. arXiv:2501.00663.

Yu, K. et al. (2024). Byte Latent Transformer: Patches Scale Better Than Tokens. arXiv:2412.09871.

DeepSeek AI (2024). DeepSeek-V3 Technical Report. arXiv:2412.19437.

DeepSeek AI (2025). DeepSeek-R1: Incentivizing Reasoning Capability in LLMs via Reinforcement Learning. arXiv:2501.12948.

MiniMax (2025). MiniMax-01: Scaling Foundation Models with Lightning Attention. arXiv:2501.09755.

DeepSeek AI (2025). Native Sparse Attention: Hardware-Aligned and Natively Trainable Sparse Attention. arXiv:2502.05171.

Jordan, K. et al. (2025). Muon is Scalable for LLM Training. arXiv:2502.16982.

Meta AI (2025). The Llama 4 Herd: The Beginning of a New Era of Natively Multimodal AI at Meta. arXiv:2505.09343.

Moonshot AI (2025). Kimi K2: Open Agentic Intelligence. arXiv:2506.01939.