Transformer Architecture

Introduction

Before 2017, the world of Natural Language Processing (NLP) was dominated by Recurrent Neural Networks (RNNs) and LSTMs. These models processed text sequentially, one word at a time, just like a human reading a book. But this sequential nature had two massive flaws: it was impossible to parallelize across GPUs (making training painfully slow), and the models suffered from "amnesia"—by the time they reached the end of a long paragraph, they forgot what happened at the beginning.

Then, Google researchers published a paper titled "Attention Is All You Need." They introduced the Transformer, an architecture that completely discarded sequential processing. Instead of reading left-to-right, the Transformer reads the entire sentence at once, using a mechanism called Self-Attention to figure out which words are related to each other, regardless of how far apart they are in the text.

This single breakthrough is the foundation of all modern AI. It is the "T" in ChatGPT (Generative Pre-trained Transformer). It powers BERT, Midjourney (Vision Transformers), and AlphaFold.

Core Concepts Explained Intuitively

1. Self-Attention: The Core Engine

Imagine the sentence: "The bank of the river was steep, so he couldn't deposit his boat." What does the word "bank" mean here? A financial institution or the side of a river? As a human, you know it's a riverbank because you look at the surrounding words ("river", "steep", "boat").

Self-Attention is the mathematical way a computer does this. For every word in a sentence, the Transformer calculates an "attention score" against every other word in the sentence.

• The word "bank" will compute a high attention score with "river" and "boat".
• The word "bank" will compute a low attention score with "The" and "was".

It does this by creating three vectors for every word: a Query, a Key, and a Value (think of a database: you use a Query to search for a Key, and retrieve a Value). The attention score is calculated via the dot product of the Query of the current word and the Keys of all other words.

2. Multi-Head Attention

Why stop at one attention mechanism? A sentence has multiple types of relationships. One attention "head" might focus on grammar (linking verbs to nouns). Another head might focus on sentiment. Another might focus on pronouns (linking "he" to "Bob").

The Transformer uses Multi-Head Attention, meaning it runs the self-attention process 8, 12, or even 96 times in parallel, allowing the model to capture incredibly complex nuances in the text.

3. Positional Encoding

If the Transformer reads all words simultaneously, how does it know the order of the words? "The dog bit the man" and "The man bit the dog" have the exact same words, but completely different meanings.

Because the model has no inherent sense of sequence, we must inject time/position information artificially. We do this via Positional Encoding. Before the text enters the neural network, we add a unique, oscillating mathematical pattern (usually sine and cosine waves) to the word embeddings. This slight numerical tweak tells the model exactly where each word sits in the sequence.

4. The Encoder-Decoder Structure

The original Transformer was designed for translation (e.g., English to French) and consisted of two halves:

• Encoder: Takes the English sentence, runs it through self-attention, and produces a deep, contextual mathematical representation of the sentence. (BERT is an Encoder-only model).
• Decoder: Takes the Encoder's representation, and generates the French sentence one word at a time, paying attention both to what it has generated so far (masked self-attention) and what the Encoder told it (cross-attention). (GPT is a Decoder-only model).

Code Examples

1. Beginner Code Example: Using a Pre-Trained Transformer (HuggingFace)

You don't need to write the math from scratch to use a Transformer. The transformers library makes it a 3-line process.

# pip install transformers torch
from transformers import pipeline

# 1. Initialize a pipeline for Sentiment Analysis
# By default, this downloads a pre-trained DistilBERT (Transformer) model
classifier = pipeline("sentiment-analysis")

# 2. Pass some text through the Transformer
texts = [
    "I absolutely loved this movie, the acting was phenomenal!",
    "The food was terrible and the service was even worse."
]

# 3. Get predictions
results = classifier(texts)

for text, result in zip(texts, results):
    print(f"Text: '{text}'")
    print(f"Prediction: {result['label']} (Confidence: {result['score']:.4f})\n")

2. Intermediate Code Example: The Math of Scaled Dot-Product Attention

This is the heart of the Transformer. We will calculate the attention scores for a sequence using PyTorch matrix multiplication.

import torch
import torch.nn.functional as F
import math

# Let's assume a sequence of 3 words, each embedded in a 4-dimensional space.
# sequence length = 3, d_model = 4
Q = torch.tensor([[1.0, 0.0, 1.0, 0.0],  # Word 1 query
                  [0.0, 1.0, 0.0, 1.0],  # Word 2 query
                  [1.0, 1.0, 0.0, 0.0]]) # Word 3 query

K = torch.tensor([[1.0, 0.0, 1.0, 0.0],  # Word 1 key
                  [0.0, 1.0, 0.0, 1.0],  # Word 2 key
                  [1.0, 1.0, 0.0, 0.0]]) # Word 3 key

V = torch.tensor([[10.0, 20.0],          # Word 1 value
                  [30.0, 40.0],          # Word 2 value
                  [50.0, 60.0]])         # Word 3 value

d_k = K.size(-1) # Dimension of keys (4)

# Step 1: Dot product of Queries and Keys (Attention Scores)
scores = torch.matmul(Q, K.transpose(-2, -1))

# Step 2: Scale down by the square root of the key dimension (prevents exploding gradients)
scaled_scores = scores / math.sqrt(d_k)

# Step 3: Apply Softmax to turn scores into probabilities (0 to 1)
attention_weights = F.softmax(scaled_scores, dim=-1)

# Step 4: Multiply the weights by the Values to get the final output
output = torch.matmul(attention_weights, V)

print("Attention Weights (Who is paying attention to whom?):")
print(attention_weights)
print("\nFinal Output for the layer:")
print(output)

3. Advanced Code Example: A PyTorch Transformer Block

