The Dog, the Man, and the Clock

In positional embeddings we walked a simple path: attention is parallel, so it is permutation-equivariant, so "the dog bit the man" and "the man bit the dog" look the same until you bolt position onto the embeddings. We tried raw indices, length normalization, then sine and cosine, and landed on sinusoidal encodings.

That article answers why order matters and how one classical fix works.

This one asks a different question.

What are all the positional encodings that could possibly exist?

Not "what have people invented so far." What is allowed at all, if you make a few assumptions so mild you would have made them without noticing.

The answer comes from group theory. The Jane Street writeup by Alok Puranik is the source. What follows is me learning that argument from the ground up: clocks, matrices, one-parameter groups, then back to attention, RoPE, decay, and ALiBi.

Scattered through this piece are boxes marked WAIT. Those are questions I actually got stuck on while reading. If a box asks the thing already forming in your head, good. That is the point.

Three pictures matter more than everything else: the circle, the fade, and the slide. Remember them. They come back wearing different clothes.

The itch

Same six words. Different meaning. Entirely because of order:

The dog bit the man.

The man bit the dog.

Now the uncomfortable fact. The core operation under every LLM cannot tell those two sentences apart by itself.

Inside a transformer each token produces three vectors. A query: what am I looking for? A key: what do I offer? A value: here is my content. To decide how much token $i$ should care about token $j$, the model computes a dot product:

$$\mathrm{score}(i,j) = q_i \cdot k_j$$

Big number, strong match. That is the engine.

Look at the formula and hunt for the indices $i$ and $j$.

They are in the labels. They are nowhere in the arithmetic.

Shuffle the words and every score comes out identical. Each one just lands in a different slot. The model sees a bag of tokens, not a sequence.

So something has to be bolted on to tell the model where each word sits.

Researchers have bolted on several things: sinusoidal embeddings, learned position vectors, ALiBi, RoPE, exponential decay. On the page they look like unrelated inventions.

They are not. They are points on one short list, and under mild assumptions the list is finished.

To see why, we leave transformers for a bit and talk about clocks.

What is a group?

Stare at a clock

Forget transformers. Here is a 12-hour clock. The only thing you are allowed to do is move the hour hand forward by some number of hours.

Call the moves $h_1$ (forward one hour), $h_5$, $h_9$, and so on.

Four things are true about this little world.

One. Do $h_5$, then $h_4$. You are nine hours ahead. So $h_9$. Combining two moves gives another move. You never fall outside the collection. Do $h_7$ then $h_8$: fifteen hours wraps to three. Still $h_3$. Still inside.

Two. There is a move that does nothing: $h_0$.

Three. Every move can be undone. Forward 5? Forward 7 more and you are back, because $5+7=12=0$. So $h_7$ undoes $h_5$.

Four. Do $h_2$, then $h_3$, then $h_4$. Brackets do not matter. Either way you moved nine hours.

That is a group. You just checked all four axioms using a clock.

The definition, after you already met it

A group is a collection of things you can do, plus a way to combine two of them, satisfying:

1. Closure: combining two gives another in the collection
2. Identity: there is a do-nothing element
3. Inverses: every element can be undone
4. Associativity: brackets do not matter

That is the bare minimum description of a system of reversible actions.

WAIT: that seems almost insultingly simple. Why bother naming it?

Weak axioms on purpose. Lots of things satisfy them: clocks, rotations, integers under addition, shuffles of a deck, and (spoiler) positional encodings.

The moment your thing satisfies these four rules, every theorem anyone ever proved about groups applies for free. Someone proved them in the 1800s for unrelated reasons.

Some families of groups are completely classified. Finished list. Nothing else can exist. If your problem lands in one of those families, the problem is not "being studied." It is done.

Hold that. Later our positional encoding problem lands exactly there.

A few more, fast

Integers under addition. Closure, identity (zero), inverses ($-5$ undoes $5$), associativity. A group.

Rotations of a square: $0^\circ, 90^\circ, 180^\circ, 270^\circ$. Same clock, four hours instead of twelve.

