Module 01: Transformer Architecture

AI Engineering Handbook
Who is this for? A Computer Science student who knows how to program but has never
studied AI infrastructure or deep learning internals. Every technical term is explained
when it first appears. Nothing is skipped.

Table of Contents
1. Chapter 1 — What is a Transformer?
2. Chapter 2 — Embeddings
3. Chapter 3 — Positional Encoding
4. Chapter 4 — Self Attention
5. Chapter 5 — Multi Head Attention
6. Chapter 6 — Feed Forward Network
7. Chapter 7 — LayerNorm
8. Chapter 8 — Residual Connections
9. Chapter 9 — Decoder Only Architecture
10. Chapter 10 — Output Layer
11. Mini Project — Build a Tiny Transformer in Python
12. Module Summary
13. Key Takeaways

Chapter 1 — What is a Transformer?

What is it?
A Transformer is a type of neural network architecture introduced by Google researchers
in 2017 in a paper titled "Attention is All You Need". It is the backbone of almost every
modern large language model (LLM) — including GPT-4, Claude, Gemini, LLaMA, Mistral,
and thousands of others.
Before Transformers, neural networks that processed text used RNNs (Recurrent Neural
Networks) and LSTMs (Long Short-Term Memory networks). These processed text one
word at a time, left to right, which made them slow and hard to train on long texts.
The Transformer completely changed this by processing all words at once using a
mechanism called attention (covered in Chapter 4). This made it much faster to train on
modern hardware (especially GPUs), and it handles long texts far better than its
predecessors.
Key idea: Instead of reading a sentence word by word like humans do, a Transformer reads
the entire sentence simultaneously and figures out which words are related to which other
words — all at once.
Why do we need it?
Before Transformers, the best models for text had serious problems:

Problem Old Approach (RNN/LSTM) Transformer Solution

Long-range Forgot earlier words in long texts Attention directly connects any two
dependencies words

Training speed Sequential — word by word, slow Parallel — all words at once, fast

GPU utilization Poor — hard to parallelize Excellent — matrix math on GPU

Scalability Couldn't scale to billions of Scales well — GPT-3 has 175B

parameters parameters

Context length Struggled past ~100 tokens Handles thousands of tokens

Real-World Analogy
Imagine you are reading a 300-page novel and someone asks: "In Chapter 20, when John
mentioned the key, what lock was he referring to?"
With an RNN, you would have to read the entire book word by word in order. By the time
you reach page 300, you have mostly forgotten what was on page 1.
With a Transformer, it is as if you laid out every page of the book on a giant table. You can
look at page 20 and page 1 at the same time, directly compare them, and immediately see
the connection — without having to "remember" anything from earlier.

Internal Architecture
The original Transformer had two main parts:
1. Encoder — reads the input and builds an understanding of it
2. Decoder — generates the output word by word
However, modern LLMs like GPT, Claude, and LLaMA use only the Decoder part
(explained in Chapter 9). For now, let us understand the full architecture.
 INPUT SIDE

Raw Text: 'The cat sat'

Tokenizer

Token IDs: 1, 42, 7, 88

Embedding Layer

Positional Encoding

ENCODER STACK (N layers)

Multi-Head Self-Attention

DECODER STACK (N layers)

Target Tokens Add & LayerNorm

Embedding + Positional Feed Forward Network

Encoding

Masked Multi-Head Add & LayerNorm

Self-Attention

Add & LayerNorm Encoder Output

 Add & LayerNorm Encoder Output

Cross-Attention with
Encoder Output

Add & LayerNorm

Feed Forward Network

Add & LayerNorm

OUTPUT SIDE

Linear Layer

Softmax

Predicted Next Token

Explanation of the diagram:

The left side (Input Side) takes raw text and converts it into numbers that the model
can work with.
The Encoder stack processes the input and builds a rich understanding of it.
The Decoder generates the output one token at a time, using both what it has
generated so far and the Encoder's understanding of the input.
The Output Side converts the decoder's output back to actual words.
For decoder-only models (like GPT), there is no Encoder. The model generates text purely
from previous tokens.

How it Works Step by Step

Let us trace the sentence "The cat sat on the mat" through a Transformer:
Step 1 — Tokenization Split text into tokens. Tokens are chunks of text — sometimes whole
words, sometimes parts of words.

"The cat sat on the mat" → [The, cat, sat, on, the, mat]
→ Token IDs: [464, 3797, 3332, 319, 262, 2603]

Step 2 — Embedding Each token ID is converted to a vector (a list of numbers) of fixed size
(e.g., 512 or 4096 numbers). This vector represents the meaning of the word in a
mathematical space.
Step 3 — Positional Encoding Since the Transformer reads all words at once, it needs to
know the order of words. Positional encoding adds position information to each word's
vector.
Step 4 — Self Attention The model looks at every word and asks: "How much should I pay
attention to every other word when trying to understand this word?"
Step 5 — Feed Forward Network Each word's representation is passed through a small
neural network to add more expressive power.
Steps 4 and 5 repeat for N layers (e.g., GPT-3 has 96 layers).
Step 6 — Output The final layer produces a probability distribution over all possible next
tokens. The highest probability token is selected as the output.

Data Flow

Text Tokens Embeddings + Position Info After Attention After FFN Output Logits Probabilities Predicted Token
'Hello world' [15496, 995] [0.2, -0.5, 0.8, ...] [0.3, -0.4, 1.1, ...] [context-aware vectors] [enriched vectors] [0.01, 0.3, 0.001, ...] [softmax applied] '!'
[0.1, 0.9, -0.3, ...] [0.4, 1.1, -0.1, ...]

Memory Flow
When processing a sentence of length L with embedding dimension d :
 Stage Memory Size Notes

Token IDs L × 4 bytes Integer IDs

Embeddings L × d × 4 bytes Float32 vectors

After Attention L × d × 4 bytes Same size, different values

Attention Scores L × L × 4 bytes The expensive part!

After FFN L × d × 4 bytes Same size

The L × L attention matrix is why long contexts are expensive. If L = 100,000 tokens (100K
context), the attention matrix is 100,000 × 100,000 = 10 billion numbers. This is one of the
biggest engineering challenges in modern LLMs.

CPU vs GPU Responsibilities

Task CPU GPU

Tokenization ✅ Yes ❌ No

Loading model weights ✅ Yes (then transfers) ✅ Stores weights

Matrix multiplication (attention, FFN) ❌ Too slow ✅ Primary job

Sampling (picking next token) ✅ Sometimes ✅ Can do both

Orchestration / scheduling ✅ Yes ❌ No

Production Examples
GPT-4: Transformer with ~1.8 trillion parameters (estimated), 128K context length
Claude 3.5 Sonnet: Decoder-only Transformer with 200K context window
LLaMA 3.1: 128K context, open-source, runs on consumer hardware
BERT: Encoder-only Transformer, used for classification tasks, not generation
Common Misconceptions
Misconception Reality

"Transformer = ChatGPT" ChatGPT is a product built on top of a Transformer

model

"Attention means the model pays attention Attention is just a weighted sum — a mathematical
like a human" operation

"Bigger = smarter" A well-trained smaller model can outperform a

poorly-trained larger one

"Transformers understand language" They learn statistical patterns; understanding is

debated

"Encoder and Decoder are always both Modern LLMs use Decoder only
present"

FAQ
Q: What is a "parameter" in a Transformer? A: A parameter is a number (specifically a
floating-point weight) that the model learned during training. GPT-3 has 175 billion such
numbers. These numbers are what make the model "smart."
Q: What is a "token"? A: A token is a chunk of text. The word "unhappiness" might be split
into ["un", "happiness"] — 2 tokens. The average English word is about 1.3 tokens.
Q: Why not just use a big lookup table instead of a Transformer? A: Language has infinite
combinations. You can't pre-store all possible sentences. A Transformer generalizes — it
learns patterns and applies them to new combinations.
Q: How long does it take to train a Transformer? A: GPT-3 took about 34 days on 1,024
A100 GPUs. LLaMA 3 took millions of GPU-hours.

Interview Questions
1. What problem did Transformers solve that RNNs couldn't? (Parallelization, long-
range dependencies)
2. What are the two main components of the original Transformer? (Encoder and
Decoder)
3. What is the time and memory complexity of self-attention? (O(n²) in sequence
length n)
4. Why do modern LLMs use only the Decoder? (For text generation, a decoder auto-
regressively generates tokens)
5. What paper introduced the Transformer? ("Attention is All You Need," Vaswani et al.,
2017)
Hands-on Exercises
1. Exercise 1: Open [Link] — an interactive 3D visualization of a small
GPT-2 model. Click through each layer and watch how data flows.
2. Exercise 2: Install the transformers library and load a model:

python

from transformers import AutoTokenizer, AutoModelForCausalLM

tokenizer = AutoTokenizer.from_pretrained("gpt2")
model = AutoModelForCausalLM.from_pretrained("gpt2")

text = "The cat sat on"

inputs = tokenizer(text, return_tensors="pt")
print(inputs) # See token IDs
outputs = [Link](**inputs, max_new_tokens=5)
print([Link](outputs[0]))

3. Exercise 3: Print the number of parameters in GPT-2:

python

total = sum([Link]() for p in [Link]())

print(f"GPT-2 has {total:,} parameters") # ~117 million

Summary
A Transformer is a neural network that processes all input tokens simultaneously using the
attention mechanism. It replaced older sequential models (RNNs/LSTMs) because it is
faster to train, scales better, and handles long texts better. Modern LLMs are decoder-only
Transformers that generate one token at a time.

Chapter 2 — Embeddings
What is it?
An embedding is a way to convert a discrete token (like the word "cat" or the number 4237)
into a continuous vector — a list of floating-point numbers.
In simpler terms: a Transformer only works with numbers. The word "cat" isn't a number.
Embeddings are the bridge between human language and math.
For example, with an embedding dimension of 4 (real models use 512 to 16384):
 "cat" → [0.25, -0.13, 0.87, 0.44]
"dog" → [0.27, -0.11, 0.85, 0.41]
"car" → [-0.72, 0.65, -0.31, 0.20]

Notice that "cat" and "dog" have similar vectors (because they are semantically related),
while "car" is very different.
A vocabulary (vocab) is the complete set of all tokens the model knows about. GPT-2 has
50,257 tokens. GPT-4 uses ~100,000+ tokens. The embedding layer is a large lookup table
(a matrix) of size vocab_size × embedding_dimension .

Why do we need it?

Neural networks are mathematical functions. They only understand real numbers — no
strings, no characters.
Furthermore, not just any numbers work. We need numbers that carry meaning. If we just
assigned "cat = 1", "dog = 2", "car = 3", the model would incorrectly think "cat" and "dog" are 1
apart while "dog" and "car" are 1 apart too — but "cat" and "dog" are more similar than "dog"
and "car."
Embeddings solve this by placing similar concepts close together in a high-dimensional
space.
Famous property: In well-trained embeddings:

vector("king") - vector("man") + vector("woman") ≈ vector("queen")

This shows that the embedding space encodes meaningful relationships.

Real-World Analogy
Think of a color picker in a design application. Colors can be described as RGB values: Red =
(255, 0, 0), Green = (0, 255, 0), Blue = (0, 0, 255). Two similar colors like Red (255, 0, 0) and
Dark Red (200, 0, 0) are numerically close to each other.
Embeddings do the same thing for words. "Cat" and "dog" are numerically close (in a 512-
dimensional space). "Cat" and "democracy" are far apart.

Internal Architecture
The embedding layer is simply a matrix — a 2D array of numbers.
 Embedding Matrix E: shape = [vocab_size × d_model]

For GPT-2:
vocab_size = 50,257
d_model = 768
E shape = [50,257 × 768]
Memory = 50,257 × 768 × 4 bytes ≈ 154 MB

When you want the embedding for token ID 4237:

You look up row 4237 of matrix E.
That row is a vector of 768 numbers.
This is called an embedding lookup.
It is literally just E[token_id] — like indexing an array. No multiplication needed. This is
why it is called a lookup table.
 Vocabulary (50,257 tokens)

0: <|endoftext|>

Embedding Matrix
50,257 × 768

1: '!' Row 0: [0.12, -0.43, 0.77, ...]

... Row 1: [-0.88, 0.21, 0.55, ...]

464: 'The' Row 464: [0.34, -0.19, 0.61,

lookup row 464 ...]

3797: 'cat' Row 3797: [0.25, -0.13, 0.87,

lookup row 3797 ...]

...

50256: last token

Explanation: Given the token "The" (ID 464), we simply retrieve row 464 from the
embedding matrix. That row's 768 numbers become the vector representation of "The" that
flows through the rest of the Transformer.

How it Works Step by Step

Step 1 — Tokenize the input

python

text = "The cat sat"

# After tokenization:
token_ids = [464, 3797, 3332]

Step 2 — Look up each token in the embedding matrix

 python

# E is the embedding matrix, shape [50257, 768]

embedding_464 = E[464] # Vector for "The", shape [768]
embedding_3797 = E[3797] # Vector for "cat", shape [768]
embedding_3332 = E[3332] # Vector for "sat", shape [768]

Step 3 — Stack into a matrix

python

# Result: shape [3, 768] — 3 tokens, each with 768 dimensions

