Story of Attention — Part 2: Attention All You Need

Part 2 of 6 in the Story of Attention series.

1. Before the Transformer — how n-gram counting and RNN memory hit their ceilings, and how additive/multiplicative attention first let a decoder look back.
2. Attention All You Need (this post) — scaled dot-product attention, multi-head attention, and the four-year search for a way to encode position.
3. KV-Cache Bottleneck: MQA and GQA — sharing and grouping key/value heads to shrink inference memory.
4. Sparse Attention, Then and Now — fixed, windowed, hashed, and learned sparsity patterns, from 2019 to 2025.
5. Linear Attention: Kernels, Decay, and Gating — replacing softmax with decomposable kernels, decay, and gated recurrent state.
6. Where Attention Stands Today — latent compression, differential attention, and what frontier LLMs actually ship.

---

In 2017, a paper titled "Attention Is All You Need" made a case that sounded too simple: remove the recurrent network entirely and build a language model purely out of attention. Although it sounded interesting, there is a gap. A model that reads words sequentially knows the order very well inherently. But attention looks at every word at once. So it has to be told the order explicitly. This part is the story of the attention mechanism inside the Transformer, and how to tell the model about the position.

By the end of the previous part, attention was already doing the work of choosing where to look, while an RNN underneath carried the sequence forward one token at a time.

---

1. Scaled Dot-Product Attention

Vaswani et al. identified the problem of recurrence in a paper titled "Attention Is All You Need" [1]. The processing is one-by-one because each hidden state waits for the previous one. The gradients fade because they travel through every intermediate step. And the whole history gets squeezed into a fixed-size vector because that vector is the only memory a recurrence has. Removing recurrence entirely is what gives the resulting architecture its name: the Transformer.

_{The Transformer architecture: stacked encoder and decoder blocks built entirely from attention and feed-forward layers, with positional encodings injected at the input. Redrawn from Vaswani et al., 2017 [1].}

The core part of the Transformer is Luong's dot product [20] from the previous part, in a more generalized form. Every token gets projected into three vectors, a Query (what am I looking for?), a Key (what do I contain?), and a Value (what do I offer?). The relevance between two tokens is the dot product between the query of one and the key of the other, which is basically a similarity score between tokens. Then softmax turns these relevance scores into weights, and the output is the weighted mix of the values:

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

Here $Q \in \mathbb{R}^{n \times d_k}$, $K \in \mathbb{R}^{n \times d_k}$, $V \in \mathbb{R}^{n \times d_v}$, and $n$ is the sequence length. The product $QK^T$ is an $n \times n$ matrix of raw similarity scores, every token scored against every other token, in one big matrix multiplication. These raw scores are called logits, the unnormalized numbers that go into softmax. Now why divide by $\sqrt{d_k}$? Without the division, the variance of the logits grows with $d_k$. Large logits push softmax into its saturated region, where the gradients are nearly zero, and there the model stops learning. Dividing by $\sqrt{d_k}$ keeps the logit variance near 1 at any dimensionality. One more step happens before softmax, and it is called masking. The scores for the positions a token should not see are set to $-\infty$, hence softmax gives them exactly zero weight. In a decoder this is the causal mask. It blocks each token from attending to anything to its right, because without it the model could read the future token it is being trained to predict. Applied to the score matrix, the mask fills the upper triangle with $-\infty$:

$\text{mask}(QK^T) = \text{mask}\left(\begin{bmatrix} e_{11} & e_{12} & \cdots & e_{1n} \\ e_{21} & e_{22} & \cdots & e_{2n} \\ \vdots & \vdots & \ddots & \vdots \\ e_{n1} & e_{n2} & \cdots & e_{nn} \end{bmatrix}\right) = \begin{bmatrix} e_{11} & -\infty & \cdots & -\infty \\ e_{21} & e_{22} & \cdots & -\infty \\ \vdots & \vdots & \ddots & \vdots \\ e_{n1} & e_{n2} & \cdots & e_{nn} \end{bmatrix}$

