AI Engineer Exhaustive
MediumPlusgpt-oss

YaRN RoPE

gpt-oss

Medium

NTK-by-parts inverse frequency scaling that extends RoPE to long contexts.

Independent study note. Written from the public paper and official code. This is not TensorTonic Plus and does not reproduce their exercises, starter code, or tests. For the official version, subscribe on TensorTonic.

Overview

gpt-oss applies rotary position embeddings to every attention layer and stretches the dense layers to a 131,072-token context with YaRN (Peng et al., 2023). The model card names RoPE and YaRN and gives the target length. It does not give (\theta), the scale factor, or the NTK cutoffs. Those live in ModelConfig and RotaryEmbedding._compute_concentration_and_inv_freq in the official reference model.

The trained rotary basis is the 4,096-token “initial context.” At inference the inverse frequencies of the low-frequency pairs are interpolated toward a (32\times) longer horizon ((4096\times 32=131072)), high-frequency pairs keep their original (\theta), and a YaRN concentration multiplies (\cos) and (\sin). Sliding-window layers still use the same rotary module; they simply mask keys. Positions stay absolute.

How it works

Head width is (d_h=64), so there are (d_h/2=32) frequency pairs. The unscaled RoPE base is (\theta_{\mathrm{base}}=150{,}000):

f_i = \theta_{\mathrm{base}}^{2i/d_h},\qquad i=0,\ldots,31,

and the ordinary inverse frequency is (1/f_i). When the YaRN scale (s=32>1), the code forms two candidates

\omega^{\mathrm{interp}}_i=\frac{1}{s\,f_i},\qquad \omega^{\mathrm{extra}}_i=\frac{1}{f_i}

and blends them with an NTK-by-parts ramp. Let (L=4096), (\alpha=1), (\beta=32). Dimension-space cutoffs are

r_{\mathrm{low}}=\frac{d_h}{2}\frac{\log\bigl(L/(\beta\cdot 2\pi)\bigr)}{\log\theta_{\mathrm{base}}},\qquad r_{\mathrm{high}}=\frac{d_h}{2}\frac{\log\bigl(L/(\alpha\cdot 2\pi)\bigr)}{\log\theta_{\mathrm{base}}}.

A linear ramp (u_i=(i-r_{\mathrm{low}})/(r_{\mathrm{high}}-r_{\mathrm{low}})) is clipped to ([0,1]). The mix mask is (m_i=1-\mathrm{clip}(u_i)):

\tilde{\omega}_i = (1-m_i)\,\omega^{\mathrm{interp}}_i + m_i\,\omega^{\mathrm{extra}}_i.

Small (i) (high-frequency pairs) have (m_i=1) and keep the original (\theta). Large (i) (low-frequency pairs) have (m_i=0) and use (\theta/s). The mid band interpolates. That is the “NTK-by-parts” recipe the YaRN paper describes and that the official comment points at (https://arxiv.org/abs/2309.00071).

YaRN’s attention-temperature / concentration factor is

c = 0.1\log s + 1.

For token index (t) and pair (i),

\cos_{t,i}=c\cos(t\tilde{\omega}_i),\qquad \sin_{t,i}=c\sin(t\tilde{\omega}_i).

Application is a half-rotation, not the interleaved even/odd form. Split the last axis of a head in half, (x=(x_1,x_2)), then

x'_1 = x_1\cos - x_2\sin,\qquad x'_2 = x_2\cos + x_1\sin.

The same ((\cos,\sin)) tables rotate (Q) and (K). (V) is not rotated. If (s=1), concentration is 1 and (\tilde{\omega}_i=1/f_i) (plain RoPE). The public checkpoint uses (s=32).

The paper specifies only “YaRN to 131,072.” The numbers (150{,}000), (4096), (32), (\alpha=1), (\beta=32), and the half-split rotation are inferred from official code.

Official code

gpt_oss/torch/model.py: ModelConfig (rope_theta, initial_context_length, rope_scaling_factor, rope_ntk_alpha, rope_ntk_beta). RotaryEmbedding._compute_concentration_and_inv_freq implements the ramp; _compute_cos_sin applies (c); _apply_rotary_emb does the half-split rotate. AttentionBlock constructs one RotaryEmbedding per layer with those config values.

Watch-outs

Sources