The Building Blocks of Transformers: A Complete Guide

The Building Blocks of Transformers: A Complete Guide

A beginner-friendly guide to understanding the parts of the Transformer architecture.

1. Introduction: From Attention to Transformers

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:

What if we remove the RNN entirely?

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.

pratik1.png

2. The Big Picture: Encoder and Decoder

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:

  • Encoder = Listener: It listens to what you say and creates an understanding.
  • Decoder = Speaker: It uses that understanding to produce a response.

The Encoder Stack:

  • Self-Attention

  • Multi-Head Attention

  • Positional Encoding

  • Feed-Forward Network (FFN)

The Decoder Stack:

  • Masked Multi-Head Attention

  • Cross-Attention

  • Feed-Forward Network (FFN)

In this article, we will explain each of these building blocks in detail.

image (10).png

PART 1: THE ENCODER

3. Self-Attention: Words Talking to Each Other

What is Self-Attention?

In a normal sentence, the meaning of a word depends on the words around it.

Example:

  • "The bank of the river was beautiful."
  • "I went to the bank to deposit money."

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.

The Three Players: Query, Key, Value

In Self-Attention, every word plays three roles:

RoleAnalogyPurpose
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

How it works:

  1. Each word creates a Query: "Who should I pay attention to?"
  2. Each word creates a Key: "Here's my label, see if you need me."
  3. Each word creates a Value: "Here's my content if you pick me."

Step-by-Step Process

Step 1: Create Q, K, V Vectors

For each word embedding x, we multiply by three weight matrices (W^Q, W^K, W^V) to get:

  • q = x × W^Q (Query)
  • k = x × W^K (Key)
  • v = x × W^V (Value)

pratik2.png

Step 2: Calculate Attention Scores

We calculate how much each word should "attend" to every other word using the dot product:

Attention Score = Q × K^T / √d_k

Where:

  • Q × K^T gives us the similarity between words
  • √d_k is a scaling factor to keep values stable

pratik3.png

Step 3: Apply Softmax

We convert scores to probabilities using softmax:

Attention Weights = softmax(Attention Score)

Step 4: Weighted Sum of Values

Finally, we compute the output as a weighted sum:

Output = Attention Weights × V

pratik4.png

Why Divide by √d_k?

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.

Self-Attention Formula

The complete Self-Attention formula is:

pratik5.png

pratik6.png

Code Example

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)

4. Multi-Head Attention: Looking from Multiple Perspectives

Why Multiple Heads?

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.

What is Multi-Head Attention?

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:

  • Head 1 might focus on subject-verb relationships
  • Head 2 might focus on adjective-noun relationships
  • Head 3 might focus on long-distance dependencies

How it Works

Step 1: Split the embedding dimension into h heads

  • If d_model = 512 and h = 8, each head gets d_k = 512/8 = 64 dimensions

Step 2: Each head performs Self-Attention independently

Step 3: Concatenate all head outputs

Step 4: Apply a final linear transformation

The Formula

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)

pratik7.png

Code Example

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)

5. Positional Encoding: Teaching Order to Transformers

The Problem

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.

The Solution: Positional Encoding

image (11).png

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 Sinusoidal Formula

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

Why Sine and Cosine?

  1. Unique patterns: Each position gets a unique pattern

  2. Relative positions: The model can learn relative positions (e.g., "3 words apart")

  3. Generalization: Works for any sequence length, even longer than training

Visualization

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)

pratik8.png

Code Example

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)

image (12).png

PART 2: THE DECODER

pratik9.png

6. Masked Multi-Head Attention: No Peeking at the Future!

The Problem in Decoding

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!

The Solution: Masking

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)

image (13).png

When we apply softmax to (-∞), it becomes 0, effectively blocking those positions.

The Mask Matrix

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]

pratik10.png

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

pratik11.png

Code Example

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)

7. Cross-Attention: Encoder Meets Decoder

What is Cross-Attention?

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.

How is it Different from Self-Attention?

Attention TypeQuery comes fromKey & Value come from
Self-AttentionSame sequenceSame sequence
Cross-AttentionDecoderEncoder

In Cross-Attention:

  • Query (Q) comes from the Decoder (what am I translating now?)
  • Key (K) and Value (V) come from the Encoder (what did the source say?)

The Formula

CrossAttention(Q_decoder, K_encoder, V_encoder) = 
   
   softmax(Q_decoder × K_encoder^T / √d_k) × V_encoder

pratik12.png

Code Example

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)

8. Feed-Forward Network (FFN): The Thinking Layer

What is the FFN?

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 Structure

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

pratik14.png

Key Details

  1. Expansion then compression: The hidden layer is typically 4× larger than the input

    • Input: 512 dimensions
    • Hidden: 2048 dimensions (512 × 4)
    • Output: 512 dimensions
  2. Position-wise: The same FFN is applied to each position independently

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

10. Summary: The Transformer Recipe

Let's summarize what we learned:

ComponentLocationPurpose
Self-AttentionEncoder & DecoderWords attending to each other in the same sequence
Multi-Head AttentionEverywhereMultiple perspectives on relationships
Positional EncodingInputTeaching word order
Masked Multi-Head AttentionDecoderPreventing future peeking during generation
Cross-AttentionDecoderDecoder attending to Encoder output
Feed-Forward NetworkAfter each attentionIndependent processing at each position

Key Insights

  1. Self-Attention allows every word to see every other word

  2. Multi-Head provides multiple relationship perspectives

  3. Positional Encoding adds order information

  4. Masking enables autoregressive generation

  5. Cross-Attention connects encoder and decoder

  6. FFN adds non-linear processing capacity

11. What's Next?

In the next article, we will cover:

  1. The Complete Transformer Architecture: How to stack these layers

  2. Training Transformers: Loss functions, backpropagation, and attention loss

  3. Fine-tuning: How to adapt pre-trained models

  4. Practical Implementation: Building a character-level GPT from scratch using PyTorch on the Shakespeare dataset

Stay tuned for the complete Transformer journey!

References

https://infominez.com/blog/vanilla_rnns_struggle

About the Author

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.