RMS Normalization - PyTorch Basics

RMS Normalization

What you'll build: rmsnorm(x, gamma, eps), which rescales each feature vector by its root mean square and applies a learnable per-feature gain.

Background

RMSNorm (Zhang & Sennrich, 2019) is LayerNorm without the mean subtraction and without the bias. It is used in LLaMA, Mistral, Gemma and most recent LLMs: it is cheaper (one reduction instead of two) and works as well in practice, because the re-scaling, not the re-centring, is what stabilises training.

Formula

For each vector x along the last dimension (size d):

rms(x) = sqrt( mean(x²) + eps ) = sqrt( (1/d) * sum_i x_i² + eps )
RMSNorm(x) = x / rms(x) * gamma

eps goes inside the square root, as in PyTorch's nn.RMSNorm.

Comparison with LayerNorm

Property	LayerNorm	RMSNorm
Subtracts the mean	yes	no
Divides by	standard deviation	root mean square
Learnable scale (gamma)	yes	yes
Learnable shift (beta)	yes	no
Output mean is 0	yes	not in general

Inputs

• x: float array of shape (..., d) (any number of leading dimensions).
• gamma: float array of shape (d,).
• eps: small positive float (default 1e-6).

Expected Return Format

An array of the same shape as x.

Edge Cases and Conventions

• Normalise over the last axis only; every leading index is an independent vector.
• An all-zero vector stays all zero (the eps prevents division by zero).

Constraints

• Vectorised; no loops over rows.
• PyTorch track: use tensor ops; don't call F.rms_norm or nn.RMSNorm.

Follow-up

• Why can models get away with dropping the mean subtraction?
• Where does the norm sit in a pre-norm transformer block, and why did pre-norm replace post-norm?

Input: rmsnorm(np.array([[1.0, 2.0, 3.0]]), np.array([1.0, 1.0, 1.0]))
Output: array([[0.4629, 0.9258, 1.3887]])

Input: rmsnorm(np.array([[3.0, 4.0], [1.0, 1.0]]), np.ones(2))
Output: array([[0.8485, 1.1314], [1.    , 1.    ]])

import numpy as np

def rmsnorm(x: np.ndarray, gamma: np.ndarray, eps: float = 1e-6) -> np.ndarray:
    """
    Apply RMS Normalization over the last axis.

    Args:
        x: Input array of shape (..., features)
        gamma: Learnable scale of shape (features,)
        eps: Small constant added inside the square root

    Returns:
        Normalized array, same shape as x
    """
    # Your code here
    pass