Skip to main content

Chapter 7.2 - 1st Layer Normalization

After adding the Residual Connection, we get our combined matrix. Now we need to normalize it to keep the values stable for the next layer!

1. Residual Output Matrix​

This is the input matrix we will be normalizing:

[  
  [ 1.133, -0.299,  0.503],  
  [-0.176,  0.911,  0.143],  
  [ 0.766,  0.781, -0.518],  
  [ 1.601, -0.669,  0.368],  
  [ 0.663,  0.779,  0.648]  
]

2. Calculating Mean​

Layer Normalization works row by row. Let's walk through the math for the first row: [1.133, -0.299, 0.503].

First, we calculate the mean (average) of the row.

mean = (1.133 - 0.299 + 0.503) / 3  
mean ≈ 0.446

3. Calculating Variance & Standard Deviation​

Next, we subtract the mean from each value in the row to find the difference:

1.133 - 0.446 =  0.687
-0.299 - 0.446 = -0.745
0.503 - 0.446 =  0.057

# Row Difference
[0.687, -0.745, 0.057]

We square these differences to calculate the Variance:

variance = average((x - mean)²)

[0.687², (-0.745)², 0.057²]
= [0.472, 0.555, 0.003]

variance = (0.472 + 0.555 + 0.003) / 3  
variance = 0.343
The Epsilon Constant (ε)

To prevent division by zero in the next step, we add a tiny constant called epsilon (ε) to the variance. Since ε is extremely small (like 1e-5), 0.343 + ε ≈ 0.343.

Finally, we find the Standard Deviation by taking the square root:

standard_deviation = √(variance + ε) 
standard_deviation = √0.343 ≈ 0.586

4. Final Normalization​

Now, we divide each value from our "Row Difference" by the standard deviation.

# [Difference / Standard Deviation]
[0.687 / 0.586, -0.745 / 0.586, 0.057 / 0.586]

= [1.17, -1.27, 0.10]
Gamma (γ) and Beta (β)

The final step is LayerNorm = (γ × normalized) + β. In most default setups before training, Gamma is 1 and Beta is 0. So multiplying by 1 and adding 0 leaves our values exactly the same!

Final Result for Top Row:

[1.17, -1.27, 0.10]

5. Complete LayerNorm Output​

After repeating this exact math for every row in the matrix, our final LayerNorm Output is ready!

[  
  [ 1.40, -1.52,  0.12],  
  [-1.15,  1.29, -0.14],  
  [ 0.71,  0.70, -1.41],  
  [ 1.31, -1.14, -0.17],  
  [-0.16,  0.91, -0.75]  
]

💻 Code Implementation​

Here is the exact PyTorch implementation for the concepts discussed above:

import torch
import torch.nn as nn
class LayerNorm(nn.Module):
    def __init__(self, emb_dim):
        super().__init__()
        self.eps = 1e-5
        self.scale = nn.Parameter(torch.ones(emb_dim))
        self.shift = nn.Parameter(torch.zeros(emb_dim))

    def forward(self, x):
        mean = x.mean(dim=-1, keepdim=True)
        var = x.var(dim=-1, keepdim=True, unbiased=False)
        norm_x = (x - mean) / torch.sqrt(var + self.eps)
        return self.scale * norm_x + self.shift