RoPE scaling: extending a model past its trained context length

The problem, from first principles

[04-seq-modeling/02-embeddings/06-rope]'s RoPE encodes position as a ROTATION angle, growing linearly with position index. A model trained ONLY on sequences up to some trained_max_len never sees rotation angles beyond whatever that length produces; when run on a LONGER sequence at inference time, positions past trained_max_len produce rotation angles the model's learned weights have simply never encountered during training, and quality typically degrades sharply, exactly the practical problem this question's title names.

Chen et al. (2023, "Position Interpolation") proposed a strikingly simple, training-free (or cheap-to-fine-tune) fix: instead of letting positions run up to the new, longer target_max_len directly, COMPRESS them first, dividing every position by scale_factor = target_max_len / trained_max_len before computing RoPE's angles. The position at the very end of the EXTENDED sequence then produces the exact SAME angle the position at the very end of the ORIGINAL trained range used to produce, so every angle the model now encounters falls back WITHIN the range it was actually trained on, even though the sequence itself is now much longer. The tradeoff: positions are now more densely packed into that same angular range (less angular separation between adjacent tokens than during training), which is why Position Interpolation-style scaling often benefits from a short period of further fine-tuning at the new length, even though it can work reasonably well with NO fine-tuning at all.

From theory to code

Implement compute_rope_angles_scaled(seq_len, dim, scale_factor): [06-rope]'s compute_rope_angles, with every position divided by scale_factor before the angle computation.

Constraints

• scale_factor=1.0 must reduce EXACTLY to [06-rope]'s unscaled compute_rope_angles.
• Positions are DIVIDED by scale_factor (compressing the range), never multiplied (which would do the opposite, expanding it).
• The per-dimension frequency formula (10000^(-2i/dim)) is unchanged from [06-rope]; only the POSITION fed into it is scaled.
• The relative-position dot-product invariance property [06-rope] established must still hold after scaling: two pairs of vectors with the same relative offset must still produce the same dot product, regardless of their absolute positions.

Hints