Long-Context Extrapolation Positional Encoding RoPE ALiBi and Context-Window Extension
# Long-Context Extrapolation & Positional Encoding: RoPE, ALiBi, and Context-Window Extension
## Introduction & Motivation
Transformer self-attention has no inherent notion of token order: the operation $ ext{softmax}(QK^T/\sqrt{d_k})V$ is a permutation-equivariant function of its input set, meaning that without some additional mechanism, shuffling the input tokens would produce a correspondingly shuffled output with no change in the underlying computation. Positional encoding is the mechanism that injects order information into an otherwise order-blind architecture, and the specific choice of positional encoding scheme has turned out to have an outsized influence on one of the most commercially and technically important properties of a modern large language model: how well it generalizes to sequences longer than anything it saw during training.
Early Transformers used absolute sinusoidal or learned positional embeddings, added once to the token embeddings at the input layer. These schemes work well within the trained context length but degrade sharply, sometimes catastrophically, when a model is asked to process sequences beyond that length, because the model has simply never observed the positional signal patterns that appear past its training boundary. As applications increasingly demand processing of long documents, entire codebases, extended conversations, and large retrieved-context windows in retrieval-augmented generation pipelines, the ability to either train efficiently at long context lengths or to cheaply extend a model trained at a shorter length to handle longer sequences at inference time has become a central engineering concern.
Two ideas dominate the modern landscape: relative positional encodings that are computed as a function of the distance between two tokens rather than their absolute positions, most notably Rotary Position Embedding (RoPE) and Attention with Linear Biases (ALiBi); and a family of post-hoc context-extension techniques, including position interpolation, NTK-aware scaling, and YaRN, that allow a model trained at one context length to be adapted, often with little or no additional fine-tuning, to reliably handle sequences several times longer. Understanding why absolute positional schemes fail to extrapolate, why relative schemes fare better but are not immune to their own failure modes, and how the interpolation-based extension techniques work mathematically, is essential for anyone deploying, fine-tuning, or evaluating long-context language models.
## Core Concepts & Theory
Absolute positional embeddings, whether the fixed sinusoidal functions used in the original Transformer or a learned embedding table as used in early BERT and GPT variants, assign a distinct vector to each position index from 0 up to some maximum trained length. This works cleanly within that range but has an obvious structural limitation: a learned embedding table simply has no entry for position 5000 if it was only ever trained with positions 0 through 511, and even the fixed sinusoidal scheme, while mathematically defined for any position, produces high-frequency patterns at large positions that the model's attention and feedforward layers were never exposed to during training, so there is no guarantee, and empirically very little evidence, that the model can meaningfully use them.
Relative positional encodings sidestep the absolute-position problem by making the positional signal a function of the offset between a query position and a key position rather than either position individually. Rotary Position Embedding achieves this elegantly by rotating the query and key vectors in each attention head by an angle proportional to their absolute position, using a set of geometrically decreasing rotation frequencies across the embedding dimension, such that the dot product between a rotated query and a rotated key depends mathematically only on their relative offset, not on their absolute positions. This means that, in principle, the same relative-offset pattern seen during training at positions 10 and 15 produces an identical attention-score contribution as the same offset seen at positions 10,010 and 10,015, a property RoPE has by construction that absolute embeddings lack entirely.
ALiBi takes an even more direct route: it adds no positional information to the token embeddings at all, and instead directly biases the raw attention logits before the softmax with a term that is linearly proportional to the negative distance between query and key positions, scaled by a small, head-specific slope. Because this bias is a simple, unbounded linear function of relative distance rather than a lookup into any finite table, it is trivially well-defined for any sequence length, and the near-term recency bias it imposes (very distant tokens receive a strongly negative, softmax-suppressing bias) has been empirically found to generalize to much longer sequences than were seen in training, even without any specialized modification, though it does this by making the model behave, in the long-range limit, increasingly like a fairly short local-attention window rather than by giving the model genuine long-range positional resolution.
## Mathematical Formulation
Rotary Position Embedding partitions each head's dimension $d$ into $d/2$ two-dimensional subspaces and assigns each subspace $i$ a rotation frequency
$$ heta_i = ext{base}^{-2i/d}, \quad i = 0, 1, \ldots, d/2 - 1 $$
where $ ext{base}$ is a hyperparameter (commonly 10,000). For a token at position $p$, the query or key vector is rotated within each subspace by angle $p \cdot heta_i$, so that the rotated vector at position $p$ is
$$ R_p x = \begin{pmatrix} \cos(p heta_i) & -\sin(p heta_i) \\ \sin(p heta_i) & \cos(p heta_i) \end{pmatrix} \begin{pmatrix} x_{2i} \\ x_{2i+1} \end{pmatrix} ext{ for each subspace } i $$
The key algebraic property that makes RoPE a relative encoding is that the inner product between a query at position $m$ and a key at position $n$ depends only on their offset,
$$ (R_m q)^T (R_n k) = q^T R_{n-m} k $$
since rotation matrices compose additively, $R_m^T R_n = R_{n-m}$.
ALiBi instead modifies the raw attention logits directly. For query position $i$ and key position $j$ in a given attention head $h$ with head-specific slope $m_h$, the attention score before softmax becomes
$$ ext{score}(i, j) = \frac{q_i^T k_j}{\sqrt{d_k}} - m_h \cdot |i - j| $$
where the slopes $m_h$ across heads are typically set as a geometric sequence, $m_h = 2^{-8h/H}$ for $H$ total heads, giving different heads different effective attention ranges, from very short-range (large slope) to comparatively long-range (small slope).
Position interpolation, a post-hoc context-extension technique, addresses RoPE's extrapolation failure by rescaling positions before rotation rather than modifying the rotation frequencies themselves: given a model trained at length $L_{ ext{train}}$ that must now handle length $L_{ ext{new}} > L_{ ext{train}}$, the effective position fed into the rotation is
$$ p' = p \cdot \frac{L_{ ext{train}}}{L_{ ext{new}}} $$
which guarantees that the maximum rotation angle encountered at the new, extended length never exceeds the maximum angle the model saw during training, at the cost of uniformly reducing the angular resolution available to distinguish nearby positions. NTK-aware scaling instead rescales the RoPE base itself,
$$ ext{base}' = ext{base} \cdot \left( \frac{L_{ ext{new}}}{L_{ ext{train}}} ight)^{d/(d-2)} $$
which, because the exponent $-2i/d$ in the frequency formula makes high-frequency (small $i$) components nearly insensitive to a change in base while low-frequency (large $i$) components are stretched substantially, achieves a non-uniform effect: local, high-resolution positional discrimination near $i=0$ is largely preserved while the slowest-rotating, longest-range components are compressed enough to stay within their trained angular range.
## Advanced Theory & Extensions
YaRN (Yet another RoPE extensioN method) generalizes and improves on both position interpolation and NTK-aware scaling by treating different frequency bands differently and explicitly rather than through a single global formula: very high-frequency dimensions (which complete many full rotation cycles even within the original training length) are left essentially untouched, since these are already the components most robust to extrapolation and most responsible for fine-grained local position discrimination; very low-frequency dimensions are interpolated fully, following the position-interpolation formula, since these are the components that would otherwise be pushed furthest outside their trained angular range; and a smooth ramp function interpolates between these two regimes for intermediate frequencies. YaRN additionally introduces a temperature adjustment to the attention softmax itself, compensating for a subtle but measurable entropy increase in attention distributions that interpolation-based scaling methods otherwise introduce, which further improves perplexity at extended lengths relative to interpolation alone.
Dynamic NTK scaling addresses a practical limitation of static NTK-aware scaling, namely that a fixed base rescaling computed for one specific target length $L_{ ext{new}}$ is not optimal for shorter sequences processed by the same deployed model; the dynamic variant instead recomputes the effective base on the fly, as a function of the actual current sequence length being processed at each inference call, so that short sequences within the model's original trained range see no distortion at all, while the degree of frequency rescaling smoothly increases only as the actual input grows beyond the original training length. Sliding-window and local-attention mechanisms, orthogonal to the choice of positional encoding, offer an entirely different strategy for handling long sequences: rather than requiring the model to maintain meaningful attention over the full extended range, they restrict most attention computation to a bounded local window (optionally combined with a small number of global or retrieved tokens), trading long-range dense attention for computational and memory efficiency, an approach exemplified by models using windowed or dilated attention patterns and often combined with ALiBi's naturally local-favoring bias.
## Computational Considerations
RoPE's rotation can be applied to queries and keys with negligible additional compute relative to the attention operation itself, since it is a fixed, input-independent linear transformation computed once per position and cached; most implementations precompute the cosine and sine tables for the maximum expected sequence length once and reuse them, though context extension via NTK-aware or YaRN-style rescaling requires either recomputing these tables when the effective base changes or maintaining separate cached tables per supported context length, a minor but non-zero engineering overhead in serving systems that must support variable-length requests efficiently.
ALiBi's linear bias term must be computed for every query-key pair at every attention computation, but since it depends only on the integer distance $|i-j|$, it can be precomputed as a single, reusable bias matrix (or even a single row, exploiting the Toeplitz structure of the bias matrix, since the bias depends only on relative offset) and added to the raw attention logits at negligible marginal cost relative to the $QK^T$ matrix multiplication that dominates attention's computational profile.
Extending context length at inference time, regardless of the positional-encoding-level technique used, does not by itself solve the quadratic scaling of attention compute and memory with sequence length; interpolation and rescaling techniques address whether the model's positional signal remains meaningful at longer lengths, not the underlying $O(n^2)$ cost of full attention, which is why long-context deployment in practice is usually combined with efficient attention implementations such as FlashAttention (which reduces memory traffic without changing the asymptotic compute) or with sparse and local attention patterns that reduce the effective $n^2$ term directly, and with key-value cache memory management strategies, since a long context directly and linearly increases the KV-cache memory footprint that must be held during autoregressive generation.
## Practical Implementation Strategies
The most reliable way to obtain strong long-context performance remains training or continuing to pretrain directly at the target context length on a data mixture containing genuinely long-range dependencies, since interpolation-based extension techniques, however well designed, are fundamentally working around the fact that a model's attention and feedforward layers were optimized for a particular distribution of positional signals and were never directly trained to exploit information at the extended range; continued pretraining at the longer length, even for a comparatively small number of additional tokens relative to the original pretraining budget, consistently yields better long-context task performance than training-free extension alone.
When training-free extension is necessary, for instance to serve an existing model at a longer context length without the cost of additional pretraining, a short fine-tuning phase at the target extended length, applied after position interpolation or NTK-aware rescaling has been applied to the positional encoding, substantially closes the gap to models natively trained at that length, and is standard practice rather than relying purely on zero-shot extrapolation. Evaluation of long-context capability should never rely solely on perplexity measured on long documents, since low perplexity can be achieved largely from strong local (short-range) language modeling even when a model is effectively ignoring distant context; targeted retrieval-style evaluations, such as inserting a specific fact at a controlled position within a long context and querying for it afterward (needle-in-a-haystack style tests), are necessary to directly measure whether information at a given distance is actually being used by the model's attention rather than merely tolerated without causing a quality collapse.
Choice between RoPE-based and ALiBi-based architectures should account for the intended extension strategy: RoPE's rich ecosystem of interpolation and rescaling techniques (linear interpolation, NTK-aware, YaRN, dynamic NTK) gives practitioners fine-grained, well-studied control over the extrapolation-versus-resolution tradeoff, while ALiBi's simplicity, requiring no positional encoding modification at all to process longer sequences, comes with less flexibility to recover fine-grained long-range positional resolution beyond its inherent recency bias, making the two schemes suited to somewhat different deployment priorities.
## Benchmark Datasets & Evaluation
The needle-in-a-haystack evaluation methodology, in which a specific short fact is inserted at a controlled, varying depth within a long, otherwise irrelevant context and the model is queried to retrieve it, has become a standard diagnostic for long-context models specifically because it isolates positional retrieval capability from general long-document comprehension, and is commonly reported as a two-dimensional heatmap of retrieval accuracy across both context length and insertion depth, revealing failure patterns (such as degraded retrieval specifically in the middle of long contexts, popularly termed "lost in the middle") that aggregate perplexity metrics obscure entirely. LongBench and ZeroSCROLLS aggregate a diverse suite of realistic long-context tasks, including long-document question answering, summarization, and few-shot learning with many in-context examples, providing broader coverage of practical long-context use cases than synthetic retrieval tests alone.
RULER extends the needle-in-a-haystack paradigm with a more demanding and diverse set of synthetic long-context tasks, including multi-needle retrieval (finding several inserted facts simultaneously), variable tracking, and aggregation tasks that require combining information distributed across the entire context rather than retrieving a single localized fact, and has become an important benchmark specifically because many models that perform well on single-needle retrieval degrade sharply on RULER's more demanding variants, revealing that basic retrieval capability at long range does not imply robust long-range reasoning capability. Perplexity on long-document corpora remains a useful, cheap-to-compute secondary signal for detecting gross positional-encoding failures, such as a sharp perplexity spike exactly at the boundary of a model's originally trained context length, but should always be reported alongside a targeted retrieval or reasoning benchmark rather than as a standalone measure of long-context quality.
## Key Challenges & Limitations
The "lost in the middle" phenomenon, in which models exhibit a pronounced U-shaped accuracy curve on long-context retrieval tasks, performing well when the relevant information is near the beginning or end of the context but substantially worse when it is buried in the middle, remains only partially explained and is not fully resolved by any current positional encoding scheme or extension technique; it appears related both to the statistical structure of training data (where relevant information is disproportionately near document boundaries) and to more fundamental attention-allocation dynamics that current architectures do not fully overcome regardless of positional encoding choice.
Extension techniques that rescale RoPE's frequencies inherently trade off extrapolation range against positional resolution: any given rescaling factor that successfully keeps the lowest-frequency components within their trained angular range at some target extended length necessarily reduces the angular resolution available for discriminating nearby positions, and pushing the extension factor too far (attempting to extend context length by, for instance, 32 times or more with a single static rescaling) causes measurable degradation even in local, short-range language modeling quality, which is why very large context extensions in practice are typically achieved through a combination of moderate rescaling and genuine continued pretraining rather than rescaling alone.
Computational cost remains an obstinate limitation independent of positional encoding quality: even a positional encoding scheme with theoretically perfect long-range extrapolation does not reduce the quadratic attention compute and linearly growing KV-cache memory cost of processing longer sequences, meaning that positional-encoding advances alone cannot make arbitrarily long context economically practical without complementary efficient-attention and memory-management techniques. Evaluation itself remains an open challenge, since synthetic retrieval benchmarks can be gamed or saturated by models specifically tuned to that benchmark's structure without necessarily reflecting genuine improvement on the messier, less clearly delineated long-range reasoning tasks that motivate long-context deployment in the first place.
## Hyperparameter Tuning
The RoPE base value, conventionally 10,000, directly determines the range of rotation frequencies across the embedding dimension and therefore the tradeoff between local resolution and the position at which the slowest-rotating dimensions begin to wrap into angular ranges not distinguishable from smaller positions; some long-context-native pretraining efforts deliberately increase the base value from the outset (rather than relying on post-hoc rescaling) specifically to push this wrap-around point beyond the intended training and deployment context length. The position interpolation or NTK-aware scaling factor must be matched to the actual ratio between the target extended length and the original trained length; under-scaling leaves some low-frequency dimensions still outside their trained angular range, while over-scaling unnecessarily sacrifices positional resolution that could otherwise have been preserved for shorter, more common request lengths.
For ALiBi, the per-head slope geometric sequence base (commonly $2^{-8/H}$) determines how sharply different heads specialize into short-range versus longer-range attention patterns; steeper slope ratios across heads increase the diversity of effective attention ranges available to the model at the cost of any individual head's long-range resolution, and this ratio is typically left at published default values rather than extensively retuned, since ALiBi's original design already spans a wide range of effective receptive fields by construction. When performing a short fine-tuning phase after context extension, the learning rate and number of fine-tuning steps should be kept modest relative to original pretraining, since the goal is to adapt the model to the rescaled positional signal rather than to substantially alter the already-learned attention and feedforward representations, and excessive fine-tuning at this stage risks degrading short-context performance that was already strong.
## Real-World Applications & Case Studies
LLaMA and its successors originally shipped with a RoPE-based architecture trained at a comparatively modest context length (2,048 to 4,096 tokens in earlier releases), and the broad open-source ecosystem's rapid adoption of position interpolation and later NTK-aware and YaRN-style scaling techniques, applied to these released model weights without requiring access to the original full pretraining pipeline, is a widely cited case study in how effective training-free context extension can be, enabling community-extended variants supporting context lengths many times the originally trained length with comparatively modest additional fine-tuning cost.
GPT-NeoX and other open pretraining efforts have used ALiBi specifically to sidestep the need for any context-extension technique altogether, training at a manageable length while relying on ALiBi's inherent extrapolation behavior to handle longer sequences at inference time, an explicit design tradeoff favoring deployment simplicity over the finer-grained resolution control that RoPE-based interpolation techniques offer. Production long-context model releases, particularly those advertising context windows in the hundreds of thousands of tokens, are understood in the technical community to virtually always combine native long-context continued pretraining with efficient attention implementations and careful KV-cache memory management, rather than relying on any single positional-encoding trick in isolation, underscoring that long-context capability at genuine production quality is a systems-level achievement rather than a single-technique achievement.
## Integration with Other Methods
Long-context positional encoding techniques interact directly with retrieval-augmented generation system design, since a model's actual usable context length, as distinct from its nominally advertised maximum, determines how many retrieved passages can be included in a prompt before quality degrades; RAG system designers frequently rely on needle-in-a-haystack-style empirical testing of their specific deployed model rather than trusting an advertised maximum context length figure, precisely because positional-encoding extrapolation quality varies substantially and unpredictably near the edges of a model's effective range.
Speculative decoding and KV-cache compression techniques both interact with context length in ways that compound with positional encoding choices: a longer effective context increases the KV-cache that must be held and searched, directly affecting the memory and bandwidth budget available for speculative decoding's draft-and-verify pipeline, making positional-encoding-driven context extension and KV-cache-efficiency techniques complementary components of a single long-context serving system rather than independent concerns. Mixture-of-Experts architectures, which typically apply positional encoding only within their shared, non-expert attention sublayers, inherit whatever extrapolation properties their chosen positional scheme provides without requiring any MoE-specific modification, since sparsity is applied to the feedforward computation rather than to the positional or attention mechanism itself.
## Future Research Directions
Learned or input-adaptive positional encoding schemes, in which the effective frequency allocation across dimensions is not fixed in advance but is itself optimized or dynamically adjusted based on the actual sequence length or content being processed, represent an active alternative to the current landscape of hand-designed interpolation formulas, aiming to remove the need for practitioners to manually select among linear interpolation, NTK-aware, and YaRN-style schedules for a given deployment scenario. Theoretical work continues to investigate exactly why relative encodings such as RoPE extrapolate substantially better than absolute encodings but still degrade at sufficiently extreme extension factors, seeking a more principled account of the interaction between rotation frequency distribution, attention head specialization, and the "lost in the middle" phenomenon than current largely empirical explanations provide.
Architectural alternatives that reduce reliance on positional encoding extrapolation altogether, including recurrent-memory-augmented Transformers, retrieval-based external memory mechanisms that avoid the need to hold an entire long context within the attention window at all, and state-space-model hybrids with fundamentally different long-range dependency mechanisms than dot-product attention, continue to be explored as complementary or alternative paths to long-context capability, motivated partly by the persistent quadratic compute cost that even a perfectly extrapolating positional encoding scheme cannot itself eliminate. Standardized, harder-to-game long-context evaluation protocols, extending beyond current needle-in-a-haystack and RULER-style synthetic benchmarks toward tasks requiring genuine long-range multi-step reasoning across realistic document distributions, remain an acknowledged priority for the field to keep pace with rapidly increasing advertised context-window sizes.
## Summary & Key Takeaways
Positional encoding determines whether a Transformer-based language model can meaningfully generalize beyond the sequence lengths it was trained on, and the shift from absolute positional embeddings, which have no defined behavior beyond their trained range, to relative encodings such as RoPE and ALiBi, which express positional information as a function of the offset between tokens, is the foundational reason modern long-context models are feasible at all. RoPE's rotation-based relative encoding pairs with a rich family of training-free context-extension techniques, including position interpolation, NTK-aware base rescaling, and YaRN's frequency-band-aware refinement, each trading off extrapolation range against positional resolution in different ways, while ALiBi achieves extrapolation through an inherently local-favoring linear attention bias that requires no rescaling at all but offers correspondingly less fine-grained control. None of these techniques eliminate the quadratic compute cost of full attention or fully resolve known failure modes such as the "lost in the middle" phenomenon, meaning that genuine production-grade long-context capability requires combining positional-encoding choices with continued pretraining at extended length, efficient attention implementations, careful KV-cache management, and rigorous retrieval-style evaluation rather than relying on any single technique in isolation.
Keywords: positional encoding, RoPE, rotary position embedding, ALiBi, position interpolation, NTK-aware scaling, YaRN, context window extension, long-context language models, needle in a haystack, RULER benchmark, lost in the middle, relative position encoding, KV-cache, context length extrapolation
---
## Appendix: Practical Labs
### Lab 1: RoPE Encodes Relative, Not Absolute, Position
This lab implements the core RoPE rotation and verifies its defining algebraic property: the attention score between a rotated query and key depends only on their relative offset, not on their absolute positions.
import numpy as np
def rope_inv_freq(dim, base=10000.0):
"""Rotation frequency for each of the dim/2 two-dimensional subspaces."""
return 1.0 / (base ** (np.arange(0, dim, 2) / dim))
def rope_rotate(x, pos, inv_freq):
"""Apply RoPE rotation to vector x at absolute position `pos`."""
angles = pos * inv_freq
cos = np.cos(angles)
sin = np.sin(angles)
x1 = x[0::2]
x2 = x[1::2]
out = np.empty_like(x)
out[0::2] = x1 * cos - x2 * sin
out[1::2] = x1 * sin + x2 * cos
return out
def attn_score(q, k):
return np.dot(q, k)
def test_rope_relative_position_invariance():
rng = np.random.default_rng(0)
dim = 16
inv_freq = rope_inv_freq(dim)
q0 = rng.normal(size=dim)
k0 = rng.normal(size=dim)
# score at absolute positions (5, 8), (105, 108), (505, 508) --
# same relative offset of 3 in every case
scores = []
for base_pos in [0, 100, 500]:
q = rope_rotate(q0, base_pos + 5, inv_freq)
k = rope_rotate(k0, base_pos + 8, inv_freq)
scores.append(attn_score(q, k))
print("scores at different absolute positions, same relative offset:", scores)
assert max(scores) - min(scores) < 1e-6, (
"RoPE attention score should depend only on relative offset"
)
# a different relative offset should generally give a different score
q = rope_rotate(q0, 5, inv_freq)
k_diff_offset = rope_rotate(k0, 20, inv_freq)
score_diff_offset = attn_score(q, k_diff_offset)
assert abs(score_diff_offset - scores[0]) > 1e-3
print("RoPE relative-position invariance test passed.")
if __name__ == "__main__":
test_rope_relative_position_invariance()### Lab 2: Position Interpolation Bounds the Out-of-Distribution Rotation Angle
This lab demonstrates the core justification for position interpolation: naive extrapolation pushes the slowest-rotating RoPE dimension far outside the angle range seen during training, while linearly rescaling positions keeps that angle exactly within the trained bound.
import numpy as np
def rope_inv_freq(dim, base=10000.0):
return 1.0 / (base ** (np.arange(0, dim, 2) / dim))
def position_interpolation_scale(pos, L_train, L_new):
"""Compress extrapolated positions back into the trained range."""
return pos * (L_train / L_new)
def test_position_interpolation_bounds_rotation_angle():
dim = 64
base = 10000.0
inv_freq = rope_inv_freq(dim, base)
L_train = 512
L_new = 4096 # 8x context extension
# the lowest-frequency (slowest-rotating) dimension pair is the hardest
# to extrapolate, since it has barely rotated by the end of training
slow_freq = inv_freq[-1]
max_angle_train = L_train * slow_freq
print("max rotation angle (slowest dim) seen during training:", max_angle_train)
naive_angle_at_Lnew = L_new * slow_freq
print("naive extrapolation angle at L_new:", naive_angle_at_Lnew)
assert naive_angle_at_Lnew > max_angle_train * 4, (
"naive extrapolation should push the slow dimension far outside its trained range"
)
pi_pos = position_interpolation_scale(L_new, L_train, L_new)
pi_angle_at_Lnew = pi_pos * slow_freq
print("position-interpolation angle at L_new:", pi_angle_at_Lnew)
assert abs(pi_angle_at_Lnew - max_angle_train) < 1e-9
positions = np.linspace(0, L_new, 200)
naive_angles = positions * slow_freq
pi_angles = position_interpolation_scale(positions, L_train, L_new) * slow_freq
frac_naive_oob = np.mean(naive_angles > max_angle_train + 1e-9)
frac_pi_oob = np.mean(pi_angles > max_angle_train + 1e-9)
print(f"fraction of range with out-of-trained-distribution angle: "
f"naive={frac_naive_oob:.2f} PI={frac_pi_oob:.2f}")
assert frac_naive_oob > 0.8
assert frac_pi_oob < 1e-6
print("Position-interpolation angle-bounding test passed.")
if __name__ == "__main__":
test_position_interpolation_bounds_rotation_angle()### Lab 3: NTK-Aware Scaling Preserves High-Frequency Resolution
This lab compares linear position interpolation against NTK-aware base rescaling, showing that NTK-aware scaling leaves the high-frequency (local resolution) component nearly untouched while still bounding the low-frequency component's out-of-range angle, unlike linear interpolation, which distorts every frequency uniformly.
import numpy as np
def rope_inv_freq(dim, base=10000.0):
return 1.0 / (base ** (np.arange(0, dim, 2) / dim))
def ntk_aware_base(base, L_train, L_new, dim):
scale = (L_new / L_train) ** (dim / (dim - 2))
return base * scale
def test_ntk_aware_scaling_preserves_high_freq_resolution():
dim = 64
base = 10000.0
L_train = 512
L_new = 4096
inv_freq_orig = rope_inv_freq(dim, base)
# linear position interpolation uniformly scales ALL frequencies by L_train/L_new
pi_scale = L_train / L_new
inv_freq_pi = inv_freq_orig * pi_scale
# NTK-aware rescales the base, stretching low-frequency dims far more
# than high-frequency ones
ntk_base = ntk_aware_base(base, L_train, L_new, dim)
inv_freq_ntk = rope_inv_freq(dim, ntk_base)
high_freq_orig = inv_freq_orig[0]
high_freq_pi = inv_freq_pi[0]
high_freq_ntk = inv_freq_ntk[0]
print("high-freq component: orig=%.6f linear-PI=%.6f NTK=%.6f" %
(high_freq_orig, high_freq_pi, high_freq_ntk))
rel_error_pi = abs(high_freq_pi - high_freq_orig) / high_freq_orig
rel_error_ntk = abs(high_freq_ntk - high_freq_orig) / high_freq_orig
print(f"high-freq relative distortion: linear-PI={rel_error_pi:.3f} NTK={rel_error_ntk:.3f}")
assert rel_error_ntk < rel_error_pi * 0.1, (
"NTK-aware scaling should distort high-frequency resolution far less than linear PI"
)
low_freq_orig = inv_freq_orig[-1]
max_angle_train = L_train * low_freq_orig
naive_angle = L_new * low_freq_orig
ntk_angle = L_new * inv_freq_ntk[-1]
print(f"trained max angle={max_angle_train:.4f} naive@L_new={naive_angle:.4f} "
f"NTK@L_new={ntk_angle:.4f}")
assert naive_angle > max_angle_train * 4
assert ntk_angle < naive_angle * 0.5
print("NTK-aware scaling test passed.")
if __name__ == "__main__":
test_ntk_aware_scaling_preserves_high_freq_resolution()### Lab 4: ALiBi Extrapolates Smoothly While a Learned Position Table Saturates
This lab contrasts a fixed-size learned absolute positional embedding table, which cannot represent any position beyond its trained range, with ALiBi's linear distance-based bias, which remains well-defined and strictly monotonic for arbitrarily large relative distances.
import numpy as np
def learned_pe_lookup(positions, table):
"""Simulate a learned absolute positional embedding table with clipping
for out-of-range positions (a common naive fallback)."""
L = table.shape[0]
clipped = np.clip(positions, 0, L - 1)
return table[clipped]
def alibi_bias(rel_distance, slope):
return -slope * np.abs(rel_distance)
def test_alibi_extrapolates_smoothly_while_learned_pe_saturates():
rng = np.random.default_rng(0)
L_train = 512
table = rng.normal(size=(L_train, 8))
far_positions = np.arange(L_train, L_train + 500)
learned_embeds = learned_pe_lookup(far_positions, table)
unique_rows = np.unique(learned_embeds, axis=0)
print("distinct learned-PE embeddings for 500 out-of-range positions:", len(unique_rows))
assert len(unique_rows) == 1, (
"a fixed-size learned absolute PE table cannot distinguish positions beyond L_train"
)
slope = 0.05
rel_distances = np.arange(0, 5000)
biases = alibi_bias(rel_distances, slope)
diffs = np.diff(biases)
print("ALiBi bias at distance 0, 512, 4999:", biases[0], biases[512], biases[4999])
assert np.all(diffs < 0), "ALiBi bias should be strictly monotonically decreasing"
assert len(np.unique(biases)) == len(biases), "every distance should map to a distinct bias"
print("ALiBi long-range extrapolation vs. learned-PE saturation test passed.")
if __name__ == "__main__":
test_alibi_extrapolates_smoothly_while_learned_pe_saturates()