Deep Learning
Transformer Architecture
Follow attention, feed-forward layers, normalization, and residual connections through a Transformer block.
Jump to a section
Why this matters
A Transformer combines attention, position-wise feed-forward layers, normalization, and residual connections. Attention mixes information across positions; the feed-forward network transforms each position's representation.
Follow the residual stream through one block before thinking about a whole stack. The ordering of normalization and residual addition matters, and encoder, decoder, and encoder-decoder models use different attention masks and connections.
Build the mental model
One block does two things in sequence. Attention lets every token gather context from every other token: it's the only place information moves between positions. The feed-forward network (FFN) then processes each position independently, adding capacity and non-linearity. Residual connections (x + sublayer(x)) give gradients a clean path and let the network learn refinements rather than full remappings; LayerNorm keeps activations well-scaled.
Read one pre-norm Transformer block
The residual stream adds each sublayer’s output back to its input. The sequence below groups each normalization with its corresponding sublayer.
The visual loads as you reach this section.
Read the connections
Residual input -> Norm + attention : sublayer 1 Residual input -> Identity A : bypass Norm + attention -> Add attention : attention output Identity A -> Add attention : unchanged input Add attention -> Norm + MLP : sublayer 2 Add attention -> Identity B : bypass Norm + MLP -> Add MLP : MLP output Identity B -> Add MLP : unchanged input Add MLP -> Next block : residual stream
Because attention is permutation-invariant, a Transformer has no inherent notion of order: positional encoding is injected at the input so the model knows token positions.
Work through the math
A modern pre-norm Transformer block, for input :
The feed-forward network is a two-layer MLP applied position-wise, usually expanding to :
where is a non-linearity (ReLU, or GELU in most LLMs), , .
Pre-norm vs. post-norm: the original paper put LayerNorm after the residual (). Pre-norm (LN inside, before the sublayer) trains far more stably at depth because the residual path stays an identity: this is why nearly all large models use pre-norm. The FFN holds roughly parameters per block and dominates the parameter count; attention contributes from the Q/K/V/O projections.
Key takeaway: A Transformer block is "move information between positions (attention), then transform each position on its own (FFN)," with residuals keeping the gradient highway open. Everything else (pre/post-norm, encoders/decoders) is configuration around that skeleton.
The three configurations at a glance:
| Configuration | Attention masking | Sees | Example models | Built for |
|---|---|---|---|---|
| Encoder-only | Bidirectional | Full context | BERT, RoBERTa | Understanding, classification |
| Decoder-only | Causal | Past tokens only | GPT, Llama | Autoregressive generation |
| Encoder-decoder | Bidirectional enc, causal dec | Cross-attention links them | T5, original Transformer | Seq2seq (translation) |
Read the implementation
1import torch
2from torch import Tensor, nn
3
4
5class TransformerBlock(nn.Module):
6 """Pre-norm decoder block: causal self-attention + position-wise FFN."""
7
8 def __init__(self, d_model: int, n_heads: int, mlp_ratio: int = 4, p: float = 0.1) -> None:
9 super().__init__()
10 self.ln1 = nn.LayerNorm(d_model)
11 self.attn = nn.MultiheadAttention(d_model, n_heads, dropout=p, batch_first=True)
12 self.ln2 = nn.LayerNorm(d_model)
13 self.ffn = nn.Sequential(
14 nn.Linear(d_model, mlp_ratio * d_model),
15 nn.GELU(),
16 nn.Linear(mlp_ratio * d_model, d_model),
17 nn.Dropout(p),
18 )
19
20 def forward(self, x: Tensor, attn_mask: Tensor | None = None) -> Tensor:
21 h = self.ln1(x)
22 # causal mask makes this a decoder block (each token sees only the past)
23 attn_out, _ = self.attn(h, h, h, attn_mask=attn_mask, need_weights=False)
24 x = x + attn_out # residual 1
25 x = x + self.ffn(self.ln2(x)) # residual 2 (pre-norm)
26 return x
27
28
29block = TransformerBlock(d_model=512, n_heads=8)
30seq = torch.randn(2, 16, 512)
31causal = torch.triu(torch.full((16, 16), float("-inf")), diagonal=1)
32assert block(seq, attn_mask=causal).shape == (2, 16, 512)Questions and trade-offs
- Conceptual: What is the difference between encoder-only, decoder-only, and encoder-decoder Transformers, and give one model for each. (BERT = encoder/bidirectional; GPT = decoder/causal; T5 = encoder-decoder/seq2seq.)
- Implementation: Why are residual connections essential in a deep Transformer? (They give gradients an identity path, preventing vanishing gradients and letting each block learn a refinement.)
- Applied: Why pre-norm over post-norm in large models? (Pre-norm keeps the residual stream an identity, making deep stacks trainable without careful warmup; post-norm is unstable at depth.)
- Systems-level: Where do most of a Transformer's parameters and FLOPs live? (The FFN (~8d² params per block) dominates parameters; attention's quadratic cost in sequence length dominates compute for long contexts.)
- Failure modes: Why does a Transformer need positional encoding at all? (Self-attention is permutation-invariant, without positional information it can't distinguish token order.)
Check your understanding
From memory: draw one pre-norm Transformer block, write the two residual equations, state what the FFN does that attention doesn't, and name the three encoder/decoder configurations with an example each. Check against Stages 2–3.
Oliver Perrin
Machine Learning Engineer · Founder, LiminalML
LiminalML brings together concept explanations, mathematical examples, and working visuals. Find more writing by Oliver on the LiminalML Substack.
Keep exploring the idea.
Ask a follow-up in a guided session, work through the self-check, or open the visual in Studio to build on it.
Related concepts
Deep Learning
Attention Mechanisms
Turn compatibility scores into attention weights, then combine values. Explore the geometry, normalization, and limits of the mechanism.
LLMs
KV Cache
See which token projections caching saves, what attention must still read, and how memory use grows with context.
Deep Learning
Batch Normalization
Normalize activations across a batch, apply a learned scale and shift, and understand why training and evaluation behave differently.