Phase 5 • EduArtha

Large Language Models (LLMs)

This is the core of how modern AI works — Transformer architecture, pre-training on text, and alignment techniques. Every chatbot, code assistant, and AI agent is built on these foundations.

⏱ 6–12 months | 14 Chapters | 50+ Exercises | Industry Problems

Part I

Transformer Architecture

The architecture that changed everything

Chapter 1

Self-Attention & Multi-Head Attention

Learning Objectives

Implement scaled dot-product attention from scratch
Understand queries, keys, values — the information retrieval analogy
Build multi-head attention and understand why multiple heads help
Compute attention complexity and memory requirements

Attention(Q, K, V) = softmax(QKᵀ / √dₖ) · V

Python
import torch
import torch.nn as nn
import math

class MultiHeadAttention(nn.Module):
    def __init__(self, d_model=512, n_heads=8, dropout=0.1):
        super().__init__()
        assert d_model % n_heads == 0
        self.d_k = d_model // n_heads
        self.n_heads = n_heads
        self.W_q = nn.Linear(d_model, d_model)
        self.W_k = nn.Linear(d_model, d_model)
        self.W_v = nn.Linear(d_model, d_model)
        self.W_o = nn.Linear(d_model, d_model)
        self.dropout = nn.Dropout(dropout)

    def forward(self, x, mask=None):
        B, T, C = x.shape
        # Project and reshape to [B, n_heads, T, d_k]
        Q = self.W_q(x).view(B, T, self.n_heads, self.d_k).transpose(1,2)
        K = self.W_k(x).view(B, T, self.n_heads, self.d_k).transpose(1,2)
        V = self.W_v(x).view(B, T, self.n_heads, self.d_k).transpose(1,2)

        # Scaled dot-product attention
        scores = (Q @ K.transpose(-2,-1)) / math.sqrt(self.d_k)
        if mask is not None:
            scores = scores.masked_fill(mask == 0, float('-inf'))
        attn = self.dropout(torch.softmax(scores, dim=-1))
        out = (attn @ V).transpose(1,2).contiguous().view(B, T, C)
        return self.W_o(out)

# Causal mask for autoregressive (GPT-style)
def causal_mask(T):
    return torch.tril(torch.ones(T, T)).unsqueeze(0).unsqueeze(0)

Industry Problem: Quadratic Memory in Long Documents

Problem: Self-attention is O(n²) in sequence length. Processing a 100K-token legal contract requires 100K × 100K = 10 billion attention scores per layer per head — impossible to fit in memory.

Solutions: (1) Flash Attention — fuses operations, reduces memory from O(n²) to O(n). (2) Sliding window attention (Mistral) — attend to local windows. (3) Ring attention — distributes across GPUs. (4) Sparse attention (BigBird) — attend to only important positions.

Exercises

Exercise 1.1: Why scale by √dₖ and what happens without it?

Without scaling, dot products grow with dimension dₖ (variance ≈ dₖ for random vectors). Large values push softmax into saturation — one position gets ~100% attention, gradients vanish. Scaling by √dₖ keeps variance at ~1, ensuring softmax outputs are smooth and informative. For d_k=64: scores ÷ 8.

Exercise 1.2: Compute memory for MHA with d_model=4096, n_heads=32, seq_len=8192

Attention matrix per head: 8192 × 8192 × 4 bytes (FP32) = 256 MB. With 32 heads: 8 GB. For one layer! A 32-layer model needs 256 GB just for attention matrices. This is why Flash Attention (which never materializes the full matrix) is essential for long contexts.

Exercise 1.3: Why use multiple heads instead of one large attention?

Different heads learn different relationship types: head 1 might attend to syntactic neighbors, head 2 to semantic relationships, head 3 to positional patterns. This is like having multiple "perspectives" on the same data. Empirically, 8-64 heads consistently outperform single-head attention of the same total dimension.

Chapter Summary

Self-attention computes relevance between all position pairs — O(n²) but powerful
Multi-head attention learns diverse relationship types in parallel subspaces
Causal masking enables autoregressive generation (GPT-style LLMs)
Industry challenge: quadratic scaling → solved by Flash Attention and sparse methods