Batch normalisation averages over the examples in a batch. Layer normalisation averages over the features of one example instead, so it behaves the same for a batch of one and in training and evaluation. Transformers use it everywhere.

import torch
from torch import nn

x = torch.randn(2, 6)
mean = x.mean(dim=-1, keepdim=True)
var = x.var(dim=-1, keepdim=True,
            unbiased=False)
mine = (x - mean) / torch.sqrt(var + 1e-5)
theirs = nn.LayerNorm(6)(x)
print(torch.allclose(mine, theirs,
                     atol=1e-5))

nn.LayerNorm then multiplies by a learned weight and adds a learned bias, both starting at 1 and 0.

A residual connection computes x + f(x) instead of f(x). The gradient flows through the plus sign unchanged, so even a stack of forty blocks passes a useful gradient back to its first layer, and each block only has to learn a correction.

Write layer_norm(x, weight, bias, eps=1e-5) normalising over the last dimension by hand, with the biased variance. Write Residual(block), an nn.Module whose forward(x) returns x + self.block(x). Write deep_mlp(width, depth, residual) returning an nn.Sequential of depth blocks, each nn.Sequential(nn.LayerNorm(width), nn.Linear(width, width), nn.Tanh()), wrapped in Residual when residual is true.

The explanation opens here once every check passes.