Lesson 11

Residual connections, normalization, and feed-forward layers

Introduction

A residual branch should modify an existing feature rather than erase it. If the old feature is 3 and the branch contributes 0.4, the residual result is 3.4 before later processing.

Learning goal

Trace one complete transformer block through normalization, attention, residual paths, and position-wise feed-forward layers.

Before you start

Attention, tensor shapes, linear layers, activation functions, averages, and basic normalization.

Lesson plan

  1. Name the four block parts and explain the job each one performs.
  2. Run one encoder layer while tracking sequence, batch, and feature axes.
  3. Add token and position information, then compare bidirectional and causal masking.

Why combine these parts?

Attention lets one token position collect information from other positions. That is only half of a transformer block. After information moves between positions, a feed-forward network transforms each position independently. Residual connections preserve an easier path for old information and gradients. Layer normalization keeps each token's feature values on a controlled scale. Stacking blocks repeats the pattern: communicate, transform, and stabilize.

The four parts

Self-attentionMix useful information across positions in the same sequence.
Feed-forward networkApply the same small MLP to every token position.
Residual connectionAdd a sublayer's input to its output.
Layer normalizationNormalize features within each token representation.
Model dimensionThe feature width of every token vector, written D.
Attention headOne learned query-key-value matching process.

With batch-first layout, a block accepts and returns shape (B,T,D). Keeping this shape unchanged makes blocks easy to stack. The model dimension must divide evenly across the number of attention heads. For D=8 and two heads, each head works with four features.

Worked example: one encoder layer

Predict the shape: a residual addition combines a block input with the block result. Both must have shape (B,T,D). Attention may mix information across T, but it must return the same width D so element-by-element addition is defined.

nn.TransformerEncoderLayer provides the complete block. The word “encoder” here means it is bidirectional by default: every position may attend to every other position. We set dropout to zero so repeated evaluation is stable and easy to inspect.

import torch
from torch import nn

torch.manual_seed(11)
layer = nn.TransformerEncoderLayer(
    d_model=8,
    nhead=2,
    dim_feedforward=16,
    dropout=0.0,
    batch_first=True,
)

x = torch.randn(2, 4, 8)  # B=2, T=4, D=8
output = layer(x)

print(tuple(x.shape))
print(tuple(output.shape))
print(torch.isfinite(output).all().item())

# Expected:
# (2, 4, 8)
# (2, 4, 8)
# True

The block changes the values but preserves the shape. Internally, self-attention mixes positions, while the feed-forward network expands from 8 to 16 features and projects back to 8. Residual additions require matching shapes, which is why the external width remains 8.

Add token meaning and position

For one scalar feature, a residual step can look like 3 + 0.4 = 3.4. Removing the 3 forces the next layer to reconstruct the old signal from 0.4. If a residual addition fails, print both shapes before changing the model width.

Attention alone does not know that token 0 came before token 1. A transformer therefore adds positional information to token embeddings. The example below uses a learned position table. Its token vectors and position vectors both have shape (B,T,D) after broadcasting.

vocabulary_size = 6
maximum_length = 4
token_embedding = nn.Embedding(vocabulary_size, 8)
position_embedding = nn.Embedding(maximum_length, 8)

token_ids = torch.tensor([[0, 1, 2, 3], [3, 2, 1, 0]])
positions = torch.arange(token_ids.shape[1])
x = token_embedding(token_ids) + position_embedding(positions)
output = layer(x)
print(tuple(output.shape))  # Expected: (2, 4, 8)

The one-dimensional position tensor has shape (T,). Its embedding has shape (T,D), and PyTorch broadcasts it across the batch. Every sequence receives positions 0 through 3.

Bidirectional or causal?

An encoder-style task such as masked-token prediction can use both earlier and later context. A GPT-style next-token model must hide future positions with a causal mask. Without the mask, training becomes invalid because a position can read the answer from the target side.

causal_mask = nn.Transformer.generate_square_subsequent_mask(4)
causal_output = layer(x, src_mask=causal_mask)
print(tuple(causal_mask.shape))    # Expected: (4, 4)
print(tuple(causal_output.shape))  # Expected: (2, 4, 8)

A padding mask solves a different problem: it hides storage positions added to make variable-length sequences rectangular. Do not confuse future masking with padding masking.

Common pitfalls

  • d_model must be divisible by nhead.
  • Token embeddings need positional information; otherwise order is not represented directly.
  • A GPT training block needs a causal mask. An encoder task usually does not.
  • batch_first=True means (B,T,D); the default in some APIs may differ.
  • Dropout behaves differently in training and evaluation, so call model.eval() for inference.

Try it

Create a layer with model dimension 12, three heads, and feed-forward width 24. Pass five-token sequences for a batch of four. What are the input, output, and causal-mask shapes?

Reveal the worked answer 
layer = nn.TransformerEncoderLayer(
    d_model=12, nhead=3, dim_feedforward=24,
    dropout=0.0, batch_first=True,
)
x = torch.zeros(4, 5, 12)
mask = nn.Transformer.generate_square_subsequent_mask(5)
y = layer(x, src_mask=mask)
print(tuple(x.shape))     # Expected: (4, 5, 12)
print(tuple(y.shape))     # Expected: (4, 5, 12)
print(tuple(mask.shape))  # Expected: (5, 5)

Three heads are valid because 12 divides evenly by 3. The feed-forward width is internal and does not change the returned model dimension.

Recap

A transformer block combines self-attention, a position-wise MLP, residual paths, and layer normalization. Its external shape stays (B,T,D), which allows stacking. Token and position embeddings provide content and order. Choose masking based on the task: bidirectional context for encoder-style understanding, causal context for next-token generation.

Reference: PyTorch TransformerEncoderLayer documentation.