Residual Connections

Your program will wire a sublayer into a network the way a transformer does: by adding its output to its input instead of replacing the input. You will build the three arrangements side by side (pre-norm, post-norm, and none) and watch what repeated layers do to a vector.

The idea

A plain layer computes x -> f(x). A residual layer computes

x -> x + f(x)

The sublayer f (attention or the FFN) now only has to produce a correction. If the best thing a layer can do is nothing, it learns f = 0 and passes x through untouched. Without the skip it would have to learn the identity map, which is hard to represent exactly.

Why deep stacks need it

Backpropagation through y = x + f(x) gives dy/dx = 1 + df/dx. The 1 is a direct path from the loss back to every earlier layer, so even if df/dx is tiny the gradient survives 24 or 96 layers. Remove the skip and the gradient is a product of 96 factors, each of them possibly small, which is the vanishing gradient problem the ResNet paper (He et al., 2016) fixed.

It also gives a useful picture: the vector flowing down the stack is a residual stream. Every attention layer and every FFN reads from it and writes an addition back; nothing ever overwrites it.

Where the normalisation goes

LayerNorm (lesson 3) can go in two places.

post-norm (Vaswani 2017):   x -> LN( x + f(x) )
pre-norm  (GPT-2 on):       x -> x + f( LN(x) )
plain (no skip, for contrast):  x -> f( LN(x) )

In pre-norm the stream itself is never normalised: the skip path is an untouched wire from input to output, so gradients flow unscaled and deep stacks train stably. The price is that the stream's scale grows as each layer adds to it, which is why GPT-2 adds one more LayerNorm after the last block. In post-norm every layer re-normalises the stream, so the scale stays fixed, but the gradient passes through a LayerNorm at every layer and very deep models become hard to train without careful warm-up.

Worked trace, D = 3, three layers, sublayer f(h) = W h with W = [[1,0,1],[0,1,0],[1,-1,0]], input [4, 0, 2]:

pre   (grows):        [5.2247, -1.2247, 4.4495] -> [6.6304, -2.6304, 6.6927] -> [8.0445, -4.0445, 8.8069]
post  (norm fixed):   [1.1355, -1.2978, 0.1622] -> [0.6732, -1.4137,  0.7405] -> [0.5544, -1.4039,  0.8495]
plain (input erased): [1.2247, -1.2247, 2.4495] -> [1.3363, -1.3363,  1.6036] -> [1.4094, -1.4094,  2.0135]

Look at the first layer. Pre-norm's output is still recognisably the input plus a correction (4 became 5.22). Post-norm and plain have already forgotten the input's scale: after one layer the output is on the LayerNorm's scale whatever went in. Only pre-norm kept a path for information to travel through unchanged.

Shapes must match

x + f(x) is only defined when f returns the same shape as x. That is why attention ends with the W_O projection back to width D and the FFN ends by contracting H back to D.

Common mistakes

• Normalising the wrong thing: in pre-norm the LayerNorm feeds f only. The skip connection carries the raw x, not LN(x).
• Writing post-norm as x + LN(f(x)) or LN(x) + f(x). The norm wraps the sum.
• Dropping the skip in "plain" mode and then being surprised that f = 0 zeroes the whole vector; with the skip, f = 0 is the identity.
• Updating the stream in place while still reading it for the sublayer's input.

Your exercise

Implement one block for each of the three modes. The starter supplies LayerNorm (gamma 1, beta 0, eps 1e-5), the matrix-vector sublayer, and the loop that applies your block LAYERS times, printing the stream after every layer.