Modernize: RoPE and RMSNorm

The model from the previous lesson uses a learned position embedding and LayerNorm, the 2017 choices. Current LLMs (LLaMA, Mistral, Qwen) use rotary position embeddings (RoPE) and RMSNorm. This lesson makes both swaps and explains why each is a strict simplification, not just a substitution.

1. RMSNorm

LayerNorm centers and scales: it subtracts the mean and divides by the standard deviation, then applies a learned gain and bias. RMSNorm drops the centering and the bias, normalizing only by the root mean square:

where $g \in \mathbb{R}^d$ is a learned gain and $\varepsilon$ guards against division by zero. It has fewer operations and parameters than LayerNorm and matches its quality in practice, which is why it became standard[1].

import torch
import torch.nn as nn

class RMSNorm(nn.Module):
    def __init__(self, dim, eps=1e-5):
        super().__init__()
        self.eps = eps
        self.gain = nn.Parameter(torch.ones(dim))

    def forward(self, x):
        rms = x.pow(2).mean(dim=-1, keepdim=True).add(self.eps).rsqrt()
        return x * rms * self.gain

Replace every nn.LayerNorm(n_embd) in the block and the final norm with RMSNorm(n_embd).

2. Rotary position embeddings

A learned position embedding stores one vector per absolute position and adds it once at the input. RoPE instead injects position inside attention by rotating each query and key by an angle proportional to its position. The dot product of a query at position $m$ and a key at position $n$ then depends only on the relative offset $m - n$, which is what gives RoPE its length-generalization behavior[2], derived in positional encodings.

Group the head dimension into pairs and rotate each pair $(x_{2i}, x_{2i+1})$ by angle $m\,\theta_i$ with $\theta_i = 10000^{-2i/d_h}$:

def build_rope_cache(seq_len, head_dim, base=10000.0, device="cpu"):
    theta = 1.0 / (base ** (torch.arange(0, head_dim, 2, device=device) / head_dim))
    pos = torch.arange(seq_len, device=device)
    freqs = torch.outer(pos, theta)            # (T, head_dim/2)
    return torch.cos(freqs), torch.sin(freqs)

def apply_rope(x, cos, sin):
    # x: (B, n_head, T, head_dim); rotate even/odd pairs
    x1, x2 = x[..., 0::2], x[..., 1::2]
    T = x.shape[-2]
    c, s = cos[:T].unsqueeze(0).unsqueeze(0), sin[:T].unsqueeze(0).unsqueeze(0)
    out1 = x1 * c - x2 * s
    out2 = x1 * s + x2 * c
    return torch.stack((out1, out2), dim=-1).flatten(-2)

In the attention head, apply apply_rope to the query and key after their projections and before the score, and delete the position_embedding table from the model. Position now enters through the rotation, not through an additive vector.

3. Why these are simplifications

Both changes remove parameters and operations. RMSNorm drops the mean subtraction and the bias; RoPE removes the learned position table and the $T \times d$ addition, replacing them with a fixed, parameter-free rotation. The model gets smaller and generalizes to longer contexts than it trained on, while the training loop is unchanged.

Exercises

Exercise 1. RMSNorm drops LayerNorm's mean subtraction. Show that RMSNorm is invariant to scaling its input by a positive constant: $\text{RMSNorm}(cx) = \text{RMSNorm}(x)$ for any $c > 0$ (ignoring the $\varepsilon$ guard). Then explain why LayerNorm, which subtracts the mean, is not invariant to adding a constant to every coordinate in the same way, contrasting the two normalizers' symmetries.

Exercise 2 (code). Implement RMSNorm in numpy and verify both its definition and the scale-invariance from Exercise 1 on a random vector. Use a unit gain so the output is the pure normalizer.

import numpy as np

rng = np.random.default_rng(0)

def rmsnorm(x, gain, eps=1e-5):
    rms = np.sqrt((x ** 2).mean(axis=-1, keepdims=True) + eps)
    return x / rms * gain

d = 8
x = rng.standard_normal(d)
gain = np.ones(d)

y = rmsnorm(x, gain)

# Definition: each coordinate is x_i divided by the root-mean-square of x.
rms = np.sqrt((x ** 2).mean())
assert np.allclose(y, x / rms, atol=1e-4)

# Scale-invariance: scaling the input by c > 0 leaves the output fixed.
y_scaled = rmsnorm(3.0 * x, gain)
assert np.allclose(y, y_scaled, atol=1e-4)

print("rms of x:", round(float(rms), 4))
print("max |y - y_scaled|:", float(np.abs(y - y_scaled).max()))

Exercise 3 (code). The point of RoPE is that the dot product of a rotated query at position $m$ and a rotated key at position $n$ depends only on the relative offset $m - n$. Verify this in numpy: rotate the same query and key vectors by two different absolute-position pairs that share an offset, and check the dot products agree. Use a single pair of head dimensions (one rotation angle) for clarity.

import numpy as np

rng = np.random.default_rng(1)

def rotate(v, angle):
    # rotate a 2D vector (one RoPE pair) by `angle` radians
    c, s = np.cos(angle), np.sin(angle)
    R = np.array([[c, -s], [s, c]])
    return R @ v

theta = 0.7                      # the per-pair frequency
q = rng.standard_normal(2)
k = rng.standard_normal(2)

def rope_score(q, k, m, n):
    return rotate(q, m * theta) @ rotate(k, n * theta)

# Two absolute-position pairs with the same offset m - n = 2.
s1 = rope_score(q, k, m=5, n=3)
s2 = rope_score(q, k, m=9, n=7)
assert np.isclose(s1, s2)

# A different offset gives a different score.
s3 = rope_score(q, k, m=5, n=4)   # offset 1
assert not np.isclose(s1, s3)

print("offset-2 scores:", round(float(s1), 4), round(float(s2), 4))
print("offset-1 score:", round(float(s3), 4))

The architecture is now current. The next lesson takes this pretrained base and runs a minimal preference-tuning pass, the alignment step from E4 at toy scale.