Here is how you define a single Transformer block (which consists of Multi-Head Attention + a Feed Forward Network + Layer Normalization) using PyTorch's built-in modules.

import torch
import torch.nn as nn

class TransformerBlock(nn.Module):
    def __init__(self, embed_dim, num_heads, ff_dim, dropout=0.1):
        super(TransformerBlock, self).__init__()
        
        # 1. Multi-Head Attention layer
        self.attention = nn.MultiheadAttention(embed_dim, num_heads, dropout=dropout, batch_first=True)
        
        # 2. Feed Forward Network (MLP)
        self.ffn = nn.Sequential(
            nn.Linear(embed_dim, ff_dim),
            nn.ReLU(),
            nn.Linear(ff_dim, embed_dim)
        )
        
        # 3. Layer Normalization
        self.norm1 = nn.LayerNorm(embed_dim)
        self.norm2 = nn.LayerNorm(embed_dim)
        
        # 4. Dropout for regularization
        self.dropout = nn.Dropout(dropout)
        
    def forward(self, x):
        # Step A: Self-Attention with Residual Connection
        # Query, Key, and Value are all 'x' (Self-Attention)
        attn_output, _ = self.attention(x, x, x)
        x = self.norm1(x + self.dropout(attn_output)) # Add & Norm
        
        # Step B: Feed Forward with Residual Connection
        ffn_output = self.ffn(x)
        x = self.norm2(x + self.dropout(ffn_output)) # Add & Norm
        
        return x

# Test the block
# Batch size: 2, Sequence length: 10 words, Embedding dimension: 64
dummy_input = torch.randn(2, 10, 64)
block = TransformerBlock(embed_dim=64, num_heads=8, ff_dim=256)
output = block(dummy_input)

print(f"Input shape: {dummy_input.shape}")
print(f"Output shape: {output.shape}") # Shape remains unchanged!

Common Mistakes

1. Confusing Encoders and Decoders:

   • Mistake: Using an Encoder-only model (like BERT) to write a poem, or a Decoder-only model (like GPT) to perform complex bidirectional sentiment analysis.
   • Fix: Encoders look at the whole sentence at once (bidirectional). They are great for classification, search, and understanding. Decoders can only look at past words (unidirectional/masked). They are designed strictly for text generation.

2. Ignoring the Quadratic Complexity of Attention:

   • Mistake: Trying to pass a 100,000-word document into a standard Transformer and watching your GPU crash due to Out of Memory (OOM) errors.
   • Fix: Standard self-attention compares every word to every other word. The memory required scales quadratically ($O(N^2)$) with sequence length. 10 times the text means 100 times the memory. You must chunk your text, or use specialized architectures (like Longformer or Ring Attention) for massive contexts.

3. Not Using Residual Connections:

   • Mistake: Building a deep transformer from scratch and forgetting the x = x + layer(x) step.
   • Fix: Without residual (skip) connections, the gradients will vanish rapidly during backpropagation, and a 12-layer transformer will simply fail to learn.

Interview Q&A

Q1: Why did the Transformer replace LSTMs/RNNs? A1: Two main reasons. First, RNNs process data sequentially, making them impossible to parallelize on GPUs, resulting in slow training. Transformers process everything simultaneously, allowing massive parallelism. Second, RNNs suffer from vanishing gradients over long sequences (they forget early context). Transformers use self-attention, meaning the distance between word 1 and word 1000 is just one calculation step, completely solving the long-term dependency issue.

Q2: What is the purpose of the Query, Key, and Value vectors in Self-Attention? A2: They are abstractions borrowed from retrieval systems. The Query represents what the current word is "looking for" in the context. The Key represents what a word "contains" or "offers". The dot product of a Query and a Key determines the attention score. The Value is the actual semantic content of the word that is multiplied by the attention score to form the final output.

Q3: Explain "Masked" Self-Attention in the Decoder. A3: In a decoder (like GPT) generating text, the model is trained to predict the next word. If we allowed standard self-attention during training, the model would just "cheat" by looking at the future words already present in the training sequence. Masking artificially sets the attention scores of future words to negative infinity, ensuring the model can only attend to past and current words.

Q4: Why do we divide by $\sqrt{d_k}$ in the scaled dot-product attention? A4: When the dimension of the key/query vectors ($d_k$) gets very large, the dot products can result in massive numbers. When these massive numbers are passed into the Softmax function, it pushes the function into regions where gradients are extremely small (vanishing gradient). Dividing by the square root of the dimension scales the variance back down to 1, keeping gradients healthy.

Q5: What are Residual Connections (Add & Norm) and why are they used in Transformers? A5: Residual connections bypass a layer and add the original input back to the output of that layer. They provide a "highway" for gradients to flow backward through the network unimpeded, preventing the vanishing gradient problem in deep networks of 12, 24, or 96 layers.

Practice Problems

1. Problem: Use the HuggingFace pipeline to instantiate a text-generation model (like gpt2). Provide it with the prompt "The secret to machine learning is" and let it generate 50 words.
2. Problem: In the PyTorch math example, modify the Q, K, and V tensors so that Word 1 pays exactly 100% attention to Word 3, and 0% to the others. Print the resulting attention weights.
3. Problem: Explain why positional encoding is added to the embeddings rather than concatenated to them. (Hint: Think about parameter counts and the mathematical properties of high-dimensional space).

Summary

• Transformers process entire sequences simultaneously, allowing for massive GPU parallelization and enabling the current era of LLMs.
• Self-Attention allows words to dynamically calculate their relevance to all other words in a sentence using Query, Key, and Value vectors.
• Multi-Head Attention allows the model to analyze different types of relationships simultaneously.
• Positional Encoding injects necessary sequence/order information into the data.
• Encoders (BERT) are for understanding text, while Decoders (GPT) are for generating text.