Sliding along a line: slide right by any real $x$ (negative means left). Slide $2.5$ then $1.5$ is slide $4$. Identity is slide $0$. Slide $-3$ undoes slide $3$. A group. Unlike the clock, there is an action for every real number. That difference matters later.

Two things that are not groups

Positive whole numbers under addition. No zero, no negatives. Not a group.

All square matrices under multiplication. Some matrices cannot be undone. Not a group. Fixing that is the next section.

Groups made of matrices

A matrix is a verb

For the rest of this article, a matrix is something you do to an arrow.

$$\begin{pmatrix}2 & 0 \\ 0 & 2\end{pmatrix} \;\longrightarrow\; \text{double everything}$$ $$\begin{pmatrix}0 & -1 \\ 1 & 0\end{pmatrix} \;\longrightarrow\; \text{rotate a quarter turn}$$ $$\begin{pmatrix}1 & 0 \\ 0 & 1\end{pmatrix} \;\longrightarrow\; \text{do nothing}$$ $$\begin{pmatrix}1 & 0 \\ 0 & 0\end{pmatrix} \;\longrightarrow\; \text{squash flat onto the x-axis}$$

Three of those can be undone. The squash cannot. Once $(5,3)$ and $(5,99)$ both become $(5,0)$, you cannot tell which was which.

Matrix multiplication means "do one, then the other." $MN$ means apply $N$ first, then $M$.

Throw out the broken ones

Keep only matrices you can undo:

$$GL(n) = \{\, n \times n \text{ matrices that can be undone}\,\}$$

"GL" is General Linear. Invertible. Determinant not zero. Same idea.

$GL(n)$ is a group. This is the room we live in from here on.

Homomorphism: the slide rule idea

One more concept. It turns out to be the answer.

Before calculators, engineers multiplied with a slide rule. The trick is one equation:

$$\log(a \times b) = \log a + \log b$$

Multiplication goes in. Addition comes out. A faithful translator between two worlds.

That is a homomorphism: a translation between two systems that respects how things combine.

$$\varphi(a \circ b) = \varphi(a) \ast \varphi(b)$$

Left side: combine first, then translate. Right side: translate each, then combine. A homomorphism is a map where those two routes always agree.

Flip the slide rule:

$$e^{a+b} = e^a \cdot e^b$$

Memorize the shape. Addition on the left, multiplication on the right. You will see this shape again and again.

WAIT: is a homomorphism just a function?

Most functions are not homomorphisms. Try $\varphi(x)=x^2$ from addition to addition.

$\varphi(a+b)=(a+b)^2 = a^2+2ab+b^2$, but $\varphi(a)+\varphi(b)=a^2+b^2$.

The cross term ruins it. Squaring mangles the structure.

The special thing is carrying the structure across intact.

One-parameter groups

The definition in English

A one-parameter group is: for every real number $t$, a matrix $A(t)$ such that

$$A(0) = I \qquad\text{and}\qquad A(u+v) = A(u)\,A(v)$$

There is an action for every moment in time. Doing nothing takes zero time. Doing the $u$-action then the $v$-action is the same as doing the $(u+v)$-action.

Look at the second equation. Addition on the left, multiplication on the right. It is the shape. A one-parameter group is a homomorphism from time-under-addition into matrices-under-multiplication.

Picture one: the circle

$$A(t) = \begin{pmatrix}\cos t & -\sin t \\ \sin t & \cos t\end{pmatrix}$$

Rotate by $t$. Nothing grows. Nothing shrinks. It just turns.

Check: $A(0)=I$. $A(2)A(3)=$ turn 3 then turn 2 $=$ turn 5 $= A(5)$.

Picture two: the fade

$$A(t) = e^{-\gamma t}$$

A dimmer switch. Brightness dies smoothly and never quite reaches zero.

Check: $A(0)=1$. $A(2)A(3)=e^{-2\gamma}e^{-3\gamma}=e^{-5\gamma}=A(5)$.

Picture three: the slide

$$A(t) = \begin{pmatrix}1 & t \\ 0 & 1\end{pmatrix}$$

Applied to $(x,y)$:

$$\begin{pmatrix}1 & t \\ 0 & 1\end{pmatrix}\begin{pmatrix}x \\ y\end{pmatrix} = \begin{pmatrix}x + ty \\ y\end{pmatrix}$$

