Chapter 15: Transformers and Attention Mechanisms

Bloom's Taxonomy Map for This Chapter

Learning Objectives

By the end of this chapter, you will be able to:

1. Derive the attention mechanism from first principles — starting from the intuition of "soft dictionary lookup" and arriving at Attention(Q,K,V) = softmax(QK^T/√d_k)V
2. Prove why the √d_k scaling factor is necessary by computing the variance of dot products in high-dimensional spaces
3. Implement single-head attention, multi-head attention, and a complete Transformer block from scratch in NumPy
4. Explain sinusoidal positional encoding — derive the formulas, understand why sine and cosine encode relative positions, and compare with learned positional embeddings
5. Diagram the complete Transformer architecture (encoder + decoder, 6 layers each) including masked multi-head attention, cross-attention, feed-forward networks, and residual connections
6. Compare encoder-only (BERT), decoder-only (GPT), and encoder-decoder (T5) architectures — their pre-training objectives, strengths, and applications
7. Analyze self-attention complexity O(n²d) and explain efficient attention variants: Linformer, Performer, FlashAttention
8. Build a character-level mini-GPT in PyTorch that generates text
9. Discuss real-world deployments: AI4Bharat IndicBERT for 11 Indian languages and the OpenAI GPT-1→4 evolution
10. Solve GATE-style numerical problems on attention computation and Transformer parameter counting

Opening Hook

The Intuition First — Why Attention?

The Library Analogy 📚

Imagine you walk into a massive library looking for information about "climate change effects on agriculture." Here's what you do:

1. Your Query (Q): You have a question in your mind — "climate change effects on agriculture." This is your query.
2. Book Labels / Keys (K): Every book on the shelf has a title and index — "Organic Farming Techniques," "Climate Science 101," "History of Ancient Rome." These are keys that you compare against your query.
3. Book Contents / Values (V): Each book contains actual content — facts, figures, analysis. These are the values you want to retrieve.
4. Matching: You mentally compare your query against each book's key. "Climate Science 101" gets high relevance. "History of Ancient Rome" gets near-zero relevance.
5. Weighted Retrieval: You don't just read one book — you read a weighted combination, spending 60% of your time on the climate book, 30% on the farming book, and 10% skimming others.

This is exactly what the attention mechanism does. Every token in a sentence generates a query ("What am I looking for?"), a key ("What do I contain?"), and a value ("What information do I carry?"). The magic: every token gets to "look up" every other token and retrieve a weighted combination of their information.

Why RNNs Failed: The Bottleneck Problem 🔬

In Chapter 14, you learned that LSTMs and GRUs mitigate the vanishing gradient problem. But they still have a fundamental limitation: information must flow sequentially. If word 1 needs information from word 100, that information must survive passage through 99 intermediate states. Even with gating, this bottleneck causes information loss.

The "Aha!" Question

"What if, instead of processing a sequence step by step, we let every position in the sequence attend to every other position simultaneously — and used nothing but this attention operation to build the entire model?"

15.1 The Attention Mechanism — Derived from Scratch

Step 1: The Simplest Attention — Dot-Product Similarity

You have a sequence of n tokens, each represented as a d-dimensional vector. Let's call these vectors x₁, x₂, ..., xₙ, each ∈ ℝ^d. Stack them into a matrix X ∈ ℝ^n×d.

The most basic question of attention: "How much should token i pay attention to token j?"

The simplest answer? Take the dot product of their vectors:

If two vectors point in similar directions, the dot product is large (high attention). If they're orthogonal, the dot product is zero (no attention). This is just cosine similarity scaled by magnitude.

Step 2: From Raw Scores to Probabilities — Softmax

The raw scores e_ij can be any real number — positive, negative, huge, tiny. You need them to be non-negative and sum to 1 (a probability distribution). The softmax function does exactly this:

Now α_ij represents "the fraction of attention that token i pays to token j." These are called attention weights.

Step 3: Weighted Sum — The Output

Once you know how much token i attends to every other token, you compute a weighted sum of all token representations:

This gives token i a new representation that's a blend of all other tokens, weighted by relevance. Beautiful, isn't it? But there's a problem...

Step 4: The Q-K-V Projection — Why We Need Three Separate Roles

