RoFormer: Enhanced Transformer with Rotary Position Embedding (RoPE)¶

Transformer attention is permutation-invariant:

  attention("the cat sat")  ==  attention("sat cat the")

Without position information, language models can't distinguish:
  "Dog bites man"  vs  "Man bites dog"

The model needs SOME signal that token i comes before token j.
But how do you encode position so the model can use it effectively?

Original Transformer added a fixed sinusoidal vector to each token:

  PE(pos, 2i)   = sin(pos / 10000^(2i/d))
  PE(pos, 2i+1) = cos(pos / 10000^(2i/d))

  x_pos = token_embedding + PE(pos)

Pros:
  - No learned parameters
  - Theoretically extrapolates to longer sequences

Cons:
  - Position info gets diluted as it flows through layers
  - Adding to embeddings mixes "what" and "where" channels
  - In practice, extrapolation BEYOND training length fails badly
  - Encodes ABSOLUTE position, but attention cares about RELATIVE distance

Replace the sinusoidal formula with a learned table:

  PE = nn.Embedding(max_seq_len, d_model)
  x_pos = token_embedding + PE[pos]

Pros:
  - Model learns task-specific position patterns
  - Often performs slightly better than sinusoidal at training length

Cons:
  - HARD CAP at max_seq_len (BERT: 512, GPT-2: 1024)
  - Cannot extrapolate AT ALL beyond training length
  - Each position is a separate parameter, no smooth interpolation
  - Wastes parameters that scale with max sequence length

Add a bias to attention scores based on relative distance:

  attention(q_i, k_j) = (q_i . k_j) + bias(i - j)

Pros:
  - Directly encodes what matters: relative offsets
  - Generalizes better to longer sequences

Cons:
  - Requires modifying the attention computation itself
  - Bucket-based variants (T5) lose precision at long distances
  - Extra learned parameters per attention head
  - Still struggles to extrapolate cleanly

What we want from position encoding:

  1. Relative-position-aware (attention should care about i - j, not absolute i)
  2. Smoothly extrapolates to unseen lengths
  3. Decays with distance (far tokens should attend less)
  4. No extra parameters
  5. No invasive changes to the attention kernel

Sinusoidal:        relative no   extrapolates partial  decay no   params 0     invasive no
Learned absolute:  relative no   extrapolates no       decay no   params O(L)  invasive no
T5 relative bias:  relative yes  extrapolates partial  decay yes  params O(buckets)  invasive yes

RoPE: relative yes  extrapolates yes  decay yes  params 0  invasive (small)

Instead of ADDING a position vector to embeddings:
  x_pos = embed + PE(pos)        (old way)

ROTATE the query and key vectors by a position-dependent angle:
  q_pos = R(pos) . q              (new way)
  k_pos = R(pos) . k

Where R(pos) is a rotation matrix that depends on the position.

Why? When you compute attention dot product:
  q_m^T R(m)^T R(n) k_n  =  q_m^T R(n - m) k_n

The absolute rotations CANCEL and only the relative offset (n - m) remains.

This is the magic: an ABSOLUTE position encoding that produces
RELATIVE position behavior for free, just from how dot products work.

For a 2D vector q = (q_x, q_y), rotation by angle theta is:

  R(theta) = | cos(theta)  -sin(theta) |
             | sin(theta)   cos(theta) |

  R(theta) . q = | q_x cos(theta) - q_y sin(theta) |
                 | q_x sin(theta) + q_y cos(theta) |

Property: R(a) . R(b) = R(a + b) (rotations compose by adding angles)

For RoPE, set theta_pos = pos * frequency.

Then:  q_m = R(m * f) . q
       k_n = R(n * f) . k

  q_m . k_n = q^T R(m*f)^T R(n*f) k
            = q^T R((n-m)*f) k

The dot product depends ONLY on (n - m), the relative distance!

For a d-dimensional vector (d even), split into d/2 pairs of 2D subspaces.
Each pair gets its OWN frequency.

Frequencies follow the sinusoidal pattern from Vaswani:
  theta_i = 10000^(-2i/d)   for i in [0, d/2)

Frequency spectrum:
  - Low-index pairs: HIGH frequency (rotate fast with position)
  - High-index pairs: LOW frequency (rotate slowly with position)

This gives the model a rich multi-scale view of position:
  Some dimensions distinguish neighbors (i vs i+1)
  Other dimensions distinguish distant chunks (i vs i+1000)

Position 0: identity (no rotation)
  q_0 = [a, b, c, d, e, f, ...]

Position 1: rotate each (x, y) pair by its frequency
  pair 0 (high freq):  rotate by theta_0 = 1.0      (a, b) -> (a', b')
  pair 1 (med freq):   rotate by theta_1 = 0.1      (c, d) -> (c', d')
  pair 2 (low freq):   rotate by theta_2 = 0.01     (e, f) -> (e', f')

