Shortcut Connections & the Transformer Block

One building block remains before assembly: shortcut connections, the deceptively simple trick that makes 12-layer (and 48-layer) networks trainable at all. With it in hand, this unit assembles LayerNorm, GELU and the FFN plus chapter 3’s multi-head attention into the transformer block — the unit repeated 12× in GPT-2 small.

4.4 Adding shortcut connections

A shortcut connection (also skip or residual connection) feeds a layer’s input around the layer and adds it to the output: x = x + layer(x). Originally proposed for deep computer-vision networks (residual networks), it mitigates the vanishing gradient problem: gradients that shrink as they propagate backward through many layers, making early layers nearly untrainable.

The book makes this concrete with a 5-layer toy network run twice from the same seed — once without shortcuts, once with:

Recreation of the book’s p. 56 figure with the real measured gradients: without shortcuts, layer 1’s mean gradient is 0.0002 — with them, 0.22. Same layers, same seed, ~1000× more learning signal where it was scarcest.

</> A configurable deep network — shortcuts on or off

class ExampleDeepNeuralNetwork(nn.Module):
  def __init__(self, layer_sizes, use_shortcut):
      super().__init__()
      self.use_shortcut = use_shortcut
      self.layers = nn.ModuleList([
          nn.Sequential(nn.Linear(layer_sizes[0], layer_sizes[1]), GELU()),
          nn.Sequential(nn.Linear(layer_sizes[1], layer_sizes[2]), GELU()),
          nn.Sequential(nn.Linear(layer_sizes[2], layer_sizes[3]), GELU()),
          nn.Sequential(nn.Linear(layer_sizes[3], layer_sizes[4]), GELU()),
          nn.Sequential(nn.Linear(layer_sizes[4], layer_sizes[5]), GELU())
      ])

  def forward(self, x):
      for layer in self.layers:
          layer_output = layer(x)
          if self.use_shortcut and x.shape == layer_output.shape:  #A
              x = x + layer_output
          else:
              x = layer_output
      return x

def print_gradients(model, x):
  output = model(x)              #B
  target = torch.tensor([[0.]])

  loss = nn.MSELoss()
  loss = loss(output, target)    #C

  loss.backward()                #D

  for name, param in model.named_parameters():
      if 'weight' in name:
          print(f"{name} has gradient mean of {param.grad.abs().mean().item()}")

#A The shortcut is only legal when input and output shapes match — the same reason the transformer block preserves its shape everywhere.
#B Forward pass on the sample input.
#C A dummy objective: mean-squared error against a target of 0 — enough to produce a loss to differentiate.
#D loss.backward() computes gradients for every parameter; we then print the mean absolute gradient of each weight matrix — a proxy for how much learning signal reaches that layer.

</> The experiment: same seed, shortcuts off vs on

layer_sizes = [3, 3, 3, 3, 3, 1]
sample_input = torch.tensor([[1., 0., -1.]])

torch.manual_seed(123)
model_without_shortcut = ExampleDeepNeuralNetwork(layer_sizes, use_shortcut=False)
print_gradients(model_without_shortcut, sample_input)
# layers.0.0.weight has gradient mean of 0.00020173587836325169
# layers.1.0.weight has gradient mean of 0.0001201116101583466
# layers.2.0.weight has gradient mean of 0.0007152041653171182
# layers.3.0.weight has gradient mean of 0.001398873864673078
# layers.4.0.weight has gradient mean of 0.005049646366387606

torch.manual_seed(123)
model_with_shortcut = ExampleDeepNeuralNetwork(layer_sizes, use_shortcut=True)
print_gradients(model_with_shortcut, sample_input)
# layers.0.0.weight has gradient mean of 0.22169792652130127
# layers.1.0.weight has gradient mean of 0.20694105327129364
# layers.2.0.weight has gradient mean of 0.32896995544433594
# layers.3.0.weight has gradient mean of 0.2665732502937317
# layers.4.0.weight has gradient mean of 1.3258541822433472

Sort the table by either column and note the pattern: without shortcuts the gradient decays toward the early layers; with them it stays healthy everywhere. Click a row for the interpretation:

