Attention And Transformers
The Full Transformer Architecture
Every component is now in place: multi-head self-attention (Chapter 2), positional encoding (Chapter 3), residual connections and layer normalization (Module 3,
Jr Codex Deep Learning Notes
Level: Advanced Prerequisites: Chapter 3: Positional Encoding Time to complete: ~30 minutes
Table of Contents
- Assembling Every Piece
- The Encoder Block
- The Decoder Block
- Cross-Attention — Connecting Encoder and Decoder
- The Feed-Forward Sub-Layer
- Stacking Blocks: the Full Architecture
- Encoder-Only, Decoder-Only, and Full Encoder-Decoder Variants
- A Complete Transformer in PyTorch
- Summary & Next Steps
1. Assembling Every Piece
Every component is now in place: multi-head self-attention (Chapter 2), positional encoding (Chapter 3), residual connections and layer normalization (Module 3, Module 4). This chapter assembles them into the complete Transformer architecture from "Attention Is All You Need."
2. The Encoder Block
One Transformer Encoder Block
─────────────────────────────────────────
Input
│
├──────────────────┐
▼ │ (residual connection,
Multi-Head Module 4 Ch.2 / Module
Self-Attention 3 Ch.6)
(Chapter 2) │
│ │
▼ │
[+]◄──────────────────┘
│
▼
Layer Norm (Module 3, Ch.3)
│
├──────────────────┐
▼ │
Feed-Forward │ (another residual
Network (Section 5) │ connection)
│ │
▼ │
[+]◄──────────────────┘
│
▼
Layer Norm
│
▼
Output (fed to the NEXT encoder block, or the decoder)
─────────────────────────────────────────
import torch
import torch.nn as nn
class EncoderBlock(nn.Module):
def __init__(self, d_model, num_heads, d_ff, dropout=0.1):
super().__init__()
self.self_attn = nn.MultiheadAttention(d_model, num_heads, batch_first=True)
self.norm1 = nn.LayerNorm(d_model)
self.norm2 = nn.LayerNorm(d_model)
self.feed_forward = nn.Sequential(
nn.Linear(d_model, d_ff), nn.ReLU(), nn.Linear(d_ff, d_model)
)
self.dropout = nn.Dropout(dropout) # Module 3, Ch.4
def forward(self, x):
attn_output, _ = self.self_attn(x, x, x) # self-attention (Ch.2)
x = self.norm1(x + self.dropout(attn_output)) # RESIDUAL connection + LayerNorm
ff_output = self.feed_forward(x)
x = self.norm2(x + self.dropout(ff_output)) # ANOTHER residual connection + LayerNorm
return xBoth residual connections directly reuse Module 3, Chapter 6's fix for vanishing gradients — they let gradients bypass the attention and feed-forward sub-layers entirely if needed, which is essential for training the very deep stacks of blocks (Section 6) used in real models.
3. The Decoder Block
The decoder block is similar, with two important additions: masked self-attention (Chapter 2, Section 6, preventing attending to future positions), and a cross-attention sub-layer connecting to the encoder's output.
One Transformer Decoder Block
─────────────────────────────────────────
Input (previously generated tokens)
│
▼
MASKED Multi-Head Self-Attention (Ch.2, Sec.6)
│
[+ residual, LayerNorm]
│
▼
Cross-Attention (Section 4) ◄──── Encoder's output
│
[+ residual, LayerNorm]
│
▼
Feed-Forward Network (Section 5)
│
[+ residual, LayerNorm]
│
▼
Output
─────────────────────────────────────────
class DecoderBlock(nn.Module):
def __init__(self, d_model, num_heads, d_ff, dropout=0.1):
super().__init__()
self.masked_self_attn = nn.MultiheadAttention(d_model, num_heads, batch_first=True)
self.cross_attn = nn.MultiheadAttention(d_model, num_heads, batch_first=True)
self.norm1 = nn.LayerNorm(d_model)
self.norm2 = nn.LayerNorm(d_model)
self.norm3 = nn.LayerNorm(d_model)
self.feed_forward = nn.Sequential(
nn.Linear(d_model, d_ff), nn.ReLU(), nn.Linear(d_ff, d_model)
)
self.dropout = nn.Dropout(dropout)
def forward(self, x, encoder_output, causal_mask):
attn1, _ = self.masked_self_attn(x, x, x, attn_mask=causal_mask) # Section 6, Ch.2
x = self.norm1(x + self.dropout(attn1))
attn2, _ = self.cross_attn(query=x, key=encoder_output, value=encoder_output) # Section 4
x = self.norm2(x + self.dropout(attn2))
ff_output = self.feed_forward(x)
x = self.norm3(x + self.dropout(ff_output))
return x4. Cross-Attention — Connecting Encoder and Decoder
Cross-attention is a direct, elegant generalization of Chapter 1's attention mechanism: the query comes from the decoder (what am I generating right now?), while the keys and values come from the encoder's output (what's in the source sequence?).
Cross-Attention vs Self-Attention
─────────────────────────────────────────
Self-attention (Ch.2): Q, K, V all come from the SAME
sequence
Cross-attention: Q comes from the DECODER;
K and V come from the
ENCODER's output
This is EXACTLY the mechanism that replaced the single
fixed-size context vector from Module 5, Chapter 4's basic
encoder-decoder RNN model — instead of one compressed
summary, the decoder can directly attend to EVERY encoder
position, at every decoding step.
─────────────────────────────────────────
5. The Feed-Forward Sub-Layer
Each block also includes a simple, position-wise feed-forward network — a small MLP (Module 2, Chapter 1) applied identically and independently to every position.
Why a Feed-Forward Layer, After Attention
─────────────────────────────────────────
Attention (Ch.1-2) computes a WEIGHTED COMBINATION of
values — this is fundamentally a LINEAR operation with
respect to the values. The feed-forward sub-layer adds
NON-LINEARITY and additional representational capacity,
processing each position's attention output further —
directly echoing Module 2, Chapter 2's point that
non-linearity is what gives depth its power.
─────────────────────────────────────────
Typically, the feed-forward layer's hidden dimension (d_ff) is significantly larger than the model's main dimension (d_model) — commonly 4x larger (e.g. d_model=512, d_ff=2048) — making this sub-layer responsible for a large fraction of the Transformer's total parameters.
6. Stacking Blocks: the Full Architecture
The Complete Transformer
─────────────────────────────────────────
Input tokens
│
▼
Token Embedding + Positional Encoding (Ch.3)
│
▼
Encoder Block × N (N is often 6, 12, 24, or more)
│
▼
Encoder output ──────────────┐
│
Output tokens (shifted right)│
│ │
▼ │
Token Embedding + │
Positional Encoding │
│ │
▼ │
Decoder Block × N ◄───────────┘ (cross-attends to encoder output)
│
▼
Final Linear layer + Softmax (Module 2, Ch.2/3)
│
▼
Output probability distribution over the vocabulary
─────────────────────────────────────────
7. Encoder-Only, Decoder-Only, and Full Encoder-Decoder Variants
The full encoder-decoder Transformer above was designed for translation. Most modern large language models use only half of this architecture:
Three Architectural Families
─────────────────────────────────────────
Encoder-only: uses ONLY the encoder stack (Section 2) —
well suited to tasks needing a rich
understanding of a FULL input at once
(classification, extracting information)
— covered in the NLP & LLM Notes' chapter
on BERT-style models
Decoder-only: uses ONLY the (masked) decoder stack,
WITHOUT cross-attention (since there's
no separate encoder) — well suited to
GENERATION tasks, predicting the next
token autoregressively (Module 5, Ch.4)
— covered in the NLP & LLM Notes'
chapter on GPT-style models; this
family underlies virtually all modern
large language models
Full encoder-decoder: the ORIGINAL design (Sections 2-4) —
still used for tasks with a clear
INPUT and OUTPUT sequence, like
translation or summarization
─────────────────────────────────────────
This distinction is the direct entry point into the NLP & LLM Notes, which covers BERT (encoder-only), GPT (decoder-only), and their applications in full depth.
8. A Complete Transformer in PyTorch
import torch
import torch.nn as nn
class SimpleTransformer(nn.Module):
def __init__(self, vocab_size, d_model=512, num_heads=8, num_layers=6, d_ff=2048, max_len=5000):
super().__init__()
self.embedding = nn.Embedding(vocab_size, d_model)
self.pos_encoding = PositionalEncoding(d_model, max_len) # Chapter 3
self.encoder_blocks = nn.ModuleList([
EncoderBlock(d_model, num_heads, d_ff) for _ in range(num_layers)
])
self.output_layer = nn.Linear(d_model, vocab_size)
def forward(self, x):
x = self.embedding(x)
x = self.pos_encoding(x)
for block in self.encoder_blocks:
x = block(x)
return self.output_layer(x)
# PyTorch ALSO provides a full built-in implementation directly
builtin_transformer = nn.Transformer(
d_model=512, nhead=8, num_encoder_layers=6, num_decoder_layers=6,
dim_feedforward=2048, batch_first=True,
)9. Summary & Next Steps
Key Takeaways
- Each Transformer block combines multi-head self-attention and a feed-forward network, each wrapped in a residual connection and layer normalization, directly reusing Module 3's training stability techniques.
- Decoder blocks add masked self-attention (preventing attending to future tokens) and cross-attention (attending to the encoder's output).
- Cross-attention generalizes attention to let a decoder's query attend to a different sequence's (the encoder's) keys and values, replacing the fixed-size context vector bottleneck entirely.
- Most modern large language models use only an encoder-only (BERT-style) or decoder-only (GPT-style) architecture, rather than the full original encoder-decoder design.
Module 6 Complete — What's Next
You now understand the complete architecture behind virtually every modern state-of-the-art AI system for language, and increasingly for vision and other domains. Module 7 covers how to actually use these (often enormous) architectures practically — through transfer learning, rather than training from scratch.
Concept Check
- What is the purpose of the two residual connections in an encoder block?
- How does cross-attention differ from self-attention, in terms of where Q, K, and V come from?
- Why do most modern large language models use only an encoder-only or decoder-only architecture, rather than the full encoder-decoder design?
Next Module
→ Module 7: Transfer Learning & Practical Training Strategies
Jr Codex — 1-on-1 Personalized Coaching | Back to Module Index | Back to DL Index