Positional encoding
Attention is permutation-equivariant: shuffle the input and the output shuffles identically. Language is not. This chapter covers how position information enters the transformer: sinusoidal encoding (the original), learned embeddings (BERT, GPT-2), Rotary Positional Embedding (the modern default in Llama, Mistral, and most open-source LLMs), and ALiBi (length extrapolation).
The transformer block from Chapter 5 has a curious property: it does not know the order of its inputs. Shuffle the input sequence, “cat the on mat sat the” instead of “the cat sat on the mat”, and the block’s output is shuffled identically. Multi-head attention, the FFN, the residuals, the layer norms: none of them encode position. Each row of the output is determined by which other rows interact with it, not by where they sit in the sequence.
Language does not have that symmetry. “The dog bit the man” and “The man bit the dog” share the same multiset of tokens; they do not share a meaning. Word order encodes grammatical role, scope, dependency. A model that processes the two sequences as permutations of the same bag will produce essentially the same output for both, and that is the wrong behavior for every language task there is.
Position information has to enter the model somewhere. Four approaches dominate. The original transformer (Vaswani et al. 2017) added precomputed sinusoidal waves to the input embeddings. BERT and GPT-2 learned a position-specific embedding for each slot, treating “position 5” like a vocabulary word. Modern open-source LLMs (Llama, Mistral, Qwen, most of the post-2021 field) use Rotary Positional Embedding (RoPE), which rotates the query and key vectors inside each attention layer. A fourth approach, ALiBi, skips positional encoding entirely and adds a distance-based bias to attention scores; this turns out to extrapolate cleanly to sequences much longer than those seen during training.
Why position matters
Self-attention is permutation-equivariant. Stated formally: for any permutation acting on the row indices of an input matrix and any self-attention layer ,
Shuffle the input rows in any order. The output rows come out in the same shuffled order, with the same values they would have had if the rows had not been shuffled at all. Self-attention sees the input as a set of feature vectors, not a sequence.
The property falls directly out of the attention formula. Look at one output row. The output at position is
a weighted sum over all positions where the weights are computed from the dot products of with every . Relabel the positions (call position 3 “position 7” and vice versa), and exactly the same dot products get computed, the same weights, the same weighted sum. The output rows simply get the new labels. Nothing about the operation depends on what integer is attached to a row; it depends only on which feature vector lives there.
The FFN does not break the symmetry either. The FFN is position-wise: the same MLP runs independently on each row. Permuting the rows permutes the outputs identically. Residual connections add the input back, which preserves the permutation. Layer norm operates on each row independently. Every component of the Chapter 5 block treats the input as a set.
This is fine for some tasks. Set transformers and graph neural networks exploit permutation invariance directly. It is not fine for language. “The dog bit the man” has a subject (dog) and an object (man) determined by position, not by token identity. Switching the positions of “dog” and “man” produces a syntactically valid sentence with the opposite meaning. SVO English distinguishes itself from OSV Yoda-speech by word order alone. Any language model has to encode position.
There are two natural places to inject position. One is to modify the input: add a position-dependent vector to each token’s embedding before the first block runs. The token embedding says which word; the added vector says where in the sequence. The block downstream operates on a representation that already carries both pieces of information. Sinusoidal and learned positional encodings work this way.
The other is to modify the operation: change the attention computation itself so the dot product depends on position. The Q and K vectors are rotated or biased in a position-dependent way; the value vectors and the rest of the block stay the same. RoPE and ALiBi work this way. They sit inside each attention layer rather than at the input, and they re-inject position at every block rather than letting one initial injection propagate through.
The next four sections cover the four main ways the field has chosen to tell it: sinusoidal, learned, RoPE, ALiBi. Two add to the input; two modify attention. The chapter’s destination is RoPE, which has effectively won.
Sinusoidal positional encoding
The formula
The original transformer’s positional encoding is a precomputed matrix of sines and cosines. For position and dimension-pair index (with ),
The encoding lives in the same space as the input embeddings, shape for a sequence of length , and is added element-wise to the token embeddings before the first block. Adjacent dimensions are paired: dimension holds the sine and dimension holds the cosine, both at the same frequency .
The encoding has no parameters. The values are determined entirely by , , and the constant . A model using sinusoidal PE has the same parameter count it would have without any positional encoding; the matrix is a fixed function of position, not a learned table.
The frequency schedule
The choice of frequencies is what makes the design work. The frequency produces a geometric sequence: at , and the period is positions; at , and the period is positions. The low- pairs oscillate rapidly with position, distinguishing positions 5 and 6 by a large change in phase. The high- pairs oscillate slowly, distinguishing positions 5 and 6 by almost nothing, but distinguishing positions 5 and 5000 by a small but resolvable change in phase.
A useful way to think about it: each dimension pair encodes position at a particular resolution. The fast pairs give the model precise position information about nearby tokens. The slow pairs give the model coarse position information about far-apart tokens. Together they form a fingerprint with components at many scales, so the model can learn to attend to whichever scale is useful for a given relationship.
The base controls the range covered. With , the slowest frequency has period roughly positions, which is more than enough for a model trained at sequence lengths of or . Modern long-context models sometimes use larger bases ( or even ) when targeting context lengths of , to give the slow pairs the headroom needed to distinguish far-apart positions in the very-long-context regime.
How it’s used
Sinusoidal PE is precomputed once for the maximum sequence length the model will ever see, then sliced and added to the token embeddings at the start of every forward pass. If is the matrix of token embeddings for the current sequence, the input to the first transformer block is simply
with PE sliced down to the current sequence length. The block stack then runs on this combined representation. The position information enters once, at the input, and propagates through the network via the residual stream, every block sees a sum of “what” (the embedding contribution) and “where” (the PE contribution), and is free to use either.
The numpy implementation is a few lines. The matrix is built by broadcasting a column of positions against a row of inverse frequencies, then filling the even columns with sine and the odd columns with cosine.
A footnote in Vaswani’s paper claimed a subtle property of the sinusoidal design, and subsequent work made it rigorous.
The implicit relative-position argument
The original transformer paper made a striking claim about sinusoidal PE in a footnote of Section 3.5: “We chose this function because we hypothesized it would allow the model to easily learn to attend by relative positions, since for any fixed offset , can be represented as a linear function of .” The footnote made the choice sound principled rather than arbitrary. It is worth working through the math, because the property is real but its practical impact has been overstated for years.
Consider one frequency pair in the sinusoidal PE. At position this pair has the value , a point on the unit circle at angle . At position this same pair has the value , the same point on the unit circle, rotated by an additional angle . Standard angle-addition identities give the rotation explicitly:
The matrix on the right depends only on , the relative offset between the two positions. It does not depend on , the absolute position. So for each frequency pair is a rotation of for that pair, and the rotation is purely a function of the offset.
What this implies: a model could in principle learn linear projections of the PE that extract relative-position information. Run and through some learned linear map, and the resulting attention scores can be made to depend only on rather than on and separately. The information is representationally there.
What this does not imply: that the trained model actually learns to use this property. The model has to discover the right linear combinations on its own, and there is no architectural pressure toward them. Empirically, sinusoidal PE often underperforms learned PE on in-distribution evaluation, and the relative-position structure does not seem to be reliably exploited. The footnote’s claim is true in principle and somewhat true in practice, but only somewhat.
Shaw, Uszkoreit & Vaswani (2018, arxiv.org/abs/1803.02155) made the relative-position idea explicit a year after the original transformer: instead of adding absolute PE to the input, modify attention to include learned relative-position offsets at each layer. The Shaw-style scheme was influential conceptually but never dominant in code; modern usage moved past it in favor of RoPE, which captures the relative-position property without learned parameters and without the per-attention-pair bookkeeping Shaw’s formulation required.
The simplest alternative is to learn the position encoding instead of computing it.
Learned positional embeddings
A simpler approach drops the sinusoidal formula entirely and treats position 0, position 1, position 2, all the way up to position , as a vocabulary. Each “position word” gets its own embedding vector, a learnable parameter of the model, initialized small and updated by gradient descent like any other embedding. The input to the first block is, as before, the sum of a token embedding and a position embedding, but the position embedding is now looked up from a table instead of computed.
The approach was used by BERT, GPT-2, and GPT-3 (the original GPT-3 paper used learned positional embeddings, though variants since have moved to other schemes). Mechanically it is identical to token embeddings, a matrix of learnable parameters, with the maximum sequence length the model will ever process. To embed a sequence of length , you slice rows through out of the table.
Two things stand out about the design. First, it adds parameters: per model. For GPT-2 small (, ), that is about parameters, small next to the M total, but nonzero. Second, and more importantly, it does not extrapolate. If you train at , there is no embedding for position ; the embedding table simply does not have that row. The model cannot process longer sequences without either retraining or replacing the position-embedding mechanism.
The structural inductive bias is also weak. Sinusoidal PE has continuity built in: neighboring positions have similar encodings because the underlying functions are continuous. Learned PE has no such guarantee. Two adjacent positions get two adjacent rows in the table, but those rows are initialized independently and updated independently. Whether the trained embeddings end up smooth in position space is a function of what gradients happened to do during training, not of the architecture.
The trade-off is straightforward. Learned PE has the most flexibility (the model can learn whatever positional representation works best for the task) and the least structure. It performs well in distribution and fails completely out of distribution. Sinusoidal PE has more structure (the multi-scale frequency design) and somewhat better extrapolation, at the cost of being less expressive when the model has plenty of training data at every position it will see. The two schemes pick different points on the same trade-off curve.
That regime, context-length extension and clean relative-position structure, is exactly what RoPE was designed for.
Rotary Positional Embedding (RoPE): the modern default
The setup
Sinusoidal and learned PE share an architectural assumption: position is injected at the input, once, before the first transformer block. The injection propagates through the network via the residual stream, every block sees a mix of token content and position content, and is free to use either. The choice is parsimonious: one injection, downstream effects everywhere.
RoPE (Su et al. 2021, arxiv.org/abs/2104.09864) rejects this. Instead of adding position to the input, it modifies the Q and K vectors inside every attention layer. The token embeddings are not touched. The V vectors are not touched. Only Q and K, and only at the moment they enter the attention computation, get rotated by an angle that depends on their position.
The motivation is direct. Sinusoidal PE allows the model to learn relative-position attention (the implicit relative-position argument above) but does not enforce it. The trained model has to discover the relative-position structure through the gradients. RoPE asks: can we build the relative-position property into the attention computation itself, so the model gets it for free regardless of how training goes?
The answer turns out to be yes, and it requires no learned parameters. The rotation angles are determined by position and dimension; nothing is learned beyond what is already in the standard Q and K projection matrices. A RoPE-equipped attention layer has exactly the same parameter count as a vanilla attention layer.
The rotation
The construction works dimension-pair by dimension-pair, just like sinusoidal PE. For a Q or K vector of dimension , treat it as pairs of 2D vectors. Each pair will be rotated independently.
For position and pair index , the rotation angle is
The frequency schedule is identical to sinusoidal PE, same progression, but the role is different. In sinusoidal PE, controlled the period of an additive sine. In RoPE, controls the angular speed of a rotation. At pair the rotation is fast (large angle per position); at the highest pair the rotation is slow (small angle per position).
The 2D rotation matrix at angle is the standard
The RoPE-transformed query and key, for the -th pair, are
where and are the absolute positions of the query and the key. The first line is the rotation. The second line is the property that makes the construction work.
The relative-position property
Read the second line of (6.rope) carefully. The dot product of the rotated query at position and the rotated key at position equals the dot product of the unrotated query and a key rotated by angle . The dependence on and separately has collapsed; only the difference matters.
Why this is true: rotation matrices compose. , because the transpose of a rotation is its inverse. So . Rotating both vectors by the same amount cancels out, what remains is the rotation by the difference of the two amounts.
The consequence: attention scores in a RoPE-equipped layer depend only on the relative position of the query and the key, not on their absolute positions. The same query-key pair will produce the same attention score whether it sits at positions , , or . Translation invariance is built into the operation, exactly the property the sinusoidal scheme had to hope the model would learn.
A useful framing: in attention, the query at position asks a question and the key at position provides a partial answer. RoPE rotates the query by an angle proportional to and the key by an angle proportional to . When the dot product is computed, the absolute rotations cancel and only the difference matters, as if both the question and the answer were rotated by the same amount, leaving the relative orientation between them unchanged. Position ‘s query at offset sees the same content as position ‘s query at offset . The model is forced to think in terms of offsets, not positions.
Implementation
The implementation is short and mostly about getting the dimension pairing right. The standard convention pairs adjacent dimensions: dimensions and form one 2D pair, dimensions and form another, and so on. The rotation is applied independently to each pair.
The same demo serves a second purpose. The numerical check at the end verifies the relative-position property directly: for several pairs with the same offset, the dot product after rotation comes out to the same value.
Four query-key pairs at four different absolute-position pairs, all with offset . The four dot products come out to the same number. That is the relative-position property in numerical form: the model’s attention score for “this query, that key, five positions apart” is a single number, regardless of where in the sequence the pair sits.
RoPE has effectively won. Llama (all generations), Mistral, Qwen, PaLM, GPT-NeoX, and most open-source LLMs released since 2022 use RoPE. Production code wraps the rotation into the attention kernel; the math is the same as the numpy implementation above. Anyone building a transformer from scratch today picks RoPE by default. The remaining variants are mostly there for special cases, and ALiBi, the next section, is the most useful of those special cases.
ALiBi: linear attention biases
ALiBi (Press, Smith & Lewis 2021, arxiv.org/abs/2108.12409) takes a different route. Instead of adding positional information to the inputs or rotating the Q and K vectors, it adds a linear bias to the attention scores, with the bias proportional to the distance between the query and the key. Tokens close to the query keep their unbiased scores; tokens far from the query have their scores reduced by an amount that grows linearly with the distance.
The formula is short. For a query at position and a key at position ,
where is a head-specific slope. The first term is the standard attention score from Chapter 4. The second term, the ALiBi correction, subtracts a positive quantity proportional to the distance between the two positions. The further apart the query and the key, the smaller the (post-bias) attention score, the lower the weight the softmax assigns.
The head-specific slopes are chosen so different heads decay at different rates. For a model with heads, the slopes are for . The first head () gets the largest slope and decays the fastest, its attention is concentrated near the query. The last head () gets the smallest slope and decays the slowest, its attention spans long distances. Different heads end up with different “receptive fields,” set by the slopes rather than learned by the model.
The key property of ALiBi is the one its paper title advertises: “train short, test long.” The bias is a function of distance, computed at inference time, with no learned per-position parameters. At evaluation time, even if you are processing sequences longer than anything seen during training, the bias formula keeps working. Position at offset from the query gets a bias of , exactly the same bias position at offset would have gotten. ALiBi-trained models can be evaluated at context lengths several times longer than their training length with minimal quality loss, a property neither learned PE nor (untreated) sinusoidal PE shares.
The trade-off with RoPE is real but narrow. RoPE is somewhat more expressive in-distribution (it lets the model use the full relative-position information rather than just distance) and matches or beats ALiBi on most standard benchmarks. ALiBi is simpler, has cleaner length-extrapolation behavior out of the box, and works without the position-interpolation gymnastics RoPE requires for very long contexts. Modern open-source LLMs are dominated by RoPE; ALiBi has found a niche in BLOOM, some MosaicML models, and projects where long-context extrapolation without retraining is the explicit goal.
ALiBi’s existence is the proof that “position encoding” need not mean “position embedding.” The model can be made position-aware by modifying its attention scores rather than by modifying its inputs or its Q/K projections. RoPE and ALiBi together demonstrate that the right level to inject position is inside the attention computation, not at the input layer. They disagree about how, rotation versus bias, but they agree that’s where.
Length extrapolation: extending context at inference time
A recurring theme across the chapter has been what each PE scheme does to sequences longer than the training context. The point deserves its own section, because the choice of PE is increasingly the most consequential architectural decision for production LLMs targeting long contexts.
The problem: a model trained at context length usually struggles to process sequences longer than . The exact failure mode depends on which PE scheme the model used:
- Learned PE fails immediately and completely. There is no embedding row for position ; the lookup throws an error. Without retraining, the model cannot process anything past its trained range.
- Sinusoidal PE keeps producing values for any position, the formula has no cutoff, but the model has only seen certain ranges of those values during training. Sequences beyond the trained length produce PE values that look unfamiliar to the network, and empirical quality degrades fairly quickly.
- ALiBi extrapolates cleanly out of the box. The linear-distance bias is meaningful at any length; the model trained at generalizes to with minimal quality loss, as the original paper demonstrated.
- RoPE sits in between. The rotation angles are well-defined for any position, but for positions far outside the training range the rotation patterns are unfamiliar. Quality degrades, but RoPE-trained models can be extended through inference-time tricks that no other scheme supports as cleanly.
The extension techniques for RoPE are what made the 100K+ context LLMs possible. Position Interpolation (Chen et al. 2023, arxiv.org/abs/2306.15595) was the first widely-used method: at inference time, divide all position indices by some factor before applying RoPE. A model trained at and evaluated at uses position indices divided by , so the rotation angles stay in the same range the model saw during training. The model “sees” compressed position information, but in a familiar range, and quality holds up well with a short fine-tuning pass on the longer context. NTK-aware scaling (a community contribution, bloc97, 2023) and YaRN (Peng et al. 2023, arxiv.org/abs/2309.00071) refine the idea further: rather than rescaling position indices like PI, NTK-aware scaling adjusts the base itself so different frequencies get stretched by different amounts (the slow frequencies get compressed more, the fast frequencies less), and YaRN combines both ideas. These frequency-dependent schemes have become the standard for production long-context RoPE models.
The practical implication is direct. If you are designing a model that needs long context, the PE choice matters. RoPE plus position interpolation is the dominant pattern: train at a moderate context length (a few thousand to a few tens of thousands), then extend at inference time through PI or YaRN to reach 32K, 128K, or longer. Llama 3.1 trained at 8K and extended to 128K; many production models follow similar recipes. Sinusoidal PE and learned PE do not support these extensions, which is much of why RoPE has displaced both.
The full story of position interpolation, NTK-aware scaling, and YaRN is more involved than this section can fit. The conceptual point is: extrapolation isn’t a property of “good PE”, it’s a property of specific PE schemes plus specific extension techniques. The pairing matters.
What we have: everything except training
Six chapters in, the architectural picture is complete.
Chapter 1 introduced the neural network primitives, forward pass, gradients, autograd. Chapter 2 covered embeddings: how discrete tokens become continuous vectors. Chapter 3 covered tokenization: how text becomes tokens. Chapter 4 derived scaled dot-product attention from first principles. Chapter 5 generalized attention to multi-head and wrapped it in the transformer block, multi-head attention, FFN, residuals, layer norms, the unit that stacks times. This chapter, Chapter 6, fills in the last piece of the architecture: position. The block as written in Chapter 5 is position-blind; sinusoidal, learned, RoPE, or ALiBi inject the information the block needs to distinguish “dog bites man” from “man bites dog.”
Tokens in. Embedded. Position-encoded. Passed through twelve or twenty-four or ninety-six transformer blocks, each block doing multi-head attention then FFN with residuals and layer norms keeping training stable. The last block’s output projected back into vocabulary space and softmaxed to produce a probability distribution over the next token. Sample, append, repeat. That is the architecture of every modern decoder-only LLM. It is also, modulo specific component choices, RMSNorm versus LayerNorm, SwiGLU versus GELU, RoPE versus sinusoidal, what GPT, Llama, Mistral, and Claude do at the level of the forward pass.
The architecture is defined. The weights start as random noise. What remains is training, the process that turns those random weights into something that can write, reason, and converse. Loss functions, optimizers, data pipelines, distributed training, scaling laws, the engineering of taking 175 billion parameters from random initialization to a model that knows the periodic table and can explain it. The next several chapters are about that process.
Exercises
The exercises build on the chapter. Each is a self-contained problem with a starting template. Hints are collapsed by default; try the problem first.
Exercise 1 (easy): Compute sinusoidal PE and verify the structure
Compute sinusoidal positional encoding for several positions. Verify that (a) PE values are in [-1, 1], (b) adjacent dimensions form sin/cos pairs at the same frequency, (c) the leftmost dimensions oscillate quickly and the rightmost slowly.
Hint
PE(p, 2k) = sin(p / 10000^(2k/d_model)); PE(p, 2k+1) = cos(p / 10000^(2k/d_model)). Verify pair structure by checking that PE[p, 2k]² + PE[p, 2k+1]² ≈ 1 for all p, k (since (sin θ)² + (cos θ)² = 1).
Solution
Build the (max_len, d_model) matrix by broadcasting position indices against inverse frequencies, then filling even columns with sin and odd columns with cos.
import numpy as np
def sinusoidal_pe(max_len, d_model, base=10000.0):
pe = np.zeros((max_len, d_model))
positions = np.arange(max_len)[:, None]
k_indices = np.arange(0, d_model, 2)
inv_freq = 1.0 / (base ** (k_indices / d_model))
angles = positions * inv_freq
pe[:, 0::2] = np.sin(angles)
pe[:, 1::2] = np.cos(angles)
return pe
PE = sinusoidal_pe(max_len=100, d_model=64)
print(f"PE range: [{PE.min():.3f}, {PE.max():.3f}]")
# PE range: [-1.000, 1.000]
pair_sumsq = PE[:, 0] ** 2 + PE[:, 1] ** 2
print(pair_sumsq.min(), pair_sumsq.max())
# 1.0 1.0 (sin^2 + cos^2 = 1 at every position)
print(abs(PE[0, 0] - PE[1, 0])) # 0.841 -- dim 0 moves a lot between p=0,1
print(abs(PE[0, 62] - PE[1, 62])) # 0.000133 -- dim 62 barely movesDim 0 has omega = 1 (period ~6.3), dim 62 has omega = 1/10000^(62/64) (period ~47,000) — confirming fast-left, slow-right.
Exercise 2 (medium): Why learned PE doesn’t extrapolate
Train a tiny learned PE (random init) on positions 0-31 and try to use it at position 64. Demonstrate that there’s no meaningful “position 64 embedding”; the model has no information about positions it wasn’t trained on. Then compare to sinusoidal PE at the same position, sinusoidal gives a well-defined value, even though it wasn’t seen during “training” (a learned PE wasn’t trained on it).
Hint
For learned PE, position 64 has no entry in the 32-position embedding table. Any value you put there is arbitrary (typically: out-of-bounds error, or zero-padded, or random). For sinusoidal PE, the formula gives a value for position 64 even if you never used it during a hypothetical training pass, the formula is deterministic.
Solution
learned_pe is a (32, 32) array; indexing past its first axis is simply out of bounds. sinusoidal_pe_at is a closed-form function of p, so it’s defined for any p.
import numpy as np
d_model = 32
trained_len = 32
rng = np.random.default_rng(42)
learned_pe = rng.normal(0, 0.02, (trained_len, d_model))
def sinusoidal_pe_at(p, d_model, base=10000.0):
pe = np.zeros(d_model)
k_indices = np.arange(0, d_model, 2)
inv_freq = 1.0 / (base ** (k_indices / d_model))
pe[0::2] = np.sin(p * inv_freq)
pe[1::2] = np.cos(p * inv_freq)
return pe
try:
learned_pe[64]
except IndexError as e:
print("learned_pe[64] raises IndexError:", e)
# learned_pe[64] raises IndexError: index 64 is out of bounds for axis 0 with size 32
print(sinusoidal_pe_at(64, d_model)[:5].round(3))
# [ 0.92 0.392 -0.99 -0.138 0.984] -- well-defined, no errorThere is no principled “position 64” for the learned table: it was never fit, so any value returned there (crash, zero-pad, wraparound) is arbitrary. The sinusoidal formula has no such gap.
Exercise 3 (medium): Implement RoPE and verify the relative-position property
Implement RoPE: rotate Q and K vectors by position-dependent angles before computing attention scores. Verify the key property: for fixed offset n - m, the dot product RoPE(q, m) · RoPE(k, n) is the same regardless of m.
Hint
Treat the d_k-dimensional Q vector as d_k/2 pairs of 2D vectors. For position p and pair k, rotate the k-th pair by angle p × omega_k where omega_k = 1 / 10000^(2k/d_k). Apply the same to K. Then compute the dot product. Test at offsets (0, 5), (3, 8), (10, 15), all have relative offset 5, so dot products should be equal.
Solution
Split x into even/odd dims, treat each pair as a 2D vector, and rotate by position * omega_k using the standard 2D rotation formula.
import numpy as np
def rope(x, position_ids, base=10000.0):
"""Apply RoPE to (seq_len, d_k) tensor."""
seq_len, d_k = x.shape
k_indices = np.arange(0, d_k, 2)
inv_freq = 1.0 / (base ** (k_indices / d_k))
angles = position_ids[:, None] * inv_freq[None, :] # (seq_len, d_k/2)
cos, sin = np.cos(angles), np.sin(angles)
x_even, x_odd = x[:, 0::2], x[:, 1::2]
rot_even = x_even * cos - x_odd * sin
rot_odd = x_even * sin + x_odd * cos
out = np.empty_like(x)
out[:, 0::2] = rot_even
out[:, 1::2] = rot_odd
return out
d_k = 8
rng = np.random.default_rng(42)
q = rng.normal(0, 1, (1, d_k))
k = rng.normal(0, 1, (1, d_k))
pairs = [(0, 5), (3, 8), (10, 15), (50, 55)]
for m, n in pairs:
q_rot = rope(q, np.array([m]))
k_rot = rope(k, np.array([n]))
dot = (q_rot @ k_rot.T)[0, 0]
print(f" q@pos{m}, k@pos{n}: dot = {dot:.6f}")
# all four print dot = 0.164293Exercise 4 (hard): ALiBi and length extrapolation
Implement ALiBi attention scores: standard scaled dot-product attention plus a linear-distance bias. Demonstrate that the bias pattern is well-defined at any sequence length, even lengths far beyond training. Compute the attention pattern on a 16-position sequence with the bias; verify the bias is monotonically negative with distance.
Hint
ALiBi adds -m_h * |query_pos - key_pos| to each attention score, where m_h is a head-specific slope (typically 2^(-8h/n_heads)). For seq_len=16 and 4 heads, build the bias matrix as (n_heads, seq_len, seq_len). Add it to your scaled dot-product attention scores. After softmax, attention should be concentrated near the query position with smooth decay.
Solution
Slopes are 2^(-8h/n_heads) for h = 1..n_heads; the bias is -slope * |m - n|, broadcast over heads. Attention is standard scaled dot-product with that bias added before softmax.
import numpy as np
def softmax(x, axis=-1):
x = x - x.max(axis=axis, keepdims=True)
return np.exp(x) / np.exp(x).sum(axis=axis, keepdims=True)
def alibi_bias(seq_len, n_heads):
heads = np.arange(1, n_heads + 1)
slopes = 2.0 ** (-8.0 * heads / n_heads) # (n_heads,)
positions = np.arange(seq_len)
distance = np.abs(positions[:, None] - positions[None, :]) # (seq_len, seq_len)
return -slopes[:, None, None] * distance[None, :, :]
def attention_with_alibi(Q, K, V, n_heads):
seq_len, d_head = Q.shape[1], Q.shape[2]
scores = Q @ K.transpose(0, 2, 1) / np.sqrt(d_head)
scores = scores + alibi_bias(seq_len, n_heads)
weights = softmax(scores, axis=-1)
return weights @ V
seq_len, n_heads = 16, 4
bias = alibi_bias(seq_len, n_heads)
print(bias.shape) # (4, 16, 16)
print(bias[0, 8].round(3))
# [-2. -1.75 -1.5 -1.25 -1. -0.75 -0.5 -0.25 -0. -0.25 -0.5 -0.75 -1. -1.25 -1.5 -1.75]
# V-shape centered at 0 (position 8), monotonically more negative with distance
bias_long = alibi_bias(1024, n_heads)
print(bias_long.shape) # (4, 1024, 1024) -- same formula, no retraining neededHead 1’s slope is 2^-2 = 0.25 (steepest decay); head 4’s slope is 2^-8 = 0.0039 (gentlest) — ALiBi gives each head a different attention “radius” for free.
A modern LLM is roughly seven architectural decisions, each of which falls out of the material so far: the layer count , the model dimension , the number of heads per layer , the positional encoding scheme, the context length , the vocabulary size , and the activation function inside the FFN. Pick those, and you have an architecture. Train it, and you have a model. The chapter completes the first half. Chapter 7 starts the second.