How the AI-Engineering-From-Scratch Curriculum Teaches Transformer Architectures from Self-Attention to Full GPT

The AI-Engineering-From-Scratch curriculum teaches transformer architectures through a bottom-up pedagogical approach spanning 10 progressive lessons, starting with NumPy-based self-attention mathematics and culminating in a trainable decoder-only GPT model with automated testing.

The rohitg00/ai-engineering-from-scratch repository provides a comprehensive educational framework for understanding modern transformer architectures by building every component from first principles. Each lesson in the deep-dive phase pairs theoretical explanations in docs/en.md files with executable PyTorch implementations in code/main.py, enabling learners to verify theoretical understanding through practical code execution and unit testing.

The 10-Step Progressive Learning Path

The curriculum follows a strict bottom-up methodology where each module depends only on previously mastered concepts. The journey begins with mathematical foundations and ends with a capstone training project:

• Self-Attention Theory – Deriving query, key, and value operations with scaled dot-product attention and causal masking. Documentation resides in phases/07-transformers-deep-dive/01-self-attention/docs/en.md with implementation in phases/07-transformers-deep-dive/01-self-attention/code/main.py.

• Multi-Head Attention – Parallel attention heads with independent projection matrices (W_Q, W_K, W_V), concatenation, and output projection.

• LayerNorm & Residual Connections – Pre-normalization strategies to stabilize deep stacks, implemented in phases/07-transformers-deep-dive/03-layernorm-residual/code/main.py.

• Transformer Block – Combining attention, feed-forward networks (MLP), and normalization into a single composable unit.

• Token & Positional Embeddings – Converting discrete tokens to vectors and injecting positional information through learned embeddings.

• Stacking Blocks – Building deep transformer bodies using nn.ModuleList to sequentially process embedded inputs.

• Language-Model Head – Projecting final hidden states back to vocabulary logits for next-token prediction.

• Full GPT-Style Model – Complete assembly of token embeddings, stacked blocks, and LM head into a unified GPTModel class.

• Training Loop – Implementing masked language-model loss, AdamW optimizer with weight decay, and cosine learning-rate schedules.

• Capstone Project – End-to-end training of a tiny decoder-only transformer on a synthetic corpus, located in phases/19-capstone-projects/35-gpt-model-assembly/.

Self-Attention: The Mathematical Foundation

The curriculum begins with the scaled dot-product attention mechanism implemented from scratch. In phases/07-transformers-deep-dive/01-self-attention/code/main.py, the implementation emphasizes causal masking to ensure decoder-only models attend only to previous tokens.

The core attention calculation follows the standard transformer formula:

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

This implementation demonstrates how the query, key, and value matrices interact through matrix multiplication, with the scaling factor math.sqrt(d_k) preventing softmax saturation. The causal mask ensures autoregressive behavior by setting future positions to negative infinity before the softmax operation.

Multi-Head Attention Parallelization

Step 2 extends single-head attention to multi-head parallelism, where each head learns different representation subspaces. The implementation in phases/07-transformers-deep-dive/02-multihead-self-attention/code/main.py wraps the attention mechanism in a module that handles projection matrices and parallel computation:

class MultiHeadAttention(nn.Module):
    def __init__(self, n_heads, d_model):
        super().__init__()
        self.n_heads = n_heads
        self.d_k = d_model // n_heads
        self.qkv_proj = nn.Linear(d_model, 3 * d_model)
        self.out_proj = nn.Linear(d_model, d_model)

    def forward(self, x, mask=None):
        B, T, D = x.size()
        qkv = self.qkv_proj(x).reshape(B, T, self.n_heads, 3 * self.d_k)
        Q, K, V = qkv.chunk(3, dim=-1)
        # transpose for batched matmul

        Q, K, V = (Q.transpose(1, 2), K.transpose(1, 2), V.transpose(1, 2))
        attn = attention(Q, K, V, mask)
        attn = attn.transpose(1, 2).contiguous().view(B, T, D)
        return self.out_proj(attn)

The code implements pre-normalization patterns where LayerNorm is applied before attention and MLP sub-layers, a design choice that improves gradient flow in deep networks. The qkv_proj layer combines the three linear projections for efficiency, splitting them into individual Q, K, V tensors within the forward pass.

Stabilizing Deep Architectures: LayerNorm and Residuals

Before stacking multiple transformer blocks, the curriculum addresses training stability through pre-normalization and residual connections. The implementation in phases/07-transformers-deep-dive/03-layernorm-residual/code/main.py teaches the x = x + sublayer(x) pattern critical for gradient highway preservation.

The transformer block implementation combines these stabilization techniques with the attention and feed-forward components:

class TransformerBlock(nn.Module):
    def __init__(self, d_model, n_heads, d_ff):
        super().__init__()
        self.ln1 = nn.LayerNorm(d_model)
        self.attn = MultiHeadAttention(n_heads, d_model)
        self.ln2 = nn.LayerNorm(d_model)
        self.mlp = nn.Sequential(
            nn.Linear(d_model, d_ff),
            nn.GELU(),
            nn.Linear(d_ff, d_model)
        )

    def forward(self, x, mask=None):
        x = x + self.attn(self.ln1(x), mask)   # residual + attention

        x = x + self.mlp(self.ln2(x))          # residual + MLP

        return x