In the naive version above, every token uses the same vector for both "querying" and "being queried." This is like a library where the same text serves as both the book's title and its content — confusing!

The solution: project each token into three different spaces using learned weight matrices:

Step 5: Putting It All Together — Raw Attention

Now combine everything. For each query, compute similarity with all keys, normalize with softmax, and retrieve weighted values:

Matrix dimensions check:

• Q ∈ ℝ^n×d_k, K^T ∈ ℝ^d_k×n → Q·K^T ∈ ℝ^n×n (attention scores matrix)
• After softmax: still ℝ^n×n (each row sums to 1)
• V ∈ ℝ^n×d_v → output ∈ ℝ^n×d_v ✓

This is almost the final formula. But there's one critical missing piece...

15.2 Scaled Dot-Product Attention — The √d_k Derivation

The Problem: Why Naive Dot Products Blow Up

Here's a subtle but critical issue. Consider a query vector q and a key vector k, each with d_k independent components drawn from a standard normal distribution (mean 0, variance 1).

The dot product is:

The Solution: Scale by √d_k

To bring the variance back to 1 (well-behaved for softmax), divide by √d_k:

After scaling:

• Each element of QK^T/√d_k has variance ≈ 1
• Softmax operates in a region with informative gradients
• The model can learn nuanced attention distributions instead of "all or nothing"

Masked Attention: Preventing the Future from Leaking

In language generation (like GPT), when predicting the next word, you must NOT let the model peek at future words. This is enforced by masking:

15.3 Multi-Head Attention — Why Multiple Heads?

The Problem with Single-Head Attention

A single attention head computes one set of attention weights. But a word can have multiple types of relationships simultaneously:

• "The cat sat on the mat because it was tired" — "it" needs to attend to "cat" (coreference resolution)
• But "it" also needs to attend to "sat" (syntactic dependency) and "tired" (semantic link)

A single attention head must average all these relationships into one set of weights — a lossy compromise.

The Solution: Multiple Parallel Heads

Instead of one attention function with d_model-dimensional keys, queries, and values, run h parallel attention heads, each with d_k = d_model/h dimensions:

Parameter Count Check

With h=8 heads and d_model=512 (as in the original Transformer):

• d_k = d_v = 512/8 = 64 per head
• Per head: W^Q(512×64) + W^K(512×64) + W^V(512×64) = 3 × 32,768 = 98,304
• All 8 heads: 8 × 98,304 = 786,432
• Output projection W^O: 512 × 512 = 262,144
• Total: 1,048,576 ≈ 1M parameters (same as one large single-head attention!)

What Do Different Heads Learn?

Research has shown that different heads naturally specialize:

15.4 Positional Encoding — Teaching Position to a Permutation-Invariant Model

The Problem: Self-Attention Ignores Order

Here's a subtle but devastating issue. Consider the sentences:

• "The dog chased the cat" (meaning A)
• "The cat chased the dog" (meaning B — completely different!)

Self-attention computes pairwise similarities between tokens. Since these sentences have exactly the same tokens (just reordered), self-attention produces identical outputs. The attention mechanism is permutation-invariant — it doesn't care about word order!

This is clearly unacceptable. We must inject positional information somehow.

Solution: Sinusoidal Positional Encoding

Wavelength Intuition

Each dimension pair (2i, 2i+1) creates a sine-cosine wave with a different wavelength:

• Dimension 0,1: wavelength = 2π ≈ 6.28 positions (fastest oscillation)
• Dimension 2,3: wavelength = 2π × 10000^2/512 ≈ 7.4 positions
• ...
• Dimension 510,511: wavelength ≈ 2π × 10000 ≈ 62,832 positions (slowest)

Think of it like a binary clock: the least significant bit oscillates fastest, the most significant bit oscillates slowest. Each position gets a unique "binary-like" representation from combining all these waves.

Learned vs. Sinusoidal Positional Encoding

15.5 The Full Transformer Architecture

The Big Picture

Encoder Block (in detail)

Each of the 6 encoder layers contains exactly two sub-layers:

Decoder Block (in detail)

Each of the 6 decoder layers has three sub-layers:

Transformer Hyperparameters (Original Paper)

15.6 Layer Normalization and Residual Connections

Residual Connections: The Gradient Highway