Layer⇅	Without shortcuts⇅	With shortcuts⇅
layers.0 (earliest)	0.0002017	0.2217	+
layers.1	0.0001201	0.2069	+
layers.2	0.0007152	0.329	+
layers.3	0.0013989	0.2666	+
layers.4 (last)	0.0050496	1.3259	+

Mean absolute weight gradient per layer — the same network, ±shortcuts.

Key idea — a gradient highway

Backpropagating through

x + f(x)

always includes the identity’s derivative (exactly 1), so the error signal can flow around every layer instead of only through it. GPT-2 small stacks 12 transformer blocks — 24 sublayers; without this trick, the early blocks would barely train.

4.5 Connecting it all: the transformer block

Now the assembly. A transformer block wires the four components together — twice over the same pattern: normalize → transform → drop → add the input back:

</> The TransformerBlock

from previous_chapters import MultiHeadAttention

class TransformerBlock(nn.Module):
  def __init__(self, cfg):
      super().__init__()
      self.att = MultiHeadAttention(
          d_in=cfg["emb_dim"],
          d_out=cfg["emb_dim"],  #A
          context_length=cfg["context_length"],
          num_heads=cfg["n_heads"],
          dropout=cfg["drop_rate"],
          qkv_bias=cfg["qkv_bias"])
      self.ff = FeedForward(cfg)
      self.norm1 = LayerNorm(cfg["emb_dim"])
      self.norm2 = LayerNorm(cfg["emb_dim"])
      self.drop_shortcut = nn.Dropout(cfg["drop_rate"])

  def forward(self, x):
      # Shortcut connection for attention block
      shortcut = x
      x = self.norm1(x)  #B
      x = self.att(x)    # Shape [batch_size, num_tokens, emb_size]
      x = self.drop_shortcut(x)
      x = x + shortcut   #C

      # Shortcut connection for feed forward block
      shortcut = x
      x = self.norm2(x)
      x = self.ff(x)
      x = self.drop_shortcut(x)
      x = x + shortcut

      return x

#A d_in = d_out = 768: chapter 3's MultiHeadAttention at full width, 12 heads × head_dim 64. The causal mask and attention dropout live inside it.
#B PRE-LayerNorm: normalize at the sublayer's entrance (before attention / before the FFN), which gives more stable training than the original transformer's post-norm placement.
#C The shortcut add — the raw input x (saved before normalization) joins the transformed path. Requires the shapes to match: everything in the block is (b, num_tokens, 768).

Step through the block’s dataflow with the book’s p. 58 figure — one sublayer at a time:

Input: embedded tokens

Four 768-dimensional vectors — "Every effort moves you" after ch2's token+positional embedding. The block will return exactly this shape.

1 / 8

One transformer block processing the embedded "Every effort moves you" (4 tokens × 768). Follow the two normalize→transform→drop→add cycles.

Why shape preservation matters

The block maps (batch, num_tokens, 768) → (batch, num_tokens, 768) by design: the one-to-one correspondence between input and output vectors is what lets blocks stack arbitrarily (§4.6 stacks 12), lets the residual additions work, and suits sequence-to-sequence tasks where each output position corresponds to an input position — while each output vector is nevertheless a context vector encapsulating information from the entire (visible) input.

Key idea — global gather, local process

Self-attention provides global information — it captures relationships between sequence elements so the model understands context and structure. The feed-forward network strengthens local information — a nonlinear per-position transformation that refines each vector’s own features. In a translation, attention understands the sentence structure and word relationships; the FFN adjusts and optimizes each word’s representation. The combination handles complex patterns neither could alone — and it is the entire recipe: §4.6 just stacks it 12 times.

📝 Check yourself: shortcuts & the transformer block

0 / 5

1.What exactly does a shortcut (residual) connection compute?
2.In the book's 5-layer demo, the first layer's gradient mean jumps from 0.00020 (no shortcuts) to 0.222 (with shortcuts) — roughly 1000×. Why do shortcuts help so much?
3.In the transformer block, where does LayerNorm sit relative to attention and the FFN?
4.What is the division of labor between multi-head attention and the feed-forward network in a block?
5.Predict: a (2, 4, 768) tensor enters a TransformerBlock. What comes out?