Normalization

Normalization refers to the transformation of intermediate representations so that their numerical scale can be controlled. Gradient based updates may require small updates along some directions while still making large errors along others. Example:

Here the features have very different variances hence the contours of the loss are highly elongated. In such scenario, Gradient descent tends to oscillate or descent slowly because the curvature is very different along different directions. After normalizing:

After standardizing the inputs, the loss contours become more balanced and closer to circular. Normalization improves the conditioning of the optimization problem that makes Gradient-based optimization more efficient $\to$ increase the speed of convergence.

Batch Normalization (BatchNorm)

We normalize the activations using the examples of a mini-batch. Suppose a particular feature $i$ has values $x_i^{(1)},x_i^{(2)},x_i^{(3)},\dots ,x_i^{(B)}$ for the $B$ examples in the batch. BatchNorm computes mean and variance as follows:

$${\mu }_i=\frac{1}{B}\sum _{b=1}^B{x}_i^{(b)}$$ $${\sigma }_i^2=\frac{1}{B}\sum _{b=1}^B(x_i^{(b)}-{\mu }_i{)}^2$$

each activation is then standardized:

$${\hat{x}}_i^{(b)}=\frac{x_i^{(b)}-{\mu }_i}{\sqrt{{\sigma }_i^2+ϵ}}$$

$ϵ$ is added to avoid division by zero error. After this we do:

$$y_i^{(b)}={\gamma }_i{\hat{x}}_i^{(b)}+{\beta }_i$$

where ${\gamma }_i\&{\beta }_i$ both are learned parameters (scale and shift).

Layer Normalization (LayerNorm)

Unlike BatchNorm, in Layer Normalization, the normalization is done within each individual example, across its feature dimensions. So we don’t look at any other examples in the batch.

For Transformer with $d_{model}$ and processing a sequence $X\in {\mathbb{R}}^{n\times d_{model}}$, and each token being represented as $x_1,x_2,x_3,\dots ,x_n$ we compute mean and variance:

$${\mu }_i=\frac{1}{d_{model}}\sum _{j=1}^{d_{model}}{x}_{ij}$$ $${\sigma }_i^2=\frac{1}{d_{model}}\sum _{j=1}^{d_{model}}(x_{ij}-{\mu }_i{)}^2$$

and then transform each feature as:

$${\hat{x}}_{ij}=\frac{x_{ij}-{\mu }_i}{\sqrt{{\sigma }_i^2+ϵ}}$$

to not restrict the features with same normalized statistics, we use two learned vectors:

$$\gamma ,\beta \in {\mathbb{R}}^{d_{model}}$$

and the final output is:

$$y_{ij}={\gamma }_j{\hat{x}}_{ij}+{\beta }_j$$

Implementation:

class LayerNormalization(nn.Module):
    
    def __init__(self,d_model, ep: float = 10**-6):
        super().__init__()
        self.ep = ep ## to avoid division by zero
        self.alpha = nn.Parameter(torch.ones(d_model)) # parameter to multiply
        self.bias = nn.Parameter(torch.zeros(d_model)) # parameter to add
        
    def forward(self,x):
        mean = x.mean(dim=-1,keepdim=True) # mean of embedding dimension + keepdim to preserve dimension 
        var = ((x-mean)**2).mean(dim=-1,keepdim=True)
        x_hat = (x-mean)/torch.sqrt(var+self.ep)
        return self.bias + self.alpha*x_hat # layer normalization

Root Mean Square Normalization (RMSNorm)

RMSNorm is a simplified normalization that removes the mean-centering operation from LayerNorm. Instead of normalizing by the variance around the mean, it normalizes the input by its root mean square.

$$y_{ij}={\gamma }_j\frac{x_{ij}}{\sqrt{{\sigma }^2}}+{\beta }_j$$

Implementation:

class RMSNormalization(nn.Module):
 
    def __init__(self, dim: int, eps: float = 1e-6):
        super().__init__()
        self.eps = eps
        self.weight = nn.Parameter(torch.ones(dim))
 
    def forward(self, x):
        rms = torch.sqrt(torch.mean(x.float() ** 2, dim=-1, keepdim=True) + self.eps)
        x_hat = x / rms.to(x.dtype)
        return self.weight * x_hat

Pre vs Post Norm

(a) Post-LN Transformer Layer; (b) Pre-LN Transformer Layer ¹

“The two main variants of residual connections, Pre-Norm and Post-Norm, each make distinct trade-offs between gradient vanishing and representation collapse. Pre-Norm applies normalization operations to the input before each residual block, effectively addressing the problem of gradient vanishing. However, it can also lead to the issue of collapse in deep representations, where hidden features in deeper layers become highly similar, diminishing the contribution of additional layers as their number increases. In contrast, Post-Norm applies normalization after the output of each residual block, reducing the influence of a hidden state on subsequent layers. This approach can alleviate the issue of representation collapse but also reintroduces the problem of vanishing gradients. The vanishing gradient and the representation collapse are like two ends of a seesaw, with these two variants making respective trade-offs between these issues. The key issue is that residual connections, including both Pre-Norm and Post-Norm variants, predefine the strength of connections between the output and input within a layer.” ²

Implementation:

class ResidualConnection(nn.Module):
    
    def __init__(self, d_model: int, dropout: float):
        super().__init__()
        self.dropout = nn.Dropout(dropout)
        self.norm = Normalization(d_model)
    
    def forward(self, x, sublayer):
        return x + self.dropout(sublayer(self.norm(x))) ## Pre-Norm
        return self.norm(self.dropout(sublayer(x))+x) ## Post-Norm

Footnotes

1. On Layer Normalization in the Transformer Architecture ↩

2. Hyper Connections ↩