A beginner-friendly guide to understanding the parts of the Transformer architecture.
In the previous article, we learned about RNNs, LSTMs, and Attention mechanisms. We saw how Attention solved the "memory bottleneck" problem by allowing the model to "look back" at all words in a sequence.
But there was still one problem: RNNs are slow.
Why? Because RNNs read words one by one, like reading a book one word at a time. You can't read word 5 until you've read words 1, 2, 3, and 4. This is called sequential processing.
In 2017, a groundbreaking paper called "Attention Is All You Need" (1706.03762) made a radical proposal:
The result was the Transformer an architecture that processes all words in parallel, making it incredibly fast and powerful.
Analogy: Imagine you're a teacher grading 30 exam papers. The RNN way is to grade them one by one. The Transformer way is to have 30 teaching assistants, each grading one paper at the same time.

The Transformer has two main parts:
1. Encoder: Reads and understands the input (e.g., an English sentence). 2. Decoder: Generates the output (e.g., a French translation).
Think of it like this:
Self-Attention
Multi-Head Attention
Positional Encoding
Feed-Forward Network (FFN)
Masked Multi-Head Attention
Cross-Attention
Feed-Forward Network (FFN)
In this article, we will explain each of these building blocks in detail.

In a normal sentence, the meaning of a word depends on the words around it.
The word "bank" has different meanings based on context. Self-Attention helps the model understand this by allowing each word to "look at" every other word in the sentence.
Analogy: Imagine you're at a party. To understand who someone is, you look at the people around them, their clothes, and their behavior. Self-Attention does the same — each word looks at all other words to understand its own meaning.
In Self-Attention, every word plays three roles:
| Role | Analogy | Purpose |
|---|---|---|
| Query (Q) | "What am I looking for?" | The question a word asks |
| Key (K) | "What do I have to offer?" | The label or identifier of each word |
| Value (V) | "What information do I contain?" | The actual content of each word |
For each word embedding x, we multiply by three weight matrices (W^Q, W^K, W^V) to get:

We calculate how much each word should "attend" to every other word using the dot product:
Attention Score = Q × K^T / √d_k
Where:

We convert scores to probabilities using softmax:
Attention Weights = softmax(Attention Score)
Finally, we compute the output as a weighted sum:
Output = Attention Weights × V

As dimensions grow, dot products become very large. Large values push softmax to extreme regions (close to 0 or 1), causing tiny gradients during training. Dividing by √d_k keeps values in a healthy range.
Example:
If d_k = 64, then √d_k = 8.
This scaling keeps our values manageable.
The complete Self-Attention formula is:

python
import torch
import torch.nn as nn
import torch.nn.functional as F
class SelfAttention(nn.Module):
def __init__(self, d_model, d_k):
super().__init__()
self.d_k = d_k
# Learned weight matrices
self.W_Q = nn.Linear(d_model, d_k, bias=False)
self.W_K = nn.Linear(d_model, d_k, bias=False)
self.W_V = nn.Linear(d_model, d_k, bias=False)
def forward(self, x):
# x shape: (batch_size, seq_len, d_model)
Q = self.W_Q(x) # Query
K = self.W_K(x) # Key
V = self.W_V(x) # Value
# Attention scores
scores = torch.bmm(Q, K.transpose(-2, -1)) / (self.d_k ** 0.5)
# Softmax to get weights
weights = F.softmax(scores, dim=-1)
# Weighted sum of values
output = torch.bmm(weights, V)
return output, weights
# Example usage
d_model, d_k = 512, 64
model = SelfAttention(d_model, d_k)
x = torch.randn(1, 10, 512)
# 1 batch, 10 words, 512 dimensions
output, weights = model(x)
print(f"Output shape: {output.shape}")
# (1, 10, 64)
Single Self-Attention has a limitation: it looks at relationships from only one perspective.
Analogy: Imagine you're analyzing a crime scene. One detective might focus on fingerprints, another on footprints, and another on the timeline. Each perspective gives different insights. Multi-Head Attention is like having multiple detectives working together.
Instead of doing Self-Attention once, we do it h times in parallel (h = number of heads, typically 8 in LLM). Each head learns to focus on different types of relationships:
Step 1: Split the embedding dimension into h heads
Step 2: Each head performs Self-Attention independently
Step 3: Concatenate all head outputs
Step 4: Apply a final linear transformation
MultiHead(Q, K, V) = Concat(head_1, head_2, ..., head_h) × W^O
where head_i = Attention(Q × W^Q_i, K × W^K_i, V × W^V_i)

