You could have designed state of the art positional encoding

Gall’s Law
A complex system that works is invariably found to have evolved from a simple
system that worked
John Gall

This post walks you through the step-by-step discovery of state-of-the-art positional encoding in transformer models. We will achieve
this by iteratively improving our approach to encoding position, arriving at Rotary Postional Encoding (RoPE) used in the latest LLama 3.2 release and most modern transformers. This post intends to limit the mathematical knowledge required to follow along, but some basic linear algebra, trigonometry and understanding of self attention is expected.

Problem Statement

You shall know a word by the company it keeps
John Rupert Firth

As with all problems, it is best to first start with understanding exactly what we are trying to achieve. The self attention mechanism in transformers is utilized to understand relationships
between tokens in a sequence. Self attention is a set operation, which
means it is permutation equivariant. If we do not
enrich self attention with positional information, many important relationships are
incapable of being determined.

This is best demonstrated by example.

Motivating Example

Consider this sentence with the same word in different positions:

$$\text{The dog chased another dog}\backslash text\{The\; dog\; chased\; another\; dog\}$$

Intuitively, “dog” refers to two different entities. Let’s see what happens if we first tokenize them, map to the real token embeddings of Llama 3.2 1B and pass them through torch.nn.MultiheadAttention.

import torch
import torch.nn as nn
from transformers import AutoTokenizer, AutoModel

model_id = "meta-llama/Llama-3.2-1B"
tok = AutoTokenizer.from_pretrained(model_id)
model = AutoModel.from_pretrained(model_id)

text = "The dog chased another dog"
tokens = tok(text, return_tensors="pt")["input_ids"]
embeddings = model.embed_tokens(tokens)
hdim = embeddings.shape[-1]

W_q = nn.Linear(hdim, hdim, bias=False)
W_k = nn.Linear(hdim, hdim, bias=False)
W_v = nn.Linear(hdim, hdim, bias=False)
mha = nn.MultiheadAttention(embed_dim=hdim, num_heads=4, batch_first=True)

with torch.no_grad():
    for param in mha.parameters():
        nn.init.normal_(param, std=0.1) 

output, _ = mha(W_q(embeddings), W_k(embeddings), W_v(embeddings))

dog1_out = output[0, 2]
dog2_out = output[0, 5]
print(f"Dog output identical?: {torch.allclose(dog1_out, dog2_out, atol=1e-6)}") 

As we can see, without any positional information, the output of a (multi
headed) self attention operation is identical for the same token in
different positions, despite the tokens clearly representing distinct entities. Let’s begin designing a method of enhancing self attention with positional information, such that it can determine relationships between words encoded by
their positions.

How should an ideal positional encoding scheme behave?

Desirable Properties

Let’s try and define some desirable properties that will make the optimization
process as easy as possible.

Property 1 – Unique encoding for each position (across sequences)

Each position needs a unique encoding that remains consistent regardless of sequence length – a token at position 5 should have the same encoding whether the current sequence is of length 10 or 10,000.

Property 2 – Linear relation between two encoded positions

The relationship between positions should be mathematically simple. If we know the encoding for position $pp$, it should be straightforward to compute the encoding for position $p+kp+k$

If you think about how we represent numbers on a number line, it’s easy to understand that 5 is 2 steps away from 3, or that 10 is 5 steps from 15. The same intuitive relationship should exist in our encodings.

Property 3 – Generalizes to longer sequences than those encountered in training

To increase our models’ utility in the real world, they should generalize outside
their training distribution. Therefore, our encoding scheme needs to be
adaptable enough to handle unexpected input lengths, without
violating any of our other desirable properties.

Property 4 – Generated by a deterministic process the model can learn

It would be ideal if our positional encodings could be drawn from a
deterministic process. This should allow the model to learn the mechanism
behind our encoding scheme efficiently.

Property 5 – Extensible to multiple dimensions

