Skip to main content

Transformer Architecture

A transformer is more than attention. Token and position representations enter repeated blocks containing multi-head attention, residual connections, normalization, and position-wise feed-forward networks. Encoder-only, decoder-only, and encoder-decoder systems differ in their masks and information flow.

Prerequisites​

Complete Attention Step by Step and understand embeddings, padding, and train/evaluation modes.

Learning goals​

  • trace batch, sequence, model, and head axes through a small block.
  • distinguish bidirectional padding masks from autoregressive causal masks.
  • connect architecture choice to classification, generation, and sequence-to-sequence tasks.

Worked example​

This decoder-style block uses both a causal mask and a key padding mask. PyTorch's batch_first=True keeps tensors in (batch, sequence, model) order; per-head weights are (batch, heads, query, key).

import torch
from torch import nn

torch.manual_seed(45)
batch, length, model_width, heads = 2, 4, 4, 2
token_ids = torch.tensor([[2, 3, 4, 0], [5, 4, 0, 0]])
padding = token_ids.eq(0)
causal = torch.triu(torch.ones(length, length, dtype=torch.bool), diagonal=1)

tokens = nn.Embedding(6, model_width, padding_idx=0)
positions = nn.Embedding(length, model_width)
attention = nn.MultiheadAttention(model_width, heads, dropout=0.0, batch_first=True)
norm1, norm2 = nn.LayerNorm(model_width), nn.LayerNorm(model_width)
feed_forward = nn.Sequential(nn.Linear(model_width, 8), nn.ReLU(), nn.Linear(8, model_width))

x = tokens(token_ids) + positions(torch.arange(length))[None, :, :]
attended, weights = attention(
    x, x, x, attn_mask=causal, key_padding_mask=padding,
    need_weights=True, average_attn_weights=False,
)
x = norm1(x + attended)
x = norm2(x + feed_forward(x))
x = x.masked_fill(padding.unsqueeze(-1), 0.0)

future = causal[None, None, :, :].expand_as(weights)
padded_keys = padding[:, None, None, :].expand_as(weights)
print(f"tokens={tuple(token_ids.shape)} hidden={tuple(x.shape)}")
print(f"causal={tuple(causal.shape)} padding={tuple(padding.shape)} heads={tuple(weights.shape)}")
print(f"future_weight_max={weights.masked_select(future).max():.1f}")
print(f"padding_weight_max={weights.masked_select(padded_keys).max():.1f} padded_output_norm={x[padding].norm():.1f}")

Output:

tokens=(2, 4) hidden=(2, 4, 4)
causal=(4, 4) padding=(2, 4) heads=(2, 2, 4, 4)
future_weight_max=0.0
padding_weight_max=0.0 padded_output_norm=0.0

An encoder classifier normally permits both left and right context among non-padding tokens. A decoder language model uses a causal mask. Cross-attention in an encoder-decoder model uses decoder states as queries and encoder outputs as keys and values.

Common mistakes​

  • Calling a single attention operation a complete transformer block.
  • Omitting positional information and expecting token order to be recoverable.
  • Using a causal mask for a bidirectional classifier without intending the restriction.
  • Confusing attention-head count with independent models or guaranteed interpretability.
  • Ignoring memory growth from the sequence-by-sequence attention matrix.

Practice​

Remove causal while keeping padding. Explain which entries may become nonzero and which must remain zero. Then identify the appropriate mask for sentiment classification versus next-token generation.

Next steps​

Pretraining and Transfer Learning explains how transformer parameters acquire reusable representations.

Primary sources​

Checked 2026-10-16: Attention Is All You Need, PyTorch MultiheadAttention, PyTorch TransformerEncoderLayer, and PyTorch transformer building-block tutorial.