python
import torch
import torch.nn as nn
import torch.nn.functional as F
class MultiHeadAttention(nn.Module):
def __init__(self, d_model, num_heads):
super().__init__()
self.num_heads = num_heads
self.d_k = d_model // num_heads
self.W_Q = nn.Linear(d_model, d_model, bias=False)
self.W_K = nn.Linear(d_model, d_model, bias=False)
self.W_V = nn.Linear(d_model, d_model, bias=False)
self.W_O = nn.Linear(d_model, d_model, bias=False)
def forward(self, x):
batch_size, seq_len, d_model = x.size()
# Linear projections
Q = self.W_Q(x)
K = self.W_K(x)
V = self.W_V(x)
# Reshape for multi-head
Q = Q.view(batch_size, seq_len, self.num_heads, self.d_k).transpose(1, 2)
K = K.view(batch_size, seq_len, self.num_heads, self.d_k).transpose(1, 2)
V = V.view(batch_size, seq_len, self.num_heads, self.d_k).transpose(1, 2)
# Attention scores
scores = torch.matmul(Q, K.transpose(-2, -1)) / (self.d_k ** 0.5)
weights = F.softmax(scores, dim=-1)
# Apply attention to values
attention_output = torch.matmul(weights, V)
# Concatenate heads
attention_output = attention_output.transpose(1, 2).contiguous()
attention_output = attention_output.view(batch_size, seq_len, d_model)
# Final linear projection
output = self.W_O(attention_output)
return output
# Example
model = MultiHeadAttention(d_model=512, num_heads=8)
x = torch.randn(1, 10, 512)
output = model(x)
print(f"Multi-Head Output shape: {output.shape}")
# (1, 10, 512)
Unlike RNNs, Transformers process all words simultaneously. But this creates a problem: the model doesn't know word order!
For the Transformer, "dog bites man" and "man bites dog" look the same because both sentences have the same words.

We add special vectors to each word embedding that encode its position in the sequence.
Analogy: Imagine you're arranging seats in a theater. Each seat has a unique number (row 1 seat 5, row 3 seat 12). Positional encoding gives each word a unique "seat number" so the model knows where each word sits in the sentence.
The original Transformer uses sine and cosine functions to create positional encodings:
PE(pos, 2i) = sin(pos / 10000^(2i/d_model))
PE(pos, 2i+1) = cos(pos / 10000^(2i/d_model))
Where:
pos = position of the word (0, 1, 2, ...)
i = dimension index
d_model = embedding dimension
Unique patterns: Each position gets a unique pattern
Relative positions: The model can learn relative positions (e.g., "3 words apart")
Generalization: Works for any sequence length, even longer than training
If you plot positional encodings, you'll see wave patterns:
Lower dimensions have slow waves (capture long-range positions)
Higher dimensions have fast waves (capture fine-grained positions)

python
import torch
import torch.nn as nn
import math
class PositionalEncoding(nn.Module):
def __init__(self, d_model, max_len=5000):
super().__init__()
# Create positional encoding matrix
pe = torch.zeros(max_len, d_model)
position = torch.arange(0, max_len).unsqueeze(1).float()
div_term = torch.exp(
torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model)
)
pe[:, 0::2] = torch.sin(position * div_term) # Even indices
pe[:, 1::2] = torch.cos(position * div_term) # Odd indices
pe = pe.unsqueeze(0) # Add batch dimension
self.register_buffer('pe', pe)
def forward(self, x):
# x shape: (batch_size, seq_len, d_model)
seq_len = x.size(1)
return x + self.pe[:, :seq_len, :]
# Example
pos_encoder = PositionalEncoding(d_model=512)
x = torch.randn(1, 10, 512) # 10 words
output = pos_encoder(x)
print(f"With positional encoding: {output.shape}") # (1, 10, 512)


When generating output (like translating to French), the decoder produces words one at a time. When predicting word 3, it should only see words 1 and 2 — not word 4 or 5!
Analogy: Imagine taking an exam. You can refer to your previous answers, but you can't peek at future questions. That would be cheating!
We use a mask to hide future positions. The mask is a matrix with:
0 for positions we can see (current and past)
-∞ for positions we can't see (future)

