Michael Brenndoerfer | Data & AI, Private Equity, Technology

Michael Brenndoerfer

RMSNorm › Implementation

Michael BrenndoerferUpdated June 7, 20255 min read

Language AI HandbookAll chapters

Part 15 of 67Transformer Blocks

Chapter 105RMSNorm › Implementation

Reading section 5 of 15 · Implementation · 5 minSection 5 of 15 · 5 min

Section 5 of 15 · 5 min read

RMSNorm › ImplementationOne section · Sections in this chapter5 / 15 ↓

Sections in this chapter15

Reading tools Definitions & inline links

Choose your expertise level to adjust how many terms are explained. Beginners see more tooltips, experts see fewer to maintain reading flow. Hover over underlined terms for instant definitions.

Article links

Make inline references clickable

With the mathematical foundation established, let's translate our understanding into working code. We'll build RMSNorm from first principles, compare it with LayerNorm, and observe how the two behave on real data. The implementation is straightforward because the formula is simple: compute the RMS, divide by it, and scale by the learned parameters.

Before we start, the position of $\epsilon$ in the formula matters. Some implementations add $\epsilon$ inside the square root, as $\sqrt{\frac{1}{d}\sum x_i^2 + \epsilon}$ , while others add it outside, as $\sqrt{\frac{1}{d}\sum x_i^2} + \epsilon$ . The two approaches differ in their behavior for very small inputs. Adding inside the square root is more common because it provides smoother numerical behavior: when all inputs are near zero, the gradient of $\sqrt{\epsilon}$ is finite, whereas the gradient of $0 + \epsilon$ is simply $1/\epsilon$ , which can be large. The LLaMA implementation adds inside the square root.

Building RMSNorm Step by StepLink Copied

The implementation follows directly from the formula. For each input vector, we need to:

Square all elements to get $x_i^2$ for each element
Compute the mean of these squares: $\frac{1}{d}\sum x_i^2$
Add $\epsilon$ for stability and take the square root: $\sqrt{\frac{1}{d}\sum x_i^2 + \epsilon}$
Divide the original input by this RMS value to get unit-RMS output
Multiply by the learnable scale parameters $\gamma$ to allow the network to adjust the scale

Let's implement both RMSNorm and LayerNorm as Python classes:

In[9]:

Code

class RMSNorm:
    """
    Root Mean Square Layer Normalization.

    Normalizes inputs by their RMS value without mean centering.
    """

    def __init__(self, dim, eps=1e-6):
        """
        Initialize RMSNorm layer.

        Args:
            dim: Dimension of the input features
            eps: Small constant for numerical stability
        """
        self.eps = eps
        self.weight = np.ones(dim)  # Learnable scale parameter (gamma)

    def __call__(self, x):
        """
        Apply RMSNorm to input.

        Args:
            x: Input tensor of shape (..., dim)

        Returns:
            Normalized tensor of the same shape
        """
        # Compute RMS along last dimension
        rms = np.sqrt(np.mean(x**2, axis=-1, keepdims=True) + self.eps)

        # Normalize and scale
        return self.weight * (x / rms)

    def _compute_rms(self, x):
        """Helper to compute RMS for analysis."""
        return np.sqrt(np.mean(x**2, axis=-1, keepdims=True) + self.eps)


class LayerNorm:
    """
    Standard Layer Normalization for comparison.
    """

    def __init__(self, dim, eps=1e-6):
        self.eps = eps
        self.weight = np.ones(dim)
        self.bias = np.zeros(dim)

    def __call__(self, x):
        mean = np.mean(x, axis=-1, keepdims=True)
        var = np.var(x, axis=-1, keepdims=True)
        x_norm = (x - mean) / np.sqrt(var + self.eps)
        return self.weight * x_norm + self.bias

Notice how much shorter the RMSNorm.__call__ method is compared to LayerNorm.__call__. LayerNorm needs to compute the mean, subtract it, compute the variance, and then normalize. RMSNorm skips straight to computing the mean of squares. The code simplicity directly reflects the mathematical simplicity.

Testing the ImplementationsLink Copied

With both classes defined, let's apply them to realistic input data and examine the output statistics. We'll use a tensor shaped like typical transformer activations: batch size 16, sequence length 128, hidden dimension 512.

In[10]:

Code

# Test both implementations
dim = 512
rms_norm_layer = RMSNorm(dim)
layer_norm_layer = LayerNorm(dim)

# Create test input
x_test = np.random.randn(16, 128, dim)  # (batch, seq_len, dim)

# Apply both normalizations
rms_output = rms_norm_layer(x_test)
ln_output = layer_norm_layer(x_test)

Out[11]:

Visualization

Three overlapping histograms showing input distribution, RMSNorm output, and LayerNorm output. This shows how each normalization maps the data. — Distribution of values before and after applying RMSNorm and LayerNorm to an input with mean approximately 1.5. LayerNorm shifts the distribution to be centered at zero while RMSNorm rescales without centering, resulting in a distribution that retains a positive mean offset. For typical transformer inputs with near-zero means, the two outputs would overlap almost perfectly.

Out[12]:

Console

Output statistics after normalization:

RMSNorm output:
  Mean:     -0.000319
  Std:      0.999999
  RMS:      0.999999

LayerNorm output:
  Mean:     0.000000
  Std:      0.999999
  RMS:      0.999999

The statistics reveal the key behavioral difference between the two normalizations. Look at the RMS values: RMSNorm produces output with RMS close to 1.0, which is exactly what it's designed to do. The mean, however, is not forced to zero.

LayerNorm tells a different story. Its output has near-zero mean (by design) and unit standard deviation. Because the mean is zero, the RMS and standard deviation are approximately equal, confirming our earlier mathematical derivation.

Despite these differences, both approaches accomplish the primary goal: they normalize the magnitude of activations to a consistent scale, which is what matters for stable training. The question of whether the mean is exactly zero turns out to matter far less than the question of whether the magnitude is controlled.