Long Context — RoPE Scaling, YaRN, ALiBi

A longer input is not automatically a longer useful memory

A model can accept more positions yet fail to retrieve, combine or summarize information across them. Position handling, training distribution, attention computation, KV capacity and task reasoning are separate constraints. RoPE scaling changes one part of that system; it is not a free universal 4k→128k capability upgrade.

Prerequisites are causal attention, rotary coordinates, cache positions and fine-tuning. This lesson derives rotations, PI, a static NTK-style base rule and a version-bound YaRN construction; it contrasts ALiBi, distributed attention, streaming and recurrent hybrids. CPU programs inspect arithmetic and test a retrieval fixture. The complete pretrained extension procedure is unexecuted: no weights, long corpus, GPU kernel or service is downloaded or run.

Rotations, content and the correct frequency direction

For a two-dimensional pair, rotating a query at m and key at n gives:

The relative offset appears, but the score also depends on q and k. It is not simply cosine of distance. Orthogonal rotations preserve norms. For full rotary width d, theta_i=base^(-2*i/d) for i from0 to d/2−1. Small i is high frequency, with i0 one radian per token; large i is low frequency. At d128/base10000, pair0 has period2π≈6.283 tokens, while pair63 has period about54,410 tokens. The original text contradicted itself by reversing these bands and claiming high-index pairs see thousands of cycles.

A clock-hand analogy helps: fast hands complete many cycles within training, while slow hands explore only a small arc. New long offsets change the joint phase/content patterns seen by the network. Every sine/cosine remains in[−1,1]; extrapolation does not create a new numeric range outside those bounds. Slow phases can enter portions unobserved at a particular training length, but high-frequency behavior is not guaranteed safe and network failure is not determined by a single cosine.

Retained rotation-identity test

import numpy as np
 
def rotate(x, angle):
    c, s = np.cos(angle), np.sin(angle)
    return np.array([[c,-s],[s,c]]) @ x
 
q = np.array([1.,2.]); k = np.array([-1.,.5])
m, n, theta = 7, 12, .3
left = rotate(q,m*theta) @ rotate(k,n*theta)
right = q @ rotate(k,(n-m)*theta)
np.testing.assert_allclose(left,right,atol=1e-12)
np.testing.assert_allclose(np.linalg.norm(rotate(q,100*theta)),np.linalg.norm(q))
assert not np.isclose(q @ rotate(k,theta), np.cos(theta))
print('RoPE relative-offset identity and content dependence pass')

Random-vector rotation plots contain no learned training boundary. A curve cannot demonstrate a model “going wild past4096” merely because you label that x-coordinate. Such a plot illustrates frequencies; learned extrapolation requires a trained checkpoint and a defined evaluation task.

Position interpolation and the static base heuristic

For training length L and extension factor s, position interpolation uses effective position m/s, equivalent to replacing every frequency by theta/s. New positions below sL map inside the old coordinate span. Adjacent-token angles also shrink by s, changing local resolution. Smooth rotations do not ensure smooth task quality: the model sees more keys, different content and fractional positions it was not trained on. Evaluate original-length as well as extended-length performance.

A static NTK-inspired heuristic changes the base to base*s^(d/(d-2)), for even rotary width d greater than2. This keeps pair0 frequency unchanged and scales the lowest-frequency pair by1/s. It is not the same as dynamic NTK variants whose effective factor depends on current sequence length. For s4,d128 the base is about40,890, not a universal40,600. Use the actual rotary width if partial RoPE is used, and preserve the checkpoint's pairing convention; adjacent-pair and split-half layouts are not interchangeable without changing weights.

YaRN bands and the squared attention factor

YaRN combines interpolation and extrapolation across frequency bands. In the pinned Transformers4.57.1 implementation, a correction range is derived from the number of rotations within the original context, using beta_fast32 and beta_slow1 by default. Low-index/high-frequency pairs retain more original frequency; high-index/low-frequency pairs receive more interpolation. A linear ramp blends the transition.

The default returned attention factor is a=1+0.1*ln(s) for s greater than1 and scales cosine/sine values. Applying it to both rotated Q and K multiplies their dot product by a². The old article multiplied logits once by a and called that the same recipe. At s4, a≈1.13863 and a²≈1.29648. Implementations with explicit attention factors, partial rotary dimensions or model-specific mscale ratios need their own accounting.