Position 100:
  pair 0: rotate by 100 * theta_0 = 100 rad   (many full revolutions)
  pair 1: rotate by 100 * theta_1 = 10 rad    (about 1.5 revolutions)
  pair 2: rotate by 100 * theta_2 = 1 rad     (a fraction of a revolution)

Different frequencies capture position at different scales,
just like how Fourier features capture signals at multiple scales.

Naively, R(pos) is a d x d block-diagonal matrix. We don't materialize it.
Instead, compute element-wise:

def apply_rope(x, pos):
    # x has shape [..., d], split into d/2 pairs of 2D vectors
    # x = [x_0, x_1, x_2, x_3, ..., x_{d-2}, x_{d-1}]
    # pairs = [(x_0, x_1), (x_2, x_3), ..., (x_{d-2}, x_{d-1})]

    freqs = 1.0 / (10000 ** (torch.arange(0, d, 2) / d))   # [d/2]
    angles = pos * freqs                                    # [d/2]
    cos_a, sin_a = torch.cos(angles), torch.sin(angles)

    x_even = x[..., 0::2]
    x_odd  = x[..., 1::2]

    rotated_even = x_even * cos_a - x_odd * sin_a
    rotated_odd  = x_even * sin_a + x_odd * cos_a

    out = interleave(rotated_even, rotated_odd)
    return out

Cost: O(d) per token. No materialized rotation matrix. No extra params.
Applied to q and k BEFORE attention; v is left untouched.

Standard attention block:

  q = W_q . x        <- project token to query
  k = W_k . x        <- project token to key
  v = W_v . x        <- project token to value
  attn_score = softmax(q . k^T / sqrt(d))
  out = attn_score . v

With RoPE:

  q = W_q . x
  k = W_k . x
  v = W_v . x
  q = apply_rope(q, pos)   <-- NEW: rotate q by its position
  k = apply_rope(k, pos)   <-- NEW: rotate k by its position
  attn_score = softmax(q . k^T / sqrt(d))
  out = attn_score . v

Note: RoPE touches ONLY q and k. v is unchanged. No additive PE on the input.

The clever bit: each token's q,k is rotated by its OWN absolute position.
But because rotations compose, the dot product q_m . k_n only sees (n - m).

You get relative-position semantics with absolute-position bookkeeping.
No need for an O(L^2) relative-position table or attention bias modification.

For sufficiently random q and k, the expected dot product after RoPE
DECREASES as |n - m| grows (proven in the paper).

Distant tokens naturally attend less, which is the inductive bias we want
for language. No hand-designed decay function needed.

Different (q,k) channel pairs rotate at different frequencies.
Like a Fourier basis, this gives the model multiple resolutions of position.

The model can specialize:
  - Some heads attend based on local position (high-freq channels)
  - Other heads attend based on global structure (low-freq channels)

The base 10000 (inherited from the original Transformer) sets the longest period.
Tweaking this base is what NTK-aware RoPE and YaRN exploit for long-context.

RoPE works inside linear attention variants too:

  Linear attention: phi(q) . phi(k)^T  (kernel approximation)
  With RoPE:        R(m) phi(q) . phi(k)^T R(n)^T  =  same trick

This was a big motivator: prior position encodings were incompatible with
linear attention because they relied on the softmax denominator. RoPE bypasses
that constraint entirely.

RoPE's rotation angle is theta_pos = pos * frequency.

If a model trains on positions 0..2047, it sees angles up to 2047 * frequency
for each frequency band. At inference, position 4096 produces angles
4096 * frequency. The MATH still works (rotations are well-defined for all angles).

But: the model has never seen those rotation regimes during training,
so attention quality degrades for very long contexts.

This is why naive context extension fails.

Insight: Don't change the formula, change the BASE.

Original RoPE: theta_i = 10000^(-2i/d)
NTK-aware:     theta_i = (10000 * scale)^(-2i/d)

By scaling the base instead of the position, low-frequency dimensions
(which encode long-range structure) get stretched proportionally,
while high-frequency dimensions (which encode local structure)
stay near their training distribution.

Result: Llama-2 4K -> 8K context with NO finetuning.

YaRN refines NTK-aware with:
  1. Per-frequency scaling (different stretch for different bands)
  2. Attention temperature adjustment (softmax sharpness)
  3. Light fine-tuning (~400 steps)

Result: Llama-2 4K -> 128K with minimal compute.
Most production long-context LLMs (Mistral 32K, Llama-3 128K) use YaRN
or close variants.

Simpler trick: scale POSITIONS instead of base.

  pos_inference = pos_real * (train_len / target_len)

