§4.2–4.3Layer Normalization & the GELU Feed-Forward Network

The backbone works but its blocks are hollow. This unit builds components 2 and 3–4 of the chapter plan: layer normalization, which keeps activations statistically tame as depth grows, and the GELU feed-forward network, the transformer block’s per-position workhorse.

4.2 Normalizing activations with layer normalization

Training deep networks is fragile: as signals and gradients pass through many layers, they can shrink toward zero or blow up — vanishing vanishing gradient Gradients shrinking toward zero as they propagate backward through many layers, starving early layers of learning signal. defined in ch. 4 — open in glossary or exploding gradients. The result is unstable training dynamics: the network struggles to find weights that minimize the loss. Layer normalization layer normalization Normalizing each token's activations to mean 0 and unit variance across the feature dimension, then applying a learnable scale and shift — stabilizes deep-network training. defined in ch. 4 — open in glossary tames this by adjusting the activations of a layer to have mean 0 and variance 1 (unit variance) — every token vector, individually, at every normalized layer.

$\mathrm{LayerNorm}(x) = \gamma \odot \frac{x - \mu}{\sqrt{\sigma^2 + \epsilon}} + \beta$

$\mu = \tfrac{1}{n}\textstyle\sum_i x_i \qquad \sigma^2 = \tfrac{1}{n}\textstyle\sum_i (x_i - \mu)^2$

Watch it work on real numbers. A toy layer (Linear(5, 6) + ReLU) produced the output below; step through the normalization exactly as the book’s p. 53 figure draws it:

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

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

4.3 The feed-forward network with GELU activations

The second half of every transformer block is a small per-position MLP. Its activation is not ReLU but GELU gelu Gaussian Error Linear Unit — the smooth activation ≈ 0.5x(1+tanh(√(2/π)(x+0.044715x³))) used in GPT's feed-forward layers. defined in ch. 4 — open in glossary (Gaussian Error Linear Unit), used here via its tanh approximation:

$\mathrm{GELU}(x) \approx 0.5 \cdot x \cdot \left(1 + \tanh\!\left[\sqrt{\tfrac{2}{\pi}}\left(x + 0.044715\,x^3\right)\right]\right)$

Worked example: $\mathrm{GELU}(-1) \approx -0.159$ — small but not zero (ReLU would output exactly 0), and $\mathrm{GELU}(1) \approx 0.841$. Explore both curves — hover to compare, and note the smooth dip below zero around $x \approx -0.75$:

class GELU(nn.Module):
  def __init__(self):
      super().__init__()

  def forward(self, x):
      return 0.5 * x * (1 + torch.tanh(
          torch.sqrt(torch.tensor(2.0 / torch.pi)) *
          (x + 0.044715 * torch.pow(x, 3))
      ))

class FeedForward(nn.Module):
  def __init__(self, cfg):
      super().__init__()
      self.layers = nn.Sequential(
          nn.Linear(cfg["emb_dim"], 4 * cfg["emb_dim"]),  #A
          GELU(),
          nn.Linear(4 * cfg["emb_dim"], cfg["emb_dim"]),  #B
      )

  def forward(self, x):
      return self.layers(x)

Two of the four building blocks are done. Next: shortcut connections and the assembled transformer block.