You already met residual connections in Chapter 11 (ResNets). The Transformer uses them identically:

Why they're essential: With 6 stacked layers, each containing 2-3 sub-layers, gradients must flow through 12-18 transformations. Without skip connections, gradients would vanish or explode. The residual path provides a direct gradient highway from the output back to the input.

Layer Normalization (Not Batch Normalization!)

The Transformer uses Layer Normalization, not Batch Normalization. Here's why and how:

Why LayerNorm, Not BatchNorm?

Pre-Norm vs Post-Norm

The original Transformer uses Post-Norm: output = LayerNorm(x + SubLayer(x))

Modern Transformers (GPT-2+) use Pre-Norm: output = x + SubLayer(LayerNorm(x))

15.7 BERT vs GPT — Encoder-Only vs Decoder-Only

Three Flavors of Transformers

BERT: Bidirectional Encoder Representations from Transformers

GPT: Generative Pre-trained Transformer

Head-to-Head Comparison

15.8 Self-Attention Complexity and Efficient Attention

The Quadratic Bottleneck

Comparison: Transformer vs RNN Complexity

Efficient Attention Variants

Worked Examples

Worked Example 1: By-Hand Attention Computation

Q = [[1, 0],     K = [[1, 0],     V = [[1, 0],
     [0, 1],          [0, 1],          [0, 1],
     [1, 1]]          [1, 1]]          [0.5, 0.5]]

QKT = Q · KT = [[1·1+0·0, 1·0+0·1, 1·1+0·1],   = [[1, 0, 1],
                     [0·1+1·0, 0·0+1·1, 0·1+1·1],      [0, 1, 1],
                     [1·1+1·0, 1·0+1·1, 1·1+1·1]]      [1, 1, 2]]

QKT/√d_k = [[0.707, 0.000, 0.707],
              [0.000, 0.707, 0.707],
              [0.707, 0.707, 1.414]]

Row 1: exp([0.707, 0.000, 0.707]) = [2.028, 1.000, 2.028]
        sum = 5.056
        softmax = [0.401, 0.198, 0.401]

Row 2: exp([0.000, 0.707, 0.707]) = [1.000, 2.028, 2.028]
        sum = 5.056
        softmax = [0.198, 0.401, 0.401]

Row 3: exp([0.707, 0.707, 1.414]) = [2.028, 2.028, 4.113]
        sum = 8.169
        softmax = [0.248, 0.248, 0.503]

Attention weights:           V:
[[0.401, 0.198, 0.401],    [[1.0, 0.0],
 [0.198, 0.401, 0.401],     [0.0, 1.0],
 [0.248, 0.248, 0.503]]     [0.5, 0.5]]

Output = weights · V:
Row 1: 0.401×[1,0] + 0.198×[0,1] + 0.401×[0.5,0.5] = [0.602, 0.399]
Row 2: 0.198×[1,0] + 0.401×[0,1] + 0.401×[0.5,0.5] = [0.399, 0.602]
Row 3: 0.248×[1,0] + 0.248×[0,1] + 0.503×[0.5,0.5] = [0.500, 0.500]

Attention(Q,K,V) = [[0.602, 0.399],
                    [0.399, 0.602],
                    [0.500, 0.500]]

Worked Example 2: Indian Industry — Multilingual Attention

Attention head #3 (sentiment detection):
                Yeh  phone bahut accha hai  but  battery life  is  terrible
terrible  →     0.01  0.02  0.01  0.35  0.01 0.15  0.12   0.08  0.02  0.23
accha     →     0.02  0.05  0.20  0.15  0.03 0.25  0.08   0.05  0.02  0.15

Note: "terrible" attends strongly to "accha" (0.35) — detecting the 
sentiment contrast. "accha" attends to "but" (0.25) — the pivot word.

Worked Example 3: US Industry — GPT Text Generation

Step 1: Input: "Explain quantum computing in simple terms."
        → Model predicts next token probabilities:
          P("Quantum") = 0.12, P("Think") = 0.08, P("Imagine") = 0.15, ...
        → Sample (or greedy): "Imagine"

Step 2: Input: "Explain quantum computing in simple terms. Imagine"
        → Causal mask ensures "Imagine" can attend to all previous tokens
           but previous tokens CANNOT attend to "Imagine"
        → Predict: P("a") = 0.18, P("you") = 0.12, ...
        → Select: "a"