input_embeddings = stack([embedding_464, embedding_3797, embedding_3332])

Step 4 — Feed into next stage (Positional Encoding) The shape [sequence_length,
d_model] is the fundamental shape that flows through most of the Transformer.
Data Flow

Raw text: 'The cat sat'

Tokenizer
(splits into subword tokens)

Token IDs: [464, 3797, 3332]

Embedding Layer
(lookup table: 50257 × 768)

Token Embeddings
Shape: [3 × 768]

(Next: Positional Encoding)

Memory Flow
For a sequence of length L with model dimension d :

Input: L integer IDs → L × 4 bytes (tiny)

Output: L × d floats → L × d × 4 bytes

For GPT-2 with L = 1024, d = 768:

Output size = 1024 × 768 × 4 bytes = 3,145,728 bytes ≈ 3 MB

For LLaMA 3 with L = 4096, d = 4096:

Output size = 4096 × 4096 × 4 bytes = 67,108,864 bytes ≈ 64 MB

The embedding matrix itself (parameters stored on GPU VRAM):

GPT-2: 50,257 × 768 × 4 bytes ≈ 154 MB

LLaMA 3: 128,256 × 4096 × 2 bytes ≈ 1 GB (using bfloat16)

CPU vs GPU Responsibilities

Task CPU GPU

Tokenization (text → IDs) ✅ Yes ❌ No

Embedding matrix storage ❌ No (would be slow) ✅ Lives in VRAM

Embedding lookup ❌ No ✅ Parallel memory access

Gradient updates during training ❌ No ✅ Yes

The embedding lookup is essentially a gather operation — for each token ID, grab the
corresponding row from the matrix. GPUs are excellent at this because they have high
memory bandwidth (the ability to read large amounts of memory quickly).

Production Examples
GPT-2: d_model = 768, vocab = 50,257
GPT-3: d_model = 12,288, vocab = 50,257
LLaMA 3: d_model = 4,096 (8B model), vocab = 128,256
Claude models: Exact dimensions not public, but similar scale
Weight Tying: In many models, the embedding matrix is shared with the output layer (the
matrix that converts back to token probabilities). This saves memory because you don't
need two large matrices. GPT-2 does this.
Common Misconceptions
Misconception Reality

"Embeddings are permanent — Embeddings are learned during training and change with
they never change" every gradient update

"A larger vocabulary is always Larger vocab = larger embedding matrix = more memory.
better" Trade-offs exist.

"Embeddings capture the true They capture statistical patterns, not true meaning. Context
meaning of words" matters.

"Each token has one fixed The input embedding is fixed per token, but the contextual
embedding" representation changes at every layer

FAQ
Q: What is the difference between "embedding" and "token"? A: A token is a piece of text
(like "cat" or "##ing"). An embedding is the vector of numbers that represents that token.
Q: How are embeddings learned? A: During training, the model starts with random
embeddings and adjusts them using backpropagation — nudging the vectors of similar
words closer together and dissimilar words further apart.
Q: Can I visualize embeddings? A: Yes! Use dimensionality reduction techniques like t-
SNE or UMAP to project 768-dimensional vectors down to 2D. You'll see clusters of related
words.
Q: What is "embedding dimension" (d_model)? A: The number of numbers in each
embedding vector. Larger d_model = more expressive but more memory and computation.

Interview Questions
1. What is an embedding and why is it needed? (Converts discrete tokens to continuous
vectors that neural networks can process)
2. What is the shape of the embedding matrix for GPT-2? (50,257 × 768)
3. What does "weight tying" mean in the context of embeddings? (Sharing the
embedding matrix between input and output layers)
4. Why are embeddings better than one-hot encoding? (Dense, low-dimensional,
semantically meaningful vs sparse, high-dimensional, no semantic information)
5. What is the time complexity of an embedding lookup? (O(1) — just an array index)

Hands-on Exercises
Exercise 1: Inspect GPT-2 embeddings

python
 import torch
from transformers import GPT2Model

model = GPT2Model.from_pretrained("gpt2")
E = [Link] # Word Token Embeddings
print(f"Embedding matrix shape: {[Link]}") # [50257, 768]

# Get embedding for token 3797 ("cat")

cat_embedding = E[3797]
print(f"Embedding for 'cat': {cat_embedding[:5]}") # First 5 numbers

Exercise 2: Measure similarity between word embeddings

python

import [Link] as F

tokenizer_output = tokenizer(["cat", "dog", "car"], return_tensors="pt")

ids = tokenizer_output["input_ids"]

cat_vec = E[ids[0][0]]
dog_vec = E[ids[1][0]]
car_vec = E[ids[2][0]]

cos = F.cosine_similarity

print(f"cat vs dog: {cos(cat_vec.unsqueeze(0), dog_vec.unsqueeze(0)).item():.3f

print(f"cat vs car: {cos(cat_vec.unsqueeze(0), car_vec.unsqueeze(0)).item():.3f
# cat vs dog should be higher (more similar)
 

Summary
Embeddings are learned lookup tables that convert token IDs into dense floating-point
vectors. The embedding matrix has shape [vocab_size × d_model] and lives on the GPU.
Similar words end up with similar embeddings because they appear in similar contexts
during training. The output of the embedding layer is a matrix of shape [sequence_length ×
d_model] that flows through the rest of the Transformer.

Chapter 3 — Positional Encoding

What is it?
Positional Encoding is a mechanism that injects information about the position (order) of
each token into its embedding vector.
Here is the problem: the Transformer's attention mechanism processes all tokens
simultaneously. It has no built-in sense of order. If you gave it "cat sat the" versus "the cat
sat," the attention scores would compute identically (just rearranged) — the model couldn't
tell the difference.
Positional encoding fixes this by adding a unique "position fingerprint" to each token's
embedding, so the model knows that token 0 comes before token 1, which comes before
token 2, etc.

Why do we need it?

Consider the two sentences:
"The dog bit the man."
"The man bit the dog."
Both have the same words. They mean completely different things. The only difference is
word order. Without positional encoding, a Transformer sees both as identical bags of
tokens and produces the same output for both — which is catastrophically wrong.

Real-World Analogy
Imagine you write the following on separate sticky notes:

"Mary", "loves", "John"

and mix them up in a bowl. You can't tell what the original sentence was, because the order
is gone.
Positional encoding is like writing the note's position on each sticky note in invisible ink:

"Mary [pos=1]", "loves [pos=2]", "John [pos=3]"

Now even when mixed up, you can reconstruct the original order.

Internal Architecture
There are two main types of positional encoding used in practice:
Type 1 — Sinusoidal Positional Encoding (Original paper, still used in some models)
For each position pos and dimension i in the embedding:

PE(pos, 2i) = sin(pos / 10000^(2i/d_model))

PE(pos, 2i+1) = cos(pos / 10000^(2i/d_model))

Even dimensions get a sine wave.

Odd dimensions get a cosine wave.
 Different frequencies are used for different dimensions.
The result is added to the token embedding:

final_embedding = token_embedding + positional_encoding

Type 2 — Learned Positional Encoding (GPT-2, BERT)

Instead of using a formula, the model learns a separate vector for each position. There is a
second matrix of shape [max_sequence_length × d_model] , and during training, these
position vectors are learned just like the token embeddings.
Type 3 — RoPE (Rotary Positional Embedding) (LLaMA, GPT-NeoX, modern LLMs)
A more sophisticated method that encodes position by rotating vectors in the attention
computation rather than adding. This is now the standard for most modern LLMs because
it generalizes better to long contexts.

Sinusoidal Encoding
Learned Encoding

Position 0 Learned vector P₀

(trained parameter)

Position 1 Learned vector P₁

(trained parameter) Position 0 Position 1 Position 2

Position 2 Learned vector P₂

(trained parameter)

sin/cos at position 0 sin/cos at position 1 sin/cos at position 2

Add to Embeddings

Token Embedding T₀ Token Embedding T₁ Token Embedding T₂

+ Position P₀ + Position P₁ + Position P₂
= Final₀ = Final₁ = Final₂

How it Works Step by Step

Using sinusoidal encoding for simplicity:
Step 1: You have token embeddings:

Token 0 ("The"): [0.34, -0.19, 0.61, 0.22, ...] ← shape [d_model]

Token 1 ("cat"): [0.25, -0.13, 0.87, 0.44, ...]
Token 2 ("sat"): [0.11, 0.55, 0.29, -0.67, ...]

Step 2: Compute positional encodings:

 Position 0: PE₀ = [sin(0/1), cos(0/1), sin(0/100), cos(0/100), ...]
PE₀ = [0.0, 1.0, 0.0, 1.0, ...]

Position 1: PE₁ = [sin(1/1), cos(1/1), sin(1/100), cos(1/100), ...]

PE₁ = [0.841, 0.540, 0.010, 1.000, ...]

Position 2: PE₂ = [sin(2/1), cos(2/1), sin(2/100), cos(2/100), ...]

PE₂ = [0.909, -0.416, 0.020, 0.9998, ...]

Step 3: Add them element-wise:

Final₀ = Token₀ + PE₀ = [0.34+0.0, -0.19+1.0, 0.61+0.0, 0.22+1.0, ...] =

[0.34, 0.81, 0.61, 1.22, ...]
Final₁ = Token₁ + PE₁ = [0.25+0.841, -0.13+0.540, ...]
Final₂ = Token₂ + PE₂ = [0.11+0.909, 0.55+(-0.416), ...]

Step 4: The result (same shape as embeddings: [L × d_model] ) flows into the first
attention layer.

Data Flow

Token Embeddings
[L × d_model]
Shape: [3 × 768]
Final Input Embeddings
Element-wise Addition [L × d_model]
⊕
Shape: [3 × 768]
(carries both meaning AND
position)
Positional Encodings
[L × d_model]
Shape: [3 × 768]
(precomputed or learned)

Memory Flow
Item Size

Positional encoding matrix (sinusoidal) 0 bytes — computed, not stored

Positional encoding matrix (learned) max_len × d_model × 4 bytes

For GPT-2 (max_len=1024, d=768) 1024 × 768 × 4 = 3MB

After addition Same as input: L × d × 4 bytes

CPU vs GPU Responsibilities
Task CPU GPU

Sinusoidal PE computation ✅ Can precompute ✅ Fast on GPU too

Storing learned PE matrix ❌ ✅ In VRAM with other weights

Element-wise addition ❌ (slow) ✅ Yes — trivially parallel

RoPE computation ❌ ✅ Fused with attention kernel

Production Examples
Model PE Type Max Context

Original Transformer (2017) Sinusoidal 512 tokens

GPT-2 Learned 1,024 tokens

BERT Learned 512 tokens

GPT-NeoX RoPE 2,048 tokens

LLaMA 3 RoPE 128K tokens

Claude 3 ALiBi / custom 200K tokens

Context Length Extension: Modern models use tricks like YaRN (Yet another RoPE
extensioN) or position interpolation to extend beyond their training context length. For
example, training on 4K tokens but later serving 128K tokens by scaling the position
frequencies.

Common Misconceptions
Misconception Reality

"Without PE, the model is completely Without PE, the model treats input as a set (like a bag of
random about order" words) — not random, but order-blind

"Sinusoidal PE is still state of the art" RoPE is now standard for LLMs

"PE is added once and that's it" In RoPE, position information is injected at every
attention layer

"Larger max position = larger memory" For sinusoidal PE, no extra memory. For learned PE, yes.
FAQ
Q: What happens if input is longer than max context? A: The model simply cannot
process it. You get an error or the excess tokens are truncated. This is why "context length"
is such an important spec for LLMs.
Q: What is RoPE and why is it better? A: RoPE (Rotary Position Embedding) encodes
position by rotating query and key vectors before computing attention. It naturally extends
to longer sequences than trained on and handles relative positions more gracefully.
Q: Is positional encoding trainable? A: Sinusoidal PE is not — it's fixed by formula.
Learned PE is trainable. RoPE has no parameters itself but is applied during attention
computation.

Interview Questions
1. Why does a Transformer need positional encoding? (Attention is permutation-
invariant — it processes all tokens simultaneously with no built-in order)
2. What is the difference between absolute and relative positional encoding?
(Absolute: encodes position 0, 1, 2... Relative: encodes distance between tokens)
3. What is RoPE? (Rotary Position Embedding — encodes position by rotating Q/K
vectors)
4. What is the maximum context length of GPT-2? (1,024 tokens)
5. What happens if you remove positional encoding entirely? (Model becomes order-
blind — "cat sat" and "sat cat" look identical)

Hands-on Exercise

python
 import numpy as np
import [Link] as plt

def sinusoidal_pe(max_len, d_model):

PE = [Link]((max_len, d_model))
for pos in range(max_len):
for i in range(0, d_model, 2):
PE[pos, i] = [Link](pos / (10000 ** (i / d_model)))
PE[pos, i+1] = [Link](pos / (10000 ** (i / d_model)))
return PE

pe = sinusoidal_pe(100, 64)
[Link](figsize=(12, 6))
[Link](pe, aspect='auto', cmap='RdBu')
[Link]()
[Link]("Embedding Dimension")
[Link]("Position")
[Link]("Sinusoidal Positional Encoding")
[Link]("positional_encoding.png")
[Link]()
# Each row is a unique "fingerprint" for each position

Summary
Positional encoding solves the Transformer's order-blindness by injecting position
information into each token's embedding. The three main approaches are: sinusoidal (fixed
formula), learned (trained parameters), and RoPE (rotation-based, now standard). Without
positional encoding, the model cannot distinguish "dog bites man" from "man bites dog."

Chapter 4 — Self Attention

What is it?
Self Attention is the heart of the Transformer. It is a mechanism that lets every token in a
sequence look at every other token and decide how much to "pay attention" to each one
when building its own representation.
The word "self" means the input is attending to itself — a sentence asking: "for each word,
which other words in this same sentence are relevant?"
Formal definition: Self attention computes, for each token, a weighted sum of all token
representations, where the weights are determined by the similarity between tokens.
Let us unpack that carefully.
Why do we need it?
Consider: "The animal didn't cross the street because it was too tired."
What does "it" refer to? "animal" or "street"? As a human, you know it refers to "animal." But
how?
You used context — specifically, you looked at other words ("animal," "tired") and made a
connection. Self attention does exactly this: it allows "it" to attend strongly to "animal" and
weakly to "street" based on their relevance.
Old RNNs struggled with this because by the time the model processed "it," the information
about "animal" (earlier in the sentence) was partially lost. Self attention has direct
connections between all token pairs, so "it" can directly query "animal" regardless of
distance.

Real-World Analogy
Imagine a meeting room with 10 people, each holding a document. When Person 5
(representing "it") wants to write their report, they ask every person in the room: "How
relevant is your document to mine?" Each person responds with a score (0 to 1). Person 5
then takes a weighted combination of everyone's documents — mostly from Person 1 (who
has the "animal" document) and a little from everyone else.
That weighted combination is "it"'s new, context-enriched representation.

