Michael Brenndoerfer | Data & AI, Private Equity, Technology

Michael Brenndoerfer

RMSNorm › RMSNorm vs LayerNorm: Empirical Comparison

Michael BrenndoerferUpdated June 7, 20254 min read

Language AI HandbookAll chapters

Part 15 of 67Transformer Blocks

Chapter 105RMSNorm › RMSNorm vs LayerNorm: Empirical Comparison

Reading section 9 of 15 · RMSNorm vs LayerNorm: Empirical Comparison · 4 minSection 9 of 15 · 4 min

Section 9 of 15 · 4 min read

RMSNorm › RMSNorm vs LayerNorm: Empirical ComparisonOne section · Sections in this chapter9 / 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

The theoretical efficiency of RMSNorm is clear, but does it maintain model quality? This is the empirical question. A normalization method that is 15% faster but converges to a worse solution would be a net loss. Let's compare the two normalizations on a small learning task to see their behavior during training.

The key finding from the original paper and subsequent work is that RMSNorm does not sacrifice quality. On standard language modeling benchmarks, models trained with RMSNorm match those trained with LayerNorm in perplexity and downstream task performance. The reason is the one we identified mathematically: in well-initialized transformers, activations are approximately centered, so the centering step that LayerNorm performs but RMSNorm skips has very little effect on the final normalized output.

In[22]:

Code

import torch
import torch.nn as nn


class TorchRMSNorm(nn.Module):
    """RMSNorm implemented in PyTorch."""

    def __init__(self, dim, eps=1e-6):
        super().__init__()
        self.eps = eps
        self.weight = nn.Parameter(torch.ones(dim))

    def forward(self, x):
        rms = torch.sqrt(torch.mean(x**2, dim=-1, keepdim=True) + self.eps)
        return self.weight * (x / rms)


class SimpleTransformerBlock(nn.Module):
    """Simplified transformer block for comparison."""

    def __init__(self, dim, norm_type="rmsnorm"):
        super().__init__()

        # Choose normalization
        if norm_type == "rmsnorm":
            self.norm1 = TorchRMSNorm(dim)
            self.norm2 = TorchRMSNorm(dim)
        else:
            self.norm1 = nn.LayerNorm(dim)
            self.norm2 = nn.LayerNorm(dim)

        # Simple self-attention substitute (linear projection)
        self.attn = nn.Linear(dim, dim)

        # Feed-forward network
        self.ffn = nn.Sequential(
            nn.Linear(dim, dim * 4),
            nn.GELU(),
            nn.Linear(dim * 4, dim),
        )

    def forward(self, x):
        # Pre-norm architecture
        x = x + self.attn(self.norm1(x))
        x = x + self.ffn(self.norm2(x))
        return x


class TinyTransformer(nn.Module):
    """Small transformer for testing normalizations."""

    def __init__(self, vocab_size, dim, n_layers, norm_type="rmsnorm"):
        super().__init__()
        self.embedding = nn.Embedding(vocab_size, dim)
        self.blocks = nn.ModuleList(
            [SimpleTransformerBlock(dim, norm_type) for _ in range(n_layers)]
        )

        if norm_type == "rmsnorm":
            self.final_norm = TorchRMSNorm(dim)
        else:
            self.final_norm = nn.LayerNorm(dim)

        self.output = nn.Linear(dim, vocab_size)

    def forward(self, x):
        x = self.embedding(x)
        for block in self.blocks:
            x = block(x)
        x = self.final_norm(x)
        return self.output(x)

In[23]:

Code

# Training comparison
def train_model(model, data, targets, epochs=100, lr=1e-3):
    """Train model and return loss history."""
    optimizer = torch.optim.Adam(model.parameters(), lr=lr)
    criterion = nn.CrossEntropyLoss()

    losses = []
    for epoch in range(epochs):
        optimizer.zero_grad()
        output = model(data)
        loss = criterion(output.view(-1, output.size(-1)), targets.view(-1))
        loss.backward()
        optimizer.step()
        losses.append(loss.item())

    return losses


# Create synthetic data
torch.manual_seed(42)
vocab_size = 100
seq_len = 32
batch_size = 16
dim = 64
n_layers = 4

data = torch.randint(0, vocab_size, (batch_size, seq_len))
targets = torch.randint(0, vocab_size, (batch_size, seq_len))

# Train both models
torch.manual_seed(42)
model_rms = TinyTransformer(vocab_size, dim, n_layers, "rmsnorm")
losses_rms = train_model(model_rms, data, targets)

torch.manual_seed(42)
model_ln = TinyTransformer(vocab_size, dim, n_layers, "layernorm")
losses_ln = train_model(model_ln, data, targets)

Out[24]:

Visualization

Line plot showing training loss curves for RMSNorm and LayerNorm converging to similar values over 100 epochs. — Training loss comparison between RMSNorm and LayerNorm on a small 4-layer transformer over 100 epochs. Both normalizations achieve similar convergence rates and final loss values, empirically confirming that removing mean centering does not degrade learning dynamics when activations are approximately centered.

Out[25]:

Console

Final training loss comparison:
  RMSNorm:   1.8955
  LayerNorm: 1.8958
  Difference: 0.0003

Both normalizations converge to similar loss values, confirming that RMSNorm doesn't sacrifice model quality for efficiency. In larger-scale experiments on language modeling benchmarks, RMSNorm has been shown to match or slightly exceed LayerNorm performance. The "slightly exceed" finding is interesting: it suggests that removing the mean-centering step, far from being a neutral change, may provide a small regularization benefit by reducing the number of operations that the gradient must backpropagate through.