Step 3: Continue until <EOS> token or max_length

Attention mask (1 = attend, 0 = blocked):
              Explain quantum computing in simple terms . Imagine
Explain    [    1       0        0       0    0      0   0    0   ]
quantum    [    1       1        0       0    0      0   0    0   ]
computing  [    1       1        1       0    0      0   0    0   ]
in         [    1       1        1       1    0      0   0    0   ]
simple     [    1       1        1       1    1      0   0    0   ]
terms      [    1       1        1       1    1      1   0    0   ]
.          [    1       1        1       1    1      1   1    0   ]
Imagine    [    1       1        1       1    1      1   1    1   ]  ← can see everything before it

Python Implementation: From-Scratch NumPy

13.1 Scaled Dot-Product Attention (NumPy)

Python (NumPy)
import numpy as np

def softmax(x, axis=-1):
    """Numerically stable softmax."""
    e_x = np.exp(x - np.max(x, axis=axis, keepdims=True))
    return e_x / np.sum(e_x, axis=axis, keepdims=True)

def scaled_dot_product_attention(Q, K, V, mask=None):
    """
    Scaled dot-product attention from scratch.
    
    Args:
        Q: Queries  (n, d_k)
        K: Keys     (n, d_k)
        V: Values   (n, d_v)
        mask: Optional boolean mask (n, n), True = mask out
    
    Returns:
        output: (n, d_v) — weighted sum of values
        weights: (n, n) — attention weights (for visualization)
    """
    d_k = Q.shape[-1]
    
    # Step 1: Compute raw attention scores
    scores = Q @ K.T  # (n, n)
    
    # Step 2: Scale by sqrt(d_k)
    scores = scores / np.sqrt(d_k)
    
    # Step 3: Apply mask (for causal / padding)
    if mask is not None:
        scores = np.where(mask, -1e9, scores)
    
    # Step 4: Softmax to get attention weights
    weights = softmax(scores, axis=-1)  # (n, n), each row sums to 1
    
    # Step 5: Weighted sum of values
    output = weights @ V  # (n, d_v)
    
    return output, weights

# === DEMO ===
np.random.seed(42)
n, d_k, d_v = 4, 8, 8
Q = np.random.randn(n, d_k)
K = np.random.randn(n, d_k)
V = np.random.randn(n, d_v)

output, weights = scaled_dot_product_attention(Q, K, V)
print("Attention weights (each row sums to 1):")
print(np.round(weights, 3))
print(f"\nOutput shape: {output.shape}")

# Causal mask demo
causal_mask = np.triu(np.ones((n, n), dtype=bool), k=1)
output_masked, weights_masked = scaled_dot_product_attention(Q, K, V, mask=causal_mask)
print("\nCausal attention weights:")
print(np.round(weights_masked, 3))

13.2 Multi-Head Attention (NumPy)

Python (NumPy)
class MultiHeadAttention:
    """Multi-Head Attention from scratch in NumPy."""
    
    def __init__(self, d_model, n_heads):
        assert d_model % n_heads == 0, "d_model must be divisible by n_heads"
        self.d_model = d_model
        self.n_heads = n_heads
        self.d_k = d_model // n_heads  # dimension per head
        
        # Initialize weight matrices (Xavier initialization)
        scale = np.sqrt(2.0 / (d_model + self.d_k))
        self.W_Q = np.random.randn(d_model, d_model) * scale
        self.W_K = np.random.randn(d_model, d_model) * scale
        self.W_V = np.random.randn(d_model, d_model) * scale
        self.W_O = np.random.randn(d_model, d_model) * scale
    
    def split_heads(self, x):
        """Reshape (n, d_model) → (n_heads, n, d_k)"""
        n = x.shape[0]
        x = x.reshape(n, self.n_heads, self.d_k)  # (n, h, d_k)
        return x.transpose(1, 0, 2)  # (h, n, d_k)
    
    def forward(self, X, mask=None):
        """
        Args:
            X: Input sequence (n, d_model)
            mask: Optional mask (n, n)
        Returns:
            output: (n, d_model)
            all_weights: (n_heads, n, n) — attention weights per head
        """
        # Project to Q, K, V
        Q = X @ self.W_Q  # (n, d_model)
        K = X @ self.W_K
        V = X @ self.W_V
        
        # Split into h heads
        Q_heads = self.split_heads(Q)  # (h, n, d_k)
        K_heads = self.split_heads(K)
        V_heads = self.split_heads(V)
        
        # Compute attention for each head independently
        all_outputs = []
        all_weights = []
        for i in range(self.n_heads):
            out, w = scaled_dot_product_attention(
                Q_heads[i], K_heads[i], V_heads[i], mask
            )
            all_outputs.append(out)    # (n, d_k)
            all_weights.append(w)      # (n, n)
        
        # Concatenate heads: (n, h*d_k) = (n, d_model)
        concat = np.concatenate(all_outputs, axis=-1)
        
        # Final linear projection
        output = concat @ self.W_O  # (n, d_model)
        
        return output, np.array(all_weights)