Internal Architecture
Self attention uses three learned linear transformations applied to each token:
Q (Query): What am I looking for?
K (Key): What do I have to offer?
V (Value): What information do I actually give out?
These are all linear projections (matrix multiplications) of the token embeddings:

Q = X · W_Q (shape: [L × d_k])

K = X · W_K (shape: [L × d_k])
V = X · W_V (shape: [L × d_v])

Where:
X is the input matrix (shape: [L × d_model] )
W_Q , W_K , W_V are learnable weight matrices
d_k is the key/query dimension (typically d_model / num_heads )
The attention output is:

Attention(Q, K, V) = softmax(Q · K^T / √d_k) · V

Let us break this formula down piece by piece.

How it Works Step by Step

Step 1 — Linear Projections Project the input into Q, K, V spaces.

Input X: [L × d_model] (e.g., [4 × 8] for a tiny example)

W_Q: [d_model × d_k] (learned weights)

W_K: [d_model × d_k]
W_V: [d_model × d_v]

Q = X × W_Q: [L × d_k]
K = X × W_K: [L × d_k]
V = X × W_V: [L × d_v]

Step 2 — Compute Attention Scores For each query (row in Q), compute how much it
"matches" each key (row in K) using dot product:

Scores = Q · K^T shape: [L × L]

Each entry Scores[i][j] is the dot product between query of token i and key of token j .
A higher score means token i pays more attention to token j .
Step 3 — Scale Divide by √d_k to prevent scores from becoming too large (which would
cause gradients to vanish in softmax):

Scaled Scores = Scores / √d_k

Step 4 — Softmax Convert scores to probabilities (all positive, sum to 1):

Attention Weights = softmax(Scaled Scores) shape: [L × L]

Each row sums to 1. Row i tells you how much token i attends to every other token.
Step 5 — Weighted Sum of Values Multiply attention weights by the value matrix:

Output = Attention Weights · V shape: [L × d_v]

Each output row is a weighted average of all value vectors, where the weights come from
how much attention was paid.
 Input X
Shape: [L × d_model]

W_Q (learned) W_K (learned)

Q = X × W_Q K = X × W_K
Shape: [L × d_k] Shape: [L × d_k]

Q × K^T
Shape: [L × L]
(Raw attention scores)

÷ √d_k W_V (learned)

(Prevent gradient issues)

Softmax (row-wise)
Shape: [L × L] V = X × W_V
(Attention weights: each Shape: [L × d_v]
row sums to 1)

×V
Weighted sum of values

Output
Shape: [L × d_v]
(Context-enriched
representations)
Explanation: This diagram shows the complete single-head attention computation. The Q
and K matrices compute compatibility scores, which after scaling and softmax become
attention weights. These weights determine how much of each Value vector contributes to
each output position.

Concrete Tiny Example

Let's use a 4-token sequence with d_model=4, d_k=2:

Tokens: ["I", "love", "dogs", "."]

Token IDs: [40, 1842, 10012, 13]

After embedding (simplified):

X = [[1, 0, 1, 0], # "I"
[0, 2, 0, 2], # "love"
[1, 1, 1, 1], # "dogs"
[0, 0, 1, 0]] # "."

Suppose W_Q = W_K = W_V = identity for simplicity.

Then Q = K = V = X.

Q × K^T (dot products between all pairs):

= [[1*1+0*0+1*1+0*0, 1*0+0*2+1*0+0*2, 1*1+0*1+1*1+0*1, 1*0+0*0+1*1+0*0],
[0*1+2*0+0*1+2*0, 0*0+2*2+0*0+2*2, 0*1+2*1+0*1+2*1, 0*0+2*0+0*1+2*0],
...]

The actual numbers aren't critical for understanding. The point is: after softmax, we get a
4×4 matrix where each row shows how much each word attends to every other word.

Data Flow

3 Projections (learned)
Input Score Computation

Q: [L × d_k]
Output

QK^T: [L × L] ÷√d_k: [L × L] Softmax: [L × L]

X: [L × d_model] K: [L × d_k] Attn × V: [L × d_v]

V: [L × d_v]

Memory Flow
This is the most memory-intensive part of the Transformer:
 Tensor Shape Memory (L=2048, d=4096, d_k=128)

Q [L × d_k] 2048 × 128 × 2 = 0.5 MB

K [L × d_k] 2048 × 128 × 2 = 0.5 MB

V [L × d_v] 2048 × 128 × 2 = 0.5 MB

Attention scores (QK^T) [L × L] 2048 × 2048 × 2 = 8 MB per head

Attention weights (after softmax) [L × L] 8 MB per head

The L × L matrix is the memory bottleneck. For L = 100K:

100,000 × 100,000 × 2 bytes = 20 GB — just for one attention head!

This is why FlashAttention (covered in Module 02) was invented: it computes attention
without materializing the full L×L matrix.

CPU vs GPU Responsibilities

Operation CPU GPU

W_Q, W_K, W_V matrix multiplications ❌ Too slow ✅ Batched GEMM on Tensor Cores

QK^T computation ❌ O(L² × d_k) ✅ Highly parallelized

Softmax ❌ ✅ Fused kernel

Value weighted sum ❌ ✅ GEMM again

Masking (causal) ❌ ✅ Elementwise

All of self-attention is GPU-bound. It is pure matrix math.

Production Examples
Causal (Masked) Self Attention: In decoder-only models, a token can only attend to
previous tokens (not future ones — that would be cheating during generation). This is
implemented by setting future positions in the attention score matrix to -∞ before
softmax, so they get weight ≈ 0.
 Causal mask for 4 tokens:
I love dogs .
I [ 1, 0, 0, 0 ]
love [ 1, 1, 0, 0 ]
dogs [ 1, 1, 1, 0 ]
. [ 1, 1, 1, 1 ]

(0 means "attend to this", values will be set to -inf before softmax)

Actually:
(1 = allowed, 0 = masked → set to -inf)

FlashAttention: NVIDIA/Stanford technique that tiles the attention computation to stay

within GPU SRAM (L1 cache), avoiding writing the L×L matrix to HBM (GPU main
memory). Result: 3-8x speedup, reduced memory by O(L²) to O(L).

Common Misconceptions
Misconception Reality

"Q, K, V are different inputs" Q, K, V all come from the same input X, just transformed
differently

"Softmax always produces 0 or It produces continuous values between 0 and 1, summing to 1

1"

"Attention weights show what They're intermediate computations, not reliable explanations
the model thinks"

"Division by √d_k is optional" Without it, dot products become very large in high dimensions,
causing softmax saturation and vanishing gradients

FAQ
Q: Why use three different matrices (W_Q, W_K, W_V) instead of one? A: They serve
different roles. Q asks "what am I looking for?" K says "here's what I can match." V says
"here's what I'll give you." Using three separate transformations gives the model more
flexibility to learn different relationships.
Q: What does "attention" actually capture? A: In practice, attention heads learn various
things: some track syntactic dependencies (subject-verb), some track coreference ("it" →
"animal"), some track positional proximity, etc.
Q: Is attention always over the full sequence? A: In vanilla attention, yes. But techniques
like Sparse Attention (BigBird, Longformer) only attend to a subset of tokens for efficiency.
Interview Questions
1. Write the self-attention formula. (Attention(Q,K,V) = softmax(QK^T/√d_k)V)
2. Why do we scale by 1/√d_k? (Prevent dot products from growing too large, causing
softmax saturation)
3. What is the time complexity of self attention? (O(L² · d) in sequence length L and
dimension d)
4. What is a causal mask? (A triangular mask that prevents tokens from attending to
future tokens)
5. What is FlashAttention? (A memory-efficient attention algorithm that avoids
materializing the L×L matrix)

Hands-on Exercise

python
 import torch
import [Link] as F

def self_attention(X, W_Q, W_K, W_V, causal=False):

"""
X: [batch, seq_len, d_model]
W_Q, W_K, W_V: [d_model, d_k]
"""
L, d_model = [Link][1], [Link][2]
d_k = W_Q.shape[1]

Q = X @ W_Q # [batch, L, d_k]

K = X @ W_K # [batch, L, d_k]
V = X @ W_V # [batch, L, d_k]

# Attention scores
scores = (Q @ [Link](-2, -1)) / (d_k ** 0.5) # [batch, L, L]

# Causal mask
if causal:
mask = [Link]([Link](L, L), diagonal=1).bool()
scores = scores.masked_fill(mask, float('-inf'))

# Softmax
weights = [Link](scores, dim=-1) # [batch, L, L]

# Weighted sum
output = weights @ V # [batch, L, d_k]
return output, weights

# Test
batch, L, d_model, d_k = 1, 4, 8, 4
X = [Link](batch, L, d_model)
W_Q = [Link](d_model, d_k)
W_K = [Link](d_model, d_k)
W_V = [Link](d_model, d_k)

out, attn = self_attention(X, W_Q, W_K, W_V, causal=True)

print(f"Output shape: {[Link]}") # [1, 4, 4]
print(f"Attention weights:\n{attn[0]}") # 4×4 matrix, rows sum to 1

Summary
Self attention lets every token attend to every other token by computing query-key dot
products to get attention weights, then using those weights to compute a weighted sum of
value vectors. This produces a new, context-enriched representation for each token. The
main bottleneck is the O(L²) attention matrix. Modern optimizations like FlashAttention
reduce this cost significantly.

Chapter 5 — Multi Head Attention

What is it?
Multi Head Attention (MHA) runs self attention multiple times in parallel, each with
different learned weight matrices (W_Q, W_K, W_V). Each parallel run is called a "head."
The outputs are concatenated and projected back to the original dimension.
Think of it as: instead of one question-answering session, you run 8 (or 16 or 32)
simultaneous sessions, each focusing on different aspects of the relationships between
tokens.
In GPT-2 small (d_model=768): 12 heads, each with d_k = 768/12 = 64 In GPT-3
(d_model=12288): 96 heads, each with d_k = 128 In LLaMA 3 8B (d_model=4096): 32 heads,
each with d_k = 128

Why do we need it?

Single-head attention can only capture one type of relationship per layer. But language has
many simultaneous relationships:
Head 1 might track syntactic dependencies (subject → verb)
Head 2 might track coreference (pronoun → antecedent)
Head 3 might track positional proximity (nearby words)
Head 4 might track semantic similarity (synonyms)
Head 5 might track negation (not → modified word)
Multiple heads let the model capture all of these simultaneously.

Real-World Analogy
Imagine a team of 8 detectives investigating the same crime scene. Each detective
specializes in different evidence:
Detective 1 looks only at fingerprints.
Detective 2 looks only at footprints.
Detective 3 looks only at financial records.
Detective 8 looks only at phone records.
At the end, all their reports are combined into one comprehensive report. This combined
report is far richer than any single detective's findings alone.
Multi-head attention does the same: 8 (or more) attention "detectives" each find different
relationships, and their findings are combined.
Internal Architecture

Input X
[L × d_model]

Head 1 Head 2 Head h (nth head)

Q₁ = X·W_Q₁ K₁ = X·W_K₁ V₁ = X·W_V₁ Q₂ = X·W_Q₂ K₂ = X·W_K₂ V₂ = X·W_V₂ Q_h K_h V_h

[L × d_k] [L × d_k] [L × d_k]

Attention₁ Attention₂ Attention_h

[L × d_k] [L × d_k] [L × d_k]

Concatenate all heads

[L × (h × d_k)] = [L ×
d_model]

× W_O (output projection)

[d_model × d_model]

Multi-Head Output
[L × d_model]

