File size: 1,276 Bytes
d7228c8 | 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 35 36 37 38 39 40 41 42 43 44 45 | """
Layer Normalization (LayerNorm) implemented from scratch.
"""
import torch
import torch.nn as nn
class CustomLayerNorm(nn.Module):
"""
Standard Layer Normalization implemented from PyTorch primitives.
"""
def __init__(self, dim: int, eps: float = 1e-5, bias: bool = True) -> None:
"""
Initializes LayerNorm parameters.
Args:
dim: Feature hidden dimension.
eps: Epsilon denominator value for stability.
bias: Whether to include learnable bias parameter beta.
"""
super().__init__()
self.eps = eps
self.weight = nn.Parameter(torch.ones(dim))
if bias:
self.bias = nn.Parameter(torch.zeros(dim))
else:
self.register_parameter("bias", None)
def forward(self, x: torch.Tensor) -> torch.Tensor:
"""
Forward pass applying mean-centering, variance scaling, and affine transform.
"""
mean = x.mean(-1, keepdim=True)
var = x.var(-1, keepdim=True, unbiased=False)
x_norm = (x - mean) / torch.sqrt(var + self.eps)
output = x_norm * self.weight
if self.bias is not None:
output = output + self.bias
return output
|