The following NumPy implementation reproduces the pinned default parameter construction in float64 for inspection, with the stated finite, moderate configuration values. It is an arithmetic teaching function, not a complete validator for NaN, infinity or arbitrary array shapes. It is not a drop-in model conversion: no weights, caches or tokenizer are changed.

import math
 
def frequencies(d=128,base=10000.,factor=4.,length=4096):
    if type(d) is not int or d <= 2 or d%2 or base <= 1 or factor < 1 or length <= 0:
        raise ValueError('invalid rotary configuration')
    i=np.arange(d//2,dtype=float)
    original=base**(-2*i/d)
    pi=original/factor
    ntk=(base*factor**(d/(d-2)))**(-2*i/d)
    correction=lambda rotations: d*math.log(length/(rotations*2*math.pi))/(2*math.log(base))
    low=max(math.floor(correction(32)),0)
    high=min(math.ceil(correction(1)),d-1)
    if high==low: high+=.001
    ramp=np.clip((i-low)/(high-low),0,1)
    yarn=original*(1-ramp)+pi*ramp
    attention_factor=1+.1*math.log(factor)
    return original,pi,ntk,yarn,attention_factor
 
def rotary(x,positions,inv):
    x=np.asarray(x)
    if x.shape[-1] != 2*len(inv): raise ValueError('rotary width')
    angles=np.asarray(positions)[...,None]*inv
    pairs=x.reshape(x.shape[:-1]+(-1,2))
    a,b=pairs[...,0],pairs[...,1]
    return np.stack((a*np.cos(angles)-b*np.sin(angles),
                     a*np.sin(angles)+b*np.cos(angles)),axis=-1).reshape(x.shape)
 
o,pi,ntk,yarn,a=frequencies()
assert o[0]==1 and np.all(np.diff(o)<0)
assert np.isclose(ntk[0],o[0]) and np.isclose(ntk[-1],o[-1]/4)
assert np.all(yarn>=pi) and np.all(yarn<=o)
rng=np.random.default_rng(78)
x=rng.normal(size=(4,128))
p=np.array([0,4095,4096,32767])
np.testing.assert_allclose(np.linalg.norm(rotary(x,p,yarn),axis=1),np.linalg.norm(x,axis=1))
assert np.isclose((a*x[0])@(a*x[1]),a*a*(x[0]@x[1]))
print('last-pair period',2*np.pi/o[-1],'NTK base',10000*4**(128/126),'YaRN a squared',a*a)

For every changed positional scheme, check full-prefix versus cached decoding. If a dynamic frequency rule changes between prefill and decode, old cached keys may carry incompatible rotations. Invalidate/recompute or use the model's documented consistent cache scheme; do not silently wrap positions modulo a precomputed table. Test offsets near the old boundary, new boundary and explicit overflow, including multi-token cached chunks.

ALiBi: a bias, not a learned rotary clock

ALiBi adds a head-specific distance penalty to attention scores instead of rotating Q/K. For causal attention where j≤i, a head with slope m_h uses score_ij -= m_h*(i-j), with future positions masked. In the original method the slopes are fixed according to a construction, not learned per head. Different heads encode different recency biases.

For raw allowed scores all zero and slope .5 at query position2, keys0,1,2 get biases[−1,−.5,0]; softmax favors recent keys. This is an inductive bias, not proof that distant information becomes inaccessible or that long-context quality has “almost no degradation.” The paper's extrapolation experiment has specific model/data/length settings. Replacing a pretrained RoPE checkpoint with ALiBi at inference changes the learned operator and is not a supported extension recipe.

scores=np.array([-1.,-.5,0.])
p=np.exp(scores-scores.max());p/=p.sum()
assert p[2]>p[1]>p[0] and np.isclose(p.sum(),1)
# Causal future positions must remain impossible regardless of bias.
masked=np.array([-1.,-.5,0.,-np.inf])
prob=np.exp(masked);prob/=prob.sum()
assert prob[-1]==0
print('ALiBi illustrative causal probabilities',prob)

Complete unexecuted extension and evaluation recipe

1. Freeze the baseline identity. Save model/checkpoint/tokenizer hashes, model code version, original trained length, rotary width/base/layout, cache implementation, dtype and attention kernel. Confirm the checkpoint actually uses the positional method being changed; a config key accepted by one model family may be ignored or rejected by another.
2. Establish baselines. Use authorized local text, with document-grouped train/validation/test splits. Measure token-count-weighted held-out NLL at the original length, retrieval at varied positions, multi-hop tracing, aggregation and task-specific summaries/code QA. Record template and maximum answer budget, and verify that the tokenizer did not truncate input.
3. Implement one candidate. Change PI, the exact NTK variant or version-bound YaRN parameters in the model's documented rotary module. Rebuild sine/cosine/cache state consistently. Unit-test norm/relative-offset identities without attention scaling, then separately test the intended scaled logits. Compare full and cached logits over boundary cases before training.
4. Training-free evaluation first. Run identical examples at L,2L,4L and intended target length within memory limits. Record short-context regressions, first-token latency, decode latency and memory. A failed long run/OOM is a result, not permission to truncate examples and call the target supported.
5. Fine-tune only with explicit resources. Use genuinely long authorized documents and a declared mixture that preserves short-context tasks. State context length, valid tokens/update, accumulation, optimizer, learning rate, steps and total token budget. Gradient accumulation changes update batch, not the activation memory of one long microbatch. Tiled attention and checkpointing may help memory but do not guarantee a7B model fits80GB. Save/reload a controlled checkpoint and compare next-step state.
6. Ablate and decide. Compare no change, scaling only and scaling+fine-tuning under equal evaluation settings; track token/compute cost and multiple seeds where feasible. Do not claim “500steps restores quality” without measurements. A model with good local perplexity may still miss a fact in the middle or fail global synthesis.
7. Release boundary. Publish the supported length only with model-bound task evidence and explicit failures, including the old-length regression suite. Keep a rollback to original rotary configuration/checkpoint; changing frequencies back does not undo fine-tuned weights.

This procedure retains the original GPU objective without pretending it ran. The old PI/NTK/YaRN universal extension-factor table, near-zero-PPL promises and needle-recall percentages are removed as unsupported generalizations. Paper examples are not a transferable service warranty.

RULER-style fixture and independent task axes

Single-needle retrieval can be solved without integrating the whole context. A useful evaluation crosses total length, evidence depth, distractor similarity, number of relevant spans and answer length. Include at least three axes: exact retrieval, multi-hop variable tracing and aggregation. Add naturalistic summary/code tasks rather than assuming synthetic success transfers.

import random
 
def long_case(length,depth,seed=78):
    if length < 8 or not 0 <= depth <= 1: raise ValueError('fixture geometry')
    r=random.Random(seed)
    items=[('d'+str(i),r.randrange(10000,99999)) for i in range(length)]
    position=min(length-1,int(depth*length))
    items[position]=('needle',314159)
    return items,position
 
def exact_lookup(items,key):
    values=[v for k,v in items if k==key]
    return values[0] if len(values)==1 else None
 
for length in (16,64,256):
    for depth in (0,.25,.5,.75,1):
        rows,pos=long_case(length,depth)
        assert exact_lookup(rows,'needle')==314159
        recent=rows[-8:]
        if pos < length-8: assert exact_lookup(recent,'needle') is None
links={'a':'b','b':'c','c':'answer'}
node='a'
for _ in range(3): node=links[node]
assert node=='answer'
assert sum([2,5,2,9])==18
print('fixture location, eviction, multihop and aggregation controls pass')

These are oracle/fixture tests, not an LLM benchmark. In a model test render the exact rows into text with a tokenizer-defined length budget, retain answer keys outside the prompt, ask one fixed question and score exact normalized answers. For tracing, vary hop count and place links far apart; for aggregation, vary distractor values and demand a count/sum of a declared subset. An eviction control should fail when the decisive fact is absent; if it succeeds, inspect leakage from the question or label.

Report the denominator before the headline. For the fifteen lookup fixtures above, full exact lookup finds all fifteen planted values; the last-eight-row control finds five of fifteen: depths .5, .75 and 1 at length 16, and depth 1 at each larger length. These rows measure storage/lookup behavior, not tokenized context or model recall. The depth parameter 1 selects index length-1, not an out-of-bounds position.

For a real retrieval comparison, fix authorized documents, questions, gold supporting IDs and the answer budget. Report support recall on answerable questions, answer correctness on all eligible questions, and abstention on no-answer questions separately. Missing support and failing to use present support are different defects. A denied document is not an ordinary retrieval miss: filter it by trusted access rights before ranking and score the permission decision separately. Cross length with evidence position and distractor similarity; keep short-context controls. Use the same questions for full-context and retrieved-context variants so that a cheaper prompt cannot silently become an easier task. See the capstone integration contract for the saved records that join these checks.

LongBench v2 covers real-world long-context reasoning tasks and reports lengths in words in its abstract; do not silently relabel them tokens. HELMET covers diverse application-centric categories and finds weak transfer from needle scores. Model judges require calibration and task-specific rubrics. “Effective context length” is conditional on the tasks and success threshold, not a single timeless property of Claude/Gemini/Llama.

Systems: dense attention, ring blocks and streaming

A naive attention-score tensor is quadratic in sequence length; tiled online-softmax attention avoids materializing the full matrix but retains quadratic pairwise arithmetic. KV storage grows linearly in retained tokens for ordinary cache designs. At1M tokens,32layers,8KVheads,head width128 and BF16, conventional cache payload is about131.072GB per sequence. This excludes model weights and workspace; it is not a proof that every1M-token model must allocate that exact state.

Work the resource arithmetic before attempting the length. For batch B, layers N, retained tokens T, KV heads K, head width D and bytes per scalar S, a conventional cache stores 2 × B × N × T × K × D × S bytes; the factor two is K plus V. For dense score materialization the query-head count H—not K—sets B × H × T² × S bytes per layer. GQA reduces the former through K; it does not eliminate query-head attention work. A causal mask excludes future pairs semantically but does not halve an ordinary dense tensor allocation.

# Payload arithmetic only: allocate no tensor at the proposed long lengths.
def kv_bytes(tokens, batch=1, layers=32, kv_heads=8, width=128, scalar_bytes=2):
    values=(tokens,batch,layers,kv_heads,width,scalar_bytes)
    if any(type(v) is not int or v <= 0 for v in values):
        raise ValueError('positive integer dimensions required')
    return 2*batch*layers*tokens*kv_heads*width*scalar_bytes
 
def dense_score_bytes(tokens, query_heads=32, scalar_bytes=2):
    if any(type(v) is not int or v <= 0 for v in (tokens,query_heads,scalar_bytes)):
        raise ValueError('positive integer dimensions required')
    return query_heads*tokens*tokens*scalar_bytes  # batch one, one layer
 
assert kv_bytes(1_000_000)==131_072_000_000
assert kv_bytes(8192)==4*kv_bytes(2048)
assert dense_score_bytes(8192)==16*dense_score_bytes(2048)
assert dense_score_bytes(8192)==4*1024**3
assert kv_bytes(8192,kv_heads=32)==4*kv_bytes(8192,kv_heads=8)
# The larger capstone's proposed 12-layer / 4-KV-head / width-64 design:
assert kv_bytes(2048,layers=12,kv_heads=4,width=64)==25_165_824
assert kv_bytes(8192,layers=12,kv_heads=4,width=64)==100_663_296
for invalid in (0,-1,True,1.5):
    try: kv_bytes(invalid)
    except ValueError: pass
    else: raise AssertionError('invalid cache dimension accepted')
print('payload bytes checked; no million-token model or tensor allocated')

Thus the capstone's proposed 2,048→8,192 extension changes this cache payload from 24 MiB to 96 MiB per sequence at two bytes per element, before weights, activations, allocator overhead or batching. A larger positional limit does not reserve this memory for you, and a smaller cache does not establish useful recall. Record decimal GB versus binary GiB explicitly: 131.072 GB is about 122.07 GiB. Tiled attention can avoid the dense score allocation; it does not make the dense pair count linear. The formula is not a sizing model for compressed/latent caches or recurrent state.

Ring Attention distributes sequence blocks across devices and rotates K/V blocks while each device accumulates blockwise attention. Online softmax maintains running maxima, normalization sums and weighted outputs so blocks can combine into dense attention up to floating-point effects. Communication can overlap compute under favorable conditions; eight devices are not guaranteed faster than one, and comparing against a one-device OOM has no measured speed ratio. No claim is made that proprietary Gemini or Llama training internally uses this exact implementation.

StreamingLLM retains initial sink tokens plus a recent window. It can stabilize streaming language modeling in the studied configurations without retaining every historical token. Evicting an old fact makes that fact unavailable unless some separate retained state captures it. “Infinite stream” therefore does not mean infinite recall or near-zero loss on all tasks. The number of sinks, position handling and model support are version-specific, not a mode every runtime universally implements.

LongRoPE and architectural alternatives

LongRoPE searches nonuniform rescaling across dimensions and token positions, uses progressive extension and readjustment for short-context recovery. Its report's2048k context is an experimental result on particular LLaMA2/Mistral configurations, not a promise of arbitrary million-token fine-tuning for a rounding-error cost. Distinguish supported input length, fine-tuning length and effective task length.

Mamba uses input-dependent selective state-space recurrence and hardware-aware computation. Its recurrent state can remain fixed-size in streaming inference, trading an explicit full token cache for compressed state. Jamba interleaves attention, Mamba and MoE; attention layers still require their own cache. Its published single-device statement depends on precision/configuration and cannot become an unconditional256k/80GB fit guarantee.

RWKV offers recurrent and parallel-training formulations with linear-attention-like structure; avoid calling all versions a conventional fully linear-attention transformer. Griffin combines gated linear recurrence and local attention. Compare copying/retrieval, summaries, throughput and state size under matched budgets; no family universally wins complex recall.

Llama4 iRoPE/NoPE interleaves positional treatments; the Transformers 4.57.1 implementation gates rotary application per layer. NoPE does not mean there is no positional signal anywhere: causal masking, surrounding layers and learned state remain. A model-card maximum context is not evidence of near-lossless retrieval at every position. A configuration can be read without reproducing its benchmark or inferring closed-model internals.

Tradeoffs, recall and worked answers

1. Why can naive extrapolation fail? New offsets change joint rotary phase/content distributions; low-frequency pairs may leave arcs observed during training. The formula remains defined and bounded. Correct the old “high frequencies leave the cosine range first” premise.
2. What does PI cost? Map m→m/s to compress positions, changing adjacent-token resolution. It does not keep every aspect of the model input distribution unchanged; test quality at both lengths.
3. YaRN's two contributions? Band-dependent interpolation/extrapolation and an attention magnitude/temperature policy. In the pinned default, scaling both Q/K rotary outputs by a yields a² in their fully rotary dot product.
4. Eight-device ring versus one-device million-token attention? Ring can distribute state and overlap communication; naive full-score materialization may not fit one device. Memory feasibility is not a latency proof. Specify kernel, device memory, topology and benchmark before claiming faster.
5. When retrieve instead? When a question needs a small subset of a large authorized corpus and retrieval can find it. Long context can be valuable for global cross-references, but RAG/hierarchical summaries may also work. Compare on the actual task rather than asserting constant RAG cost or guaranteed global-context reasoning.
6. Three initial evaluation tasks? A varied-depth exact lookup tests access; a three-hop chain tests composition across spans; a filtered count/sum tests aggregation. Add short-context regressions and naturalistic summary/code tasks so success is not one-dimensional.

Training-free scaling is cheap to explore but not cost-free inference; fine-tuning adds data, compute and checkpoint maintenance; retrieval adds its own indexing and support-selection failures. Choose the smallest verified change that meets the task, not the largest advertised context number. Revisit the frequency ordering tomorrow and defend the scaled-QK factor from first principles in a week.

Primary sources

RoPE, PI, YaRN, Transformers4.57.1 implementation, ALiBi, LongRoPE, Ring Attention, StreamingLLM, Mamba, Jamba, RWKV, Griffin, RULER, LongBench v2, HELMET.

Previous: MoE · Next: safety evaluation

Optional reading

Explore these supplementary resources after the exercises. Video titles and availability have not been independently verified; use the primary sources above for technical claims.

• RoPE (Rotary Positional Embeddings) Explained
• Extending Context Length in LLMs (PI, NTK, YaRN)
• Extending Context Window of LLMs via Position Interpolation
• YaRN: Efficient Context Window Extension of LLMs