# === DEMO ===
np.random.seed(42)
d_model, n_heads, seq_len = 32, 4, 6
mha = MultiHeadAttention(d_model, n_heads)

X = np.random.randn(seq_len, d_model)
output, attn_weights = mha.forward(X)

print(f"Input shape:  {X.shape}")
print(f"Output shape: {output.shape}")
print(f"Attention weights shape: {attn_weights.shape}")
print(f"  (n_heads={n_heads}, seq_len={seq_len}, seq_len={seq_len})")

13.3 Transformer Block (NumPy)

Python (NumPy)
def layer_norm(x, gamma, beta, eps=1e-6):
    """Layer Normalization."""
    mean = np.mean(x, axis=-1, keepdims=True)
    var = np.var(x, axis=-1, keepdims=True)
    x_norm = (x - mean) / np.sqrt(var + eps)
    return gamma * x_norm + beta

def relu(x):
    return np.maximum(0, x)

class TransformerBlock:
    """A single Transformer encoder block from scratch."""
    
    def __init__(self, d_model, n_heads, d_ff):
        self.mha = MultiHeadAttention(d_model, n_heads)
        
        # LayerNorm parameters
        self.gamma1 = np.ones(d_model)
        self.beta1 = np.zeros(d_model)
        self.gamma2 = np.ones(d_model)
        self.beta2 = np.zeros(d_model)
        
        # Feed-Forward Network parameters
        scale1 = np.sqrt(2.0 / (d_model + d_ff))
        scale2 = np.sqrt(2.0 / (d_ff + d_model))
        self.W1 = np.random.randn(d_model, d_ff) * scale1
        self.b1 = np.zeros(d_ff)
        self.W2 = np.random.randn(d_ff, d_model) * scale2
        self.b2 = np.zeros(d_model)
    
    def ffn(self, x):
        """Position-wise Feed-Forward Network."""
        return relu(x @ self.W1 + self.b1) @ self.W2 + self.b2
    
    def forward(self, x, mask=None):
        """
        Forward pass through one Transformer encoder block.
        Uses Post-Norm (original Transformer style).
        """
        # Sub-layer 1: Multi-Head Self-Attention + Residual + LayerNorm
        attn_output, attn_weights = self.mha.forward(x, mask)
        x = layer_norm(x + attn_output, self.gamma1, self.beta1)
        
        # Sub-layer 2: FFN + Residual + LayerNorm
        ffn_output = self.ffn(x)
        x = layer_norm(x + ffn_output, self.gamma2, self.beta2)
        
        return x, attn_weights

# === DEMO: Stack 6 layers like the original Transformer ===
np.random.seed(42)
d_model, n_heads, d_ff = 64, 8, 256
seq_len = 10

# Create 6 Transformer blocks
blocks = [TransformerBlock(d_model, n_heads, d_ff) for _ in range(6)]

# Simulate positional encoding + embedding
X = np.random.randn(seq_len, d_model) * 0.1  # simulated input

# Forward through all 6 layers
h = X
for i, block in enumerate(blocks):
    h, w = block.forward(h)
    print(f"Layer {i+1} output — mean: {h.mean():.4f}, std: {h.std():.4f}")

print(f"\nFinal output shape: {h.shape}")
print("✓ LayerNorm keeps values well-behaved across all 6 layers!")

