"""Normalization modules for X-former blocks."""
|
|
from typing import Dict, Optional, Tuple
|
|
import torch
|
|
|
def get_normalization(
|
normalization_type: str,
|
eps: Optional[float] = None,
|
partial: Optional[float] = None,
|
) -> Tuple[torch.nn.Module, Dict]:
|
"""Get normalization module and arguments given parameters.
|
|
Args:
|
normalization_type: Normalization module type.
|
eps: Value added to the denominator.
|
partial: Value defining the part of the input used for RMS stats (RMSNorm).
|
|
Return:
|
: Normalization module class
|
: Normalization module arguments
|
|
"""
|
norm = {
|
"basic_norm": (
|
BasicNorm,
|
{"eps": eps if eps is not None else 0.25},
|
),
|
"layer_norm": (torch.nn.LayerNorm, {"eps": eps if eps is not None else 1e-12}),
|
"rms_norm": (
|
RMSNorm,
|
{
|
"eps": eps if eps is not None else 1e-05,
|
"partial": partial if partial is not None else -1.0,
|
},
|
),
|
"scale_norm": (
|
ScaleNorm,
|
{"eps": eps if eps is not None else 1e-05},
|
),
|
}
|
|
return norm[normalization_type]
|
|
|
class BasicNorm(torch.nn.Module):
|
"""BasicNorm module definition.
|
|
Reference: https://github.com/k2-fsa/icefall/pull/288
|
|
Args:
|
normalized_shape: Expected size.
|
eps: Value added to the denominator for numerical stability.
|
|
"""
|
|
def __init__(
|
self,
|
normalized_shape: int,
|
eps: float = 0.25,
|
) -> None:
|
"""Construct a BasicNorm object."""
|
super().__init__()
|
|
self.eps = torch.nn.Parameter(torch.tensor(eps).log().detach())
|
|
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
"""Compute basic normalization.
|
|
Args:
|
x: Input sequences. (B, T, D_hidden)
|
|
Returns:
|
: Output sequences. (B, T, D_hidden)
|
|
"""
|
scales = (torch.mean(x.pow(2), dim=-1, keepdim=True) + self.eps.exp()) ** -0.5
|
|
return x * scales
|
|
|
class RMSNorm(torch.nn.Module):
|
"""RMSNorm module definition.
|
|
Reference: https://arxiv.org/pdf/1910.07467.pdf
|
|
Args:
|
normalized_shape: Expected size.
|
eps: Value added to the denominator for numerical stability.
|
partial: Value defining the part of the input used for RMS stats.
|
|
"""
|
|
def __init__(
|
self,
|
normalized_shape: int,
|
eps: float = 1e-5,
|
partial: float = 0.0,
|
) -> None:
|
"""Construct a RMSNorm object."""
|
super().__init__()
|
|
self.normalized_shape = normalized_shape
|
|
self.partial = True if 0 < partial < 1 else False
|
self.p = partial
|
self.eps = eps
|
|
self.scale = torch.nn.Parameter(torch.ones(normalized_shape))
|
|
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
"""Compute RMS normalization.
|
|
Args:
|
x: Input sequences. (B, T, D_hidden)
|
|
Returns:
|
x: Output sequences. (B, T, D_hidden)
|
|
"""
|
if self.partial:
|
partial_size = int(self.normalized_shape * self.p)
|
partial_x, _ = torch.split(
|
x, [partial_size, self.normalized_shape - partial_size], dim=-1
|
)
|
|
norm_x = partial_x.norm(2, dim=-1, keepdim=True)
|
d_x = partial_size
|
else:
|
norm_x = x.norm(2, dim=-1, keepdim=True)
|
d_x = self.normalized_shape
|
|
rms_x = norm_x * d_x ** (-1.0 / 2)
|
x = self.scale * (x / (rms_x + self.eps))
|
|
return x
|
|
|
class ScaleNorm(torch.nn.Module):
|
"""ScaleNorm module definition.
|
|
Reference: https://arxiv.org/pdf/1910.05895.pdf
|
|
Args:
|
normalized_shape: Expected size.
|
eps: Value added to the denominator for numerical stability.
|
|
"""
|
|
def __init__(self, normalized_shape: int, eps: float = 1e-5) -> None:
|
"""Construct a ScaleNorm object."""
|
super().__init__()
|
|
self.eps = eps
|
self.scale = torch.nn.Parameter(torch.tensor(normalized_shape**0.5))
|
|
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
"""Compute scale normalization.
|
|
Args:
|
x: Input sequences. (B, T, D_hidden)
|
|
Returns:
|
: Output sequences. (B, T, D_hidden)
|
|
"""
|
norm = self.scale / torch.norm(x, dim=-1, keepdim=True).clamp(min=self.eps)
|
|
return x * norm
|