Explanation: Each head independently performs self attention with its own set of Q, K, V
weight matrices. The individual outputs (each of shape [L × d_k] ) are concatenated along
the last dimension, giving [L × d_model] . A final linear projection W_O brings it back to
shape [L × d_model] .

How it Works Step by Step

Step 1 — Split into heads
For h heads each with dimension d_k = d_model / h:

python

# Weight matrices for all heads (often stored as one big matrix and split)
W_Q: [d_model × d_model] # = h × (d_model × d_k)
W_K: [d_model × d_model]
W_V: [d_model × d_model]

# Project input
Q_all = X @ W_Q # [L × d_model]
K_all = X @ W_K
V_all = X @ W_V

# Reshape to separate heads

# [L × d_model] → [L × h × d_k] → [h × L × d_k]
Q = Q_all.view(L, h, d_k).transpose(0, 1) # [h × L × d_k]
K = K_all.view(L, h, d_k).transpose(0, 1)
V = V_all.view(L, h, d_k).transpose(0, 1)

Step 2 — Attention on each head independently

python
 # Batched attention across all heads simultaneously
# [h × L × d_k] @ [h × d_k × L] = [h × L × L]
scores = Q @ [Link](-2, -1) / (d_k ** 0.5)
weights = softmax(scores, dim=-1) # [h × L × L]
attended = weights @ V # [h × L × d_k]

Step 3 — Concatenate and project

python

# [h × L × d_k] → [L × h × d_k] → [L × d_model]

attended = [Link](0, 1).reshape(L, d_model)

# Output projection
output = attended @ W_O # [L × d_model]

Grouped Query Attention (GQA) and Multi Query Attention (MQA)

Modern LLMs have evolved beyond standard MHA to save memory during inference:

Type Q heads K heads V heads Used In

MHA (Multi-Head) h h h GPT-2, BERT

MQA (Multi-Query) h 1 1 GPT-3, early fast models

GQA (Grouped-Query) h h/g h/g LLaMA 2, LLaMA 3, Mistral

Why GQA? In MHA, during inference, you store K and V caches for all heads. With 32 heads
and 32 layers, this is huge. GQA reduces K and V to fewer groups (e.g., 8 groups instead of 32
heads), drastically cutting the KV cache size (discussed in Module 02) while retaining most
of the quality.
 Grouped Query Attention (h=4, g=2)

Q₁ Q₂ K_group1 Q₃ Q₄ K_group2

Attn₁ Attn₂ Attn₃ Attn₄

Multi-Head Attention (h=4)

Q₁ K₁ Q₂ K₂ Q₃ K₃ Q₄ K₄

Attn₁ Attn₂ Attn₃ Attn₄

Explanation: In GQA with 4 query heads and 2 groups, K and V are shared within each
group. Q₁ and Q₂ share one K/V pair, Q₃ and Q₄ share another. This halves the K/V cache
size.

Data Flow