With multimodal models becoming the norm, it is crucial that our positional
encoding scheme can naturally extend from $1D1D$ to $nDnD$. This will allow models to consume data like images or brain scans, which are $2D2D$ and $4D4D$
respectively.

Now we know the ideal properties (henceforth referred to as $P{r}_{n}Pr\_n$

Integer Position Encoding

The first approach that may jump to mind is simply to add the integer value of the token position to each component of the token embedding, with values ranging from $0\to L0\; \backslash rightarrow\; L$

In the above animation, we create our positional encoding vector for the token $\text{chased}\backslash color\{\#699C52\}\backslash text\{chased\}$ from the index and add it to our token embedding. The embedding values here are a subset of the real values from Llama 3.2 1B. We can observe that they’re clustered around 0. This
is desirable to avoid vanishing or exploding gradients during training and therefore is something we’d like to maintain throughout the model.

It’s clear that our current naïve approach is going to cause problems. The magnitude of the position value
vastly exceeds the actual values of our input. This means the signal-to-noise
ratio is very low, and it’s hard for the model to separate the semantic
information from the positional information.

With this new knowledge, a natural follow on might be to normalize the position value by $\frac{1}{N}\backslash frac\{1\}\{N\}$

Is there a better way to ensure our numbers are between 0 and 1?
If we thought really hard about this for a while, we might come up with switching
from decimal to binary numbers.

Binary Position Encoding

Instead of adding our (potentially normalized) integer position to each
component of the embedding, we could instead convert it into its binary
representation and s t r e t c h our value out to match our embedding dimension, as demonstrated below.

We’ve converted the position of interest (252) into its binary representation
(11111100) and added each bit to the corresponding component of the
token embedding. The least significant bit (LSB) will cycle between 0 and 1 for every
subsequent token, whilst the most significant bit (MSB) will cycle every $2^{n-1}2^\{n-1\}$

We’ve solved the value range problem, and we now have unique encodings that are
consistent across different sequence lengths. What happens if we plot a low dimensional version of our token embedding and visualize the addition of our binary positional vector for different values.

We can see that the result is very “jumpy” (as we might expect from the
discrete nature of binary). The optimization process likes smooth, continuous and
predictable changes. Do we know any functions with similar value ranges that are smooth and continuous?

If we looked around a little, we might notice that both $\sin \backslash sin$ and $\cos \backslash cos$ fit the bill!

Sinusoidal positional encoding

The above animation visualizes our position embedding if each component is
alternatively drawn from $\sin \backslash sin$ and $\cos \backslash cos$ with gradually increasing
wavelengths. If you compare it with the previous animation, you’ll notice a striking similarity!

We’ve now arrived at Sinusoidal embeddings; originally defined in the Attention is all you need paper.
Let’s look at the equations:

$$P{E}_{(pos,2i)}=\sin \left(\frac{pos}{{10000}^{2i\mathrm{/}d}}\right)P{E}_{(pos,2i+1)}=\cos \left(\frac{pos}{{10000}^{2i\mathrm{/}d}}\right)PE\_\{(pos,2i)\}\; =\; \backslash color\{\#58C4DD\}\backslash sin\backslash left(\backslash color\{black\}\backslash frac\{pos\}\{10000^\{2i/d\}\}\backslash color\{\#58C4DD\}\backslash right)\backslash color\{black\}\; \backslash \backslash \; \backslash quad\; \backslash \backslash \; PE\_\{(pos,2i+1)\}\; =\; \backslash color\{\#FC6255\}\backslash cos\backslash left(\backslash color\{black\}\backslash frac\{pos\}\{10000^\{2i/d\}\}\backslash color\{\#FC6255\}\backslash right)\backslash color\{black\}\; \backslash \backslash$$

where $pospos$ is the tokens position index, $ii$ is the component index
in the positional encoding vector, and $dd$ is the model dimension. $10,00010,000$

There’s a few parts of this equation that are confusing at first glance. How did the
authors choose $10,00010,000$

It seems that using $10,00010,000$

Consider a sequence of sine and cosine pairs, each associated with a frequency ${\omega }_i\backslash omega\_i$

$$\mathbf{M}\cdot \left[\begin{array}{c}\sin ({\omega }_{i}p)\\ \cos ({\omega }_{i}p)\end{array}\right]=\left[\begin{array}{c}\sin ({\omega }_i(p+k))\\ \cos ({\omega }_i(p+k))\end{array}\right]\backslash mathbf\{M\}\; \backslash cdot\; \backslash begin\{bmatrix\}\; \backslash sin(\backslash omega\_i\; p)\; \backslash \backslash \; \backslash cos(\backslash omega\_i\; p)\; \backslash end\{bmatrix\}\; =\; \backslash begin\{bmatrix\}\; \backslash sin(\backslash omega\_i(p\; +\; k))\; \backslash \backslash \; \backslash cos(\backslash omega\_i(p\; +\; k))\; \backslash end\{bmatrix\}$$

The frequencies ${\omega }_i\backslash omega\_i$

$${\omega }_i=\frac{1}{{10000}^{2i\mathrm{/}d}}\backslash omega\_i\; =\; \backslash frac\{1\}\{10000^\{2i/d\}\}$$

To find this transformation matrix, we can express it as a general 2×2 matrix with unknown coefficients $u_{1}u\_1$

By applying the trigonometric addition theorem to the right-hand side, we can expand this into:

This expansion gives us a system of two equations by matching coefficients:

By comparing terms with $\sin ({\omega }_{i}p)\backslash sin(\backslash omega\_i\; p)$

These solutions give us our final transformation matrix ${\mathbf{M}}_{\mathbf{k}}\backslash mathbf\{M\_k\}$

If you’ve done any game programming before, you might notice that the
result of our derivation is oddly familiar. That’s right, it’s the Rotation Matrix! ${[}^3][^3]$

So the encoding scheme designed by Noam Shazeer in Attention is all you need was already encoding relative position as a rotation back in 2017! It took another 4 years to go from Sinusoidal Encoding to RoPE, despite rotations already being on the table…

Absolute vs Relative Position Encoding

With the knowledge in hand that rotations are important here, let’s
return to our motivating example and try to discover some intuitions for our next iteration.

Above we can see the absolute positions of our tokens, and the relative
positions from $\text{chased}\backslash color\{\#699C52\}\backslash text\{chased\}$ to every other token. With Sinusoidal Encoding, we
generated a separate vector which represents the absolute position,
and using some trigonometric trickery we were able to encode relative positions.

When we’re trying to understand these sentences, does it matter that this word is the 2157th word in this blog post? Or do we care about its relationship to the words around it? The absolute position of a word rarely matters for meaning – what matters is how words relate to each other.

Positional encoding in context

From this point on, it’s key to consider positional encoding in the context of
self attention. To reiterate, the self-attention mechanism enables the model to weigh the importance of different elements in an input sequence and dynamically adjust their influence on the output.

$$\text{Attn}(Q,K,V)=\text{softmax}\left(\frac{Q{K}^T}{\sqrt{d_k}}\right)V\backslash text\{Attn\}(Q,\; K,\; V)\; =\; \backslash text\{softmax\}\backslash left(\backslash frac\{QK^T\}\{\backslash sqrt\{d\_k\}\}\backslash right)V$$

In all our previous iterations, we’ve generated a separate positional encoding
vector and added it to our token embedding prior to our $QQ$, $KK$ and $VV$ projections.
By adding the positional information directly to our token embedding, we are
polluting the semantic information with the positional information. We should
be attempting to encode the information without modifying the norm. Shifting to multiplicative is the
key.

Using the dictionary analogy, when looking up a word (query) in our dictionary (keys), nearby words should have more influence than distant ones. The influence of one token upon another is determined by the $Q{K}^{T}QK^T$

$$\vec{a}\cdot \vec{b}=\mathrm{\mid }\vec{a}\mathrm{\mid }\mathrm{\mid }\vec{b}\mathrm{\mid }\cos \theta \backslash vec\{a\}\; \backslash cdot\; \backslash vec\{b\}\; =\; |\backslash vec\{a\}|\; |\backslash vec\{b\}|\; \backslash cos\; \backslash theta$$

The geometric interpretation of the dot product shown above gives us a magnificent insight.
We can modulate the result of our dot product of two vectors purely by
increasing or decreasing the angle between them. Furthermore, by rotating the
vector, we have absolutely zero impact on the norm of the vector, which encodes
the semantic information of our token.

So now we know where to focus our attention, and have seen from another angle why a
rotation might be a sensible “channel” in which to encode our positional
information, let’s put it all together!

Rotary Postional Encoding

Rotary Postional Encoding or RoPE was defined in the
RoFormer paper (Jianlin Su designed it independently on his blog here and here).
While it may seem like voodoo if you skip to the end result, by thinking about Sinusoidal Encoding in the
context of self attention (and more specifically dot products), we can see how
it all comes together.

Much like in Sinusoidal Encoding, we decompose our vectors $\mathbf{q}\backslash mathbf\{q\}$ or $\mathbf{k}\backslash mathbf\{k\}$, instead of pre-projection $\mathbf{x}\backslash mathbf\{x\}$) into 2D pairs/chunks. Rather than encoding absolute position directly by adding a vector we drew from sinusoidal functions of slowly decreasing frequencies, we cut to the chase and encode relative position by multiplying each pair with the rotation matrix.

Let $\mathbf{q}\backslash mathbf\{q\}$ or $\mathbf{k}\backslash mathbf\{k\}$ be our input vector at position $pp$. We create a block diagonal matrix
where ${\mathbf{M}}_{\mathbf{i}}\backslash mathbf\{M\_i\}$

Much the same as Sinusoidal Encoding, ${\mathbf{M}}_{\mathbf{i}}\backslash mathbf\{M\_i\}$

In practice, we don’t use a matrix multiplication to compute RoPE as it would be
computationally inefficient with such a sparse matrix. Instead, we can directly apply the rotations to pairs of elements independently, taking advantage of the regular pattern in the computation:

$$R_{\mathrm{\Theta },p}^{d}q=\left(\begin{array}{c}{q}_1\\ q_2\\ q_3\\ q_4\\ \mathrm{⋮}\\ q_{d-1}\\ q_d\end{array}\right)\otimes \left(\begin{array}{c}\cos p{\theta }_1\\ \cos p{\theta }_1\\ \cos p{\theta }_2\\ \cos p{\theta }_2\\ \mathrm{⋮}\\ \cos p{\theta }_{d\mathrm{/}2}\\ \cos p{\theta }_{d\mathrm{/}2}\end{array}\right)+\left(\begin{array}{c}-q_2\\ q_1\\ -q_4\\ q_3\\ \mathrm{⋮}\\ -q_d\\ q_{d-1}\end{array}\right)\otimes \left(\begin{array}{c}\sin p{\theta }_1\\ \sin p{\theta }_1\\ \sin p{\theta }_2\\ \sin p{\theta }_2\\ \mathrm{⋮}\\ \sin p{\theta }_{d\mathrm{/}2}\\ \sin p{\theta }_{d\mathrm{/}2}\end{array}\right)R\_\{\backslash Theta,p\}^d\; q\; =\; \backslash begin\{pmatrix\}\; q\_1\; \backslash \backslash \; q\_2\; \backslash \backslash \; q\_3\; \backslash \backslash \; q\_4\; \backslash \backslash \; \backslash vdots\; \backslash \backslash \; q\_\{d-1\}\; \backslash \backslash \; q\_d\; \backslash end\{pmatrix\}\; \backslash otimes\; \backslash begin\{pmatrix\}\; \backslash cos\; p\backslash theta\_1\; \backslash \backslash \; \backslash cos\; p\backslash theta\_1\; \backslash \backslash \; \backslash cos\; p\backslash theta\_2\; \backslash \backslash \; \backslash cos\; p\backslash theta\_2\; \backslash \backslash \; \backslash vdots\; \backslash \backslash \; \backslash cos\; p\backslash theta\_\{d/2\}\; \backslash \backslash \; \backslash cos\; p\backslash theta\_\{d/2\}\; \backslash end\{pmatrix\}\; +\; \backslash begin\{pmatrix\}\; -q\_2\; \backslash \backslash \; q\_1\; \backslash \backslash \; -q\_4\; \backslash \backslash \; q\_3\; \backslash \backslash \; \backslash vdots\; \backslash \backslash \; -q\_d\; \backslash \backslash \; q\_\{d-1\}\; \backslash end\{pmatrix\}\; \backslash otimes\; \backslash begin\{pmatrix\}\; \backslash sin\; p\backslash theta\_1\; \backslash \backslash \; \backslash sin\; p\backslash theta\_1\; \backslash \backslash \; \backslash sin\; p\backslash theta\_2\; \backslash \backslash \; \backslash sin\; p\backslash theta\_2\; \backslash \backslash \; \backslash vdots\; \backslash \backslash \; \backslash sin\; p\backslash theta\_\{d/2\}\; \backslash \backslash \; \backslash sin\; p\backslash theta\_\{d/2\}\; \backslash end\{pmatrix\}$$

That’s all there is to it! By artfully applying our rotations to 2D chunks of $\mathbf{q}\backslash mathbf\{q\}$ and $\mathbf{k}\backslash mathbf\{k\}$ prior to their dot product, and switching from additive to
multiplicative, we can gain a big performance boost in evaluations ${[}^4][^4]$

Extending RoPE to $nn$-Dimensions

We’ve explored the $1D1D$ case for RoPE and by this point I hope you’ve gained an
intuitive understanding of an admittedly unintuitive component of transformers.
Finally, let’s explore extending it to higher dimensions, such as images.

A natural first intuition could be to directly use the $\left[\begin{array}{c}x\\ y\end{array}\right]\backslash begin\{bmatrix\}\; x\; \backslash \backslash \; y\; \backslash end\{bmatrix\}$

In the $1D1D$ case, we encode the relative position $m-nm\; –\; n$

The future of positional encoding

Is RoPE the final incarnation of positional encoding? This recent paper from DeepMind deeply analyses RoPE and highlights some fundamental problems. TLDR: RoPE isn’t a perfect solution, and the models mostly focus on the lower frequencies, but the paper shows that removing (not rotating) the lowest frequencies improves performance on Gemma 2B!

I anticipate some future breakthroughs, perhaps taking inspiration from
signal processing with ideas like wavelets or hierarchical implementations. As models
are increasingly quantized for deployment, I’d also expect to see some
innovation in encoding schemes that remain robust under low-precision arithmetic.

Conclusion

Positional encoding has and continues to be treated as an after thought in
transformers. I believe we should view it differently – self attention has an
Achilles heel that has been repeatedly patched.

I hope this blog post showed you that you too could have discovered state of the
art positional encoding, despite it being unintuitive at first. In a follow up
post I’d love to explore practical implementation details for RoPE in order to
maximise performance.

This post was originally published here.

References

[^1]: Binary and Sinusoidal animations are reproductions of animations contained
in this video.

[^2]: Using $\theta =10000\backslash theta\; =\; 10000$

[^3]: Pieces of this post are based on this fantastic
post
by Amirhossein Kazemnejad.

[^4]: For empirical evidence, see this great post by EleutherAI.