Softmax is then applied row-wise, so each token ends up with a probability distribution over the tokens it is allowed to read from.

_{Scaled dot-product attention: queries scored against keys, scaled by the square root of the key dimension, softmaxed into weights, and applied to the values. Redrawn from Vaswani et al., 2017 [1].}

Key result: Attention works like a dictionary lookup, and it is differentiable. Each query retrieves a mixture of values, weighted by how similar the keys are. And the entire operation is a single matrix-multiply pipeline, with no recurrence, no convolution, and $O(n^2 \cdot d_k)$ complexity.

Worked example

Sequence length $n = 4$, head dimension $d_k = 4$, so $\sqrt{d_k} = 2$, with the causal mask applied, as in a decoder.

$Q = \begin{bmatrix} 1 & 0 & 1 & 0 \\ 0 & 2 & 0 & 1 \\ 1 & 1 & 0 & 0 \\ 0 & 0 & 1 & 2 \end{bmatrix}, \quad K = \begin{bmatrix} 1 & 1 & 0 & 0 \\ 0 & 1 & 1 & 0 \\ 0 & 0 & 2 & 1 \\ 1 & 0 & 0 & 1 \end{bmatrix}, \quad V = \begin{bmatrix} 2 & 0 & 1 & 1 \\ 1 & 3 & 0 & 0 \\ 0 & 1 & 4 & 2 \\ 1 & 2 & 0 & 5 \end{bmatrix}$

Step 1 — Raw scores ($QK^T$):

$QK^T = \begin{bmatrix} 1 & 1 & 2 & 1 \\ 2 & 2 & 1 & 1 \\ 2 & 1 & 0 & 1 \\ 0 & 1 & 4 & 2 \end{bmatrix}$

Step 2 — Scale (divide by 2):

$\text{Scaled} = \begin{bmatrix} 0.5 & 0.5 & 1.0 & 0.5 \\ 1.0 & 1.0 & 0.5 & 0.5 \\ 1.0 & 0.5 & 0.0 & 0.5 \\ 0.0 & 0.5 & 2.0 & 1.0 \end{bmatrix}$

Step 3 — Mask (causal). Every position above the diagonal is set to $-\infty$, so each token can only score the tokens at or before its own position:

$\text{mask}(\text{Scaled}) = \begin{bmatrix} 0.5 & -\infty & -\infty & -\infty \\ 1.0 & 1.0 & -\infty & -\infty \\ 1.0 & 0.5 & 0.0 & -\infty \\ 0.0 & 0.5 & 2.0 & 1.0 \end{bmatrix}$

Step 4 — Softmax (row-wise). Masked entries contribute $e^{-\infty} = 0$. Let's take row 3 explicitly. The exponentials are $e^{1.0} \approx 2.718$, $e^{0.5} \approx 1.649$, $e^{0.0} = 1.000$, and $0$ for the masked position. The sum is $5.367$. Dividing gives $[0.506,\; 0.307,\; 0.186,\; 0]$.

Full attention weight matrix:

$A \approx \begin{bmatrix} 1.000 & 0 & 0 & 0 \\ 0.500 & 0.500 & 0 & 0 \\ 0.506 & 0.307 & 0.186 & 0 \\ 0.078 & 0.129 & 0.579 & 0.213 \end{bmatrix}$

Step 5 — Output ($A \times V$). Element (2,1): $(0.500 \cdot 2) + (0.500 \cdot 1) = 1.500$.

$\text{Output} \approx \begin{bmatrix} 2.000 & 0.000 & 1.000 & 1.000 \\ 1.500 & 1.500 & 0.500 & 0.500 \\ 1.320 & 1.108 & 1.252 & 0.879 \\ 0.499 & 1.393 & 2.395 & 2.302 \end{bmatrix}$

You can see the effect of the mask in the output. Row 1 can see nothing but itself, so its output is exactly $V_1 = [2,\; 0,\; 1,\; 1]$. Row 4 sees the whole sequence, and it places 57.9% of its weight on position 3, the highest score in its row, so its output is dominated by $V_3 = [0,\; 1,\; 4,\; 2]$.