The second number sits still. The first absorbs the second at rate $y$ per unit time.

Physical picture: a hockey puck on frictionless ice. $x$ is position, $y$ is velocity. After time $t$, position is $x+ty$, velocity is still $y$.

Check: $A(2)A(3)=A(5)$.

Stop. Say the outrageous thing.

Turn. Fade. Slide.

Those are not three tidy examples I picked. Those three, and combinations of them, are everything a continuous one-parameter group of matrices can do. There is no fourth behaviour.

The theorem

Every continuous one-parameter group of matrices is

$$A(t) = e^{tX}$$

for exactly one matrix $X$, called the generator.

The matrix exponential is the same series as for numbers:

$$e^{tX} = I + tX + \frac{(tX)^2}{2} + \frac{(tX)^3}{6} + \cdots$$

Why it is true, without the full proof: if I tell you only the puck's velocity right now, I can predict its position forever. The spinning hand: tell me angular speed and I know every future angle. The dimmer: tell me the decay rate and I know every future brightness.

The generator $X$ is what happens in one infinitesimal tick. The exponential adds up infinitely many tiny ticks into duration $t$. Formally $X = A'(0)$.

Before: $A(t)$ was an unknown function. Infinite possibilities.

After: $A(t)$ is pinned by one matrix. In 2D, four numbers. An infinite search became a finite one.

WAIT: where did continuous come from?

Without continuity the theorem is false. You can construct solutions to $A(u+v)=A(u)A(v)$ that are monstrous: unbounded on every interval, impossible to plot, requiring the axiom of choice just to prove they exist. The original article calls them "far too deranged to implement on a physical computer."

Continuous means: it has to be something a GPU can actually compute. It costs you nothing you wanted.

Back to transformers

Toolbox built. Actual problem.

Relative, not absolute

We need the score to depend on position. Absolute position means learning what "position 1742" means. Relative means learning what "3 tokens back" means.

Relative wins for length generalization. Train up to 2000 tokens, then meet 10000. Under absolute encoding, position 8300 is a place the model never saw. Under relative encoding, "key is 3 behind query" looks the same at 8300 as at 17.

Where the fix has to live

You could correct every pair of tokens. With $n$ tokens that is $n^2$ pairs. Unaffordable, and it fights every optimisation that makes attention fast.

So touch each query and each key once, using only its own position:

$$\tilde q(t) = F(t)\,q(t), \qquad \tilde k(s) = G(s)\,k(s)$$

Then

$$\mathrm{score} = \tilde q(t)\cdot\tilde k(s) = q(t)^\top \underbrace{F(t)^\top G(s)}_{M(t,s)}\, k(s)$$

Everything positional sits in that one matrix $M$. Choosing a positional encoding is choosing $M$.

The hard constraint: the query only knows $t$. The key only knows $s$. Neither knows the other exists. Yet the result must depend on $t-s$.

The bridge

This is the section. If you read one part twice, read this.

Two inputs collapse into one

We demanded that only the gap matters. So $M(5,2)$, $M(9,6)$, $M(100,97)$ must all be the same matrix. Gap 3 each time.

Write one variable:

$$M(t,s) = A(t-s)$$

$A(3)$ means: the matrix you use when the key is 3 steps behind the query.

The lookup table

Three-word sentence. Nine pairs. Each cell says which $A$ to use:

	key $s=0$	key $s=1$	key $s=2$
query $t=0$	$A(0)$	$A(-1)$	$A(-2)$
query $t=1$	$A(1)$	$A(0)$	$A(-1)$
query $t=2$	$A(2)$	$A(1)$	$A(0)$

WAIT: a matrix of matrices? What does that mean?

It is not a matrix. It is a lookup table. Like a mileage chart. Rows are query positions. Columns are key positions. The cell tells you which correction to use. Nothing is being multiplied.

Row $t=2$, column $s=0$: key is 2 behind. Cell says $A(2)$. Carry the key forward 2.

Row $t=0$, column $s=2$: key is 2 in the future. Gap $-2$. Cell says $A(-2)$. (In a causal LM this upper triangle is masked away.)

