Transformers and Attention Questions

The setup

Single-head attention for 2 tokens with d_k = 2. Each token has a query (what it is looking for), a key (what it offers for matching) and a value (the content it passes on). Given:

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

Row i of each matrix belongs to token i. In a real model these three matrices are not typed in: they come from the token embeddings X (one row of numbers per token) multiplied by three learned weight matrices, Q = X W_Q, K = X W_K, V = X W_V. Here they are simply handed to you so the attention step can be done on its own. The formula is

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

where softmax is applied separately to each row.

Step by step

1. Score matrix Q K^T (2 by 2): entry (i, j) is the dot product of query i with key j.
   Row 0: (1,0) dot (1,0) = 1 and (1,0) dot (0,1) = 0. Row 1: (1,1) dot (1,0) = 1 and (1,1) dot (0,1) = 1.
2. Scaling by the square root of d_k = 2, which is 1.4142: 1 becomes 0.7071.
3. Row-wise softmax: for row 0 the scaled scores are 0.7071 and 0; the weights are e^0.7071 over (e^0.7071 + e^0) = 0.6698 and 0.3302. For row 1 the two scores are equal, so each weight is 0.5.
4. Output = weights times V: each output row is a weighted average of the rows of V. Row 0: 0.6698 times (1,2) plus 0.3302 times (3,4) = (1.6605, 2.6605). Row 1: 0.5 times (1,2) plus 0.5 times (3,4) = (2, 3).

The program does the same computation, adds the unscaled variant and compares with the library implementation.

import numpy as np
import torch
import torch.nn.functional as F

torch.set_num_threads(1)

Q = np.array([[1.0, 0.0], [1.0, 1.0]])
K = np.array([[1.0, 0.0], [0.0, 1.0]])
V = np.array([[1.0, 2.0], [3.0, 4.0]])
dk = Q.shape[1]

scores = Q @ K.T                      # (2 queries, 2 keys): row i, column j = q_i . k_j
scaled = scores / np.sqrt(dk)
e = np.exp(scaled - scaled.max(axis=1, keepdims=True))
weights = e / e.sum(axis=1, keepdims=True)    # softmax across each ROW
out = weights @ V

np.set_printoptions(precision=4, suppress=True)
print("scores Q K^T:\n", scores)
print("scaled by sqrt(2):\n", scaled)
print("weights (row-wise softmax):\n", weights)
eu = np.exp(scores - scores.max(axis=1, keepdims=True))
print("weights without the 1/sqrt(d_k) scaling:\n", eu / eu.sum(axis=1, keepdims=True))
print("row sums:", weights.sum(axis=1))
print("output weights @ V:\n", out)
print("check row 0 by hand-style sum:", weights[0, 0] * V[0] + weights[0, 1] * V[1])
ref = F.scaled_dot_product_attention(*(torch.tensor(a, dtype=torch.float64) for a in (Q, K, V)))
print("matches library:", np.allclose(out, ref.numpy()))

Output:

scores Q K^T:
 [[1. 0.]
 [1. 1.]]
scaled by sqrt(2):
 [[0.7071 0.    ]
 [0.7071 0.7071]]
weights (row-wise softmax):
 [[0.6698 0.3302]
 [0.5    0.5   ]]
weights without the 1/sqrt(d_k) scaling:
 [[0.7311 0.2689]
 [0.5    0.5   ]]
row sums: [1. 1.]
output weights @ V:
 [[1.6605 2.6605]
 [2.     3.    ]]
check row 0 by hand-style sum: [1.6605 2.6605]
matches library: True

What each attention row means

Row i of the weight matrix is how token i divides one unit of attention across all tokens: entry (i, j) is how much token i reads from token j. Each row sums to 1 (the printed row sums), the columns do not have to. Row 0 puts 0.6698 on token 0 because its query (1,0) matches key 0 better than key 1; row 1's query (1,1) is equally close to both keys, so it takes a plain average of the two values. The output row is the resulting blend of the value vectors; a different row of the same matrix gives a different blend, which is why each token gets its own context-aware output.

What the scaling did

Without the 1/sqrt(d_k) division, row 0's weights would be 0.7311 and 0.2689 instead of 0.6698 and 0.3302: scaling softens the distribution. With only d_k = 2 the effect is small; with d_k = 64 or more the unscaled scores become large and the softmax saturates.

Causal variant and pitfalls

• A causal mask is a rule that stops each token from attending to later positions: before the softmax, the scores of those positions are set to minus infinity, so their weights become exactly 0. With a causal mask, token 0 may not see token 1, so row 0 is forced to (1, 0) and its output is V's row 0, (1, 2). Row 1 is unchanged here because it may see both.
• Softmax across rows, not columns.
• Scale by the square root of d_k, not by d_k, and not by the number of tokens.
• Keys and values have the same number of rows (one per token), but the values can have a different width than the keys.

Running the code

Save the program as demo.py and run python demo.py. The output above was produced with Python 3.12, numpy 2.1.3 and torch 2.5.1 (CPU). The program uses no random numbers, so the output is the same on every run.

The short version

