Layer Norm & Residuals

Why Do We Need These?

As we stack more layers, two problems can make training difficult:

1. Activation scale: Values can grow or shrink as each layer transforms them.
2. Gradient flow: Repeated transformations can shrink the learning signal that reaches early layers.

Residual connections provide a direct path for the signal. Layer normalization controls the values presented to a transformation. Both help deep networks train.

Residual Connections

The problem: gradients vanish in deep networks

In a deep network, the gradient must flow backward through every layer’s transformation. If each layer shrinks the gradient even slightly (say by a factor of 0.9), the signal disappears fast. It’s the same vanishing gradient problem we saw with RNNs, but across layers instead of across time.

Toggle between “With Residuals” and “Without Residuals” below. Try increasing the layer count to see how the problem compounds:

What to look for: In “Without Residuals” mode, watch the bars shrink from right to left. Now switch to “With Residuals” and compare the earlier layers. The shortcut gives gradients a direct route around a transformation; this helps preserve the signal, but does not guarantee that every deep network trains successfully.

The solution: add a shortcut

The fix is surprisingly simple. Instead of replacing the input with the layer’s output, we add the layer’s output to the original input:

Analogy: Think of a highway with an express lane. Normal traffic (the gradient) must pass through every toll booth (layer transformation), losing speed each time. A residual connection adds an express lane that bypasses all booths. Even if the toll booths grind traffic to a halt, the express lane guarantees flow. Where the analogy breaks: unlike a real highway with separate lanes, the residual adds the express-lane signal and the transformed signal together, so the layer’s contribution isn’t lost.

Concrete example: Suppose the input vector is [1.0, 2.0] and the layer transforms it to [0.1, 0.3]:

• Without residual: output = [0.1, 0.3] (layer’s transformation replaced the input)
• With residual: output = [1.0 + 0.1, 2.0 + 0.3] = [1.1, 2.3] (original input preserved, layer adds a small adjustment)

The gradient through x + f(x) is 1 + f'(x). Even if f'(x) is tiny, the gradient is still approximately 1.

Layer Normalization

The problem: activations drift and explode

Residual connections fix gradient flow backward. But there’s a separate problem going forward: as activations pass through layer after layer, their distribution shifts. Early layers might output values centered around 0 with a standard deviation of 1. But by layer 10, the mean could be 87 and the standard deviation 4312. Later layers must constantly adapt to a moving target, and eventually the numbers overflow entirely.

Toggle “No Normalization” vs “With LayerNorm” below. Increase the layer count and watch what happens to the standard deviation:

What to look for: Without normalization, the σ bars grow exponentially. With 15+ layers, values explode past 100. Now switch to “With LayerNorm” and notice every layer stays near σ = 1.0, no matter how many layers you add.

See the normalization steps

Now that you’ve seen why we need normalization, here’s how it works. Click through the tabs to see each step. Adjust γ and β to see how the learned parameters affect the final output:

What to look for: Watch the μ and σ stats at the bottom. After step 2 (“Centered”), μ drops to 0. After step 3 (“Normalized”), σ becomes 1. The final step lets the network learn whatever scale and shift it needs, but starting from that stable baseline.

How it works

Layer normalization takes each vector and re-centers it to have a mean of 0 and a standard deviation of 1. Then it applies a learned scale (γ) and shift (β) so the network can still represent whatever distribution it needs.

Concrete example: Suppose a token’s activation vector after layer 5 is [2, −3, 1, 4].

1. Compute mean: (2 + (−3) + 1 + 4) / 4 = 1.0
2. Compute std: √(((2−1)² + (−3−1)² + (1−1)² + (4−1)²) / 4) = √(1 + 16 + 0 + 9)/4 = √6.5 ≈ 2.55
3. Normalize: [(2−1)/2.55, (−3−1)/2.55, (1−1)/2.55, (4−1)/2.55] = [0.39, −1.57, 0, 1.18]
4. Scale + shift: The network learns γ and β. If γ = [2, 2, 2, 2] and β = [0.5, 0.5, 0.5, 0.5]: output = [1.28, −2.64, 0.5, 2.86]

No matter how wild the input was (mean 87, std 4312 at layer 10), normalization resets it to a predictable range. The learned scale and shift give the network flexibility to represent whatever distribution it actually needs, but starting from a stable baseline.

class LayerNorm:
    def __init__(self, d_model, eps=1e-6):
        self.gamma = np.ones(d_model)   # Learnable scale
        self.beta = np.zeros(d_model)   # Learnable shift
        self.eps = eps

    def forward(self, x):
        mean = x.mean(axis=-1, keepdims=True)
        var = x.var(axis=-1, keepdims=True)

        x_norm = (x - mean) / np.sqrt(var + self.eps)
        return self.gamma * x_norm + self.beta

Why not batch normalization?

You might know Batch Normalization from CNNs. It normalizes across the batch dimension (across different examples). But for transformers, that doesn’t work:

The key reason: transformers process sequences of different lengths. For a batch where some sequences are 10 tokens and others are 100, what does “position 50” even mean for the short sequences? LayerNorm sidesteps this entirely by normalizing each token independently, within its own dimensions.

class TransformerBlock:
    def __init__(self, d_model, n_heads, d_ff):
        self.attn = MultiHeadAttention(d_model, n_heads)
        self.ffn = FeedForward(d_model, d_ff)
        self.ln1 = LayerNorm(d_model)
        self.ln2 = LayerNorm(d_model)

    def forward(self, x, mask=None):
        # Sub-layer 1: Attention with residual + norm
        x = x + self.attn(self.ln1(x), mask)

        # Sub-layer 2: FFN with residual + norm
        x = x + self.ffn(self.ln2(x))

        return x

class RMSNorm:
    def __init__(self, d_model, eps=1e-6):
        self.gamma = np.ones(d_model)
        self.eps = eps

    def forward(self, x):
        rms = np.sqrt(np.mean(x**2, axis=-1, keepdims=True) + self.eps)
        return self.gamma * x / rms

Summary

Test Your Understanding

What’s Next

We have vectors, sequence memory, and ways to keep deep transformations trainable. Next, each token will gather information directly from other tokens to update its representation.

Next: Self-Attention →