When we apply softmax to (-∞), it becomes 0, effectively blocking those positions.
For a sequence of 5 words, the mask looks like:
Word1 Word2 Word3 Word4 Word5
Word1 [0, -∞, -∞, -∞, -∞]
Word2 [0, 0, -∞, -∞, -∞]
Word3 [0, 0, 0, -∞, -∞]
Word4 [0, 0, 0, 0, -∞]
Word5 [0, 0, 0, 0, 0]

This means:
Word 1 can only see Word 1
Word 2 can see Words 1 and 2
Word 3 can see Words 1, 2, and 3
And so on...

python
import torch
import torch.nn as nn
import torch.nn.functional as F
class MaskedMultiHeadAttention(nn.Module):
def __init__(self, d_model, num_heads):
super().__init__()
self.num_heads = num_heads
self.d_k = d_model // num_heads
self.W_Q = nn.Linear(d_model, d_model, bias=False)
self.W_K = nn.Linear(d_model, d_model, bias=False)
self.W_V = nn.Linear(d_model, d_model, bias=False)
self.W_O = nn.Linear(d_model, d_model, bias=False)
def create_mask(self, seq_len, device):
# Create causal mask to prevent attending to future positions
mask = torch.triu(torch.ones(seq_len, seq_len, device=device), diagonal=1)
mask = mask.masked_fill(mask == 1, float('-inf'))
return mask
def forward(self, x):
batch_size, seq_len, d_model = x.size()
Q = self.W_Q(x)
K = self.W_K(x)
V = self.W_V(x)
# Reshape for multi-head
Q = Q.view(batch_size, seq_len, self.num_heads, self.d_k).transpose(1, 2)
K = K.view(batch_size, seq_len, self.num_heads, self.d_k).transpose(1, 2)
V = V.view(batch_size, seq_len, self.num_heads, self.d_k).transpose(1, 2)
# Attention scores
scores = torch.matmul(Q, K.transpose(-2, -1)) / (self.d_k ** 0.5)
# Apply mask - THIS IS THE KEY DIFFERENCE!
mask = self.create_mask(seq_len, x.device)
scores = scores + mask
weights = F.softmax(scores, dim=-1)
attention_output = torch.matmul(weights, V)
# Concatenate and project
attention_output = attention_output.transpose(1, 2).contiguous()
attention_output = attention_output.view(batch_size, seq_len, d_model)
output = self.W_O(attention_output)
return output
# Example
model = MaskedMultiHeadAttention(d_model=512, num_heads=8)
x = torch.randn(1, 10, 512)
output = model(x)
print(f"Masked output shape: {output.shape}")
# (1, 10, 512)
Cross-Attention connects the Encoder and Decoder. It allows the Decoder to "look at" the Encoder's output while generating each word.
Analogy: Imagine an interpreter translating a speech. The interpreter (Decoder) must constantly refer back to what the speaker said (Encoder output) while producing the translation.
| Attention Type | Query comes from | Key & Value come from |
|---|---|---|
| Self-Attention | Same sequence | Same sequence |
| Cross-Attention | Decoder | Encoder |
CrossAttention(Q_decoder, K_encoder, V_encoder) =
softmax(Q_decoder × K_encoder^T / √d_k) × V_encoder