Empirical Results

The original paper evaluated the model on WMT 2014, the 2014 edition of the Workshop on Machine Translation's shared task. It is a public benchmark of parallel sentence pairs, here English-German and English-French, and competing translation systems train and test on the same data, so the scores are directly comparable across papers. The quality is measured in BLEU, the same 0 to 100 score we met in Part 1, where higher means closer to the human reference translations [1]:

The quality numbers are good, but the training-cost column matters even more, because that is what changed the field. Recurrence forces sequential processing, token 4 cannot be computed before token 3. Attention has no such dependency. It is a matrix multiply, and matrix multiplies are exactly what GPUs are built to parallelize. The base model trained for just 12 hours, and the big model for 3.5 days, both on a single machine with 8 P100 GPUs. So the Transformer beat every CNN, RNN, and hybrid of its time while training in a fraction of the time. After this, scaled dot-product attention became the basic building block of everything that followed, GPT-1 [2], BERT [3], GPT-2 [4], T5 [5], and every Transformer variant since.

But attention solved the sequence-dependency problems by multiplying large matrices, so it is still expensive. The $n \times n$ score matrix has to be computed in full. At a sequence length of 4,096 it holds 16.7 million entries, and at 128,000 tokens it holds 16.4 billion. The compute scales as $O(n^2 \cdot d_k)$ and the memory as $O(n^2)$. This is known as the quadratic wall, and it has been the central problem of the field since 2017. The rest of this series, Parts 3 through 6, is in one way or another about making attention cheaper. Either you cache less, or you compute fewer of the scores, or you replace the softmax entirely.

And this is the cost of just one attention operation. The Transformer does not run one, it runs several of them side by side. Why is one head not enough, and what does running many of them cost? That is the next part of the story.

---

2. Multi-Head Attention

Let's see why a single attention head is not enough. The output of a head for each token is one weighted average over the whole sequence. Now suppose token 5 needs two different things at once, its subject sitting nearby and a pronoun reference sitting forty tokens back. One set of weights has to compromise between the two. Multi-head attention solves this by running $h$ attention operations in parallel. Each head has its own learned projections, so each head can learn its own idea of relevance. The results are then concatenated and linearly mixed:

$\text{MultiHead}(Q, K, V) = \text{Concat}(\text{head}_1, \ldots, \text{head}_h)\, W^O$

$\text{head}_i = \text{Attention}(QW_i^Q,\; KW_i^K,\; VW_i^V)$

where $W_i^Q \in \mathbb{R}^{d_\text{model} \times d_k}$, $W_i^K \in \mathbb{R}^{d_\text{model} \times d_k}$, $W_i^V \in \mathbb{R}^{d_\text{model} \times d_v}$, $W^O \in \mathbb{R}^{hd_v \times d_\text{model}}$. $W^O$ is the piece that actually lets the heads talk to each other, it projects their concatenated outputs back down to a single $d_\text{model}$-dimensional vector the next layer can use.

At training time this extra machinery costs nothing. Each head operates in a smaller subspace, $d_k = d_v = d_\text{model} / h$. So the heads split the embedding among themselves instead of duplicating it, and the total compute per layer stays identical to a single full-width head.

Inference Cost and the KV Cache

The cost shows up at generation time. Generation is autoregressive, each new token attends to all the previous tokens, so the keys and values of every past token have to be kept around. They live in a memory called the KV cache. Training never sees this cost, but at inference it grows with the number of heads, the number of layers, and the length of the context:

$\text{KV cache size} = 2 \times h \times n \times d_k \times \text{bytes per element}$

