Connecting Attention and Linear Layers in a Transformer Block

When a deep network trains poorly, one of the first questions to ask is whether the learning signal is reaching the early layers at all. Training works by sending a loss value backward through the network so that every weight can learn how to change. If that backward signal fades as it travels, the early layers hardly learn, no matter how long training runs.

This article takes a measurement approach. Instead of starting with the mathematics of the chain rule, you will run one small helper function on two versions of the same deep network: one built without shortcut connections and one built with them. The helper reports a single gradient number per layer, so the two runs can be compared directly. Keep one question in mind as you read: if the signal shrinks as it moves backward, which layer would you expect to show the smallest gradient?

The helper is called print_gradients. It takes a model and a sample input, passes the input through the model to get an output, and compares that output with a target value. In this setup the target is simply zero, chosen for simplicity. The comparison is done with a mean squared error loss, which measures how far the output is from the target.

The step that matters most is calling loss.backward(). This asks PyTorch to compute the gradient for each layer's parameters. You never write the gradient formulas yourself. PyTorch does that work, which is why training deep networks is practical.

A layer's weight matrix holds many values, and each one has its own gradient. A 3 by 3 weight matrix, for example, produces nine gradient values. Looking at nine numbers per layer is hard to compare across a deep network, so the helper reduces each layer to one summary value.

To make that summary, the helper takes the absolute value of every gradient in a layer and then averages those absolute values. The absolute value removes the sign, so positive and negative gradients both count toward the size of the signal. The result is one number per layer, which is easy to compare across the depth of the network.

The baseline run applies the helper to a model built without shortcut connections, using the same sample input. The report lists one mean gradient for each weight-bearing layer. Read from the last layer to the first, the pattern is consistent: the gradients get smaller as the signal travels backward. This is the phenomenon called the vanishing gradient problem.

The problem matters because a layer whose gradient is tiny receives almost no useful correction during training. The early layers are where many useful features get built, so when they barely learn, the whole network suffers. The source ties this problem to deep networks without shortcut connections, and the baseline run is the case where it appears.

A shortcut connection is a link that lets the signal bypass part of the network. Information and gradients can then travel across layers through a second route in addition to the usual layer path. The source describes shortcut connections as connections that allow gradients to flow more consistently across layers. In this course's model, the shortcut behavior is switched on with a single argument, use_shortcut=True, when the model is created.

Once you have both versions, read the comparison in a fixed order. Start at the last layer, where the loss is computed, and follow the gradient backward to the first layer. Without shortcut connections, the gradient shrinks along that path. With shortcut connections, the last layer still carries the largest gradient, but the value stabilizes as you move toward the first layer instead of dropping away to almost nothing.

A fair comparison fixes the random starting point and reads both reports in the same direction. The source seeds the random number generator with torch.manual_seed(123) before it creates the shortcut model. Seeding makes the random starting weights repeatable, so when two runs differ, you can attribute the difference to the architecture and not to the random starting point.