GenAI Course · Module 6 · Session 21

Inside the Encoder and Decoder Blocks

⏱ 12 min read🎯 Intermediate🖼 2 diagrams🧪 1 live demo🐍 Python code

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:

One encoder block. The original transformer stacks six of these.

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).

Pre-LN vs Post-LN: the original paper normalised after adding the residual. Most modern LLMs normalise before each sub-layer (Pre-LN), which trains more stably in very deep models.

02Inside the decoder block

The decoder block has three sub-layers instead of two:

A decoder block (Add & LayerNorm after each sub-layer not shown). The final block feeds the output head.

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.

Demo · Translating at inference time, one token per step

04Code: one block and a causal mask

blocks.py

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"

torch.Size([1, 4, 512]) tensor([[0., -inf, -inf, -inf], [0., 0., -inf, -inf], [0., 0., 0., -inf], [0., 0., 0., 0.]])

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

Created by Ajeetkumar