Inside the Encoder and Decoder Blocks
Let us open the boxes. In this session we look inside one encoder block and one decoder block, understand each small part, and see why the model behaves differently during training and when you actually use it.
In this session
- The four parts of an encoder block
- The decoder’s extra parts: masked attention, cross-attention and the output head
- Training with teacher forcing vs generating at inference time
- PyTorch code for a block and a causal mask
01Inside the encoder block
Every encoder layer has the same four parts, in this order:
| Part | Job in simple words |
|---|---|
| Self-attention | each token collects information from the other tokens |
| Feed-forward network (FFN) | a small 2-layer network applied to each token separately; it “thinks” about what was collected |
| Residual connection | adds the block’s input back to its output, so information and gradients flow easily through many layers |
| Layer normalization | keeps the numbers in a stable range so training doesn’t blow up |
Like a group study session. First everyone discusses together (attention). Then each student goes home and thinks it over alone (feed-forward). Their notes from before are never thrown away (residual).
02Inside the decoder block
The decoder block has three sub-layers instead of two:
- Masked self-attention: while writing word 3, the decoder may only look at words 1 and 2. The causal mask (Session 18) hides the future.
- Cross-attention: the queries come from the decoder, but the keys and values come from the encoder’s output. This is how the decoder “looks at” the input sentence.
- Output head: a linear layer turns the final vector into one score per vocabulary word, and softmax turns scores into probabilities.
03Training vs inference
The same model behaves differently when learning and when being used:
| Training | Inference (using the model) | |
|---|---|---|
| Decoder input | the correct previous words (teacher forcing) | its own previous predictions |
| Positions processed | all at once, in parallel | one new token per step |
| Speed | fast | slower: one step per token |
| Dropout | on | off |
| KV cache | not needed | used to avoid recomputing old tokens |
Teacher forcing is like learning to ride with support wheels. During training, the model is always given the correct previous word, even if it guessed wrong. At inference, the support wheels come off and it must use its own guesses.
04Code: one block and a causal mask
import torch import torch.nn as nn layer = nn.TransformerEncoderLayer(d_model=512, nhead=8, dim_feedforward=2048, batch_first=True) x = torch.randn(1, 4, 512) # 1 sentence, 4 tokens print(layer(x).shape) # same shape in and out mask = nn.Transformer.generate_square_subsequent_mask(4) print(mask) # -inf = "cannot look here"
Notice that a block keeps the shape the same: 4 tokens of 512 numbers in, 4 tokens of 512 numbers out. That is why we can stack as many blocks as we like.
05Summary
- Encoder block = self-attention → Add & Norm → FFN → Add & Norm.
- Decoder adds masked self-attention and cross-attention to the encoder.
- The output head turns the final vector into probabilities over the vocabulary.
- Training uses teacher forcing in parallel; inference generates one token at a time.
06Quick quiz
1. In cross-attention, where do the keys and values come from?
2. What does the residual connection do?
3. Teacher forcing means during training the decoder gets…
4. Why can’t the decoder see future tokens?
5. The feed-forward network works…
07References
- Attention Is All You Need (2017)
- The Annotated Transformer (Harvard NLP)
- PyTorch nn.TransformerEncoderLayer