components/normalization.py

# coding=utf-8
"""WiolaRMSNorm: RMS normalisation with a learned per-dimension input offset.

Implements Eq. (3) of the Wiola paper:

    WiolaRMSNorm(x) = gamma * (x + delta) / sqrt(mean((x + delta)^2) + eps)

Setting ``delta = 0`` recovers standard RMSNorm exactly, so this strictly
generalises RMSNorm. The offset shifts the *input before* normalisation,
changing the normalisation target itself rather than adding a post-norm bias.
"""

import torch
import torch.nn as nn


class WiolaRMSNorm(nn.Module):
    def __init__(self, hidden_size: int, eps: float = 1e-6):
        super().__init__()
        self.weight = nn.Parameter(torch.ones(hidden_size))  # gamma
        self.offset = nn.Parameter(torch.zeros(hidden_size))  # delta
        self.variance_epsilon = eps

    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
        input_dtype = hidden_states.dtype
        # Compute the normalisation in fp32 for numerical stability.
        z = hidden_states.to(torch.float32) + self.offset.to(torch.float32)
        variance = z.pow(2).mean(-1, keepdim=True)
        z = z * torch.rsqrt(variance + self.variance_epsilon)
        return (self.weight.to(torch.float32) * z).to(input_dtype)

    def extra_repr(self) -> str:
        return f"{tuple(self.weight.shape)}, eps={self.variance_epsilon}"