File size: 1,349 Bytes
2db32a1 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 | # 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}"
|