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 each token gather context from the positions permitted by its mask: 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.
The original post-norm Transformer block
Each attention or feed-forward sublayer is followed by a residual addition and LayerNorm.
Exact inputs and equations appear with the figure.
Preparing the scientific figure…
The encoder-block ordering from the 2017 paper. Dropout, positional embeddings, and stacked blocks are omitted. This is post-norm; the optional legacy Studio diagram is a separately labelled pre-norm variant.
Vaswani et al. · Attention Is All You Need§3.1, Figure 1
Schematic adaptation of the original post-norm ordering, not a claim that modern Transformers all use it.
Compare the pre-norm Studio variant
The residual stream adds each sublayer’s output back to its input. The sequence below groups each normalization with its corresponding sublayer.
A pre-norm decoder block overview. Attention masking, multi-head reshaping, and the MLP’s internal projections are omitted.
Read the pre-norm Studio 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
Unmasked self-attention without positional signals is permutation-equivariant: permuting inputs permutes outputs. Positional information can be added to embeddings or built into attention, as with RoPE. A fixed causal mask also breaks arbitrary permutation symmetry.
Work through the math
A modern pre-norm Transformer block, for input :
In the original architecture, the feed-forward network is a two-layer MLP applied position-wise, expanding to :
where is a non-linearity such as ReLU or GELU, , .
Pre-norm vs. post-norm: the original paper put LayerNorm after the residual (). Pre-norm places normalization before each sublayer, leaving an identity residual path and often simplifying optimization. Post-norm can also train successfully with appropriate initialization and schedules. For the dense FFN above, weights total , versus for standard Q/K/V/O projections, excluding biases. Gated FFNs, GQA, mixture-of-experts, and vocabulary embeddings change these counts; RMSNorm is another common normalization choice.
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 | Current and earlier input tokens | 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 input token sees itself and earlier input tokens)
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, helping gradient propagation and letting each block learn a refinement.)
- Applied: Why pre-norm over post-norm in large models? (Pre-norm preserves an identity residual path and can simplify optimization; warmup and initialization still depend on the architecture. Post-norm is not inherently untrainable.)
- Systems-level: Where do most of a Transformer's parameters and FLOPs live? (A standard dense 4d FFN has ~8d² weights versus ~4d² for MHA projections. Counts vary by architecture; dense attention arithmetic can dominate at sufficiently long context lengths.)
- Failure modes: Why does a Transformer need positional encoding at all? (Unmasked self-attention without positional signals is permutation-equivariant. Position encodings or attention modifications supply order; a causal mask also constrains it.)
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.
Sources and further reading
- Vaswani et al., Attention Is All You Need: original blocks, masks, and feed-forward dimensions.
- Xiong et al., On Layer Normalization in the Transformer Architecture: pre-norm, post-norm, and warmup.
- Zhang and Sennrich, Root Mean Square Layer Normalization: an alternative normalization method.
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.