I’ve recently been trying to understand the inner workings of LLMs. Some days ago, I wrote a detailed walkthrough of how Llama 3.1 generates the next token, that follows the computation from tokenization all the way through to sampling. Writing it helped me make sense of the mechanical process, but many parts of the pipeline still felt mysterious to me, one of them being the positional encoding through which the transformer blocks get information about where tokens are relative to each other.
In Llama 3.1, rotary positional embeddings (RoPE) takes pairs of components from the query and key vectors inside each transformer block and rotates each pair by an angle that depends on the token’s position. Somehow, this lets the attention block know the relative position between tokens. But why pairs? How do those rotations make relative position show up in the attention score? And what is the mathematical structure underneath it all?
After seeing all kinds of positional encodings in papers, I often find the choices feel ad hoc, almost pulled out of the ass. Why this particular formula? Why sines and cosines? Why rotations? Even when the results look good, I still want to know: what is the first principle behind putting position information into attention? Can we start with the properties we want and arrive at different positional encodings in a principled way?
While trying to understand this and reading some of the surrounding material, I came across Alok Puranik’s Using group theory to explore the space of positional encodings for attention on the Jane Street blog. This article is my attempt to work through some of those ideas and explain them as I understand them, starting with the assumptions and then coming back to RoPE.
Where Does Position Enter the Computation?
Let’s zoom into a single transformer block in Llama. The residual stream carries one vector for each token down through the block. Two branches read from that stream and add their results back into it: the attention block, then an MLP block.
The attention branch starts by normalizing the residual stream, then applying three learned projections to produce query, key, and value vectors for each token position. RoPE rotates the query and key vectors within each head, using each token’s position to set the rotation angle. The rotated queries and keys are then compared through dot products to produce the attention scores.
After scaling, a causal mask prevents each token from attending to future tokens, and softmax is applied to turn the scores into weights that multiply into the value vectors to produce a weighted sum for each query. The head outputs are concatenated and projected back to the residual stream’s width before being added to it. This follows Meta’s reference implementation.
So the place I want to focus on is those two RoPE operations. How do they actually inject information about position into the computation?
Let’s focus on just one attention head. Let be the query at position and the key at position . The diagram places RoPE after these vectors have been calculated but before their dot product. Writing for the rotation at position , the score before masking and softmax is
where is the dimension of a head. Each vector gets its own position-dependent rotation.
What Should Attention Know About Position?
Before choosing a rotation, or any other encoding, it helps to ask what we want that comparison to tell us. Does attention need to know that the query is at position 12 and the key is at position 9? Or would it be enough to know that the key is three tokens behind the query?
For a relative encoding, we want our model to know the second kind of information. If we held the query and key vectors fixed and moved their positions from to , the positional effect on their score should stay the same. The signed offset is still . Holding the vectors fixed matters here: we’re testing the encoding, rather than claiming that moving words around in a sentence leaves the model’s representations unchanged.
The Assumptions Behind the Jane Street Analysis
The Jane Street article by Alok Puranik asks the more general question that I had in my mind: can we make some assumptions about the properties of the kind of transformations queries and keys need to undergo to inject position information? And can we use these assumptions to derive different positional encodings from first principles? Puranik’s idea was to assume four properties:
- Linearity in the vectors. At fixed positions, the transformations are matrices: and . Their positional contribution to the score is therefore .
- Translation invariance. Shifting both positions equally leaves that contribution unchanged. Equivalently, for some matrix-valued function .
- Identity at zero offset. Require : comparing vectors at the same position preserves their original dot product.
- Continuity. Extend positions to real numbers and require to vary continuously with the offset . This is an additional assumption beyond integer token indices.
His next step is to show that these assumptions force the relative-offset matrices to form a mathematical object called a one-parameter group. In particular, they obey
The intuition is that moving through an offset of two tokens and then another three should have the same effect as moving through five in one go. The matrices compose in the same way that offsets add. An offset of zero does nothing, and a negative offset undoes the corresponding positive one: . The “one parameter” is the single real number specifying the offset, and continuity means the transformation changes continuously as that number varies. Rotating a plane at a fixed angular speed is one example: two successive rotations add their angles, and reversing the angle undoes the rotation. Puranik’s result lets him use the established mathematics of these groups to work out which positional encodings satisfy his assumptions.
Digging into Group Theory
Before coming back to RoPE, I want to unpack that group theory bit. What is a group here, and what does it mean for a matrix to generate one?
Take all the rotations of a two-dimensional plane. Each rotation is an element of a group called , and the operation that combines two elements is matrix multiplication. Rotating through an angle and then an angle gives another rotation, through . There’s an identity element, , which does nothing, and every rotation has an inverse: undoes . Matrix multiplication is associative, so it doesn’t matter how we bracket a sequence of these operations. These are the properties that make this collection a group.
Now imagine turning the angle up continuously from zero. At zero we’re at the identity matrix. The generator tells us how the rotation matrix starts changing as we increase the angle:
For a tiny angle , we have . Applied to a vector , the term gives its initial direction of motion as we begin to rotate it. The generator describes a rate of change; it isn’t itself the small rotation. In this particular example, also happens to be the matrix for a quarter-turn, but its role here is as the derivative at zero.
The next question is how to get from that initial change to a rotation through a finite angle. That’s what the exponential map does: . This is a matrix exponential, defined by the same power series as the ordinary exponential, with matrix powers and the identity matrix in place of the number one:
Since , the powers cycle through . Collecting the even and odd powers gives , which is exactly the usual rotation matrix. We can also think of the exponential as the limit of many tiny steps: .
The matrices form a line through zero called the Lie algebra of : the space of possible initial changes at the identity. The exponential map takes us from this space into the group. For a continuous one-parameter matrix group, the corresponding result is that for a fixed generator .
The earlier assumptions don’t require to generate rotations. A general real matrix can generate transformations involving scaling, rotation, or shear. If we additionally require to be skew-symmetric, meaning , then is an orthogonal rotation for every . For a RoPE pair rotating at frequency , the angle is , so its generator with respect to position is . This is a skew-symmetric generator, one particular choice within the more general matrix-group framework.
Digging into RoPE
With that structure in mind, let’s come back to how the rotations in Llama make the positional effect depend on the relative offset.
RoPE achieves this by splitting each query and key vector into pairs of components and rotating each pair in its own two-dimensional plane. At position , pair is rotated through an angle , where is a fixed frequency for that pair. Queries and keys use the same frequencies, but each token’s position sets its rotation angles. These rotations make the positional effect on their dot product depend on the relative offset, as described in the RoFormer paper.
But why does rotating pairs of components work? The query and key are vectors in a much higher-dimensional space. Why split them into two-dimensional pieces instead of applying one big rotation to the whole vector? What does an -dimensional rotation matrix even look like?
An important property of a real skew-symmetric generator is that an orthogonal change of basis puts it into independent blocks of the form , with zeros for any remaining directions. Taking the matrix exponential turns those blocks into two-dimensional rotations and leaves the remaining directions fixed. So a continuous one-parameter family of rotations in dimensions uses the same independent planes at every position, with a fixed frequency in each plane. This decomposition applies to rotation groups, rather than to every matrix group allowed by the earlier assumptions. RoPE (Rotary Positional Embedding) chooses its generator in this block-diagonal form, applying a rotation to each pair of coordinates in the query and key vectors. Each pair rotates at a different frequency, and the angle of rotation depends on the token’s position :
where is the even vector dimension, indexes the coordinate pairs, and is a fixed frequency-base hyperparameter. Writing and , the full rotation matrix is
There are rotation blocks along the diagonal, one for each coordinate pair; all entries outside those blocks are zero. As we move along the vector dimensions, the rotation frequencies decrease geometrically; as we move through token positions, each pair’s rotation angle increases linearly at its own frequency. We can think of this as a multi-frequency positional fingerprint, somewhat analogous to Gray code: each position corresponds to a particular configuration of rotations, with fast-changing and slow-changing components jointly representing position. Unlike Gray code, however, these are continuous, periodic signals rather than discrete bits.
The particularly useful property of these rotations comes from their group structure: composing two rotations adds their angles, so . Consequently, when RoPE rotates a query at position and a key at position , their attention score becomes:
Although RoPE rotates the query and key according to their absolute positions, the positional contribution to their dot product depends only on the relative displacement. This gives the transformer access to relative positional information at multiple frequencies without explicitly computing or storing pairwise positional embeddings. Mathematically, these rotations form a one-parameter group generated by a block-diagonal skew-symmetric matrix , with . RoPE is thus one example of the one-parameter groups in Puranik’s analysis: here, the generator produces rotations in independent two-dimensional planes.
Adding decay through the generator
Puranik’s analysis also considers generators that combine rotation with exponential decay. For one coordinate pair, choose
Because and commute, the relative-offset matrix becomes
Our convention is , key position minus query position. In causal attention, the key is at or before the query, so the lag is nonnegative. The pair’s contribution to the unscaled dot product is then
The rotation encodes the offset, while the exponential gives the contribution a decaying envelope for fixed content vectors. This doesn’t guarantee that the softmax attention weight decreases, since the score also depends on content and competes with other scores.
Unlike RoPE, this construction uses different query and key transformations. Taking and gives and preserves the dot product at zero offset. Damping both vectors identically would instead introduce dependence on the sum of their positions. The decay is directional, too. Reversing the offset produces growth, so causal masking matters here.
The diagram compares three choices of generator. Each row separates the position embeddings applied to individual queries and keys, and , from the relative encoding that appears in their dot product, . The paths show what happens to a fixed vector as we increase the causal lag , rather than the absolute position .
The damped-rotation construction appears as xPos in Sun et al.’s paper. In our notation, its per-pair scale is , so the query transform is and the key transform is . Their product gives . Setting the rotation frequency to zero gives the pure-decay ablation. These formulas express the paper’s construction using our column-vector and offset conventions. The generator description is our mathematical interpretation of those transforms. Sun et al., §3.2 and §4.4.1.
Putting it together
So, putting it together, and trying to answer the questions I had about positional embeddings:
- Where does position enter the computation? RoPE rotates queries and keys before their dot product, making position affect the attention scores that determine how values are combined.
- What should attention know about position? For a relative encoding, the positional effect on the score should depend only on the signed offset between tokens, with the content vectors held fixed.
- Can we derive positional encodings from first principles? Under Puranik’s assumptions of linearity, translation invariance, identity at zero offset, and continuity, the relative transformations form a one-parameter group, .
- What is a group here, and what generates it? The relative transformations compose as offsets add, and a fixed matrix describes their initial rate of change and generates every finite transformation through the matrix exponential.
- Why rotations, sines, and cosines? Choosing a skew-symmetric generator gives rotations, whose two-dimensional blocks have sine and cosine entries because .
- Why rotate pairs instead of the whole vector at once? In a suitable basis, any continuous one-parameter rotation in higher dimensions decomposes into independent two-dimensional rotation planes, with any remaining directions fixed.
- How do absolute rotations reveal relative position? The identity makes the positional contribution to the query–key dot product depend only on their offset.
- Does this explain every choice in RoPE’s formula? The group structure explains the independent rotation planes and angles linear in position, but leaves choices such as the frequency schedule open.
- Can the same framework produce other encodings? Adding a positive scalar term to the rotation generator gives an exponential decay envelope over causal lag when queries and keys use the paired transforms derived above.
Correspondence
Reader Comments & Discussion
No comments yet. Be the first to share your thoughts.
Join the Discussion