Key result (a 70B-scale model at LLaMA 2's dimensions, using full multi-head attention): With $h = 64$ heads, $d_k = 128$, $n = 4{,}096$, in float16:

$2 \times 64 \times 4{,}096 \times 128 \times 2 = \mathbf{134\text{ MB per layer}}$

With 80 layers: 10.7 GB just for the KV cache, while the model weights themselves are 140 GB. Doubling the context window doubles this figure, since the formula is linear in $n$. Real LLaMA 2-70B never pays this cost, it uses Grouped-Query Attention (GQA) with only 8 KV heads, which shrinks the cache 8x. That fix is exactly where Part 3 picks up.

_{Multi-head attention: h parallel scaled dot-product heads, each with its own learned projections, concatenated and linearly mixed by W^O. Redrawn from Vaswani et al., 2017 [1].}

Was the multi-head design worth it? Yes. Multiple small heads beat one large head of the same total dimension, in the paper's own ablation on the dev set [1]:

Every Transformer model since 2017 uses it. MHA become the default baseline, and every mechanism in the rest of this series defines itself by which part of MHA it changes.

But MHA carries two costs. The compute is quadratic in the sequence length, because the $n \times n$ matrix now sits inside every head. And the KV cache grows linearly with the head count, since each head stores its own keys and values for every past token. In practice the second cost is the more serious one. At each decoding step, the GPU has to reload the entire cache from memory just to compute attention for one new query. That is a huge memory read for a tiny amount of arithmetic. So attention at inference time is limited by memory bandwidth, not by compute, and this fact drives most of Part 3.

But the cache problem came later. First the Transformer had to deal with a more basic problem, one it created the moment it dropped the recurrence. A model that processes all the tokens in parallel has no idea what order they came in. Without positional information, "the cat sat on the mat" and "mat the on sat cat the" produce identical outputs.

---

3. Positional Encoding

In Part 1 we saw that recurrence handles position on its own, token 5 is simply whatever gets processed after token 4. The Transformer's parallelism threw that away. So now the position has to be put back in explicitly, and it took the field four years and four mechanisms to settle on how to do so.

3a. Sinusoidal Positional Encoding

The first approach came from Vaswani et al. themselves [1]. Give each position a unique vector built from sines and cosines at different frequencies, and add it to the token embedding before the first attention layer. For position $\text{pos}$ and dimension index $i$ in a model of width $d_\text{model}$:

$PE(\text{pos},\, 2i) = \sin\!\left(\frac{\text{pos}}{10000^{2i/d_\text{model}}}\right)$

$PE(\text{pos},\, 2i{+}1) = \cos\!\left(\frac{\text{pos}}{10000^{2i/d_\text{model}}}\right)$

Each dimension pair oscillates at its own wavelength, and the wavelengths form a geometric progression from $2\pi$ to $10000 \cdot 2\pi$. This works like a clock. The fast hands track the fine position, and the slow hands track the coarse position. There was also a deeper reason behind this choice. The encoding at position $\text{pos} + k$ is a linear function of the encoding at position $\text{pos}$, so in principle the model gets a handle on relative offsets. And the periodic structure would extrapolate past the training lengths.

Let's see this with a small example. The sinusoid's argument factors as $\text{pos} \cdot \omega_i$, where $\omega_i = 1/10000^{2i/d_\text{model}}$ is the frequency of dimension pair $i$. The $\text{pos}$ part is the position pulled out front, and $i$ enters only through the exponent. With $d_\text{model} = 4$ and dimension indices $i = 0, 1$, the frequencies are $\omega_0 = 1/10000^{0/4} = 1.0$ and $\omega_1 = 1/10000^{2/4} = 0.01$:

$PE = \begin{bmatrix} 0.000 & 1.000 & 0.000 & 1.000 \\ 0.841 & 0.540 & 0.010 & 1.000 \\ 0.909 & -0.416 & 0.020 & 1.000 \\ 0.141 & -0.990 & 0.030 & 1.000 \end{bmatrix}$

Each row is a unique encoding for its position. The fast pair (columns 1 and 2) swings a lot between adjacent positions. The slow pair (columns 3 and 4) barely moves here, but this is the pair that would separate position 10 from position 1,000.

The original Transformer used this scheme [1]. GPT-2 [4] and BERT [3] used a related scheme, learned absolute position embeddings. But the extrapolation hope did not work out in practice. For example, a model trained up to length 512 has received no gradient signal for the encoding values at position 1,024. The sinusoidal structure guarantees that those values exist, but it does not guarantee that the model knows what to do with them. So in practice, the quality degrades quickly past the training length.

There is also a deeper problem, absolute position itself may be the wrong thing to encode. Language rarely cares that a token sits at position 7. What matters is that this token comes three positions after that one.

3b. Relative Position Representations: Shaw et al. and the T5 Bias

Shaw et al. [6] came up with a better approach. They moved the position out of the input and into the attention computation itself, so the question changes from "where am I?" to "how far apart are we?". The model learns embeddings indexed by the relative distance, and the attention logit becomes:

$e_{ij} = \frac{x_i W^Q (x_j W^K + a_{ij}^K)^T}{\sqrt{d_k}}$

where $a_{ij}^K \in \mathbb{R}^{d_k}$ is looked up by the clipped relative distance:

$a_{ij}^K = w^K_{\text{clip}(j-i,\, -k,\, k)}$

Clipping puts a limit on how many distances the model has to learn. The model learns $2k+1$ distinct embeddings, and everything farther than $k$ apart looks the same. For $n = 4$ and clip distance $k = 2$, the raw distance matrix has entries $d_{ij} = j - i$, position $j$ minus position $i$:

$\begin{bmatrix} 0 & 1 & 2 & 3 \\ -1 & 0 & 1 & 2 \\ -2 & -1 & 0 & 1 \\ -3 & -2 & -1 & 0 \end{bmatrix}$

Clipping then applies $\text{clip}(d,-k,k) = \max(-k, \min(d,k))$ entrywise. Only the two entries outside $[-2,2]$ actually move: $\text{clip}(3,-2,2)=2$ at the top-right corner, and $\text{clip}(-3,-2,2)=-2$ at the bottom-left corner. Every other entry is already inside $[-2,2]$ and passes through unchanged:

$\xrightarrow{\text{clip}} \begin{bmatrix} 0 & 1 & 2 & 2 \\ -1 & 0 & 1 & 2 \\ -2 & -1 & 0 & 1 \\ -2 & -2 & -1 & 0 \end{bmatrix}$

Position 3 looking back at position 0 is truly $-3$ away, but it gets clipped to $-2$. So "two apart" and "three or more apart" look the same to the model.

T5 [5] then asked the same kind of question Luong asked about Bahdanau's attention, how much of this machinery do we really need? Instead of a $d_k$-dimensional embedding per distance, T5 learns a single scalar bias per bucketed distance, and adds it to the logit:

$e_{ij} = \frac{q_i \cdot k_j}{\sqrt{d_k}} + b(i - j)$

The buckets are linear up close and logarithmic far away. Small distances each get their own bucket, for $d$ below half the bucket count $B$, $\text{bucket}(d) = d$, so the nearby offsets keep full resolution. Past that point the bucket edges grow geometrically until a maximum distance $D_{\max}$, and every distance beyond $D_{\max}$ falls into the last bucket. This spacing makes sense for language. Telling "1 apart" from "2 apart" matters a lot, while telling "50 apart" from "60 apart" mostly does not.

Shaw et al.'s scheme fed into Transformer-XL [7] with some modifications, and the scalar bias shipped in T5 and mT5 [5]. But both of them inherit the clip, past the maximum distance the position information simply stops. Both carry learned parameters, and T5 shares its scalar bias across layers, though each head within a layer still learns its own distinct bias. So the next question was, can relative position come out of the math itself, with no learned embeddings, no clipping, and no buckets?

3c. RoPE — Rotary Position Embeddings

Rotary Position Embedding (RoPE), introduced by Su et al. [18], answers this with geometry. The idea is to encode position by rotating the query and key vectors, in 2D subspaces of the embedding, by an angle proportional to the absolute position. When a rotated query meets a rotated key in a dot product, the absolute angles subtract, and only the difference survives. And that difference is exactly the relative position.

Concretely, split the $d$-dimensional query and key into $d/2$ consecutive pairs, and rotate the $i$-th pair of a token at position $m$ by the angle $m\theta_i$:

$R_{\Theta, m} = \begin{bmatrix} \cos m\theta_1 & -\sin m\theta_1 & & \\ \sin m\theta_1 & \cos m\theta_1 & & \\ & & \cos m\theta_2 & -\sin m\theta_2 \\ & & \sin m\theta_2 & \cos m\theta_2 \\ & & & & \ddots \end{bmatrix}$

with $\theta_i = 10000^{-2(i-1)/d}$, the same base frequencies as the sinusoidal encoding, applied as rotations rather than additions:

$\tilde{q}_m = R_{\Theta, m}\, q_m, \quad \tilde{k}_n = R_{\Theta, n}\, k_n$

Key result: $\tilde{q}_m^T \tilde{k}_n = q_m^T R_{\Theta, m}^T R_{\Theta, n}\, k_n = q_m^T R_{\Theta, n-m}\, k_n$

The property $R_{\Theta, m}^T R_{\Theta, n} = R_{\Theta, n-m}$ means the score depends only on relative position $n - m$. No clipping, no buckets, no learned position parameters.

Worked example

$d = 4$ (2 rotation pairs), position $m = 3$.

Frequencies: $\theta_1 = 10000^{0/4} = 1.0$, $\theta_2 = 10000^{-2/4} = 0.01$.

Rotation angles at position 3: $m\theta_1 = 3.0$ rad, $m\theta_2 = 0.03$ rad.

$\cos(3.0) \approx -0.990$, $\sin(3.0) \approx 0.141$; $\cos(0.03) \approx 1.000$, $\sin(0.03) \approx 0.030$. These four numbers are exactly the rotation matrix entries:

$R_{\Theta, 3} = \begin{bmatrix} -0.990 & -0.141 & 0 & 0 \\ 0.141 & -0.990 & 0 & 0 \\ 0 & 0 & 1.000 & -0.030 \\ 0 & 0 & 0.030 & 1.000 \end{bmatrix}$

Apply to $q_3 = [1,\; 0,\; 1,\; 2]$. Each pair rotates by the standard 2D rotation formula, $x_1' = x_1\cos\theta - x_2\sin\theta$, $x_2' = x_1\sin\theta + x_2\cos\theta$:

• Pair 1: $[1, 0]$ at $\theta = 3.0$ rad $\to [1 \cdot (-0.990) - 0 \cdot 0.141,\;\; 1 \cdot 0.141 + 0 \cdot (-0.990)] = [-0.990,\; 0.141]$
• Pair 2: $[1, 2]$ at $\theta = 0.03$ rad $\to [1 \cdot 1.000 - 2 \cdot 0.030,\;\; 1 \cdot 0.030 + 2 \cdot 1.000] = [0.940,\; 2.030]$

$\tilde{q}_3 = [-0.990,\; 0.141,\; 0.940,\; 2.030]$

The fast pair has swung nearly half a turn while the slow pair has barely moved. So the same vector carries the fine local position and the coarse global position, purely through geometry.

RoPE is the positional scheme of LLaMA 1/2/3 [8, 9, 10], Mistral [11], Qwen [12], Gemma [13], and DeepSeek [14] (via a decoupled RoPE variant, needed to keep RoPE compatible with its later latent-compression scheme), which is essentially every major open-weight model from 2023 onward. It also turned out to be easy to edit after training. The position lives in the rotation frequencies, so Neural Tangent Kernel (NTK)-aware scaling and YaRN (Yet another RoPE extensioN) interpolation [15] can stretch a trained model's positional range just by adjusting the base $\theta$. This change can extend a 4K-trained model to a 128K context without retraining from scratch.

But RoPE has its own limitations. The position enters only through the Q·K dot product, so the value pathway, and therefore the output representation, carries no explicit position signal. And far beyond the training length, the high-frequency pairs have rotated so many times that the dot products become unpredictable. This is exactly why the interpolation tricks exist.

3d. ALiBi — Attention with Linear Biases

While more models were adopting RoPE, ALiBi (Press et al. [19]) took the opposite approach and used no positional embedding at all. Instead, it subtracts a penalty proportional to the distance directly from the attention logits, so the closer tokens win by default:

$\text{Attention}(Q, K, V) = \text{softmax}\!\left(\frac{QK^T}{\sqrt{d_k}} - m \cdot D\right) V$

where $D_{ij} = i - j$ is how far back key $j$ sits from query $i$, and $m$ is a head-specific slope. ALiBi is defined for causal decoders. The causal mask still applies on top of the penalty, so only the lower triangle ($j \le i$) ever reaches softmax. The slopes form a geometric sequence across the $h$ heads:

$m_1 = 2^{-8/h},\quad m_2 = 2^{-16/h},\quad \ldots,\quad m_h = 2^{-8}$

so some heads are sharply local and others nearly global.

Worked example

$n = 4$, $h = 4$ heads, $d_k = 4$. Head 1 has slope $m_1 = 2^{-2} = 0.25$. The causal mask applies, as in Section 1's example, so only the lower triangle matters (masked entries marked $\cdot$).

Distance matrix and penalty:

$D = \begin{bmatrix} 0 & \cdot & \cdot & \cdot \\ 1 & 0 & \cdot & \cdot \\ 2 & 1 & 0 & \cdot \\ 3 & 2 & 1 & 0 \end{bmatrix}, \quad 0.25 \cdot D = \begin{bmatrix} 0 & \cdot & \cdot & \cdot \\ 0.25 & 0 & \cdot & \cdot \\ 0.50 & 0.25 & 0 & \cdot \\ 0.75 & 0.50 & 0.25 & 0 \end{bmatrix}$

Assume scaled logits:

$\text{Logits} = \begin{bmatrix} 1.00 & 0.50 & 0.80 & 0.30 \\ 0.60 & 1.00 & 0.40 & 0.70 \\ 0.90 & 0.50 & 1.00 & 0.60 \\ 0.40 & 0.80 & 0.50 & 1.00 \end{bmatrix}$

After subtracting the penalty and applying the causal mask:

$\text{Biased} = \begin{bmatrix} 1.00 & -\infty & -\infty & -\infty \\ 0.35 & 1.00 & -\infty & -\infty \\ 0.40 & 0.25 & 1.00 & -\infty \\ -0.35 & 0.30 & 0.25 & 1.00 \end{bmatrix}$

In row 4, the last token, the score for the first token drops from 0.40 to $-0.35$, so after softmax the first token contributes very little. In other words, head 1 mostly sees only its nearby tokens. Head 4, with slope $m_4 = 2^{-8} \approx 0.0039$, subtracts almost nothing, so it can still attend across the whole visible context. The slopes are fixed rather than learned, so the model always has some heads that focus locally and some heads that keep the long range in view.

BLOOM [16] and MPT [17] shipped ALiBi, and the length extrapolation worked as designed. But the field still chose RoPE. The reason is what ALiBi's penalty assumes, that relevance always decays with distance. This assumption is usually true, but it fails exactly when a long context matters the most. Think of a token at position 5,000 that needs to attend strongly to a definition back at position 50, the penalty pushes that connection down. RoPE encodes relative position without penalizing distance, so the model stays free to learn whatever position-dependent pattern the data demands. And the field valued that freedom more than the extrapolation.

By 2023 the position question was more or less settled. RoPE became the consensus, with the interpolation tricks handling length. But the bottleneck this part started with was still standing, the $n \times n$ matrix, and the KV cache that multi-head attention piles on top of it. The rest of this series is the field working through these two costs. Sparse and linear attention go after the quadratic matrix, while Multi-Query Attention (MQA), GQA, and latent compression go after the cache.

---

What comes next

In this part, we saw attention mechanisms used inside the transformer, scaled dot-product and multi-head, and four ways of putting position back in, sinusoidal encodings, relative biases, RoPE, and ALiBi. But the costs we found along the way are still standing. The KV cache stores one key and one value per head, per past token, per layer. And the attention matrix is quadratic in the sequence length.

The next part starts with the cache, and with the most direct way possible. If every head storing its own keys and values is what makes the cache big, then let the heads share. Multi-Query Attention (2019) gives all the heads one shared set. Grouped-Query Attention (2023) shares within groups instead, and that compromise became the field's default approach. Underneath both of them sits a single question, how many key/value heads does a model actually need?

Continue to Part 3: KV-Cache Bottleneck →

---

References

[1] Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L., Gomez, A. N., Kaiser, L., & Polosukhin, I. (2017). Attention Is All You Need. NeurIPS 2017. arXiv:1706.03762.

[2] Radford, A., Narasimhan, K., Salimans, T., & Sutskever, I. (2018). Improving Language Understanding by Generative Pre-Training. OpenAI.

[3] Devlin, J., Chang, M.-W., Lee, K., & Toutanova, K. (2018). BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding. NAACL 2019. arXiv:1810.04805.

[4] Radford, A., Wu, J., Child, R., Luan, D., Amodei, D., & Sutskever, I. (2019). Language Models are Unsupervised Multitask Learners. OpenAI.

[5] Raffel, C., Shazeer, N., Roberts, A., Lee, K., Narang, S., Matena, M., Zhou, Y., Li, W., & Liu, P. J. (2020). Exploring the Limits of Transfer Learning with a Unified Text-to-Text Transformer. JMLR 21. arXiv:1910.10683.

[6] Shaw, P., Uszkoreit, J., & Vaswani, A. (2018). Self-Attention with Relative Position Representations. NAACL 2018. arXiv:1803.02155.

[7] Dai, Z., Yang, Z., Yang, Y., Carbonell, J., Le, Q., & Salakhutdinov, R. (2019). Transformer-XL: Attentive Language Models Beyond a Fixed-Length Context. ACL 2019. arXiv:1901.02860.

[8] Touvron, H., Lavril, T., Izacard, G., et al. (2023). LLaMA: Open and Efficient Foundation Language Models. arXiv:2302.13971.

[9] Touvron, H., Martin, L., Stone, K., et al. (2023). Llama 2: Open Foundation and Fine-Tuned Chat Models. arXiv:2307.09288.

[10] Grattafiori, A., et al. (2024). The Llama 3 Herd of Models. arXiv:2407.21783.

[11] Jiang, A. Q., Sablayrolles, A., Mensch, A., et al. (2023). Mistral 7B. arXiv:2310.06825.

[12] Bai, J., Bai, S., Chu, Y., et al. (2023). Qwen Technical Report. arXiv:2309.16609.

[13] Gemma Team (2024). Gemma: Open Models Based on Gemini Research and Technology. arXiv:2403.08295.

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

[15] Peng, B., Quesnelle, J., Fan, H., & Shippole, E. (2023). YaRN: Efficient Context Window Extension of Large Language Models. arXiv:2309.00071.

[16] Scao, T. L., Fan, A., Akiki, C., et al. (2022). BLOOM: A 176B-Parameter Open-Access Multilingual Language Model. arXiv:2211.05100.

[17] MosaicML (2023). Introducing MPT-7B. MosaicML Blog.

[18] Su, J., Lu, Y., Pan, S., Murtadha, A., Wen, B., & Liu, Y. (2021). RoFormer: Enhanced Transformer with Rotary Position Embedding. arXiv:2104.09864.

[19] Press, O., Smith, N. A., & Lewis, M. (2021). Train Short, Test Long: Attention with Linear Biases Enables Input Length Extrapolation. ICLR 2022. arXiv:2108.12409.

[20] Luong, M.-T., Pham, H., & Manning, C. D. (2015). Effective Approaches to Attention-based Neural Machine Translation. EMNLP 2015. arXiv:1508.04025.