Member-only story
RMS Norm Explained
The math and code
The Root Mean Square Layer Normalization (RMSNorm) is a common type of normalization layer in Deep Learning. In the original paper the is defined as:
Where x is a vector of D dimension and x_i is the i-th component of vector x and γ is a learnable parameter.
However this normalization is applied over the layer dimension. If we are operating with computer vision then we would need to adapt it to compute the normalization over the channel dimension C of the x input tensor now with dims : [B, H, W, C].
The equation would change and the RMS term would only be computed over the C dimension:
We could interpret this visually as an image with 3 colors channels and a normalization tensor 1/RMS(x):
And then, the learnable parameter γ would only learn one value per channel:
This was all for the math and the conceptual understanding in visual terms. Now show me the code:
import torch
from torch import nn
class RMSNorm(nn.Module):
"""Root Mean Square Layer Normalization over channel dimension"""
def __init__(self, dim: int):
super().__init__()…