python
import torch
import torch.nn as nn
import torch.nn.functional as F
class CrossAttention(nn.Module):
def __init__(self, d_model, num_heads):
super().__init__()
self.num_heads = num_heads
self.d_k = d_model // num_heads
# Query from decoder
self.W_Q = nn.Linear(d_model, d_model, bias=False)
# Key and Value from encoder
self.W_K = nn.Linear(d_model, d_model, bias=False)
self.W_V = nn.Linear(d_model, d_model, bias=False)
self.W_O = nn.Linear(d_model, d_model, bias=False)
def forward(self, decoder_input, encoder_output):
"""
decoder_input: (batch, decoder_seq_len, d_model)
encoder_output: (batch, encoder_seq_len, d_model)
"""
batch_size = decoder_input.size(0)
dec_seq_len = decoder_input.size(1)
enc_seq_len = encoder_output.size(1)
# Q from decoder, K and V from encoder
Q = self.W_Q(decoder_input)
K = self.W_K(encoder_output)
V = self.W_V(encoder_output)
# Reshape for multi-head
Q = Q.view(batch_size, dec_seq_len, self.num_heads, self.d_k).transpose(1, 2)
K = K.view(batch_size, enc_seq_len, self.num_heads, self.d_k).transpose(1, 2)
V = V.view(batch_size, enc_seq_len, self.num_heads, self.d_k).transpose(1, 2)
# Attention: decoder attends to encoder
scores = torch.matmul(Q, K.transpose(-2, -1)) / (self.d_k ** 0.5)
weights = F.softmax(scores, dim=-1)
attention_output = torch.matmul(weights, V)
# Concatenate and project
attention_output = attention_output.transpose(1, 2).contiguous()
attention_output = attention_output.view(batch_size, dec_seq_len, -1)
output = self.W_O(attention_output)
return output
# Example
cross_attn = CrossAttention(d_model=512, num_heads=8)
encoder_output = torch.randn(1, 20, 512)
decoder_input = torch.randn(1, 15, 512)
output = cross_attn(decoder_input, encoder_output)
print(f"Cross-Attention output shape: {output.shape}")
# (1, 15, 512)
After each attention layer, the Transformer has a Feed-Forward Network. This is a simple two-layer neural network that processes each position independently.
Analogy: If attention is like a group discussion (words talking to each other), the FFN is like individual thinking time (each word processing its own information).
The FFN has two linear layers with a ReLU activation in between:
FFN(x) = max(0, x × W₁ + b₁) × W₂ + b₂
Or in simpler terms:
Linear → ReLU → Linear

Expansion then compression: The hidden layer is typically 4× larger than the input
Position-wise: The same FFN is applied to each position independently
No interaction between positions: Unlike attention, positions don't communicate here
Code Example
python
import torch
import torch.nn as nn
class FeedForwardNetwork(nn.Module):
def __init__(self, d_model, d_ff, dropout=0.1):
super().__init__()
self.linear1 = nn.Linear(d_model, d_ff)
self.linear2 = nn.Linear(d_ff, d_model)
self.relu = nn.ReLU()
self.dropout = nn.Dropout(dropout)
def forward(self, x):
# x shape: (batch_size, seq_len, d_model)
x = self.linear1(x) # (batch_size, seq_len, d_ff)
x = self.relu(x) # Apply ReLU
x = self.dropout(x) # Regularization
x = self.linear2(x) # (batch_size, seq_len, d_model)
return x
# Example
ffn = FeedForwardNetwork(d_model=512, d_ff=2048)
x = torch.randn(1, 10, 512)
output = ffn(x)
print(f"FFN output shape: {output.shape}") # (1, 10, 512)
Let's summarize what we learned:
| Component | Location | Purpose |
|---|---|---|
| Self-Attention | Encoder & Decoder | Words attending to each other in the same sequence |
| Multi-Head Attention | Everywhere | Multiple perspectives on relationships |
| Positional Encoding | Input | Teaching word order |
| Masked Multi-Head Attention | Decoder | Preventing future peeking during generation |
| Cross-Attention | Decoder | Decoder attending to Encoder output |
| Feed-Forward Network | After each attention | Independent processing at each position |
Self-Attention allows every word to see every other word
Multi-Head provides multiple relationship perspectives
Positional Encoding adds order information
Masking enables autoregressive generation
Cross-Attention connects encoder and decoder
FFN adds non-linear processing capacity
In the next article, we will cover:
The Complete Transformer Architecture: How to stack these layers
Training Transformers: Loss functions, backpropagation, and attention loss
Fine-tuning: How to adapt pre-trained models
Practical Implementation: Building a character-level GPT from scratch using PyTorch on the Shakespeare dataset
Stay tuned for the complete Transformer journey!
Attention Is All You Need → https://arxiv.org/abs/1706.03762
The Illustrated Transformer → https://jalammar.github.io/illustrated-transformer/
Why Vanilla RNNs Struggle with Sequence Problems (Previous Article) ->
About the Author

Pratik Jat
Pratik Jat is an AI Research Engineer at Infominez. He specializes in computer vision, NLP, and deep learning, with a focus on building practical AI solutions for real-world applications.
About the Author

Pratik Jat
Pratik Jat is an AI Research Engineer at Infominez. He specializes in computer vision, NLP, and deep learning, with a focus on building practical AI solutions for real-world applications.