rms_norm

rms_norm

Root Mean Square Layer Normalization.

From Zhang et al. (2019): https://arxiv.org/abs/1910.07467

RMSNorm normalizes using only the root mean square statistic, without centering (no mean subtraction). This is more efficient than LayerNorm and works well in practice for LLM-scale models.

RMSNorm

RMSNorm(normalized_shape, epsilon=1e-06, elementwise_affine=True)

Bases: Module

Root Mean Square Layer Normalization.

Parameters:

Name	Type	Description	Default
normalized_shape	int	Input shape from an expected input of size [*, normalized_shape]	required
epsilon	float	A value added to the denominator for numerical stability	1e-06
elementwise_affine	bool	If True, learns affine scaling parameters	True

Initialize RMSNorm.

Parameters:

Name	Type	Description	Default
normalized_shape	int	Feature dimension to normalize	required
epsilon	float	Small constant for numerical stability	1e-06
elementwise_affine	bool	Whether to learn scaling parameters	True

Source code in src/versatil/models/layers/normalization/rms_norm.py

def __init__(
    self,
    normalized_shape: int,
    epsilon: float = 1e-6,
    elementwise_affine: bool = True,
):
    """Initialize RMSNorm.

    Args:
        normalized_shape: Feature dimension to normalize
        epsilon: Small constant for numerical stability
        elementwise_affine: Whether to learn scaling parameters
    """
    super().__init__()
    if epsilon <= 0.0:
        raise ValueError(f"epsilon must be positive, got {epsilon}.")
    self.epsilon = epsilon
    self.elementwise_affine = elementwise_affine

    if self.elementwise_affine:
        self.weight = nn.Parameter(torch.ones(normalized_shape))
    else:
        self.register_parameter("weight", None)

forward

forward(x)

Apply RMS normalization.

Parameters:

Name	Type	Description	Default
x	Tensor	Input tensor (..., normalized_shape)	required

Returns:

Type	Description
Tensor	Normalized tensor of same shape

Source code in src/versatil/models/layers/normalization/rms_norm.py

def forward(self, x: torch.Tensor) -> torch.Tensor:
    """Apply RMS normalization.

    Args:
        x: Input tensor (..., normalized_shape)

    Returns:
        Normalized tensor of same shape
    """
    # Compute RMS: sqrt(mean(x^2))
    rms = torch.sqrt(torch.mean(x**2, dim=-1, keepdim=True) + self.epsilon)
    x_normed = x / rms

    if self.elementwise_affine:
        x_normed = x_normed * self.weight

    return x_normed