Layer normalization (LN) standardizes activations across features within each example. Each example computes its own mean and standard deviation, while all examples share learned scale γ and shift β parameters. Unlike batch normalization, LN does not depend on batch size, making it useful when batches are small or when examples must be normalized independently. section 6.2.1 section 6.2.4
Use the following implementation rules when placing LN in PyTorch models:
- Place normalization between a linear or convolutional layer and its activation; learned γ and β let training undo or adjust standardization when needed. section 6.2.1
- Construct fully connected normalization with nn.LayerNorm([D]) for tensors shaped (B, D), where D is the feature count. section 6.2.5
- Pass all feature dimensions to normalize from right to left; for (B, C, W, H), use [C, W, H]. section 6.2.5
- Track convolutional spatial dimensions after pooling, because LN’s normalized shape must match the current channel, width, and height. section 6.2.5
- Prefer LN as a reliable starting point for GANs, where batch normalization can cause problems; replacing it with BN in the generator remains an optional performance experiment. section 9.1.3
Fully connected and convolutional LN
A fully connected block can use nn.Linear(D, n), nn.LayerNorm([n]), and nn.LeakyReLU(...). For convolutional blocks, a helper can pair nn.Conv2d with nn.LayerNorm([out_filters, W//(2**pool_factor), H//(2**pool_factor)]) and an activation. Padding preserves width and height through convolutions, reducing the shape bookkeeping required after pooling. section 6.2.5
LN inside a recurrent model
A stacked GRUCell model can maintain one LayerNorm(hidden_size) per recurrent layer and normalize each updated hidden state before passing it to the next layer. A prediction head can also place LayerNorm(hidden_size) between a linear layer and LeakyReLU. section 7.5.1
B, D = X.shape #1
μ = torch.zeros((B))
σ = torch.zeros((B))
for j in range(D):
μ += X[:, j] #2
μ /= D
for j in range(D):
σ += (X[:,j]-μ)*(X[:,j]-μ) #3
σ += 1e-5
σ = torch.sqrt(σ)
Shows the per-example feature mean and standard deviation calculation underlying LN, including the numerical stabilizer added before taking the square root. section 6.2.4