flowchart TD
X["X: [batch, L, d_model]"]
W_QKV["W_Q, W_K, W_V\n[d_model × d_model each]"]
PROJ["Linear Projections"]
SPLIT["Reshape & Split Heads\n[batch, h, L, d_k]"]
ATTN["Scaled Dot-Product Attention\n(on all heads in parallel)\n[batch, h, L, d_k]"
MERGE["Concatenate Heads\n[batch, L, h × d_k] = [batch, L, d_model]"]
WO_PROJ["W_O Projection\n[d_model × d_model]"]
OUT["Output: [batch, L, d_model]"]

X --> PROJ
W_QKV --> PROJ
PROJ --> SPLIT --> ATTN --> MERGE --> WO_PROJ --> OUT

 

Memory Flow
For MHA with h=12, L=1024, d_model=768 (GPT-2 small):
d_k = 768/12 = 64
Q, K, V each: [12 × 1024 × 64] = 786,432 floats = 3 MB each
Attention matrix per head: [1024 × 1024] = 1M floats = 4 MB per head = 48 MB for all
heads
Output (after projection): [1024 × 768] = ~3 MB
For a 40B parameter model at inference with L=8192, h=64, d_k=128:
Attention matrices: 64 × 8192 × 8192 × 2 bytes ≈ 8.6 GB per layer!
This is why FlashAttention (not materializing the full matrix) is essential in production.

CPU vs GPU Responsibilities

Task GPU Operation

Q, K, V projections Batched GEMM (General Matrix-Matrix Multiply)

Attention score computation Batched BMMT (Batch Matrix-Matrix Transpose Multiply)

Softmax Custom fused kernel (row-wise normalization)

Value weighted sum Batched GEMM

Output projection GEMM

All operations are on GPU. The GPU processes all heads simultaneously using the batch
dimension.

Common Misconceptions
Misconception Reality

"Each head looks at a different part of Each head looks at the full sentence but attends to
the sentence" different relationships

"More heads = better" Diminishing returns; too many small heads lose capacity

"W_Q, W_K, W_V are separate per Usually stored as one big matrix and reshaped
head"

"Concatenating is just copying" Concatenating combines different views; the W_O

projection then mixes them

Interview Questions
1. What is the difference between single-head and multi-head attention? (MHA runs
attention multiple times with different projections, capturing diverse relationships)
2. If d_model=1024 and we have 16 heads, what is d_k? (1024/16 = 64)
3. What is GQA and why is it used? (Grouped Query Attention — shares K/V across
query groups, reducing KV cache size)
4. What is the W_O matrix in MHA? (Output projection that combines all heads' outputs
back to d_model)
 5. How does MHA improve over single-head attention? (Captures multiple different
types of relationships simultaneously)

Hands-on Exercise

python
 import torch
import [Link] as nn
import [Link] as F

class MultiHeadAttention([Link]):
def __init__(self, d_model, num_heads):
super().__init__()
assert d_model % num_heads == 0
self.num_heads = num_heads
self.d_k = d_model // num_heads

self.W_Q = [Link](d_model, d_model)

self.W_K = [Link](d_model, d_model)
self.W_V = [Link](d_model, d_model)
self.W_O = [Link](d_model, d_model)

def forward(self, x, causal=False):

B, L, d = [Link]
h = self.num_heads

Q = self.W_Q(x).view(B, L, h, self.d_k).transpose(1, 2) # [B, h, L, d_

K = self.W_K(x).view(B, L, h, self.d_k).transpose(1, 2)
V = self.W_V(x).view(B, L, h, self.d_k).transpose(1, 2)

scores = Q @ [Link](-2, -1) / (self.d_k ** 0.5) # [B, h, L, L]

if causal:
mask = [Link]([Link](L, L, device=[Link]), diagonal=1).bo
scores = scores.masked_fill([Link](0).unsqueeze(0), float('

weights = [Link](scores, dim=-1)

out = weights @ V # [B, h, L, d_k]

out = [Link](1, 2).reshape(B, L, d) # [B, L, d_model]

return self.W_O(out)

# Test
mha = MultiHeadAttention(d_model=64, num_heads=8)
x = [Link](2, 10, 64) # batch=2, seq=10, d_model=64
out = mha(x, causal=True)
print(f"Output shape: {[Link]}") # [2, 10, 64]
 
Summary
Multi-head attention runs self attention multiple times in parallel with different learned
projections. Each head captures different linguistic relationships. Outputs are concatenated
and projected. Modern variants (MQA, GQA) reduce K/V heads to save memory during
inference. This is the core computational bottleneck in Transformer inference.

Chapter 6 — Feed Forward Network

What is it?
The Feed Forward Network (FFN) (also called the Feed Forward Layer or MLP sublayer) is
the second major component in each Transformer block, applied after multi-head attention.
It is a simple two-layer fully connected neural network applied independently to each token
position:

FFN(x) = activation(x · W₁ + b₁) · W₂ + b₂

Where:
W₁ : shape [d_model × d_ff] (expands dimension)
W₂ : shape [d_ff × d_model] (compresses back)
d_ff : the "inner dimension," typically 4 × d_model

activation : usually ReLU, GELU, or SwiGLU

For GPT-2 small: d_model=768, d_ff=3072 (4× expansion) For LLaMA 3 8B: d_model=4096,
d_ff=14336 (~3.5× expansion, using SwiGLU)

Why do we need it?

Attention is great at moving information between tokens — it figures out which tokens are
relevant to each other. But it is essentially a weighted average: a linear operation over the
value vectors.
The FFN adds non-linear transformation — it can compute complex functions of each
token's representation. This is where the model's "knowledge" is primarily stored.
Research has shown that:
Attention heads store relational knowledge ("Paris is the capital of which country?")
FFN layers store factual knowledge ("The capital of France is Paris")
The FFN is sometimes called the "key-value memory" of the Transformer — facts are
stored in W₁ (keys) and W₂ (values).
Real-World Analogy
After the attention mechanism lets tokens consult each other (like team members
discussing a problem), the FFN is each team member going to their personal encyclopedia
to enrich their understanding individually.
The FFN takes each token's current representation, expands it into a larger space (d_ff),
applies a non-linearity (to make complex decisions), then compresses back. This
expansion-compression is like "thinking deeply" about the current representation.

Internal Architecture

flowchart LR
X["Token Representation\n[d_model = 768]"]

W1["W₁\n[768 × 3072]\n(Expand)"]
BIAS1["+ b₁"]
ACT["GELU Activation\n(Non-linearity)"]

W2["W₂\n[3072 × 768]\n(Compress)"]
BIAS2["+ b₂"]

OUT["Output\n[d_model = 768]"]

X --> W1 --> BIAS1 --> ACT --> W2 --> BIAS2 --> OUT

SwiGLU Variant (Used in LLaMA, PaLM, Claude):

Instead of a single activation, SwiGLU uses a gating mechanism:

FFN_SwiGLU(x) = (x · W₁ ⊙ SiLU(x · W_gate)) · W₂

This has been empirically shown to improve performance. LLaMA uses this variant.

How it Works Step by Step

For a single token representation x (shape: [d_model]):
Step 1 — Expand

python

h = x @ W1 + b1 # [d_model] → [d_ff]
# For GPT-2: [768] → [3072]
# This "expands" the representation into a higher-dimensional space

Step 2 — Non-linearity

python
 h = gelu(h) # [d_ff] → [d_ff] (same shape, but non-linear transformatio
# GELU: Gaussian Error Linear Unit
# GELU(x) ≈ x * Φ(x) where Φ is cumulative normal distribution
# Smoother than ReLU (which just clips negatives to 0)
 

Step 3 — Compress

python

out = h @ W2 + b2 # [d_ff] → [d_model]

# [3072] → [768]

Key insight: This is applied independently to each token. There is no interaction between
tokens in the FFN (unlike attention). The FFN is purely per-position.

Data Flow

flowchart TD
X["Attention Output\n[batch × L × d_model]"]
W1m["W₁: [d_model × d_ff]\nshape: [768 × 3072]"]
E1["Expand: [batch × L × d_ff]"]
ACT["Activation (GELU/SwiGLU)\n[batch × L × d_ff]"]
W2m["W₂: [d_ff × d_model]\nshape: [3072 × 768]"]
E2["Compress: [batch × L × d_model]"]
ADD["+ Residual Connection (Chapter 8)"]
OUT["FFN Output: [batch × L × d_model]"]

X --> E1
W1m --> E1
E1 --> ACT
ACT --> E2
W2m --> E2
E2 --> ADD
X --> ADD
ADD --> OUT
Memory Flow
Item Shape Memory (GPT-2, L=1024)

Input [1024 × 768] 3 MB

After W₁ [1024 × 3072] 12 MB

After activation [1024 × 3072] 12 MB

After W₂ [1024 × 768] 3 MB

W₁ weights [768 × 3072] 9.4 MB

W₂ weights [3072 × 768] 9.4 MB

The intermediate activation [L × d_ff] is 4× larger than input and is the memory peak
within the FFN.
For LLaMA 3 70B: d_model=8192, d_ff=28672, L=4096:

Peak intermediate: 4096 × 28672 × 2 bytes = 235 MB per layer

70B model has 80 layers → cannot store all simultaneously without activation
checkpointing

CPU vs GPU Responsibilities

Task GPU Operation Notes

W₁ projection GEMM Largest single matrix multiply

Activation Element-wise kernel Very fast

W₂ projection GEMM Second large matrix multiply

The FFN accounts for ~2/3 of the parameters in a standard Transformer (attention is the
other ~1/3). In a 7B parameter model:
~2.3B parameters in attention layers
~4.7B parameters in FFN layers
Activation Functions Comparison
Activation Formula Used In Notes

ReLU max(0, x) Early Transformers Simple, sharp

GELU x·Φ(x) GPT-2, BERT Smooth, empirically better

SiLU/Swish x·σ(x) EfficientNet, some LLMs Similar to GELU

SwiGLU SiLU(Wx)⊙Vx LLaMA, PaLM Current best for LLMs

Production Examples
Mixture of Experts (MoE): A powerful variant where instead of one FFN, you have many
FFNs (called "experts"). For each token, only a small number of experts (e.g., 2 out of 64) are
activated.

FFN_MoE(x) = Σ gate(x)[i] * Expert_i(x) for top-k experts

Benefit: More parameters without proportional compute increase

Used in: Mixtral 8×7B (8 experts, use 2), GPT-4 (rumored)

flowchart TD
X["Token x"]
G["Gating Network\n(Router)\nOutputs: which 2 of 8 experts to use"]
E1["Expert 1\n(FFN)"]
E2["Expert 2\n(FFN)"]
E3["Expert 3\n(FFN, not selected)"]
E8["Expert 8\n(FFN, not selected)"]
COMB["Weighted Combination\nof selected experts' outputs"]
OUT["MoE Output"]

X --> G
G -->|"weight 0.6"| E1
G -->|"weight 0.4"| E2
G -->|"0 (not selected)"| E3
G -->|"0 (not selected)"| E8
E1 --> COMB
E2 --> COMB
COMB --> OUT
Common Misconceptions
Misconception Reality

"FFN interacts with multiple tokens" FFN is purely per-position — no inter-token interaction

"The FFN is less important than FFN contains ~2/3 of parameters and stores most factual
attention" knowledge

"The expansion ratio is always exactly Varies by model and architecture choice
4×"

"ReLU is still used in LLMs" Modern LLMs use GELU or SwiGLU

Interview Questions
1. What is the FFN's role in a Transformer block? (Non-linear per-position
transformation; stores factual knowledge)
2. What are the shapes of W₁ and W₂ in GPT-2? (W₁: [768×3072], W₂: [3072×768])
3. What is MoE (Mixture of Experts)? (Multiple FFN experts per layer; only top-k are
activated per token)
4. What is SwiGLU? (A gated activation function: SiLU(Wx) ⊙ Vx, used in
LLaMA/PaLM)
5. Does the FFN see other tokens? (No — applied independently to each token)

Hands-on Exercise

python
 import torch
import [Link] as nn
import [Link] as F

class FeedForward([Link]):
def __init__(self, d_model, d_ff):
super().__init__()
self.W1 = [Link](d_model, d_ff)
self.W2 = [Link](d_ff, d_model)

def forward(self, x):

return self.W2([Link](self.W1(x)))

class FeedForwardSwiGLU([Link]):
"""LLaMA-style FFN with SwiGLU activation"""
def __init__(self, d_model, d_ff):
super().__init__()
self.W1 = [Link](d_model, d_ff, bias=False)
self.W2 = [Link](d_ff, d_model, bias=False)
[Link] = [Link](d_model, d_ff, bias=False)

def forward(self, x):

return self.W2([Link](self.W1(x)) * [Link](x))

# Test
ffn = FeedForward(d_model=768, d_ff=3072)
x = [Link](2, 10, 768) # batch=2, seq=10, d_model=768
out = ffn(x)
print(f"Output shape: {[Link]}") # [2, 10, 768]

swiglu = FeedForwardSwiGLU(d_model=768, d_ff=3072)

out2 = swiglu(x)
print(f"SwiGLU output: {[Link]}") # [2, 10, 768]

Summary
The FFN is the second sublayer in each Transformer block. It expands each token's
representation to a higher dimension, applies a non-linear activation, then compresses
back. It is applied independently per token and contains most of the model's learned factual
knowledge. Modern LLMs use SwiGLU activation. The MoE variant uses multiple sparse
FFN experts for better parameter efficiency.
Chapter 7 — LayerNorm
What is it?
Layer Normalization (LayerNorm) is a technique that normalizes the activations (the
numbers flowing through the network) within each layer to have mean ≈ 0 and variance ≈ 1.
It then applies learnable scale (γ, gamma) and shift (β, beta) parameters.
Formula:

LayerNorm(x) = γ · (x - μ) / √(σ² + ε) + β

Where:
μ = mean of x across the feature dimension
σ² = variance of x
ε = small constant (e.g., 1e-5) to prevent division by zero
γ, β = learnable parameters (initialized to 1 and 0 respectively)

Why do we need it?

During training, as data flows through many layers (Transformers have 12 to 96+ layers),
the values can become very large or very small. This is called the vanishing/exploding
gradient problem:
If values are too large: gradients explode, training becomes unstable
If values are too small: gradients vanish, learning stops
LayerNorm keeps values in a stable range at every layer, making training much more stable
and faster.

Real-World Analogy
Imagine you are grading students from 5 different countries using different grading scales:
Country A: 0–100
Country B: 0–10
Country C: 0–4.0 GPA
Country D: A, B, C, D, F
Country E: 0–20
To fairly compare them, you "normalize" all grades to a standard scale (mean=0, std=1).
Now a grade of 0 means "average for their system" and +1 means "one standard deviation
above average."
LayerNorm does this for activations within each layer — ensuring all features are on the
same scale, making the model easier to train.
Internal Architecture

flowchart LR
X["Input x\n[d_model]"]
MU["Compute Mean\nμ = mean(x)"]
SIG["Compute Variance\nσ² = var(x)"]
NORM["Normalize\nx̂ = (x - μ) / √(σ² + ε)"]
SCALE["Scale & Shift\nγ · x̂ + β\n(γ, β: learned params)"]
OUT["Normalized Output\n[d_model]"]

X --> MU
X --> SIG
X --> NORM
MU --> NORM
SIG --> NORM
NORM --> SCALE --> OUT

Pre-LN vs Post-LN
The position of LayerNorm in the Transformer block matters significantly:
Post-LN (Original paper):

x → Attention → Add → LayerNorm → FFN → Add → LayerNorm

Pre-LN (Modern practice, GPT-2, LLaMA):

x → LayerNorm → Attention → Add → LayerNorm → FFN → Add

Pre-LN is now the standard because:

More stable training (gradients flow cleanly through residual connections)
No "warm-up" period needed for learning rate
Better performance at scale
RMSNorm (Used in LLaMA, modern LLMs):

RMSNorm(x) = x / √(mean(x²) + ε) · γ

Simpler than LayerNorm (no mean subtraction, no β), computationally cheaper, similar
performance. LLaMA uses RMSNorm.

How it Works Step by Step

For a single token's representation x of shape [d_model=4] (tiny example):
 x = [2.0, -1.0, 4.0, 3.0]

Step 1 - Mean:
μ = (2.0 + (-1.0) + 4.0 + 3.0) / 4 = 8.0 / 4 = 2.0

Step 2 - Variance:
σ² = ((2-2)² + (-1-2)² + (4-2)² + (3-2)²) / 4
= (0 + 9 + 4 + 1) / 4 = 14/4 = 3.5

Step 3 - Normalize:
x̂ = (x - μ) / √(σ² + ε) (ε = 1e-5)
x̂ = ([2-2, -1-2, 4-2, 3-2]) / √3.5
x̂ = [0, -3, 2, 1] / 1.871
x̂ = [0.0, -1.604, 1.069, 0.534]

Step 4 - Scale and shift (γ=[1,1,1,1], β=[0,0,0,0] initially):

output = γ * x̂ + β = x̂ = [0.0, -1.604, 1.069, 0.534]

Now the output has mean ≈ 0 and std ≈ 1. The model can learn γ and β to rescale if needed.

Data Flow

flowchart TD
X["Input x: [batch, L, d_model]"]

subgraph LN["LayerNorm (per token)"]

M["Compute μ per token\nacross d_model dimension"]
V["Compute σ² per token\nacross d_model dimension"]
N["Normalize: (x-μ)/√(σ²+ε)"]
SS["Scale γ + Shift β\n(learned per dimension)"]
end

OUT["Output: [batch, L, d_model]\n(same shape, normalized values)"]

X --> M & V
M & V & X --> N --> SS --> OUT

style LN fill:#e8f4f8
Memory Flow
Item Size

γ parameters d_model × 4 bytes = 768×4 = 3KB (GPT-2)

β parameters d_model × 4 bytes = 3KB

Intermediate (μ, σ²) L × 4 bytes each = tiny

Output Same as input: L × d_model

LayerNorm is extremely cheap in memory. It adds negligible overhead.

CPU vs GPU Responsibilities

Task GPU

Mean computation Reduction kernel (parallel sum)

Variance computation Reduction kernel

Normalization Element-wise kernel

Scale/shift Element-wise kernel

In practice, all four steps are fused into a single GPU kernel (fused LayerNorm) for
efficiency.

Production Examples
Model Normalization Notes

Original Transformer Post-LN, LayerNorm Less stable

GPT-2 Pre-LN, LayerNorm More stable

BERT Post-LN, LayerNorm

LLaMA 1/2/3 Pre-LN, RMSNorm Fastest and stable

Mistral Pre-LN, RMSNorm

PaLM Pre-LN, RMSNorm

Common Misconceptions
Misconception Reality

"BatchNorm and LayerNorm are BatchNorm normalizes across the batch dimension; LayerNorm
the same" across feature dimension — different!

"LayerNorm learns the right It normalizes to 0 mean/1 variance, then applies learnable γ/β on
mean/variance" top

"RMSNorm is worse than In practice, RMSNorm performs similarly with less compute
LayerNorm"

"LayerNorm is expensive" It's very cheap — just addition and division

FAQ
Q: What is the difference between BatchNorm and LayerNorm? A: BatchNorm
normalizes across the batch dimension (requires large batch sizes, doesn't work with
variable-length sequences). LayerNorm normalizes across the feature dimension for each
sample independently. LayerNorm works for any sequence length and batch size.
Q: What happens without LayerNorm? A: Training becomes very unstable. Loss oscillates
wildly or diverges. Deep networks (32+ layers) are nearly impossible to train without
normalization.
Q: Why does RMSNorm omit the mean subtraction? A: The mean subtraction and bias β
add cost with little benefit. Experiments show that just the scale (γ) after RMS
normalization works just as well.

Interview Questions
1. What does LayerNorm normalize across? (The feature/embedding dimension for
each token independently)
2. What is the difference between Pre-LN and Post-LN? (Pre-LN: normalize before
attention/FFN. Post-LN: normalize after. Pre-LN is more stable.)
3. What is RMSNorm? (Simpler LayerNorm without mean subtraction; used in LLaMA)
4. What are γ and β in LayerNorm? (Learned scale and shift parameters — allow the
model to undo normalization if needed)
5. Why is BatchNorm not used in Transformers? (Requires fixed batch statistics;
doesn't work well with variable-length sequences)

Hands-on Exercise

python
 import torch
import [Link] as nn

class LayerNorm([Link]):
def __init__(self, d_model, eps=1e-5):
super().__init__()
[Link] = [Link]([Link](d_model))
[Link] = [Link]([Link](d_model))
[Link] = eps

def forward(self, x):

mean = [Link](dim=-1, keepdim=True)
var = [Link](dim=-1, keepdim=True, unbiased=False)
x_hat = (x - mean) / [Link](var + [Link])
return [Link] * x_hat + [Link]

class RMSNorm([Link]):
def __init__(self, d_model, eps=1e-8):
super().__init__()
[Link] = [Link]([Link](d_model))
[Link] = eps

def forward(self, x):

rms = [Link]([Link](2).mean(dim=-1, keepdim=True) + [Link])
return x / rms * [Link]

# Test
x = [Link]([[2.0, -1.0, 4.0, 3.0]])
ln = LayerNorm(d_model=4)
rms = RMSNorm(d_model=4)

print(f"Input: {x}")
print(f"LayerNorm:{ln(x).detach()}")
print(f"RMSNorm: {rms(x).detach()}")
print(f"Mean of LN output: {ln(x).mean().item():.4f}") # ≈ 0
print(f"Std of LN output: {ln(x).std().item():.4f}") # ≈ 1

Summary
LayerNorm stabilizes training by normalizing activations to zero mean and unit variance,
then applying learnable scale/shift. Pre-LN (normalizing before each sublayer) is the
modern standard. RMSNorm is a faster variant used in LLaMA and most modern LLMs.
Without normalization, deep Transformer training is nearly impossible.
Chapter 8 — Residual Connections
What is it?
A Residual Connection (also called a skip connection) is when you add the input of a layer
directly to its output:

output = layer(x) + x

This seems almost trivially simple — and it is! But it is one of the most impactful ideas in
deep learning. Introduced in ResNet (2015) for image classification, it was then adopted in
the Transformer architecture.
In a Transformer block:

# After Multi-Head Attention:

x = x + MultiHeadAttention(LayerNorm(x))

# After FFN:
x = x + FFN(LayerNorm(x))

The + x at the end is the residual connection.

Why do we need it?

When neural networks have many layers (deep networks), the gradient signal can vanish
as it flows backward through all the layers during training. By the time it reaches the first
layers, the gradient is essentially zero — the first layers learn nothing.
Residual connections fix this by providing gradient highways — direct paths from the
output all the way back to early layers, bypassing the individual layer transformations.
Mathematically, the gradient of the loss with respect to the early layer's input x is:

Without skip: ∂L/∂x = ∂L/∂output × ∂output/∂x

(chain of many multiplications → vanishes)

With skip: ∂L/∂x = ∂L/∂output × (1 + ∂layer(x)/∂x)

(always has a "+1" component — gradient always flows!)

The "+1" ensures gradients always flow, even if the layer's contribution vanishes.

Real-World Analogy
Imagine water flowing through a series of filters. Each filter modifies the water but also
clogs up a little. Without bypass pipes, water barely makes it through many filters. With
bypass pipes that go around each filter and rejoin the main flow, water always has a direct
path and the filters can safely do their job without bottlenecking the flow.
Residual connections are the bypass pipes for gradient flow.

Internal Architecture

flowchart TD
X["Input x\n[L × d_model]"]

subgraph BLOCK1["Transformer Block (Pre-LN style)"]

LN1["LayerNorm(x)"]
MHA["Multi-Head Attention"]
ADD1["⊕ Add (Residual)"]
LN2["LayerNorm"]
FFN["Feed Forward Network"]
ADD2["⊕ Add (Residual)"]
end

OUT["Output\n[L × d_model]"]

X --> LN1 --> MHA --> ADD1

X -->|"Residual skip\n(identity)"| ADD1
ADD1 --> LN2 --> FFN --> ADD2
ADD1 -->|"Residual skip\n(identity)"| ADD2
ADD2 --> OUT

Explanation: The key feature is the two arrows bypassing the LayerNorm+Attention and
LayerNorm+FFN respectively. These arrows carry the original x directly to the addition
operation. So the total output is x + transformation(x) , not just transformation(x) .

How it Works Step by Step

Before understanding residuals: Without them, a 12-layer Transformer would compute:

h₁ = layer₁(x)
h₂ = layer₂(h₁)
...
h₁₂ = layer₁₂(h₁₁)

If any layer produces small numbers, the information is gone. If any layer's gradient is small,
earlier layers don't learn.
With residuals: The network computes:

h₁ = x + layer₁(x)
h₂ = h₁ + layer₂(h₁)
...
h₁₂ = h₁₁ + layer₁₂(h₁₁)
Notice that h₁₂ implicitly contains x (the original input) because each step adds to the
previous:

h₁₂ = x + layer₁(x) + layer₂(h₁) + ... + layer₁₂(h₁₁)

= x + (sum of all layer contributions)

Each layer only needs to learn the residual (the small change needed), not the full
transformation. This is much easier to learn!
Residual stream view: Modern interpretability research views the Transformer as a
"residual stream" — a vector that flows through all layers. Each attention head and FFN
reads from and writes to this stream additively. No information is ever destroyed — it
accumulates.

Data Flow

flowchart LR
subgraph ONE_BLOCK["One Transformer Block"]
IN["x (input)\n[L × d_model]"]

subgraph PATH1["Attention Path"]

LN1a["LayerNorm"]
ATT["Multi-Head\nAttention"]
end

ADD1["x + Attn(LN(x))\n[L × d_model]"]

subgraph PATH2["FFN Path"]

LN2a["LayerNorm"]
FFNa["Feed\nForward"]
end

ADD2["mid + FFN(LN(mid))\n[L × d_model]"]

end

IN --> LN1a --> ATT --> ADD1

IN -->|"skip"| ADD1
ADD1 --> LN2a --> FFNa --> ADD2
ADD1 -->|"skip"| ADD2

Memory Flow
Residual connections don't add memory — the x tensor that gets added is already in
memory. The addition is element-wise and happens in-place.
However, for gradient computation during training, you need to store the intermediate
activations to compute gradients. For a model with 32 layers, storing all intermediate
activations requires enormous memory. This is why gradient checkpointing (also called
activation checkpointing) is used — it recomputes activations during the backward pass
instead of storing them.

Scenario Memory for Activations

Full storage (training) O(N × L × d) for N layers

Gradient checkpointing O(√N × L × d) — recompute

Inference Only current layer needed

CPU vs GPU Responsibilities

Task GPU

Residual addition Element-wise add — trivially fast

Gradient through skip Automatic — no special handling needed

Activation storage HBM (GPU main memory)

Residual addition is the cheapest operation in the Transformer. It's literally one GPU
instruction.

Production Examples
GPT-2: 12 transformer blocks, each with 2 residual connections (one after attention,
one after FFN)
GPT-3: 96 transformer blocks = 192 total residual connections
LLaMA 3 70B: 80 transformer blocks = 160 residual connections
Initialization: Models like GPT-2 scale the residual connection by 1/√(2N) (where N is the
number of blocks) to keep the variance stable at initialization. LLaMA uses a similar
technique.
Common Misconceptions
Misconception Reality

"Residual connections add complexity" They are literally just + — zero added complexity

"Skip connections mean the layer is The layer still trains; it just learns the residual
optional" (delta)

"You need residuals only in very deep They help even in shallow networks
networks"

"The gradient only flows through the skip" It flows through both paths — skip and through the
layer

Interview Questions
1. What is a residual connection and why is it needed? (Adding input to output: output
= layer(x) + x . Prevents vanishing gradients.)
2. What problem do residual connections solve? (Vanishing gradient problem in deep
networks)
3. What is the "residual stream" in interpretability? (The view that the residual vector
accumulates contributions from all layers additively)
4. How many residual connections does a standard Transformer block have? (2 — one
after attention, one after FFN)
5. What is gradient checkpointing? (A technique that recomputes activations during
backward pass to save memory)

Hands-on Exercise

python
 import torch
import [Link] as nn

class TransformerBlock([Link]):
def __init__(self, d_model, num_heads, d_ff):
super().__init__()
[Link] = [Link](d_model, num_heads, batch_first=True
[Link] = [Link](
[Link](d_model, d_ff),
[Link](),
[Link](d_ff, d_model)
)
self.ln1 = [Link](d_model)
self.ln2 = [Link](d_model)

def forward(self, x):

# Pre-LN + Residual around Attention
attn_out, _ = [Link](self.ln1(x), self.ln1(x), self.ln1(x))
x = x + attn_out # ← RESIDUAL CONNECTION

# Pre-LN + Residual around FFN

ffn_out = [Link](self.ln2(x))
x = x + ffn_out # ← RESIDUAL CONNECTION
return x

# Test
block = TransformerBlock(d_model=64, num_heads=8, d_ff=256)
x = [Link](2, 10, 64) # [batch=2, seq=10, d_model=64]
out = block(x)
print(f"Input shape: {[Link]}") # [2, 10, 64]
print(f"Output shape: {[Link]}") # [2, 10, 64]
# Same shape — residual connections preserve dimensions
 

Summary
Residual connections add the layer's input directly to its output ( x + layer(x) ). They solve
the vanishing gradient problem by providing direct gradient paths through deep networks.
They appear twice per Transformer block (after attention and after FFN). They are
computationally free (just addition) but enormously impactful for training stability.
Chapter 9 — Decoder Only Architecture
What is it?
A Decoder-Only Transformer is a Transformer that uses only the decoder portion of the
original encoder-decoder architecture — with one modification: the cross-attention (which
previously attended to encoder outputs) is removed. What remains is a stack of blocks, each
containing:
1. Masked Multi-Head Self-Attention (causal — can only look at past tokens)
2. Feed Forward Network
3. Layer Norm (Pre-LN in modern models)
4. Residual Connections
This is the architecture used by GPT, Claude, LLaMA, Mistral, Gemini, PaLM, and
essentially all modern text-generation LLMs.

Why do we need it?

For text generation, you need to generate text left to right — predicting the next token given
all previous tokens. The decoder's causal masking (not allowing future tokens to influence
past tokens' representations) is perfect for this.
The encoder was designed to build a bidirectional understanding of input text (every word
can see every other word). This is great for classification or translation but not for open-
ended text generation. For generation, a decoder-only approach where the model is trained
to predict the next token is simpler and scales better.
Scaling insight: Decoder-only models have been empirically found to scale better than
encoder-decoder models. GPT-3's success with a purely decoder-only approach established
this paradigm.

Real-World Analogy
A decoder-only model is like an author writing a novel word by word:
When writing word 500, you can see words 1–499.
You cannot see words 501–1000 (they haven't been written yet).
The causal mask enforces this — each word can only be informed by words before it.
An encoder-decoder model is like a translator who:
1. First reads the entire source text (encoder — bidirectional, sees everything)
2. Then generates the translation word by word (decoder — causal)
For a general-purpose AI assistant, you don't need the encoder step. You just generate
responses directly, treating the conversation history as your context.
Internal Architecture

flowchart TD
subgraph INPUT["Input Processing"]
T["Input Tokens\n[token_1, ..., token_n]"]
E["Embedding\n[n × d_model]"]
PE["+ Positional Encoding"]
end

subgraph LAYERS["N Decoder-Only Blocks\n(N = 12 for GPT-2, 80 for LLaMA 70B)"]

subgraph BLOCK["One Block"]
LN1["RMSNorm / LayerNorm"]
CMHA["Causal Multi-Head Self-Attention\n(can only attend to positions ≤ cur
ADD1["⊕ Residual Add"]
LN2["RMSNorm / LayerNorm"]
FFN["Feed Forward Network\n(SwiGLU / GELU)"]
ADD2["⊕ Residual Add"]
end
end

subgraph OUTPUT["Output"]
LNF["Final LayerNorm"]
LINEAR["Linear: [d_model → vocab_size]"]
SOFT["Softmax → probabilities"]
SAMPLE["Sample / Argmax → Next Token"]
end

T --> E --> PE --> LN1 --> CMHA --> ADD1

PE -->|"skip"| ADD1
ADD1 --> LN2 --> FFN --> ADD2
ADD1 -->|"skip"| ADD2
ADD2 -->|"...repeat N times..."| LNF --> LINEAR --> SOFT --> SAMPLE

 

The Causal Mask

The key difference from an encoder is the causal mask (also called autoregressive mask).
During attention, token at position i can only attend to positions 0, 1, ..., i :

Position: 0 1 2 3 4
Tokens: The cat sat on mat

Attention mask (1=allowed, 0=blocked):

The cat sat on mat
The [ 1, 0, 0, 0, 0 ]
cat [ 1, 1, 0, 0, 0 ]
sat [ 1, 1, 1, 0, 0 ]
on [ 1, 1, 1, 1, 0 ]
mat [ 1, 1, 1, 1, 1 ]
The 0s become -∞ before softmax, making those attention weights ≈ 0.
This mask serves two purposes:
1. During training: Allows us to train on all positions simultaneously (teacher forcing)
— each position learns to predict the next token based only on previous positions.
2. During inference: Naturally enforced — you generate tokens one at a time, so later
tokens don't exist yet.

Training vs Inference
Training (Teacher Forcing):
Feed the entire target sequence at once.
Use causal masking to prevent future token leakage.
At every position, compute the probability of the next token.
Compute cross-entropy loss across all positions simultaneously.
Very efficient — one forward pass computes loss for all positions.

flowchart LR
subgraph TRAIN["Training: Parallel Across Positions"]
direction TB
IN["Input: [The, cat, sat, on]"]
TGT["Target: [cat, sat, on, mat]"]
MASK["Causal Mask Applied"]
PRED["Predictions at all positions simultaneously"]
LOSS["Cross-Entropy Loss"]
IN & MASK --> PRED --> LOSS
TGT --> LOSS
end

Inference (Autoregressive Generation):

Start with a prompt: "The cat"
Forward pass → predict next token → "sat"
Append "sat" to input → "The cat sat"
Forward pass again → predict next token → "on"
Repeat until stop token or max length

flowchart TD
P["Prompt: 'The cat'"]
S1["Forward pass\n→ predict 'sat'"]
S2["Append: 'The cat sat'\n→ Forward pass\n→ predict 'on'"]
S3["Append: 'The cat sat on'\n→ Forward pass\n→ predict 'the'"]
SN["...repeat until &lt;eos&gt; or max length"]

P --> S1 --> S2 --> S3 --> SN

KV Cache (Preview)
During autoregressive generation, every new token requires a full forward pass. But the Key
and Value vectors for previous tokens don't change between steps! The KV Cache stores
these computed K and V vectors and reuses them.
Without KV cache: Step N requires computing attention over N tokens → O(N) work per
step With KV cache: Step N only computes Q for the new token, looks up cached K and V →
O(1) work per step (covered in detail in Module 02)

Data Flow

flowchart TD
subgraph INFERENCE["Autoregressive Inference"]
PROMPT["Prompt Tokens\n[token₁, ..., tokenₚ]"]

PREFILL["Prefill Phase\n(Process all prompt tokens at once)\n[P × d_model]"]

GEN1["Generate token₁\n→ append → new input [P+1]"]

GEN2["Generate token₂\n→ append → new input [P+2]"]
GENN["Generate tokenₙ\n→ until &lt;eos&gt;"]

OUT["Generated Text\n[token₁, ..., tokenₙ]"]

end

PROMPT --> PREFILL --> GEN1 --> GEN2 --> GENN --> OUT

Memory Flow
At inference, for a model with N layers, d_model, and current sequence length L:

Tensor Memory

Model weights Fixed: ~2 × param_count bytes (bfloat16)

KV cache N × 2 × L × d_model × num_heads × d_head bytes

Current activations O(L × d_model)

For LLaMA 3 70B at L=4096:

Model weights ≈ 140 GB (bfloat16)

KV cache = 80 layers × 2 × 4096 × 8192 × 8 heads × 128 × 2 bytes ≈ 64 GB
Total VRAM needed ≈ 200+ GB

This is why LLaMA 70B needs 4+ A100 GPUs (each with 80GB VRAM).
CPU vs GPU Responsibilities
Phase CPU GPU

Prompt parsing ✅ ❌

Tokenization ✅ ❌

Prefill (batch matrix multiplies) ❌ ✅ GPU-bound

Decode (one token per step) ❌ ✅ Memory-bandwidth bound

Sampling ❌ (sometimes) ✅

KV cache management ✅ (metadata) ✅ (actual tensors)

Prefill is compute-bound (large matrix multiplications). Decode is memory-bandwidth

bound (small matrix multiplies but lots of weight reads for each token).

Production Examples
Model Layers Heads d_model Context Vocab

GPT-2 small 12 12 768 1K 50K

GPT-3 96 96 12,288 4K 50K

LLaMA 3 8B 32 32 4,096 128K 128K

LLaMA 3 70B 80 64 8,192 128K 128K

Mistral 7B 32 32 4,096 32K 32K

Common Misconceptions
Misconception Reality

"Decoder-only means the model is a decoder It is a different, standalone architecture — not a

from encoder-decoder architecture" partial encoder-decoder

"You need an encoder to understand input" The decoder processes and understands input
and generates output in one forward pass

"Causal masking is only for generation" It's used during training too — to predict every
position's next token simultaneously

"All tokens are generated one by one" Prompt tokens are processed in parallel (prefill);
only generation is sequential
Interview Questions
1. What is a decoder-only Transformer? (A Transformer with only causal self-attention
blocks — no encoder or cross-attention)
2. What is the causal mask? (A triangular mask preventing each token from attending to
future tokens)
3. What is the difference between prefill and decode? (Prefill processes prompt in
parallel; decode generates new tokens one at a time)
4. Why does GPT use decoder-only? (For language modeling — predicting next token
— which naturally fits autoregressive generation)
5. What is teacher forcing? (During training, feeding ground-truth previous tokens as
input rather than the model's own predictions)

Hands-on Exercise

python
 import torch
import [Link] as nn
import [Link] as F

class DecoderOnlyTransformer([Link]):
def __init__(self, vocab_size, d_model, num_heads, d_ff, num_layers, max_se
super().__init__()
[Link] = [Link](vocab_size, d_model)
self.pos_embedding = [Link](max_seq_len, d_model)

[Link] = [Link]([
TransformerBlock(d_model, num_heads, d_ff)
for _ in range(num_layers)
])

self.ln_f = [Link](d_model)
[Link] = [Link](d_model, vocab_size, bias=False)
# Weight tying: share embedding and output weights
[Link] = [Link]

def forward(self, token_ids):

B, L = token_ids.shape
positions = [Link](L, device=token_ids.device)

x = [Link](token_ids) + self.pos_embedding(positions)

for block in [Link]:

x = block(x)

x = self.ln_f(x)
logits = [Link](x) # [B, L, vocab_size]
return logits

# Tiny GPT-like model

model = DecoderOnlyTransformer(
vocab_size=1000, d_model=64, num_heads=4,
d_ff=256, num_layers=4, max_seq_len=128
)
tokens = [Link](0, 1000, (2, 16)) # batch=2, seq=16
logits = model(tokens)
print(f"Logits shape: {[Link]}") # [2, 16, 1000]

# Compute loss (next-token prediction)

loss = F.cross_entropy(logits[:, :-1].reshape(-1, 1000), tokens[:, 1:].reshape(
print(f"Loss: {[Link]():.4f}")
 
Summary
Decoder-only Transformers are the backbone of all modern LLMs. They use causal self-
attention (only looking at past tokens) to enable autoregressive text generation. During
training, all positions are processed simultaneously with teacher forcing. During inference,
tokens are generated one at a time (with KV caching for efficiency). The architecture scales
extremely well, enabling models from 1B to 1T+ parameters.

Chapter 10 — Output Layer

What is it?
The Output Layer is the final component of a Transformer that converts the model's
internal representations (vectors of shape d_model ) into a probability distribution over all
possible next tokens.
It consists of:
1. Final LayerNorm: Normalize the last hidden state
2. Linear (LM Head): Project from d_model → vocab_size
3. Softmax: Convert raw scores (logits) to probabilities
After the softmax, the model picks the next token using a sampling strategy.

Why do we need it?

The Transformer's internal representations are vectors of shape [d_model] (e.g., 4096
floating-point numbers). That's not something you can directly interpret as "the next word."
The output layer bridges the gap:
Internal representation: 4096-dimensional vector (meaningless to humans)
Output: Probability over all 128,256 possible tokens (interpretable)
Then from those probabilities, a decoding strategy picks the actual output token.

Real-World Analogy
Imagine you're a judge on a quiz show. After hearing all the clues, you have a "confidence
state" in your head (complex, hard to describe). When it's time to answer, you must choose
one word from a list of 128,000 options. Your internal confidence → a ranking of all possible
answers → you choose based on that ranking.
The linear projection is "converting your confidence into rankings." Softmax is
"normalizing those rankings into percentages." The sampling strategy is "how you actually
pick the answer."
Internal Architecture

flowchart TD
HS["Last Hidden State\n[batch × L × d_model]\ne.g., [2 × 1024 × 4096]"]

LN["Final LayerNorm\n(normalize last hidden states)\n[batch × L × d_model]"]

LINEAR["Linear (LM Head)\nW: [d_model × vocab_size]\ne.g., [4096 × 128256]\nOutput

LOGITS["Logits\n[batch × L × vocab_size]\n(raw unnormalized scores)\ne.g., [2 × 102

LAST["Take last position logits\n[batch × vocab_size]\n(only care about next token)

SOFTMAX["Softmax\n[batch × vocab_size]\n(probabilities, sum to 1.0)"]

SAMPLE["Sampling Strategy\n(greedy / top-k / top-p / temperature)"]

TOKEN["Next Token ID\n(integer)"]

HS --> LN --> LINEAR --> LOGITS --> LAST --> SOFTMAX --> SAMPLE --> TOKEN

 

How it Works Step by Step

Step 1 — Final LayerNorm

python

x = LayerNorm(x) # [B, L, d_model] — stabilize before projection

Step 2 — Linear Projection (LM Head)

python

logits = x @ W_lm_head.T # [B, L, d_model] × [d_model, vocab] = [B, L, vocab]

# For LLaMA 3: [2, 1024, 4096] × [4096, 128256] = [2, 1024, 128256]

Step 3 — Extract Last Position

python

last_logits = logits[:, -1, :] # [B, vocab_size] — only the last token's pred
 

Step 4 — Temperature Scaling

python

temperature = 0.8 # < 1 makes distribution sharper, > 1 makes it flatter

last_logits = last_logits / temperature
Step 5 — Softmax → Probabilities

python

probs = softmax(last_logits, dim=-1) # [B, vocab_size], each row sums to 1

Step 6 — Sampling Strategy Choose how to pick the next token.

Sampling Strategies
This is crucial for LLM behavior. The same model with different sampling produces very
different text.
Greedy Decoding

python

next_token = argmax(probs) # Always pick highest probability token

Deterministic, but boring and often repetitive

Used when consistency matters (coding, structured outputs)
Temperature Sampling

python

# Scale logits by temperature before softmax

scaled_logits = logits / T
probs = softmax(scaled_logits)
next_token = multinomial(probs, 1) # Random sample from distribution

T < 1 (e.g., 0.3): More focused, less creative

T = 1.0: Original distribution
T > 1 (e.g., 1.5): More random, creative but can be incoherent
Top-K Sampling

python

# Keep only top K tokens, zero out the rest

top_k = 50
top_k_logits, top_k_indices = topk(logits, k=top_k)
probs = softmax(top_k_logits)
next_token = multinomial(probs, 1)
# Map back to vocab indices

Prevents sampling from very unlikely tokens

Fixed K, regardless of how peaked the distribution is
Top-P (Nucleus) Sampling

python

# Keep tokens until cumulative probability reaches p

p = 0.9
sorted_probs, sorted_indices = sort(softmax(logits), descending=True)
cumulative_probs = cumsum(sorted_probs)
# Remove tokens once cumulative prob > p
cutoff = first index where cumulative_probs > p
# Sample from remaining tokens

Dynamic K — adjusts based on how confident the model is

More commonly used in practice than top-K

flowchart LR
subgraph PROBS["Example Probabilities"]
T1["'cat' = 40%"]
T2["'dog' = 25%"]
T3["'tree' = 15%"]
T4["'car' = 10%"]
T5["'..." = 10%\n(other 128K tokens)"]
end

subgraph GREEDY["Greedy"]
G["Always pick 'cat' (40%)"]
end

subgraph TOPK["Top-K (K=3)"]

TK["Sample from: cat(40%), dog(25%), tree(15%)\nNormalize: cat=47%, dog=29%, tr
end

subgraph TOPP["Top-P (p=0.9)"]

TP["cat=40%, dog=25%, tree=15%, car=10% → 90% covered\nSample from these 4 toke
end

PROBS --> GREEDY & TOPK & TOPP

 

Beam Search
A more sophisticated generation strategy (used more in translation than LLMs):
Maintain B best partial sequences (beams) at each step
For each beam, expand with top tokens
Keep top B combined sequences
 flowchart TD
S["Start: '<bos>'"]

subgraph STEP1["Step 1 (B=3 beams)"]

B1_1["The (p=0.3)"]
B1_2["A (p=0.25)"]
B1_3["Once (p=0.2)"]
end

subgraph STEP2["Step 2"]

B2_1["The cat (0.3×0.4=0.12)"]
B2_2["The dog (0.3×0.2=0.06)"]
B2_3["A cat (0.25×0.35=0.087)"]
B2_4["A dog (0.25×0.15=0.037)"]
B2_5["Once upon (0.2×0.9=0.18)"]
end

S --> STEP1
B1_1 --> B2_1 & B2_2
B1_2 --> B2_3 & B2_4
B1_3 --> B2_5

Keep top 3: ["Once upon", "The cat", "A cat"] → continue expanding.

Data Flow

flowchart TD
H["Final Hidden States\n[B × L × d_model]"]

LN["LayerNorm → [B × L × d_model]"]

WH["W_head: [d_model × vocab_size]\nMemory: 4096×128256×2 = ~1GB for LLaMA3"]

LOGIT["Logits: [B × L × vocab_size]\nSize: 2 × 1024 × 128256 ≈ 500MB!"]

LAST_POS["Last position: [B × vocab_size]"]

TEMP["Temperature scaling"]

SF["Softmax: [B × vocab_size]"]

SAMP["Sampling: [B] (one token ID per batch item)"]

H --> LN
WH --> LOGIT
LN --> LOGIT --> LAST_POS --> TEMP --> SF --> SAMP

Memory Flow
The LM head weight matrix is one of the largest in the model:
 Model vocab_size d_model LM head size

GPT-2 small 50,257 768 ~150MB

LLaMA 3 8B 128,256 4,096 ~1GB

LLaMA 3 70B 128,256 8,192 ~2GB

With weight tying (LM head = transpose of embedding matrix), this memory is shared.
The logit tensor [B × L × vocab_size] can be huge during training. For batch=32, L=2048,
vocab=50K: 32 × 2048 × 50000 × 2 bytes = 6.5 GB . In practice, the loss is computed
directly from logits without materializing the full softmax.

CPU vs GPU Responsibilities

Task CPU GPU

LM head projection (big GEMM) ❌ ✅ Large matrix multiply

Softmax ❌ ✅ Fused kernel

Top-k selection ❌ ✅ Parallel sort

Sampling (multinomial) ❌ ✅ CUDA random

Temperature scaling ❌ ✅ Element-wise

Token ID decoding (back to text) ✅ ❌

Production Examples
Logit processors: In production (vLLM, TGI, etc.), "logit processors" are applied before
sampling:
Repetition penalty: Reduce probability of recently generated tokens
Frequency penalty: Penalize frequent tokens
Presence penalty: Penalize any token that has appeared
Stop tokens: If a stop token has high probability, halt generation
Speculative Decoding (preview of Module 02): A small "draft" model generates K tokens
quickly. The large model verifies them all in one parallel forward pass. If accepted, you
generate K tokens with the cost of ~1 forward pass of the big model. Huge speedup.
Common Misconceptions
Misconception Reality

"Softmax always picks the Softmax just converts to probabilities; sampling decides what to
highest probability" pick

"Temperature=0 means greedy" T=0 causes division by zero; instead set T→0 or just use argmax

"Top-p always selects p% of Top-p selects however many tokens are needed to reach p
tokens" cumulative probability — could be 1 or 1000 tokens

"The model knows when to It generates a special <eos> (end-of-sequence) token; the
stop" system then stops generation

FAQ
Q: What is the difference between logits and probabilities? A: Logits are raw,
unnormalized scores (can be any real number). Probabilities are after softmax — they are
positive and sum to 1.
Q: Why not always use greedy decoding? A: Greedy decoding often produces repetitive,
boring text. It always picks the most likely word, which leads to loops like "the the the the..."
Sampling with temperature produces more varied and interesting outputs.
Q: What is a good temperature for coding vs creative writing? A: Coding: 0.0–0.2
(deterministic, precise). Creative writing: 0.7–1.0 (varied, expressive). Very high
temperatures (>1.2) produce incoherent text.
Q: What is the <bos> and <eos> token? A: Beginning-of-sequence and end-of-sequence
tokens. <bos> signals the start of generation. <eos> signals the model has finished
generating. The system monitors for <eos> to stop the generation loop.

Interview Questions
1. What is the LM head? (A linear projection from d_model → vocab_size, converting
hidden states to token logits)
2. What is the difference between top-k and top-p sampling? (Top-k: fixed K
candidates. Top-p: dynamic K based on cumulative probability threshold)
3. What does temperature do to the output distribution? (T<1: sharper, more confident.
T>1: flatter, more random)
4. Why is the LM head often weight-tied with embeddings? (Saves memory; the same
semantic relationships apply — tokens close in embedding space should have similar
logit responses)
5. What is beam search? (Maintains B best partial sequences and expands them, finding
better global sequences than greedy)
Hands-on Exercise

python

import torch
import [Link] as F

def sample_next_token(logits, temperature=1.0, top_k=None, top_p=None):

"""
logits: [vocab_size] — raw model output for one position
"""
logits = logits / temperature

# Top-K filtering
if top_k is not None:
values, _ = [Link](logits, top_k)
min_val = values[-1]
logits = logits.masked_fill(logits < min_val, float('-inf'))

# Top-P (nucleus) filtering

if top_p is not None:
sorted_logits, sorted_indices = [Link](logits, descending=True)
cumulative_probs = [Link]([Link](sorted_logits, dim=-1), dim=-
# Remove tokens once cumulative prob > top_p
sorted_indices_to_remove = cumulative_probs > top_p
# Shift right to keep at least 1 token
sorted_indices_to_remove[1:] = sorted_indices_to_remove[:-1].clone()
sorted_indices_to_remove[0] = False
indices_to_remove = sorted_indices[sorted_indices_to_remove]
logits[indices_to_remove] = float('-inf')

probs = [Link](logits, dim=-1)

next_token = [Link](probs, num_samples=1)
return next_token.item()

# Simulate logits for a 10-token vocabulary

vocab_size = 10
logits = [Link]([1.0, 3.0, 0.5, 2.0, 0.1, 0.2, 1.5, 0.8, 0.3, 0.4])
probs = [Link](logits, dim=-1)
print("Probabilities:", [f"{p:.3f}" for p in [Link]()])

# Try different strategies

print("\nGreedy:", [Link]().item()) # Always picks highest
print("Temp=0.5:", sample_next_token([Link](), temperature=0.5))
print("Top-K=3:", sample_next_token([Link](), temperature=1.0, top_k=3))
print("Top-P=0.9:", sample_next_token([Link](), temperature=1.0, top_p=0.
 
Summary
The output layer converts the Transformer's final hidden states to token probabilities via a
linear projection (LM head) and softmax. The LM head is one of the largest weight
matrices. Sampling strategies (greedy, temperature, top-k, top-p) determine how the next
token is selected from the probability distribution. In production, various logit processors
apply constraints before sampling to improve output quality and safety.

Mini Project — Build a Tiny Transformer in Python

Goal
Build a working character-level language model using a decoder-only Transformer. Train it
on a small text corpus and generate new text. This combines all 10 chapters into one
working system.
Setup

python
# Install dependencies
# pip install torch

import torch
import [Link] as nn
import [Link] as F
import math

# ─────────────────────────────────────────────
# 1. DATA PREPARATION
# ─────────────────────────────────────────────

# Use a small text corpus

text = """
To be or not to be that is the question
Whether tis nobler in the mind to suffer
The slings and arrows of outrageous fortune
Or to take arms against a sea of troubles
And by opposing end them To die to sleep
No more and by a sleep to say we end
The heartache and the thousand natural shocks
""" * 50 # Repeat to have more data

# Build character vocabulary

chars = sorted(set(text))
vocab_size = len(chars)
char2idx = {c: i for i, c in enumerate(chars)}
idx2char = {i: c for c, i in [Link]()}

# Encode
data = [Link]([char2idx[c] for c in text], dtype=[Link])
print(f"Vocab size: {vocab_size}")
print(f"Data length: {len(data)} characters")

# Train/val split
n = int(0.9 * len(data))
train_data = data[:n]
val_data = data[n:]

# ─────────────────────────────────────────────
# 2. HYPERPARAMETERS
# ─────────────────────────────────────────────

BLOCK_SIZE = 64 # Context length (sequence length)

BATCH_SIZE = 32
D_MODEL = 128 # Embedding / model dimension
N_HEADS = 4 # Number of attention heads
N_LAYERS = 4 # Number of Transformer blocks
D_FF = 512 # FFN inner dimension (4× d_model)
DROPOUT = 0.1
LR = 3e-4
MAX_ITERS = 3000
DEVICE = 'cuda' if [Link].is_available() else 'cpu'
print(f"Training on: {DEVICE}")

# ─────────────────────────────────────────────
# 3. DATA LOADER
# ─────────────────────────────────────────────

def get_batch(split):
data = train_data if split == 'train' else val_data
idx = [Link](len(data) - BLOCK_SIZE, (BATCH_SIZE,))
x = [Link]([data[i:i+BLOCK_SIZE] for i in idx])
y = [Link]([data[i+1:i+BLOCK_SIZE+1] for i in idx])
return [Link](DEVICE), [Link](DEVICE)

# ─────────────────────────────────────────────
# 4. MODEL COMPONENTS
# ─────────────────────────────────────────────

class CausalSelfAttention([Link]):
"""Single attention head (we'll use [Link] for multi-head)"""
def __init__(self, d_model, num_heads):
super().__init__()
assert d_model % num_heads == 0
self.num_heads = num_heads
self.head_dim = d_model // num_heads

[Link] = [Link](d_model, 3 * d_model, bias=False)

[Link] = [Link](d_model, d_model, bias=False)
[Link] = [Link](DROPOUT)

# Causal mask (upper triangular = -inf)

mask = [Link]([Link](BLOCK_SIZE, BLOCK_SIZE), diagonal=1)
self.register_buffer('mask', [Link]())

def forward(self, x):

B, L, D = [Link]
h, d_k = self.num_heads, self.head_dim

# Project to Q, K, V
qkv = [Link](x) # [B, L, 3D]
Q, K, V = [Link](D, dim=-1) # Each [B, L, D]
 # Reshape to [B, h, L, d_k]
Q = [Link](B, L, h, d_k).transpose(1, 2)
K = [Link](B, L, h, d_k).transpose(1, 2)
V = [Link](B, L, h, d_k).transpose(1, 2)

# Attention scores
scores = Q @ [Link](-2, -1) / [Link](d_k) # [B, h, L, L]

# Apply causal mask

scores = scores.masked_fill([Link][:L, :L], float('-inf'))

# Softmax + weighted sum

weights = [Link](scores, dim=-1)
weights = [Link](weights)
out = weights @ V # [B, h, L, d_k]

# Merge heads
out = [Link](1, 2).contiguous().view(B, L, D)
return [Link](out)

class TransformerBlock([Link]):
def __init__(self, d_model, num_heads, d_ff):
super().__init__()
self.ln1 = [Link](d_model)
[Link] = CausalSelfAttention(d_model, num_heads)
self.ln2 = [Link](d_model)
[Link] = [Link](
[Link](d_model, d_ff),
[Link](),
[Link](d_ff, d_model),
[Link](DROPOUT),
)

def forward(self, x):

x = x + [Link](self.ln1(x)) # Residual + Attention
x = x + [Link](self.ln2(x)) # Residual + FFN
return x

class TinyGPT([Link]):
def __init__(self):
super().__init__()
self.token_emb = [Link](vocab_size, D_MODEL)
self.pos_emb = [Link](BLOCK_SIZE, D_MODEL)
[Link] = [Link](DROPOUT)
 [Link] = [Link](*[
TransformerBlock(D_MODEL, N_HEADS, D_FF)
for _ in range(N_LAYERS)
])
self.ln_f = [Link](D_MODEL)
self.lm_head = [Link](D_MODEL, vocab_size, bias=False)
# Weight tying
self.lm_head.weight = self.token_emb.weight

total_params = sum([Link]() for p in [Link]())

print(f"Model parameters: {total_params:,}")

def forward(self, idx, targets=None):

B, L = [Link]
positions = [Link](L, device=[Link])

# Embeddings + Positional Encoding

x = [Link](self.token_emb(idx) + self.pos_emb(positions))

# Transformer blocks
x = [Link](x)

# Output
x = self.ln_f(x)
logits = self.lm_head(x) # [B, L, vocab_size]

loss = None
if targets is not None:
loss = F.cross_entropy(
[Link](-1, vocab_size),
[Link](-1)
)
return logits, loss

@torch.no_grad()
def generate(self, idx, max_new_tokens, temperature=0.8, top_k=20):
for _ in range(max_new_tokens):
# Trim to BLOCK_SIZE
idx_cond = idx[:, -BLOCK_SIZE:]

# Forward pass
logits, _ = self(idx_cond)
logits = logits[:, -1, :] / temperature # [B, vocab_size]

# Top-K sampling
v, _ = [Link](logits, min(top_k, [Link](-1)))
logits[logits < v[:, [-1]]] = float('-inf')
 probs = [Link](logits, dim=-1)
next_token = [Link](probs, num_samples=1)
idx = [Link]([idx, next_token], dim=1)
return idx

# ─────────────────────────────────────────────
# 5. TRAINING LOOP
# ─────────────────────────────────────────────

model = TinyGPT().to(DEVICE)
optimizer = [Link]([Link](), lr=LR)

@torch.no_grad()
def estimate_loss():
[Link]()
losses = {}
for split in ['train', 'val']:
loss_vals = []
for _ in range(50):
x, y = get_batch(split)
_, loss = model(x, y)
loss_vals.append([Link]())
losses[split] = sum(loss_vals) / len(loss_vals)
[Link]()
return losses

print("Starting training...")
for step in range(MAX_ITERS):
x, y = get_batch('train')
logits, loss = model(x, y)
optimizer.zero_grad()
[Link]()
[Link].clip_grad_norm_([Link](), 1.0)
[Link]()

if step % 500 == 0:
losses = estimate_loss()
print(f"Step {step}: train_loss={losses['train']:.4f}, val_loss={losses

# ─────────────────────────────────────────────
# 6. GENERATION
# ─────────────────────────────────────────────

[Link]()
context = [Link]([[char2idx['\n']]], device=DEVICE)
generated_ids = [Link](context, max_new_tokens=300)
 generated_text = ''.join([idx2char[[Link]()] for i in generated_ids[0]])
print("\n=== GENERATED TEXT ===")
print(generated_text)

What to observe:
Loss should decrease from ~3.0 to ~1.5+ over 3000 steps
Generated text should start looking like Shakespearean phrases
Increase N_LAYERS, D_MODEL, or MAX_ITERS for better results

Module Summary
Here is a complete picture of how all 10 chapters fit together:

flowchart TD
subgraph CH1["CH1: Transformer Overview"]
ARCH["Architecture Overview\n(Encoder-Decoder / Decoder-Only)"]
end

subgraph CH2_3["CH2-3: Input Preparation"]

EMB["Embeddings\n(token → vector)"]
PE["Positional Encoding\n(add order information)"]
end

subgraph CH4_5["CH4-5: Attention"]

SA["Self Attention\n(Q·K^T/√d_k)·V"]
MHA["Multi-Head Attention\n(run attention h times in parallel)"]
end

subgraph CH6["CH6: Processing"]

FFN["Feed Forward Network\n(per-token, expand-activate-compress)"]
end

subgraph CH7_8["CH7-8: Stability"]

LN["LayerNorm / RMSNorm\n(stabilize activations)"]
RES["Residual Connections\nx + layer(x)"]
end

subgraph CH9["CH9: Architecture"]

DEC["Decoder-Only\n(causal mask, autoregressive)"]
end

subgraph CH10["CH10: Output"]

OUT["Output Layer\n(logits → probabilities → token)"]
SAMP["Sampling Strategies\n(greedy / top-k / top-p / temperature)"]
end

CH1 --> CH2_3 --> CH4_5 --> CH6 --> CH7_8 --> CH9 --> CH10
Key Takeaways
# Concept One-Line Summary

1 Transformer Neural network that processes all tokens simultaneously using

attention

2 Embeddings Lookup table converting token IDs to dense vectors; shape [vocab ×
d_model]

3 Positional Encoding Adds position information to embeddings; RoPE is modern standard

4 Self Attention Each token attends to all others: softmax(QK^T/√d_k)·V

5 Multi-Head Attention Runs attention h times in parallel; each head captures different
relationships

6 Feed Forward Per-token non-linear transformation; stores factual knowledge

Network

7 LayerNorm Normalizes activations to stabilize training; RMSNorm in modern

LLMs

8 Residual x + layer(x); prevents vanishing gradients; creates "residual stream"

Connections

9 Decoder-Only Causal masking + autoregressive generation; the backbone of all

LLMs

10 Output Layer LM head (linear) + softmax + sampling gives the next token

Training Teacher forcing: all positions learned simultaneously with causal

mask

Inference Autoregressive: generate one token at a time, with KV cache for

efficiency

Scale More layers + more heads + more d_model + more data = smarter
model

What's Next?
Type NEXT to proceed to Module 02 — LLM Inference, which covers how models like
LLaMA and GPT run efficiently at production scale: KV Cache, PagedAttention,
Continuous Batching, Speculative Decoding, and more.

AI Engineering Handbook — Module 01 of 11 "Understanding the Transformer is the

prerequisite for everything in AI Engineering."