Note the GELU activation function in the MLP, which the curriculum uses as the non-linearity between the two linear layers of the feed-forward network. The hidden dimension d_ff is typically set to four times the model dimension (4 × d_model), following the original transformer specification.

Model Assembly: From Embeddings to GPT

The upper-level architecture components transform discrete token indices into continuous representations and sequence them through the transformer stack. In phases/07-transformers-deep-dive/05-token-positional-embeddings/code/main.py, the curriculum implements learned positional embeddings rather than sinusoidal encodings, providing a simpler entry point for beginners.

The complete model assembly occurs in phases/07-transformers-deep-dive/08-full-gpt-model/code/main.py, where the GPTModel class wires together all components:

class GPTModel(nn.Module):
    def __init__(self, vocab_size, d_model, n_layers, n_heads, d_ff, max_seq):
        super().__init__()
        self.tok_emb = nn.Embedding(vocab_size, d_model)
        self.pos_emb = nn.Parameter(torch.zeros(1, max_seq, d_model))
        self.blocks = nn.ModuleList(
            [TransformerBlock(d_model, n_heads, d_ff) for _ in range(n_layers)]
        )
        self.ln_f = nn.LayerNorm(d_model)
        self.head = nn.Linear(d_model, vocab_size, bias=False)

    def forward(self, idx):
        B, T = idx.shape
        x = self.tok_emb(idx) + self.pos_emb[:, :T, :]
        mask = torch.tril(torch.ones(T, T, device=idx.device)).bool()
        for block in self.blocks:
            x = block(x, mask)
        x = self.ln_f(x)
        return self.head(x)

This implementation creates a decoder-only architecture (GPT-style) through the causal mask generated by torch.tril, ensuring each position attends only to itself and previous positions. The final LayerNorm (ln_f) stabilizes the output before the language modeling head projects hidden states to vocabulary logits.

Training Loop and Capstone Validation

The final implementation phase in phases/07-transformers-deep-dive/09-training-loop/code/main.py connects the model to optimization infrastructure. The curriculum implements cross-entropy loss over masked tokens, AdamW optimizer with weight decay, and a learning-rate schedule featuring warm-up followed by cosine decay.

The capstone project in phases/19-capstone-projects/35-gpt-model-assembly/ requires learners to train the complete model on a synthetic corpus for 20-30 steps, verifying that the full stack functions end-to-end. Unit tests in each code/tests.py file validate that components produce expected tensor shapes and maintain gradient flow through backpropagation.

How to Run the Lessons Locally

Each lesson is self-contained and executable without external dependencies beyond PyTorch:

1. Navigate to the specific lesson directory, for example: cd phases/07-transformers-deep-dive/01-self-attention

2. Execute the reference implementation: python3 code/main.py runs a forward pass and prints output shapes

3. Validate correctness: python3 -m unittest discover code/tests -v runs the test suite verifying mathematical correctness and gradient flow

All implementations are self-contained, requiring no internet connectivity or external API calls during execution.

Summary

• The curriculum employs a bottom-up pedagogical approach, teaching mathematical primitives before architectural composition.
• Each of the 10 progressive lessons pairs theoretical documentation (docs/en.md) with executable PyTorch code (code/main.py) and unit tests (code/tests.py).
• Self-attention is implemented from scratch with causal masking to ensure autoregressive behavior suitable for decoder-only language models.
• Multi-head attention uses parallel projection matrices (W_Q, W_K, W_V) followed by concatenation and output projection.
• Pre-normalization with LayerNorm and residual connections stabilizes deep transformer stacks.
• The final GPTModel class assembles token embeddings, positional embeddings, stacked transformer blocks, and a language modeling head into a trainable decoder-only architecture.
• The capstone project validates end-to-end functionality by training on synthetic data, confirming that theoretical understanding translates to working implementations.

Frequently Asked Questions

How does the curriculum implement causal masking in self-attention?

The curriculum implements causal masking by generating a lower-triangular boolean matrix using torch.tril(torch.ones(T, T)).bool(), which creates a mask where positions can only attend to themselves and previous tokens. Before applying softmax, the scores are masked with scores.masked_fill(mask == 0, -1e9), setting future positions to negative infinity to ensure zero attention weight after the softmax operation.

What is the dimensionality relationship between the model size and feed-forward networks in the transformer blocks?

According to the implementation in phases/07-transformers-deep-dive/04-transformer-block/code/main.py, the feed-forward network (MLP) uses a hidden dimension (d_ff) typically set to four times the model dimension (4 × d_model). This follows the original transformer architecture where the MLP expands the representation dimension before projecting back to the model dimension using a two-layer network with GELU activation.

How does the curriculum test that transformer components are implemented correctly?

Each lesson includes a code/tests.py file containing unit tests that verify tensor shape consistency, numerical correctness against reference implementations, and gradient flow through backpropagation. Learners run these tests using python3 -m unittest discover code/tests -v to validate that their implementations meet the expected mathematical contracts before proceeding to the next lesson.

What optimization techniques does the training loop teach for transformer models?

The training loop in phases/07-transformers-deep-dive/09-training-loop/code/main.py implements AdamW optimizer with weight decay for regularization, combined with a learning-rate schedule that includes a warm-up period followed by cosine decay. This combination helps stabilize early training while gradually reducing the learning rate to fine-tune the model parameters effectively during the capstone training phase.

{
  "mcpServers": {
    "instagit": {
      "command": "npx",
      "args": ["-y", "instagit@latest"]
    }
  }
}