13.4 Positional Encoding (NumPy)

Python (NumPy)
def sinusoidal_positional_encoding(max_len, d_model):
    """
    Generate sinusoidal positional encodings.
    
    PE(pos, 2i)   = sin(pos / 10000^(2i/d_model))
    PE(pos, 2i+1) = cos(pos / 10000^(2i/d_model))
    """
    PE = np.zeros((max_len, d_model))
    position = np.arange(max_len)[:, np.newaxis]  # (max_len, 1)
    
    # Compute the division term: 10000^(2i/d_model)
    div_term = np.exp(
        np.arange(0, d_model, 2) * (-np.log(10000.0) / d_model)
    )  # (d_model/2,)
    
    # Even dimensions: sine
    PE[:, 0::2] = np.sin(position * div_term)
    
    # Odd dimensions: cosine
    PE[:, 1::2] = np.cos(position * div_term)
    
    return PE

# === DEMO ===
PE = sinusoidal_positional_encoding(max_len=50, d_model=16)
print(f"Positional Encoding shape: {PE.shape}")
print(f"PE[0] (position 0): {np.round(PE[0], 3)}")
print(f"PE[1] (position 1): {np.round(PE[1], 3)}")
print(f"\nValues are bounded: min={PE.min():.3f}, max={PE.max():.3f}")

# Verify: dot product between nearby positions is higher
print(f"\nSimilarity PE[0]·PE[1]: {PE[0] @ PE[1]:.3f}")
print(f"Similarity PE[0]·PE[5]: {PE[0] @ PE[5]:.3f}")
print(f"Similarity PE[0]·PE[25]: {PE[0] @ PE[25]:.3f}")
print("Nearby positions have higher similarity ✓")

def broken_attention(Q, K, V):
    scores = Q @ K.T
    weights = softmax(scores, axis=0)   # <-- Bug here!
    return weights @ V

PyTorch Implementation: Mini-GPT (Character-Level)

Python (PyTorch)
import torch
import torch.nn as nn
import torch.nn.functional as F

# ── Hyperparameters ──
BLOCK_SIZE   = 64     # context window (max sequence length)
D_MODEL      = 128    # embedding dimension
N_HEADS      = 4      # number of attention heads
N_LAYERS     = 4      # number of Transformer blocks
D_FF         = 512    # feed-forward inner dimension
DROPOUT      = 0.1
LEARNING_RATE = 3e-4
MAX_ITERS    = 5000
BATCH_SIZE   = 32
DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'

# ── Single Attention Head ──
class Head(nn.Module):
    """One head of causal self-attention."""
    def __init__(self, head_size):
        super().__init__()
        self.key   = nn.Linear(D_MODEL, head_size, bias=False)
        self.query = nn.Linear(D_MODEL, head_size, bias=False)
        self.value = nn.Linear(D_MODEL, head_size, bias=False)
        self.register_buffer('tril',
            torch.tril(torch.ones(BLOCK_SIZE, BLOCK_SIZE)))
        self.dropout = nn.Dropout(DROPOUT)
    
    def forward(self, x):
        B, T, C = x.shape
        k = self.key(x)    # (B, T, head_size)
        q = self.query(x)  # (B, T, head_size)
        
        # Compute attention scores
        scores = q @ k.transpose(-2, -1) * (C ** -0.5)  # (B, T, T)
        scores = scores.masked_fill(
            self.tril[:T, :T] == 0, float('-inf'))  # causal mask
        weights = F.softmax(scores, dim=-1)
        weights = self.dropout(weights)
        
        # Weighted sum of values
        v = self.value(x)       # (B, T, head_size)
        out = weights @ v       # (B, T, head_size)
        return out

# ── Multi-Head Attention ──
class MultiHeadAttention(nn.Module):
    def __init__(self, n_heads, head_size):
        super().__init__()
        self.heads = nn.ModuleList([Head(head_size) for _ in range(n_heads)])
        self.proj = nn.Linear(D_MODEL, D_MODEL)
        self.dropout = nn.Dropout(DROPOUT)
    
    def forward(self, x):
        out = torch.cat([h(x) for h in self.heads], dim=-1)
        out = self.dropout(self.proj(out))
        return out