Both designs wrap every sublayer (one of the two blocks inside a Transformer layer: attention or feed-forward) in a residual connection (output = input + sublayer(input), so the input has a direct path forward and gradients have a direct path back; the running vector each token carries from layer to layer, which every sublayer reads and adds to, is called the residual stream) and layer normalization (LN: rescale each token's vector to zero mean and unit variance across its features, then apply a learned scale and shift). They differ in where the LN sits.

$\text{post-LN: } x_{l+1} = \mathrm{LN}\big(x_l + F(x_l)\big) \qquad \text{pre-LN: } x_{l+1} = x_l + F\big(\mathrm{LN}(x_l)\big)$

Here x_l is the activation (the vector of numbers) entering layer l and F is the attention or feed-forward sublayer. In post-LN the normalization is on the main path, so every layer's output passes through an LN. In pre-LN the main path from input to output is a pure sum, and LN only touches the branch.

At scale, prefer pre-LN (with one extra LN on the final output): it trains stably with less learning-rate warm-up and less sensitivity to hyperparameters. Post-LN is what the original Transformer used and can reach good results in shallower models when carefully tuned, but it is the fragile one.

Why gradient flow differs

Backpropagation multiplies local Jacobians (matrices of derivatives) layer by layer from the output back, so what matters is the factor each layer contributes to that product.

• Pre-LN: because x_{l+1} = x_l + F(LN(x_l)), the derivative of x_{l+1} with respect to x_l is the identity plus a branch term. The gradient reaching layer l therefore contains an untouched copy of the output gradient, plus corrections. Early layers get a gradient of a similar size to late layers.
• Post-LN: the derivative passes through an LN Jacobian at every layer (how the normalized output changes when its input changes: it divides by the token vector's spread and removes the mean direction), which rescales and projects it. There is no clean identity path, so gradient magnitudes depend on depth and on the initial scales. The analysis of Xiong et al. (2020, "On Layer Normalization in the Transformer Architecture") shows that for post-LN the expected gradients of the parameters near the output are large at initialization, which makes a large learning rate unsafe at the start and explains the need for warm-up (a learning-rate ramp from near zero). They argue pre-LN gradients behave well at initialization so warm-up can be dropped.

A numeric picture of the two paths

Treat each layer as multiplying the backward gradient by one factor. In pre-LN that factor is 1 plus a branch term (the identity path is the 1), so it stays near 1 when the branch is small. In post-LN the LN Jacobian multiplies the whole thing, so there is no guaranteed 1. The program multiplies 24 such factors, and then adds 24 equal-length random branches to a stream of length 1 to show how the pre-LN residual stream grows. These are illustrative numbers under stated assumptions, not a model of a real network.

import numpy as np

L = 24
# 1. Backward pass through L layers multiplies one factor per layer.
for name, factor in (("each layer shrinks the gradient by 10%", 0.9),
                     ("identity path plus a branch of +2%", 1.02),
                     ("identity path plus a branch of -2%", 0.98)):
    print(f"{name}: factor {factor} per layer -> gradient after {L} layers x {factor ** L:.3f}")

# 2. Pre-LN: the stream is a running sum of branch outputs.
rng = np.random.default_rng(0)
d = 64
stream = rng.standard_normal(d)
stream /= np.linalg.norm(stream)                      # input vector of length 1
for layer in range(1, L + 1):
    branch = rng.standard_normal(d)
    branch /= np.linalg.norm(branch)                  # each branch output has length 1
    stream = stream + branch
    if layer in (1, 6, 12, 24):
        print(f"layer {layer:2d}: stream length {np.linalg.norm(stream):.2f}, this branch is {1 / np.linalg.norm(stream):.0%} of it")

Output:

each layer shrinks the gradient by 10%: factor 0.9 per layer -> gradient after 24 layers x 0.080
identity path plus a branch of +2%: factor 1.02 per layer -> gradient after 24 layers x 1.608
identity path plus a branch of -2%: factor 0.98 per layer -> gradient after 24 layers x 0.616
layer  1: stream length 1.29, this branch is 77% of it
layer  6: stream length 2.51, this branch is 40% of it
layer 12: stream length 3.68, this branch is 27% of it
layer 24: stream length 4.55, this branch is 22% of it

Reading it: a factor only 10% below 1 per layer leaves 0.080 of the gradient after 24 layers, while factors of 1 plus or minus 2% give 1.608 and 0.616, so the identity path keeps the product in a sensible range. In the stream, the branch added at layer 24 is about a fifth of the stream's length, while the branch at layer 1 was more than half of it; that is the sense in which the top layers of a pre-LN model can end up changing the stream less.

A measurement at initialization

The program builds a stack of each variant (model width 64, 4 heads, feed-forward width 256), pushes random inputs through, puts a linear head and a cross-entropy loss on random labels at the end, and records the gradient norm of each layer's first feed-forward weight. It reports the ratio of the last layer's gradient norm to the first layer's, for stacks of 12, 24 and 48 layers, over 20 seeds each. Pre-LN gets a final LN, as is standard. The loss has to depend on the output: a loss such as the mean of the squared outputs of a final LN is the same number for every input (a normalized vector always has mean square close to 1), so its gradients would be rounding noise and the ratio would mean nothing.

import torch
import torch.nn as nn

torch.set_num_threads(1)

D, H, FF = 64, 4, 256

class Block(nn.Module):
    def __init__(self, pre):
        super().__init__()
        self.pre = pre
        self.attn = nn.MultiheadAttention(D, H, batch_first=True)
        self.ff = nn.Sequential(nn.Linear(D, FF), nn.ReLU(), nn.Linear(FF, D))
        self.n1, self.n2 = nn.LayerNorm(D), nn.LayerNorm(D)
    def forward(self, x):
        if self.pre:   # pre-LN: normalize the input of each sublayer, residual path stays clean
            h = self.n1(x); x = x + self.attn(h, h, h, need_weights=False)[0]
            return x + self.ff(self.n2(x))
        # post-LN: normalize after adding the residual
        x = self.n1(x + self.attn(x, x, x, need_weights=False)[0])
        return self.n2(x + self.ff(x))

def first_last_ratio(pre, depth, seed):
    torch.manual_seed(seed)
    blocks = nn.ModuleList(Block(pre) for _ in range(depth))
    final = nn.LayerNorm(D) if pre else nn.Identity()
    head = nn.Linear(D, 10)
    x = torch.randn(8, 16, D)
    for b in blocks:
        x = b(x)
    labels = torch.randint(0, 10, (8 * 16,), generator=torch.Generator().manual_seed(1000 + seed))
    loss = nn.functional.cross_entropy(head(final(x)).reshape(-1, 10), labels)
    loss.backward()
    g = [b.ff[0].weight.grad.norm().item() for b in blocks]
    return g[-1] / g[0]

for depth in (12, 24, 48):
    for name, pre in [("post-LN", False), ("pre-LN", True)]:
        r = torch.tensor([first_last_ratio(pre, depth, s) for s in range(20)])
        print(f"{depth:2d} layers {name}: last/first gradient norm, median {r.median():.2f}, "
              f"range {r.min():.2f} to {r.max():.2f} (20 seeds)")

Output:

12 layers post-LN: last/first gradient norm, median 1.15, range 0.97 to 1.47 (20 seeds)
12 layers pre-LN: last/first gradient norm, median 0.76, range 0.68 to 0.91 (20 seeds)
24 layers post-LN: last/first gradient norm, median 1.71, range 1.05 to 2.69 (20 seeds)
24 layers pre-LN: last/first gradient norm, median 0.66, range 0.52 to 0.81 (20 seeds)
48 layers post-LN: last/first gradient norm, median 2.52, range 0.86 to 6.49 (20 seeds)
48 layers pre-LN: last/first gradient norm, median 0.57, range 0.46 to 0.79 (20 seeds)

Reading it: post-LN puts more gradient on the layers nearest the output, and the imbalance grows with depth (median ratio 1.15, 1.71 and 2.52 at 12, 24 and 48 layers, with wide seed-to-seed spread), while pre-LN stays below 1 at every depth (0.57 to 0.76), so its layers get gradients of similar size. That is the direction of the analysis above, and the effect at initialization here is modest: a factor of about two at 24 layers, not an order of magnitude. Why that matters: one learning rate has to serve every layer, so the output-side layers of a post-LN stack take proportionally larger first steps than the input-side ones, and the gap widens as the stack deepens. Warm-up (small rates first) gives the weights time to settle before the full rate is applied. The ratio gives a direction, not a threshold: how large a ratio actually breaks a run depends on the optimizer, the rate and the data, and the way to find out is to log per-layer gradient norms in your own first few hundred steps. This is a measurement at random initialization on a small stack with random labels. It supports the mechanism; it does not reproduce a diverging run and does not tell you the final quality of either design after training.

Trade-offs and when to choose which

Pre-LN's weakness: because the residual stream keeps accumulating, deeper layers' contributions are small relative to it, and the top layers can end up doing less work than in post-LN. Choose post-LN only when the model is shallow, you have a tuned recipe, and you can afford restarts; choose pre-LN when a failed run costs days. What would change the call: a recipe from a published model you are reproducing, or evidence on your own evaluation set that post-LN ends higher after tuning.

Common mistakes

• Saying residuals alone fix the gradient problem: in post-LN the LN sits on the residual path and undoes part of the benefit.
• Forgetting the final LN in pre-LN.
• Treating one measured gradient ratio as proof about final accuracy.
• Thinking RMSNorm changes the placement question. RMSNorm is a lighter normalization that divides each token vector by its root mean square and skips the mean subtraction; it can replace LN in either position, and pre versus post is decided the same way.
• Mixing LN (per token, over features) with batch normalization (over the batch), which is a poor fit for variable-length sequences.

Running the code

Save each of the two programs to its own file and run it with python. The outputs above was produced with Python 3.12, numpy 2.1.3 and torch 2.5.1 (CPU). Seeds are fixed in the programs.

The idea

Self-attention treats its input as an unordered set: shuffle the tokens and, without extra information, each token's output is the same. A positional encoding is a vector for each position that is added to the token embedding so that the model can tell "first" from "fifth". The original Transformer paper used a fixed (not learned) recipe: each position gets a vector of sines and cosines at a ladder of different frequencies. The result has shape (max_len, d_model): one row per position, one column per embedding dimension.

The formula and the implementation

For position pos (0, 1, 2, ...) and pair index i from 0 to d_model/2 - 1:

Even columns hold sines and odd columns hold cosines; the two columns of a pair share one frequency. The function below builds the whole table with broadcasting (NumPy stretches an axis of size 1 to match the other array's size, so a column of positions and a row of pair indices combine into a full grid without a loop): pos has shape (max_len, 1), i has shape (1, d_model/2), their combination has shape (max_len, d_model/2), and the sine and cosine fill alternating columns. It assumes d_model is even. The program checks it against a literal double loop written from the formula.

import numpy as np

def sinusoidal_pe(max_len, d_model):
    assert d_model % 2 == 0
    pos = np.arange(max_len)[:, None]                    # (max_len, 1)
    i = np.arange(d_model // 2)[None, :]                 # (1, d_model/2)
    angle = pos / (10000.0 ** (2 * i / d_model))         # (max_len, d_model/2)
    pe = np.zeros((max_len, d_model))
    pe[:, 0::2] = np.sin(angle)
    pe[:, 1::2] = np.cos(angle)
    return pe

pe = sinusoidal_pe(50, 16)
print("shape:", pe.shape)

# reference: explicit double loop written straight from the paper's formula
ref = np.zeros((50, 16))
for p in range(50):
    for i in range(8):
        ref[p, 2 * i] = np.sin(p / 10000 ** (2 * i / 16))
        ref[p, 2 * i + 1] = np.cos(p / 10000 ** (2 * i / 16))
print("matches loop reference:", np.allclose(pe, ref))

print("position 0 row:", np.round(pe[0], 2))
print("position 1 row:", np.round(pe[1], 4)[:6], "...")
d = 16
wavelen = 2 * np.pi * 10000.0 ** (2 * np.arange(d // 2) / d)
print("wavelength per sin/cos pair:", [round(float(w), 1) for w in wavelen])
print("every row has norm sqrt(d/2):", np.allclose(np.linalg.norm(pe, axis=1), np.sqrt(d / 2)))

# shifting by k positions is a rotation of each (sin, cos) pair, the same for every p
k = 5
freq = 1.0 / 10000.0 ** (2 * np.arange(d // 2) / d)
s, c = pe[:-k, 0::2], pe[:-k, 1::2]
shift_sin = s * np.cos(k * freq) + c * np.sin(k * freq)
shift_cos = c * np.cos(k * freq) - s * np.sin(k * freq)
print("PE(p+5) is a fixed rotation of PE(p):", np.allclose(shift_sin, pe[k:, 0::2]) and np.allclose(shift_cos, pe[k:, 1::2]))
# worked case: pair 0 has rate 1, so (sin p, cos p). Rotating position 1's pair by 5 radians gives position 6's pair
sin1, cos1 = np.sin(1.0), np.cos(1.0)
print("pair 0 of PE(1) rotated by 5 rad:", np.round([sin1 * np.cos(5) + cos1 * np.sin(5), cos1 * np.cos(5) - sin1 * np.sin(5)], 4),
      " pair 0 of PE(6):", np.round(pe[6, :2], 4))
print("distinct rows among 5000 positions at d_model=64:", len({tuple(np.round(r, 6)) for r in sinusoidal_pe(5000, 64)}))

Output:

shape: (50, 16)
matches loop reference: True
position 0 row: [0. 1. 0. 1. 0. 1. 0. 1. 0. 1. 0. 1. 0. 1. 0. 1.]
position 1 row: [0.8415 0.5403 0.311  0.9504 0.0998 0.995 ] ...
wavelength per sin/cos pair: [6.3, 19.9, 62.8, 198.7, 628.3, 1986.9, 6283.2, 19869.2]
every row has norm sqrt(d/2): True
PE(p+5) is a fixed rotation of PE(p): True
pair 0 of PE(1) rotated by 5 rad: [-0.2794  0.9602]  pair 0 of PE(6): [-0.2794  0.9602]
distinct rows among 5000 positions at d_model=64: 5000

The table is added to the token embeddings, x = emb * sqrt(d_model) + pe[:T] in the original paper's setup, where T is the actual sequence length. Time and memory to build it are O(max_len x d_model) (proportional to the number of table entries, one per position and column), and it has no trainable parameters. The original paper states the sqrt(d_model) factor (in its setup the embedding weights are shared with the output projection before the final softmax) and gives no reason for it. A common explanation is about scale: embedding entries are often initialised small (of order 1/sqrt(d_model)), while every table entry is between -1 and 1, so without the factor the position signal could swamp the token identity. That is an argument, not something the paper measures.

Why the frequencies are arranged this way

The angular rate of pair i is 1 / 10000^(2i/d_model): it starts at 1 radian per position for the first pair and falls geometrically, so the wavelengths (the printed list, for d_model = 16) run from about 6.3 positions to about 19,869 positions, each roughly 3.2 times the last (for d_model = 16 the ratio is 10000^(2/16)). The paper describes the wavelengths as a geometric progression from 2 pi to 10000 x 2 pi; the longest is approached as d_model grows, and a small d_model stops short of it.

• Fast pairs separate neighbours. The first pairs change a lot from one position to the next, so nearby positions get distinguishable vectors.
• Slow pairs separate far-apart positions. A pair with wavelength near 20,000 barely changes across a few hundred positions but separates positions thousands apart. Together the pairs work like the digits of a number or the hands of a clock: the fast hand gives fine resolution, the slow hand removes ambiguity at long range, so each position gets a distinct code (5,000 distinct rows at d_model = 64 in the last printed line).
• Geometric spacing covers many scales with a modest number of dimensions. A linear ladder of frequencies would need far more dimensions to span both a few tokens and thousands of tokens.
• Sine and cosine together make offsets easy. The printed check shows that PE(p+5) is obtained from PE(p) by rotating each (sine, cosine) pair by the angle 5 x rate, the same for every p. The reason is the angle-addition identities, with a = p x rate and b = 5 x rate:

The right sides use only the old sine and cosine (at position p) and the fixed numbers cos b and sin b, so the new pair is a fixed linear mix of the old pair, whatever p is. The worked case above shows it for pair 0: rotating the pair of position 1 by 5 radians lands on the pair of position 6. A linear layer can therefore, in principle, learn to compute "the token k places away" with weights that do not depend on the absolute position. The check is on the encodings themselves; whether a trained model uses this is an empirical matter.

• Bounded and parameter-free. Every value lies in [-1, 1], every row has the same norm (the norm is the vector's length; each sine-cosine pair contributes sin^2 + cos^2 = 1 to the squared length and there are d_model/2 pairs, so the length is sqrt(d_model/2)), and the table can be generated for any length, including lengths never seen in training. Whether the model attends well at those lengths is a different matter.

Edge cases and common mistakes

• Odd d_model: the pairing breaks; pad or handle the last column separately.
• Using the wrong exponent (2i/d_model versus i/d_model) silently changes the frequency ladder; compare to a loop reference as above.
• Computing the table in half precision (16-bit floating point numbers, which keep only about three decimal digits): large positions lose resolution, so build it in float32 or float64 and cast.
• Saying sinusoidal encodings are "learned": they are fixed. Another formulation interleaves versus concatenates sines and cosines; any fixed permutation of columns is equivalent as long as it is used consistently.

Running the code

Install: pip install numpy. Save the program as sinusoidal_pe.py and run python sinusoidal_pe.py.

Why position has to be supplied

Self-attention scores every pair of tokens from their vectors alone. Nothing in "query dot key, softmax, weighted sum" knows which token came first. If you shuffle the input tokens, the outputs are the same vectors in the shuffled order: the layer is permutation equivariant. To such a layer "dog bites man" and "man bites dog" are the same bag of words. Language depends on order, so the model must be given position information, either added to the token vectors or injected into the attention scores.

Permutation equivariant means: shuffle the inputs and the outputs are shuffled in exactly the same way, so the layer cannot tell which input came first. The first check in the program further down shows this: without positions, attending to the shuffled input equals shuffling the output (it prints True); once sinusoidal positions are added to the shuffled tokens, the equality fails (it prints False), so order now changes the result.

Sinusoidal encodings (fixed)

A positional encoding is a vector of the same width d as the token embedding, added to it before the first layer. The sinusoidal version uses no learned parameters. For position pos and pair index i:

Each pair of dimensions is a sine and cosine at one frequency; the frequencies range from fast (i = 0) to very slow, so together they work like the hands of a clock turning at different speeds, and each position gets a distinct code. The fastest pair has angle pos (an angle of 1 radian per position), and a full turn is 2 x pi = 6.28 radians, so its wavelength is about 6 positions. In the program's printed table, position 0 is 0, 1, 0, 1, ... (sin 0 and cos 0 in every pair) and position 1 has a fast-changing first pair (0.841, 0.54) and a slowly changing last pair (0.001, 1.0).

The property that motivates the design: PE at position pos + k is a rotation of the (sin, cos) pair at position pos, by an angle that depends only on k and the frequency, not on pos. A rotation by a fixed angle is a linear map (the pair is multiplied by a fixed 2 by 2 matrix), so a learned linear layer can in principle read relative offsets from the encodings. One rotation with actual numbers, from the program: take the fastest pair (angle = position in radians) at position 2, which is (sin 2, cos 2) = (0.909, -0.416). Rotating it by k = 3 radians gives (-0.959, 0.284), and that is exactly the pair at position 5, (sin 5, cos 5). The program also checks this for every pair (it prints True). The dot product of two encodings, PE_i dot PE_j, also depends only on i minus j. The reason is one line: each pair contributes sin(a) sin(b) + cos(a) cos(b) = cos(a - b), and a - b is (i - j) times that pair's speed. The program checks it by taking the matrix of all PE_i dot PE_j values and looking along each diagonal (a diagonal is a set of cells with the same i minus j). The spread is the largest minus the smallest value along a diagonal, so a spread of 0 means every cell on the diagonal is equal; it prints 0.0. After query and key projections (the learned matrices W_q and W_k) are applied, this clean property is no longer guaranteed. With random, untrained W_q and W_k the spread is 3.588 for this draw; it changes with the random matrices (over 30 other draws it ran from about 1 to 9) but was never 0, whereas the unprojected spread is exactly 0 (the unprojected values lie between -4 and 4, since the encoding has squared length d/2 = 4). That number says only that projections do not automatically preserve shift-invariance; it is not evidence about a trained model, which is free to learn projections that keep the offset information. Whether a given trained model does so has to be measured.

Because the table is a formula, it can be evaluated at any position, including ones longer than seen in training. Whether the model uses those unseen positions well is a different matter that has to be tested.

Learned absolute embeddings

Here the positional encodings are an nn.Embedding(max_len, d) table: one trainable vector per position, added to the token embedding. The program counts 16 x 8 = 128 parameters. Pros: the model picks the representation it needs. Cons: a position beyond max_len has no row at all (the run shows the IndexError for position 16), and positions rarely seen in training have poorly trained vectors, so the context length is fixed by design.

Relative positional encodings

The idea: what usually matters is how far apart two tokens are, not their absolute index. A relative scheme therefore makes the attention score between tokens i and j depend on the offset j minus i. Two common routes, given as examples:

1. Add a learned vector for the offset to the key (or value) inside the attention computation.
2. Add a learned scalar bias per offset (often with far offsets grouped into buckets, so that, say, all offsets from 32 to 64 share one learned number) directly to the score before the softmax. The last lines of the output show the idea with an illustrative bias table for offsets -2 to 2 (offsets beyond are clamped to the ends): the resulting matrix has the same value along every diagonal, because the bias depends only on the offset j minus i; clamping only makes all the far-off diagonals share one learned number. Clamped means offsets beyond the table are treated as the nearest end of it: here any offset of 2 or more uses the offset-2 entry and any offset of -2 or less uses the offset -2 entry. Row 0 (offsets 0 to 5) prints [-0.5, 1.0, -1.0, -1.0, -1.0, -1.0] and row 3 (offsets -3 to 2) prints [0.0, 0.0, 0.5, -0.5, 1.0, -1.0]. Row 0's diagonal (offset 0) sits in column 0 and row 3's sits in column 3, so compare row 3 from column 3 on, [-0.5, 1.0, -1.0], with row 0 from column 0 on, [-0.5, 1.0, -1.0]: they are the same, shifted three columns to the right. The values to the left of that, 0.0, 0.0, 0.5, belong to offsets -3, -2 and -1 (offset -3 is clamped to the -2 entry). The pattern follows the offset, not the position.

Rotary embeddings are a third route: they rotate each query and key by an angle proportional to its position, so the score depends only on the offset between the two tokens.

Because only the offset enters, a pattern such as "attend to the previous token" is the same everywhere in the sequence, and sequences longer than training only need offsets that were already seen (or are clamped to the largest seen).

The program below runs four groups of checks in this order: (1) shuffling with and without positions; (2) the sinusoidal table, the rotation (a general check and the one worked pair), and the diagonal spread before and after random projections; (3) the learned table and its fixed number of rows; (4) the relative bias matrix. Its printed lines appear under the code.

import numpy as np
import torch
import torch.nn as nn

torch.manual_seed(0)
torch.set_num_threads(1)
n, d = 6, 8

# 1. Self-attention without positions treats the input as a set.
Wq, Wk, Wv = (torch.randn(d, d) * 0.4 for _ in range(3))
def attend(x):
    s = (x @ Wq) @ (x @ Wk).T / d ** 0.5
    return torch.softmax(s, -1) @ (x @ Wv)

def sinusoidal(n, d, base=10000.0):
    pos = np.arange(n)[:, None]
    i = np.arange(d // 2)[None, :]
    ang = pos / base ** (2 * i / d)                       # (n, d/2)
    pe = np.zeros((n, d))
    pe[:, 0::2], pe[:, 1::2] = np.sin(ang), np.cos(ang)
    return torch.tensor(pe, dtype=torch.float32)

x = torch.randn(n, d)
perm = torch.tensor([3, 0, 5, 1, 4, 2])
print("no positions: attention(x[perm]) equals attention(x)[perm]:",
      torch.allclose(attend(x[perm]), attend(x)[perm], atol=1e-6))
pe = sinusoidal(n, d)
print("with sinusoidal positions added to the shuffled tokens:",
      torch.allclose(attend(x[perm] + pe), attend(x + pe)[perm], atol=1e-6), "(order now matters)")

# 2. Sinusoidal table: PE[pos, 2i] = sin(pos / 10000^(2i/d)), PE[pos, 2i+1] = cos(same angle).
print("sinusoidal PE, d=8, positions 0..2:")
for p in range(3):
    print(" ", [round(float(a), 3) for a in pe[p]])
# A shift by k is a fixed rotation of each (sin, cos) pair, the same for every pos.
k, pos = 3, 2
big = sinusoidal(20, d)
ok = True
for i in range(d // 2):
    w = 1.0 / 10000.0 ** (2 * i / d)
    R = torch.tensor([[np.cos(k * w), np.sin(k * w)], [-np.sin(k * w), np.cos(k * w)]], dtype=torch.float32)
    pair = big[pos, 2 * i:2 * i + 2]                      # (sin, cos) at pos
    ok &= torch.allclose(R @ pair, big[pos + k, 2 * i:2 * i + 2], atol=1e-5)
print(f"PE[pos+{k}] is a fixed rotation of PE[pos] in every pair:", bool(ok))
# One rotation by hand: pair 0 has frequency 1, so position 2 sits at angle 2 radians and a shift of k=3 turns it by 3 radians.
a2, k3 = 2.0, 3.0
v2 = np.array([np.sin(a2), np.cos(a2)])
R3 = np.array([[np.cos(k3), np.sin(k3)], [-np.sin(k3), np.cos(k3)]])
print("pair 0 at position 2:", np.round(v2, 3).tolist(), "-> rotated by 3 radians:", np.round(R3 @ v2, 3).tolist(),
      "= pair 0 at position 5:", np.round([np.sin(5.0), np.cos(5.0)], 3).tolist())
G = pe @ pe.T                                             # PE_i . PE_j
dev = max(float(np.ptp(np.diag(G.numpy(), off))) for off in range(-(n - 1), n))
print("PE_i . PE_j depends only on i-j, max spread along a diagonal:", round(dev, 6))
Gp = (pe @ Wq) @ (pe @ Wk).T                              # same, after random projections
dev2 = max(float(np.ptp(np.diag(Gp.numpy(), off))) for off in range(-(n - 3), n - 2))
print("after random W_q, W_k the position-only score is no longer shift-invariant, spread:", round(dev2, 3))

# 3. Learned absolute positions: a table of parameters with a fixed number of rows.
max_len = 16
learned = nn.Embedding(max_len, d)
print("learned table parameters:", sum(p.numel() for p in learned.parameters()), "= max_len*d =", max_len * d)
try:
    learned(torch.tensor([max_len]))
except IndexError as e:
    print("position", max_len, "-> IndexError (no row exists)")

# 4. Relative positions: a bias per offset j-i added to the scores.
bias = torch.tensor([0.0, 0.5, -0.5, 1.0, -1.0])         # offsets -2..2 as an illustrative learned table
idx = (torch.arange(n)[None, :] - torch.arange(n)[:, None]).clamp(-2, 2) + 2
B = bias[idx]                                             # (n, n), same value along each diagonal
print("relative bias matrix row 0:", B[0].tolist(), "row 3:", B[3].tolist())

Output:

no positions: attention(x[perm]) equals attention(x)[perm]: True
with sinusoidal positions added to the shuffled tokens: False (order now matters)
sinusoidal PE, d=8, positions 0..2:
  [0.0, 1.0, 0.0, 1.0, 0.0, 1.0, 0.0, 1.0]
  [0.841, 0.54, 0.1, 0.995, 0.01, 1.0, 0.001, 1.0]
  [0.909, -0.416, 0.199, 0.98, 0.02, 1.0, 0.002, 1.0]
PE[pos+3] is a fixed rotation of PE[pos] in every pair: True
pair 0 at position 2: [0.909, -0.416] -> rotated by 3 radians: [-0.959, 0.284] = pair 0 at position 5: [-0.959, 0.284]
PE_i . PE_j depends only on i-j, max spread along a diagonal: 0.0
after random W_q, W_k the position-only score is no longer shift-invariant, spread: 3.588
learned table parameters: 128 = max_len*d = 128
position 16 -> IndexError (no row exists)
relative bias matrix row 0: [-0.5, 1.0, -1.0, -1.0, -1.0, -1.0] row 3: [0.0, 0.0, 0.5, -0.5, 1.0, -1.0]

Comparison and trade-offs

A key-value cache stores the keys and values of tokens already processed so that generation does not recompute them. Common mistakes: saying sinusoidal encodings "are learned"; claiming they let a model generalise to any length for free; adding positions to the values or to the output rather than to the inputs of attention; and forgetting that relative schemes change the attention formula, not the embeddings.

Direct answer

Causal multi-head attention with rotary position embeddings (RoPE) does five things in order: project the input into queries, keys and values (Q, K, V) and split them into heads; rotate Q and K, but not V, by an angle that depends on each token's position; compute scaled dot-product scores; mask every key that lies after the query (causal masking, so a token cannot see the future) and take a softmax; then combine the values and merge the heads. Rotation makes the dot product between a query at position m and a key at position n depend on the content and on m - n only, which gives relative-position behaviour without adding anything to the input embeddings. The cost is O(T^2 D) time for scores and values plus O(T D^2) for the projections, where T is the sequence length and D the model width.

What RoPE does

Terms: a head is one of several parallel attention computations, each working on its own slice of width Dh of the model width; softmax turns a row of scores into weights that sum to 1; and the score between a query and a key is their dot product.

Take one head with width Dh (even). Group its dimensions into Dh / 2 pairs. Pair i has a frequency

and for a token at position m the pair (x_1, x_2) is rotated by the angle m times theta_i:

Here base is a constant (10,000 in the original formulation; a design choice), and the early pairs rotate fast while the later pairs rotate slowly. The frequencies come from a clock-hand picture: pair 0 turns fast (1 radian per position), and each later pair turns more slowly, down to a tiny angle per position for the last pair, so the head sees both fine position differences (fast hands) and coarse ones (slow hands). The base 10,000 sets how slow the slowest hand is. A rotation is orthogonal, meaning it turns a vector without changing its length or the angle between two vectors that are turned together. Because of that, the dot product of a rotated query at position m with a rotated key at position n is q^T R(n - m) k, where R(n - m) is the same rotation by the angle difference. Position enters only through the offset n - m. In this implementation the pairing is "split-half": dimension i is paired with dimension i + Dh / 2 (a common alternative, called interleaved, pairs neighbouring dimensions i and i+1; both are valid, but weights must be used with the pairing they were trained with). The rotation is applied inside attention, after the head split and before the score computation, so nothing is added to the token embeddings. V is not rotated because it is the content that gets copied, not part of the matching.

The rotation on small numbers

The trick x * cos + rotate_half(x) * sin is the rotation matrix written without building the matrix. For one pair (a, b) the matrix gives (a cos - b sin, a sin + b cos). With split-half pairing the first half of the vector holds every pair's a and the second half holds every b, so x * cos supplies the a cos and b cos terms, and rotate_half(x) is the vector (-b, a) that supplies the -b sin and a sin terms. The second program below checks this on 4 numbers, where the pairs are (x0, x2) and (x1, x3), and then prints how the frequencies fall with the pair index for a head of width 8. It also counts where attention starts to dominate the cost.

import math
import torch
torch.set_num_threads(1)

def rotate_half(x):
    a, b = x[..., : x.shape[-1] // 2], x[..., x.shape[-1] // 2:]
    return torch.cat([-b, a], dim=-1)

# 1. A 2-D rotation by hand: turn (1, 0) by 90 degrees and (1, 2) by 90 degrees
def rot(v, angle):
    c, s = math.cos(angle), math.sin(angle)
    return [round(c * v[0] - s * v[1], 4), round(s * v[0] + c * v[1], 4)]
print("rotate (1, 0) by 90 degrees:", rot([1, 0], math.pi / 2))
print("rotate (1, 2) by 90 degrees:", rot([1, 2], math.pi / 2))
print("length before and after for (1, 2):", round(math.hypot(1, 2), 4), round(math.hypot(*rot([1, 2], math.pi / 2)), 4))

# 2. Split-half pairing on 4 numbers: pairs are (x0, x2) and (x1, x3)
x = torch.tensor([1., 2., 3., 4.])
theta = torch.tensor([math.pi / 2, math.pi / 4])         # one angle per pair at some fixed position
cos = torch.cat([theta, theta]).cos()
sin = torch.cat([theta, theta]).sin()
print("rotate_half([1, 2, 3, 4]) =", rotate_half(x).tolist(), "(second half negated and moved first)")
fast = x * cos + rotate_half(x) * sin
by_hand = torch.empty(4)
for i in range(2):
    a, b = x[i].item(), x[i + 2].item()
    by_hand[i] = a * math.cos(theta[i]) - b * math.sin(theta[i])
    by_hand[i + 2] = a * math.sin(theta[i]) + b * math.cos(theta[i])
print("x * cos + rotate_half(x) * sin :", [round(v, 4) for v in fast.tolist()])
print("rotation matrix applied to each pair:", [round(v, 4) for v in by_hand.tolist()])
print("same:", torch.allclose(fast, by_hand, atol=1e-6))

# 3. Frequencies for head width 8, base 10000: pair i turns theta_i radians per position
dh = 8
for i in range(dh // 2):
    th = 10000.0 ** (-2 * i / dh)
    print(f"pair {i}: {th:.4f} radians per position, one full turn every {2 * math.pi / th:.1f} positions")

# 4. When does attention dominate the FLOPs? projections 8*B*T*D^2, scores+values 4*B*T^2*D
B, D = 1, 4096
for T in [1024, 4096, 8192, 16384, 65536]:
    proj, att = 8 * B * T * D * D, 4 * B * T * T * D
    print(f"T={T:6d}: projections {proj:.3e}, scores+values {att:.3e}, ratio attention/projection {att / proj:.2f}")

Output:

rotate (1, 0) by 90 degrees: [0.0, 1.0]
rotate (1, 2) by 90 degrees: [-2.0, 1.0]
length before and after for (1, 2): 2.2361 2.2361
rotate_half([1, 2, 3, 4]) = [-3.0, -4.0, 1.0, 2.0] (second half negated and moved first)
x * cos + rotate_half(x) * sin : [-3.0, -1.4142, 1.0, 4.2426]
rotation matrix applied to each pair: [-3.0, -1.4142, 1.0, 4.2426]
same: True
pair 0: 1.0000 radians per position, one full turn every 6.3 positions
pair 1: 0.1000 radians per position, one full turn every 62.8 positions
pair 2: 0.0100 radians per position, one full turn every 628.3 positions
pair 3: 0.0010 radians per position, one full turn every 6283.2 positions
T=  1024: projections 1.374e+11, scores+values 1.718e+10, ratio attention/projection 0.12
T=  4096: projections 5.498e+11, scores+values 2.749e+11, ratio attention/projection 0.50
T=  8192: projections 1.100e+12, scores+values 1.100e+12, ratio attention/projection 1.00
T= 16384: projections 2.199e+12, scores+values 4.398e+12, ratio attention/projection 2.00
T= 65536: projections 8.796e+12, scores+values 7.037e+13, ratio attention/projection 8.00

Reading it: turning (1, 0) by 90 degrees gives (0, 1), and turning (1, 2) gives (-2, 1) with the same length, 2.2361. On the 4-number vector, rotate_half returns (-3, -4, 1, 2): the second half negated and moved to the front. Pair (1, 3) turned by 90 degrees becomes (-3, 1), and pair (2, 4) turned by 45 degrees becomes (-1.4142, 4.2426); the vectorized formula and the rotation matrix agree. The frequency lines show the clock hands: pair 0 completes a turn every 6.3 positions, pair 3 every 6283.2.

Shapes

Batching needs nothing special: the cos and sin tables depend only on position, so they broadcast across the batch and heads. For padded batches a padding mask (which hides the filler tokens added to make sequences equal length) is combined with the causal mask, and positions count from each sequence's real start; left-padding, where filler goes before the real tokens, shifts every real token's index, so positions must be counted from the first real token.

Implementation, checks and output

The program implements it, compares the attention against scaled_dot_product_attention with is_causal=True (using the same rotated Q, K and V), compares the vectorized rotation against a pair-by-pair loop written from the definition, shows that the score depends on the offset only (positions 5 and 2, and 12 and 9, have offset 3 and give the same score; offset 2 differs), and checks causality by changing only the last three tokens. The same causality test with the mask switched off is printed to show that the test can fail.

import math
import torch
import torch.nn as nn
import torch.nn.functional as F

torch.set_num_threads(1)

def rope_cos_sin(T, dh, base=10000.0):
    inv_freq = base ** (-torch.arange(0, dh, 2).float() / dh)        # (dh/2,) one frequency per pair
    ang = torch.arange(T).float()[:, None] * inv_freq[None, :]       # (T, dh/2) angle = position * frequency
    ang = torch.cat([ang, ang], dim=-1)                              # (T, dh): pair i is (dim i, dim i + dh/2)
    return ang.cos(), ang.sin()

def rotate_half(x):
    a, b = x[..., : x.shape[-1] // 2], x[..., x.shape[-1] // 2:]
    return torch.cat([-b, a], dim=-1)

def apply_rope(x, cos, sin):                                         # x: (B, H, T, Dh)
    return x * cos + rotate_half(x) * sin                            # rotate each (a, b) pair by its angle

class CausalRoPEAttention(nn.Module):
    def __init__(self, d_model, n_heads):
        super().__init__()
        assert d_model % n_heads == 0 and (d_model // n_heads) % 2 == 0
        self.h, self.dh = n_heads, d_model // n_heads
        self.qkv = nn.Linear(d_model, 3 * d_model, bias=False)
        self.out = nn.Linear(d_model, d_model, bias=False)
    def forward(self, x, return_qk=False, causal_mask=True):
        B, T, D = x.shape                                            # x: (B, T, D)
        q, k, v = (t.view(B, T, self.h, self.dh).transpose(1, 2) for t in self.qkv(x).chunk(3, -1))  # (B, H, T, Dh)
        cos, sin = rope_cos_sin(T, self.dh)                          # (T, Dh), broadcast over B and H
        q, k = apply_rope(q, cos, sin), apply_rope(k, cos, sin)      # rotate Q and K only, never V
        scores = q @ k.transpose(-2, -1) / math.sqrt(self.dh)        # (B, H, T, T)
        causal = torch.ones(T, T, dtype=torch.bool).tril()           # True where key j <= query i
        if not causal_mask:
            causal = torch.ones(T, T, dtype=torch.bool)              # only used to show what the mask prevents
        w = scores.masked_fill(~causal, float("-inf")).softmax(-1)   # (B, H, T, T)
        y = (w @ v).transpose(1, 2).reshape(B, T, D)                 # (B, T, D)
        y = self.out(y)
        return (y, q, k, v) if return_qk else y

torch.manual_seed(0)
B, T, D, H = 2, 12, 32, 4
m = CausalRoPEAttention(D, H).eval()
x = torch.randn(B, T, D)
with torch.no_grad():
    y, q, k, v = m(x, return_qk=True)
    print("shapes: x", tuple(x.shape), "q/k/v", tuple(q.shape), "output", tuple(y.shape))
    ref = F.scaled_dot_product_attention(q, k, v, is_causal=True).transpose(1, 2).reshape(B, T, D)
    print("matches F.scaled_dot_product_attention(is_causal=True):", torch.allclose(y, m.out(ref), atol=1e-5))

    # Reference rotation written pair by pair from the definition (2-D rotation by pos * freq).
    dh = D // H
    q0 = m.qkv(x).chunk(3, -1)[0].view(B, T, H, dh).transpose(1, 2)
    slow = torch.empty_like(q0)
    for pos in range(T):
        for i in range(dh // 2):
            th = pos * 10000.0 ** (-2 * i / dh)
            a, b = q0[..., pos, i], q0[..., pos, i + dh // 2]
            slow[..., pos, i] = a * math.cos(th) - b * math.sin(th)
            slow[..., pos, i + dh // 2] = a * math.sin(th) + b * math.cos(th)
    print("vectorized RoPE equals the pair-by-pair rotation:", torch.allclose(slow, q, atol=1e-5))

    # Relative-position property: the score for (m, n) depends only on m - n.
    qv, kv = torch.randn(dh), torch.randn(dh)
    def score(pm, pn):
        c, s = rope_cos_sin(max(pm, pn) + 1, dh)
        return (apply_rope(qv, c[pm], s[pm]) * apply_rope(kv, c[pn], s[pn])).sum().item()
    print("score(5,2) =", round(score(5, 2), 5), " score(12,9) =", round(score(12, 9), 5), " score(5,3) =", round(score(5, 3), 5))

    # Causality: change only the last 3 tokens; earlier outputs must not move.
    x2 = x.clone(); x2[:, -3:] = torch.randn(B, 3, D)
    print("with the mask, earlier outputs unchanged when future tokens change:", torch.allclose(m(x2)[:, :-3], y[:, :-3], atol=1e-6))
print("same test with the mask switched off (earlier outputs unchanged?):",
      torch.allclose(m(x2, causal_mask=False)[:, :-3], m(x, causal_mask=False)[:, :-3], atol=1e-6))
print("parameters:", sum(p.numel() for p in m.parameters()), "= 4 * D^2 =", 4 * D * D)

B2, T2, D2 = 8, 4096, 4096
proj = 2 * B2 * T2 * 4 * D2 * D2
attn = 2 * 2 * B2 * T2 * T2 * D2
print(f"B={B2}, T={T2}, D={D2}: projection FLOPs {proj:.3e}, score+value FLOPs {attn:.3e} (causal skipping would halve the second),"
      f" naive score tensor {B2 * 32 * T2 * T2 * 2 / 2**30:.0f} GiB at 32 heads and 2 bytes")

Output:

shapes: x (2, 12, 32) q/k/v (2, 4, 12, 8) output (2, 12, 32)
matches F.scaled_dot_product_attention(is_causal=True): True
vectorized RoPE equals the pair-by-pair rotation: True
score(5,2) = -0.89952  score(12,9) = -0.89952  score(5,3) = -1.85075
with the mask, earlier outputs unchanged when future tokens change: True
same test with the mask switched off (earlier outputs unchanged?): False
parameters: 4096 = 4 * D^2 = 4096
B=8, T=4096, D=4096: projection FLOPs 4.398e+12, score+value FLOPs 2.199e+12 (causal skipping would halve the second), naive score tensor 8 GiB at 32 heads and 2 bytes

Reading it: scores for (5, 2) and (12, 9) are both -0.89952, as the offset-only property requires, while offset 2 (positions 5 and 3) gives -1.85075. The causal test passes with the mask and fails without it. The block has 4096 parameters, exactly 4 D^2 for the three input projections and the output projection (no biases here, and RoPE itself has no learned parameters).

Complexity

• Projections: 4 D^2 weights, about 2 B T (4 D^2) FLOPs (a multiply-add counted as 2 FLOPs).
• Scores and weighted values: each is a (T by Dh) times (Dh by T) style product per head, giving 2 x 2 B T^2 D FLOPs in total across heads, so O(B T^2 D) time. With the program's example of B = 8, T = 4096, D = 4096 this is 4.398e12 for projections and 2.199e12 for scores plus values (causal skipping of the upper triangle would halve the second figure).
• Which term dominates: projections cost 8 B T D^2 FLOPs and scores plus values cost 4 B T^2 D, so they are equal at T = 2D and attention dominates only beyond that. In the program's example T = 4096 is below 2D = 8192, which is why projections (4.398e12) exceed scores plus values (2.199e12); the quadratic term is the one that grows, reaching twice the projection cost at T = 16384 and eight times at T = 65536 (second program above, D = 4096).
• RoPE: elementwise, O(B T D), negligible next to the rest.
• Memory: a naive implementation stores the (B, H, T, T) score tensor: 8 GiB at 32 heads and 2 bytes per entry for the same sizes. A fused kernel (one GPU kernel that computes attention tile by tile without writing the score tensor to memory) avoids storing it, but the time stays quadratic in T. Causal skipping means not computing the upper-triangle tiles that the mask would discard, which halves the score and value work.

Pitfalls

• Rotating V, or rotating before the head split with the wrong width (the frequency must use Dh, not D).
• Position bookkeeping: computing positions from zero at each call when using cached keys during generation gives wrong angles; the new token's position must continue where the cache ends (the cache holds the keys and values of earlier tokens, already rotated, so generation only computes the new token's).
• Mixing the split-half and interleaved pairing between code and loaded weights.
• Applying the causal mask after the softmax, or leaving the diagonal masked (a token must see itself).
• Assuming RoPE extrapolates (keeps working at lengths longer than any seen in training) for free: angles at unseen positions are out of the range the model trained on, so quality is measured, not assumed.