Row $t=1$, column $s=1$: same token. Gap zero. $A(0)=I$. Raw dot product. That is why $A(0)$ must be identity.

Look at the diagonals. Every cell on a diagonal is the same. Nine cells. Five distinct matrices. Scale up: a thousand-word document has a million cells but only about two thousand distinct $A$ values.

You thought you needed a function of two positions. You only ever needed a function of one gap.

What $A$ is actually doing

Here is the sentence that made everything click for me.

$A(u)$ is a time machine. It carries a key forward through $u$ steps of time.

The key was written three tokens ago. It is stale. Before you compare it to a query that lives now, you have to bring it forward.

Every positional encoding is a different style of that transport. RoPE turns the key around a dial. Decay dims it. ALiBi slides it. Same job.

What a time machine must obey

Three sentences.

Carry forward zero steps: nothing happens.

$$A(0) = I$$

Carry forward 3, then 2 more: you carried 5.

$$A(2)\,A(3) = A(5) \qquad\text{i.e.}\qquad A(u)A(v)=A(u+v)$$

Carry forward 3, then back 3: back to start.

$$A(3)\,A(-3) = I$$

Read those as sentences, not equations. They are what moving through time means.

Identity. Closure. Inverses. Associativity free from matrix multiplication.

Those are the four group axioms. Same four as the clock.

We asked what a time machine must do and accidentally wrote the definition of a group. And there is an $A$ for every real gap. So:

The family $\{A(t)\}$ is a one-parameter group.

WAIT: didn't you just choose assumptions that force a group?

Test: imagine violating each rule. Would you ever want to?

Shifting by zero changes something? That is a bug.

Carry 3 then 2 differs from carry 5? Internally incoherent.

Cannot reverse? You destroyed information for no reason.

Discontinuous and uncomputable? Useless.

Nobody was going to build any of those. The constraints cost nothing you wanted. That is why the conclusion bites instead of feeling rigged.

And therefore

The theorem applies immediately:

$$\boxed{A(t) = e^{tX}}$$

Every positional encoding that is linear, relative-only, and continuous is a matrix exponential. The only freedom left is the choice of one matrix $X$.

Feel the transfer

Before: a function assigning a $d\times d$ matrix to every real number. Infinite-dimensional. Unenumerable. People proposed encodings one at a time and nobody could say whether the good ones had been found.