# ── Feed-Forward Network ──
class FeedForward(nn.Module):
    def __init__(self):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(D_MODEL, D_FF),
            nn.ReLU(),
            nn.Linear(D_FF, D_MODEL),
            nn.Dropout(DROPOUT),
        )
    def forward(self, x):
        return self.net(x)

# ── Transformer Block ──
class TransformerBlock(nn.Module):
    """Transformer block: communication (attention) + computation (FFN)."""
    def __init__(self):
        super().__init__()
        head_size = D_MODEL // N_HEADS
        self.sa = MultiHeadAttention(N_HEADS, head_size)
        self.ffn = FeedForward()
        self.ln1 = nn.LayerNorm(D_MODEL)
        self.ln2 = nn.LayerNorm(D_MODEL)
    
    def forward(self, x):
        # Pre-Norm residual connections (GPT-2 style)
        x = x + self.sa(self.ln1(x))
        x = x + self.ffn(self.ln2(x))
        return x

# ── Mini-GPT Model ──
class MiniGPT(nn.Module):
    def __init__(self, vocab_size):
        super().__init__()
        self.token_embedding = nn.Embedding(vocab_size, D_MODEL)
        self.position_embedding = nn.Embedding(BLOCK_SIZE, D_MODEL)
        self.blocks = nn.Sequential(
            *[TransformerBlock() for _ in range(N_LAYERS)]
        )
        self.ln_f = nn.LayerNorm(D_MODEL)
        self.lm_head = nn.Linear(D_MODEL, vocab_size)
    
    def forward(self, idx, targets=None):
        B, T = idx.shape
        
        # Token + Position embeddings
        tok_emb = self.token_embedding(idx)         # (B, T, D_MODEL)
        pos_emb = self.position_embedding(
            torch.arange(T, device=DEVICE))          # (T, D_MODEL)
        x = tok_emb + pos_emb                       # (B, T, D_MODEL)
        
        # Transformer blocks
        x = self.blocks(x)                           # (B, T, D_MODEL)
        x = self.ln_f(x)                             # (B, T, D_MODEL)
        logits = self.lm_head(x)                     # (B, T, vocab_size)
        
        # Compute loss if targets provided
        loss = None
        if targets is not None:
            B, T, C = logits.shape
            logits_flat = logits.view(B*T, C)
            targets_flat = targets.view(B*T)
            loss = F.cross_entropy(logits_flat, targets_flat)
        
        return logits, loss
    
    def generate(self, idx, max_new_tokens):
        """Autoregressive generation."""
        for _ in range(max_new_tokens):
            # Crop to last BLOCK_SIZE tokens
            idx_cond = idx[:, -BLOCK_SIZE:]
            logits, _ = self(idx_cond)
            logits = logits[:, -1, :]  # last token's predictions
            probs = F.softmax(logits, dim=-1)
            idx_next = torch.multinomial(probs, num_samples=1)
            idx = torch.cat([idx, idx_next], dim=1)
        return idx

# ── Training Loop ──
# Load your text data
# text = open('input.txt', 'r').read()
# chars = sorted(list(set(text)))
# vocab_size = len(chars)
# stoi = {c: i for i, c in enumerate(chars)}
# itos = {i: c for c, i in stoi.items()}
# encode = lambda s: [stoi[c] for c in s]
# decode = lambda l: ''.join([itos[i] for i in l])

# model = MiniGPT(vocab_size).to(DEVICE)
# optimizer = torch.optim.AdamW(model.parameters(), lr=LEARNING_RATE)
# 
# for step in range(MAX_ITERS):
#     xb, yb = get_batch('train')  # random batch of (BATCH_SIZE, BLOCK_SIZE)
#     logits, loss = model(xb, yb)
#     optimizer.zero_grad()
#     loss.backward()
#     optimizer.step()
#     if step % 500 == 0:
#         print(f"Step {step}, Loss: {loss.item():.4f}")