Map a 8K position back into the 0..2K range the model trained on,
then RoPE-rotate normally. With short fine-tuning, this also extends context.

LLaMA-2 Long used this approach for 4K -> 32K.

Modern long-context models stack these:

  1. Pretrain with RoPE at base 10000, length 4K-8K
  2. Apply YaRN/NTK to extend rotation regime
  3. Continue-train on long documents (~1B tokens)
  4. Sometimes: interleave with short data to prevent quality drop

Llama-3.1 (128K), Mistral Large (128K), Qwen2.5 (1M) all follow this recipe.
None of it works without RoPE as the foundation.

Train length: 512 tokens. Test at varying lengths.

Length    Sinusoidal   Learned    RoPE
512       baseline     baseline   baseline
1024      +12% PPL     fails      +3% PPL
2048      +45% PPL     fails      +8% PPL
4096      +120% PPL    fails      +15% PPL

RoPE alone extrapolates 2-4x further than alternatives BEFORE
applying NTK/YaRN extensions.

ALiBi: Attention with Linear Biases

  attention(q_i, k_j) = (q_i . k_j) - m * |i - j|

  Just subtract a linear penalty proportional to distance.
  Each head gets a different slope m.

  No vector rotation. No position embedding. Pure attention bias.

ALiBi extrapolates well out-of-the-box but encodes only DISTANCE.
RoPE encodes RELATIVE OFFSET as a full rotation, preserving more structure.

Empirically:
  - At training length: comparable
  - Slight extension: RoPE wins with NTK/YaRN
  - Long-form reasoning: RoPE-based models retain better needle-in-haystack
  - Multi-needle and ordering tasks: RoPE preserves structure ALiBi loses

Once YaRN solved RoPE's extrapolation weakness in 2023, ALiBi lost its
main advantage and the field consolidated around RoPE.

Open-source LLMs using RoPE:
  - LLaMA 1, 2, 3, 4 (Meta)
  - Mistral, Mixtral (Mistral AI)
  - Qwen 1, 2, 2.5, 3 (Alibaba)
  - DeepSeek V1, V2, V3, R1 (DeepSeek)
  - Yi (01.AI)
  - GPT-NeoX, Pythia (EleutherAI)
  - Falcon (later versions)
  - Phi-2, Phi-3, Phi-4 (Microsoft)
  - Gemma 1, 2, 3 (Google)
  - Command R, R+ (Cohere)

Closed-source LLMs reported/inferred to use RoPE:
  - PaLM, PaLM 2, Gemini 1.0, 1.5, 2.0 (Google)
  - Claude (Anthropic, inferred from behavior)

That's basically the entire frontier of LLM development since 2022.

Holdouts (mostly historical):
  - GPT-2, GPT-3 (sinusoidal/learned, predates RoPE)
  - BLOOM (ALiBi)
  - Falcon-7B/40B (ALiBi, switched in later versions)

Without RoPE + YaRN:
  Models capped at training length (typically 2K-8K tokens)
  Long PDFs require chunking and retrieval

With RoPE + YaRN:
  128K-1M token contexts standard
  Whole codebase analysis (Claude, Gemini)
  Long legal/medical document review
  Book-length summarization

Code is structurally long-range:
  - Function defined on line 100, called on line 5000
  - Imports at top affect references throughout

RoPE's multi-scale frequency encoding gives code LLMs the right inductive
bias for these long dependencies. DeepSeek-Coder, Qwen-Coder, and Codestral
all rely on RoPE for their 32K-128K windows.

Long chain-of-thought outputs (10K-100K reasoning tokens) need:
  - Coherent positional reasoning across the chain
  - No degradation at the tail of the trace

RoPE-with-YaRN models maintain attention quality across long traces.
This is part of why DeepSeek-R1 and similar work at all.

RAG often stuffs many documents into the prompt:
  Top-50 retrieved chunks + query + system prompt = 30K+ tokens

RoPE-extended models can attend to ALL chunks rather than only the latest,
reducing the "lost in the middle" failure mode that plagued earlier systems.

RoPE alone extrapolates 1.5-3x beyond training length before quality drops.
Real long-context (10x+) requires NTK, YaRN, or position interpolation
plus continue-training.

Values are not rotated, so positional info reaches downstream layers
only through the attention pattern. Usually fine, but limits some niche use cases
(e.g. value-based positional reasoning).

At position 1M, angles for low-freq channels are huge.
fp16 / bf16 sin/cos can lose precision. Practical implementations
compute freqs in fp32 and cast back.

You can't take BERT-base and swap RoPE in without retraining.
The Q,K projections were trained assuming additive position info.
RoPE retrofits require at least continue-pretraining.