After: one matrix. And even that is generous, because a fixed change of coordinates does not count as a different encoding (the network's weights absorb it).

We did not invent a new encoding. We recognised which shelf to reach for.

The complete list

What can $X$ be?

Why eigenvalues

A matrix usually mixes coordinates and stretches them. Eigenvectors answer: are there directions where it only stretches?

$$Xv = \lambda v$$

If you have enough eigenvectors to form a basis, rebuild your coordinates from them. The matrix becomes diagonal. One hard $d$-dimensional problem becomes $d$ easy 1D problems.

WAIT: why is this necessary, not just convenient?

1. You cannot compute $e^{tX}$ by hand otherwise. On a diagonal matrix the series separates into $d$ scalar exponentials.
2. It splits the problem apart.
3. Eigenvalues are the only part that is real. Fixed basis changes are free. What survives conjugation? The eigenvalues. Everything else is a coordinate accident. Without this, "classify encodings" is meaningless: you drown in $d^2$ numbers that do not mean anything individually.
4. They tell you the behaviour directly. Next subsection.

Fate and rhythm

Write any eigenvalue as

$$\lambda = -\gamma + i\theta$$

Then

$$e^{\lambda t} = e^{-\gamma t}\big(\cos\theta t + i\sin\theta t\big)$$

part of $\lambda$	controls	what it does
real part $-\gamma$	magnitude	fade ($\gamma>0$), hold ($\gamma=0$), explode ($\gamma<0$)
imaginary part $\theta$	angle	rotation speed

The real part is fate. The imaginary part is rhythm.

Real eigenvalue: the fade

The coordinate is multiplied by $e^{-\gamma t}$. Figure: the fade, returning.

With $\gamma=0.1$:

gap	multiplier
0	1.000
5	0.607
10	0.368
25	0.082
50	0.0067
100	0.000045

Hard-wired recency bias. Half-life $\ln 2/\gamma \approx 7$ tokens. Choosing $\gamma$ chooses memory span.

• $\gamma>0$: sensible fading. Linear attention, SSMs.
• $\gamma=0$: encoding does nothing. NoPE. Causal models can still infer position from the mask.
• $\gamma<0$: explodes. Forbidden.

One catch: $t-s$ goes negative for future keys. With $\gamma>0$ and a key 50 ahead, multiplier $e^{5}\approx 148$. So decay only works in causal models where the future is masked.

WAIT: gated models learn a different decay per token. Are they outside this?

Better reading: they learn how far to advance the clock, not how fast to decay. Generator $X$ stays fixed. The gate controls $\Delta t$. Same object, different knob.

Complex pair: the circle. This is RoPE

Complex eigenvalues cannot be separated into two real 1D lines. They spin into each other. In real terms:

$$A(t) = e^{-\gamma t}\begin{pmatrix}\cos\theta t & -\sin\theta t \\ \sin\theta t & \cos\theta t\end{pmatrix}$$

Rotate by $\theta t$, optionally fade. The circle, returning.

This is RoPE. Llama, Qwen, Mistral. Here it does not arrive as a clever invention. It arrives as a forced consequence.

Why rotation makes relative position work

Dot product:

$$a\cdot b = |a|\,|b|\cos(\text{angle between them})$$

Two lengths and one angle. That is all it sees.

Rotation does not change lengths. It changes angles.

	before	after
query angle	$\alpha$	$\alpha + \theta t$
key angle	$\beta$	$\beta + \theta s$
angle between	$\alpha-\beta$	$(\alpha-\beta)+\theta(t-s)$
lengths	unchanged	unchanged

The angle picked up exactly $\theta(t-s)$. The gap, and nothing but the gap.

WAIT: why pairs of coordinates?

Geometric: you cannot rotate in one dimension. A line has no room to turn. Rotation needs a plane.

Algebraic: complex roots of a real characteristic polynomial come in conjugate pairs. Each pair occupies a 2D subspace.

Geometry and algebra agree. Not a coincidence.

WAIT: why many frequencies?

One hand wraps. Gap $0$ and gap $2\pi/\theta$ look identical.

Analogue clock: second hand alone cannot tell 10:00:05 from 11:00:05. Hour hand alone cannot tell 10:00:05 from 10:00:55. Together they pin time down.

Standard RoPE uses $\theta_i = 10000^{-2i/d}$: from a full turn every ~6 positions to every ~60,000. Fast planes for local structure. Slow planes for document scale. A bank of clocks.

Context extension tricks (PI, NTK-aware, YaRN) are rescaling those $\theta_i$. In this language: adjusting imaginary parts of eigenvalues.

RoPE, written three ways

If you have seen the RoPE paper, you have seen a big block-diagonal matrix and a frequency formula. Same object as above. Three levels.

Notation warning. In $w_k = 1/b^{2k/d}$, that $b$ is the base (usually 10000). In $R(\theta t)$, that $\theta$ is a frequency. Different things. Write:

$$w_k = b^{-2k/d}, \qquad k = 0,1,\ldots,\tfrac{d}{2}-1$$

Angle on plane $k$ at position $t$: $\alpha_k = t\,w_k$.

Level 1: one big matrix

$$R_t = \begin{pmatrix} \cos\alpha_0 & -\sin\alpha_0 & & \\ \sin\alpha_0 & \cos\alpha_0 & & \\ & & \ddots & \\ & & & \ddots \end{pmatrix}$$

Then $\tilde q = R_t q$, $\tilde k = R_s k$. This is $A(t)=e^{tX}$ with eigenvalues $\pm i w_k$: pure rotation, $\gamma=0$.

Level 2: the block form

Blocks do not mix. So $d/2$ independent 2D rotations:

$$\begin{pmatrix}\tilde x_{2k} \\ \tilde x_{2k+1}\end{pmatrix} = \begin{pmatrix}\cos\alpha_k & -\sin\alpha_k \\ \sin\alpha_k & \cos\alpha_k\end{pmatrix} \begin{pmatrix}x_{2k} \\ x_{2k+1}\end{pmatrix}$$

$d/2$ little arrows, each spinning on its own dial.

Level 3: the implementation

Multiply out. No matrix left:

$$\begin{aligned} \tilde x_{2k} &= x_{2k}\cos\alpha_k - x_{2k+1}\sin\alpha_k \\ \tilde x_{2k+1} &= x_{2k}\sin\alpha_k + x_{2k+1}\cos\alpha_k \end{aligned}$$

Four multiplies and two adds. At $d=128$ the Level-1 matrix has 16384 entries and only 256 nonzero. Nobody builds it. Level 3 is the same operation for cheap.

Level 1 is what it is. Level 3 is what you run. Level 2 is where you understand it.

Worked example: $d=16$, $b=10000$

Frequencies $w_k = 10000^{-k/8}$:

plane $k$	$w_k$	full turn every
0	1.0	~6 tokens
1	0.316	~20 tokens
2	0.1	~63 tokens
3	0.032	~199 tokens
4	0.01	~628 tokens
5	0.0032	~2k tokens
6	0.001	~6k tokens
7	0.00032	~20k tokens

Plane 0 is a second hand. Plane 7 barely creeps across a long document.

Let every plane of $q$ and $k$ start as $(1,0)$. At $t=5$, $s=2$ (gap 3), the score is about $5.543$. Same content at $(t,s)=(13,10)$ or $(103,100)$: same score to many decimals. Gap 7: different score.

Closed form for this special case:

$$\mathrm{score}(t,s) = \sum_{k=0}^{7}\cos\big(w_k(t-s)\big)$$

Only $t-s$ appears. Translation invariance in floating point.

Change the base and you change reach vs resolution. Fastest hand $w_0=1$ never moves with $b$. The base stretches how slow the slowest hand gets. $b=10000$ sits in the middle, which is why that number has survived since 2017.

Defective: the slide. This is ALiBi

Sometimes there are not enough eigenvectors. The matrix is defective.

Simplest example:

$$X = \begin{pmatrix}0 & 1 \\ 0 & 0\end{pmatrix}$$

Eigenvalue $0$ twice, but only one eigendirection $(1,0)$. Cannot diagonalize. Nilpotent: $X^2=0$. So the exponential stops after two terms:

$$e^{tX} = I + tX = \begin{pmatrix}1 & t \\ 0 & 1\end{pmatrix}$$

A polynomial in $t$. The slide, returning. Effect grows linearly with the gap.

After the Jane Street article shipped, someone objected that the framework seemed to exclude ALiBi:

$$\mathrm{score} = q\cdot k + m(t-s)$$

It is covered. Give the query two spare slots $(m,0)$ and the key $(0,1)$. Let the slide act:

$$\begin{pmatrix}m & 0\end{pmatrix} \begin{pmatrix}1 & t-s \\ 0 & 1\end{pmatrix} \begin{pmatrix}0 \\ 1\end{pmatrix} = m(t-s)$$

Exactly ALiBi. Produced by a defective generator.

ALiBi is the hockey puck.

The "strange theoretical" family had been shipping in BLOOM, MPT, and others for years.

The list

Eigenvalues of $X$	Behaviour	Picture	Name	Shipped in
real, positive	explodes	—	—	forbidden
real, zero	nothing	—	NoPE	some decoders
real, negative	fades	fade	decay	linear attn, SSMs
complex, $\gamma=0$	spins	circle	RoPE	Llama, Qwen, Mistral
complex, $\gamma>0$	spins + fades	both	damped RoPE	RetNet, Mamba-3
defective	grows linearly	slide	ALiBi	BLOOM, MPT

Six rows. That is the design space.

What this buys you, and what it does not

Puranik's conclusion:

It was a reassuring finding, because it means that we don't need to rack our brains to come up with some perfect positional encoding, as we are probably already using it.

No new method. A reason to stop looking along this axis. Rotation is the only option that is simultaneously non-trivial, information-preserving, and incapable of blowing up. Not "best on the leaderboard." The only one in that corner.

Be honest about the fence

Nonlinear encodings. Linearity was aesthetic, not forced. Largest open region.

General score-level maps like $q\cdot k \mapsto \phi(q\cdot k, t-s)$ for arbitrary $\phi$. ALiBi squeaked in because its bias is additive and linear in the gap.

Absolute encodings. Ruled out on purpose for length generalization. See also the positional embeddings walkthrough for why absolute hacks struggle.

Content-dependent position. Framework assumes $F$ and $G$ depend on time alone.

Live questions are inside the family: how to choose frequencies, whether to include decay, how to rescale for longer context, whether time itself should be learned.

Takeaway

One sentence:

A positional encoding is a time machine that carries stale keys forward to the present. Ask what a time machine must obey, and you have written the axioms of a group. A 200-year-old theorem then hands you the complete list.

Turn. Fade. Slide.

That is all there is.

RoPE in code

Not to ship. To read. The maths above, concrete.

import math
import torch

def build_rope_cache(seq_len, head_dim, base=10000.0):
    assert head_dim % 2 == 0
    i = torch.arange(head_dim // 2, dtype=torch.float32)
    theta = base ** (-2.0 * i / head_dim)
    t = torch.arange(seq_len, dtype=torch.float32)
    angles = torch.outer(t, theta)
    return angles.cos(), angles.sin()

def apply_rope(x, cos, sin):
    B, H, T, D = x.shape
    x = x.view(B, H, T, D // 2, 2)
    x0, x1 = x[..., 0], x[..., 1]
    c = cos[:T].view(1, 1, T, D // 2)
    s = sin[:T].view(1, 1, T, D // 2)
    out0 = x0 * c - x1 * s
    out1 = x0 * s + x1 * c
    return torch.stack((out0, out1), dim=-1).flatten(-2)

def attention_with_rope(q, k, v, cos, sin, causal=True):
    D = q.shape[-1]
    q = apply_rope(q, cos, sin)
    k = apply_rope(k, cos, sin)
    scores = (q @ k.transpose(-2, -1)) / math.sqrt(D)
    if causal:
        T = scores.shape[-1]
        mask = torch.ones(T, T, dtype=torch.bool).triu(1)
        scores = scores.masked_fill(mask, float("-inf"))
    return scores.softmax(dim=-1) @ v

Sanity check: same content, same gap, same score.

torch.manual_seed(0)
D = 8
cos, sin = build_rope_cache(64, D)
q_content = torch.randn(1, 1, 1, D)
k_content = torch.randn(1, 1, 1, D)

def score_at(t, s):
    qt = apply_rope(q_content, cos[t:t+1], sin[t:t+1])
    ks = apply_rope(k_content, cos[s:s+1], sin[s:s+1])
    return (qt * ks).sum().item()

print(score_at(3, 0))    # gap 3
print(score_at(13, 10))  # gap 3 -> same
print(score_at(43, 40))  # gap 3 -> same
print(score_at(5, 0))    # gap 5 -> different

Compute angles in fp32 even in a bf16 model. At long context, low-precision trig is a real bug source.

Glossary

Term	Plain meaning
Group	Reversible actions you can chain
Identity	Do nothing
$GL(n)$	Invertible matrices under multiplication
Homomorphism	Translation that respects combining
One-parameter group	Action for every time with $A(u)A(v)=A(u+v)$
Generator $X$	What happens in one infinitesimal tick
Matrix exponential	Turns one tick into any duration
Eigenvalue	Stretch factor along an unmixed direction
Defective	Not enough eigenvectors to diagonalize
NoPE / RoPE / ALiBi	Nothing / rotation / linear distance penalty

Where to go next

1. Alok Puranik, Using group theory to explore the space of positional encodings for attention, Jane Street Blog
2. Jane Street talk on group theory and positional embeddings (YouTube)
3. Su et al., RoFormer (original RoPE)
4. Press et al., Train Short, Test Long (ALiBi)
5. Zhang et al., GRAPE (same framework, independently)
6. Hall, Lie Groups, Lie Algebras, and Representations, ch. 2, for the theorem done properly
7. My earlier positional embeddings post, for the sinusoidal path before this classification

---

Written while learning from Puranik's Jane Street article and the questions I kept getting stuck on along the way.