# Print model size
vocab_size_demo = 65  # ASCII printable characters
model = MiniGPT(vocab_size_demo)
n_params = sum(p.numel() for p in model.parameters())
print(f"Mini-GPT Parameters: {n_params:,}")
print(f"  Token Embedding:    {vocab_size_demo * D_MODEL:,}")
print(f"  Position Embedding: {BLOCK_SIZE * D_MODEL:,}")
print(f"  Transformer Blocks: {n_params - vocab_size_demo*D_MODEL - BLOCK_SIZE*D_MODEL - D_MODEL*vocab_size_demo - D_MODEL - D_MODEL:,}")
print(f"  LM Head:            {D_MODEL * vocab_size_demo + vocab_size_demo:,}")

Visual Aids

Self-Attention Computation Flow

Multi-Head Attention Visualization

Complete Encoder Layer

Indian Industry Case Study: AI4Bharat IndicBERT

US/Global Industry Case Study: The GPT Evolution

Common Misconceptions

GATE / Exam Corner

Key Formulas for Exams

GATE-Style Problems

Prediction Table — Likely GATE Topics (2025-2027)

Interview Prep

Conceptual Questions

Coding Questions

def attention(Q, K, V, mask=None):
    d_k = Q.shape[-1]
    scores = Q @ K.transpose(-2, -1) / (d_k ** 0.5)
    if mask is not None:
        scores = scores.masked_fill(mask == 0, -1e9)
    weights = torch.softmax(scores, dim=-1)
    return weights @ V, weights

India-Specific Interview Angle

Hands-On Lab / Mini-Project

Exercises

Section A: Conceptual Questions (5)

Section B: Mathematical Questions (8)

Section C: Coding Questions (4)

Section D: Critical Thinking (3)

★ Starred Research Questions (2)

Connections

Chapter Summary

Key Intuition: The Transformer processes sequences not through step-by-step recurrence, but through layers of self-attention — letting every token directly query every other token, enabling parallelism, long-range dependencies, and the scaling that powers modern AI.

Further Reading

🇮🇳 Indian Resources

• NPTEL Course: "Deep Learning for NLP" by Prof. Mitesh Khapra (IIT Madras) — covers Transformers with Indian language examples, available free on Swayam
• AI4Bharat: ai4bharat.iitm.ac.in — IndicBERT, IndicTrans, IndicCorp resources and pre-trained models
• GATE CS/DS Prep: "Deep Learning" by GATE Wallah / Unacademy — covers attention complexity and Transformer architecture for GATE exams
• Book: "Speech and Language Processing" by Jurafsky & Martin (draft chapters free online) — Chapter 10 covers Transformers extensively

🌍 Global Resources

• Original Paper: "Attention Is All You Need" (Vaswani et al., 2017) — the foundational paper, dense but essential
• The Illustrated Transformer: Jay Alammar's visual walkthrough at jalammar.github.io/illustrated-transformer — the best visual explanation
• 3Blue1Brown: "But what is a GPT?" — visual intuition for attention, embedding, and generation
• Andrej Karpathy: "Let's build GPT from scratch" (YouTube) — builds nanoGPT step by step, the best hands-on tutorial
• Harvard NLP: "The Annotated Transformer" — line-by-line PyTorch implementation with explanations
• Lilian Weng's Blog: "The Transformer Family" — comprehensive survey of all Transformer variants (2020-2023)
• Distill.pub: "A Mathematical Framework for Transformer Circuits" (Elhage et al., 2021) — deep mechanistic interpretability of attention

📄 Key Papers (Chronological)

1. Bahdanau et al. (2014) — "Neural Machine Translation by Jointly Learning to Align and Translate" (additive attention)
2. Vaswani et al. (2017) — "Attention Is All You Need" (the Transformer)
3. Devlin et al. (2018) — "BERT: Pre-training of Deep Bidirectional Transformers" (encoder-only)
4. Radford et al. (2018, 2019) — GPT-1, GPT-2 papers (decoder-only)
5. Brown et al. (2020) — "Language Models are Few-Shot Learners" (GPT-3, scaling)
6. Dosovitskiy et al. (2020) — "An Image is Worth 16×16 Words" (Vision Transformer)
7. Su et al. (2021) — "RoFormer: Rotary Position Embedding" (RoPE)
8. Dao et al. (2022) — "FlashAttention: Fast and Memory-Efficient Exact Attention"
9. Touvron et al. (2023) — "LLaMA: Open and Efficient Foundation Language Models"
10. Gu & Dao (2023) — "Mamba: Linear-Time Sequence Modeling with Selective State Spaces"