01. Transformer Architecture Review

The transformer architecture, introduced in "Attention Is All You Need" (Vaswani et al., 2017), remains the foundation of modern large language models. Understanding its core structure is essential before modifying or optimizing it.

A transformer consists of an input embedding layer, stacked encoder or decoder blocks, and a output projection layer. The encoder processes source sequences bidirectionally; the decoder generates output autoregressively using masked self-attention.

import torch
import torch.nn as nn

class SimpleTransformer(nn.Module):
    def __init__(self, vocab_size, d_model, n_layers, n_heads):
        super().__init__()
        self.embedding = nn.Embedding(vocab_size, d_model)
        self.pos_embedding = nn.Embedding(2048, d_model)  # max sequence length
        self.blocks = nn.ModuleList([
            TransformerBlock(d_model, n_heads) 
            for _ in range(n_layers)
        ])
        self.lm_head = nn.Linear(d_model, vocab_size, bias=False)
    
    def forward(self, input_ids, attention_mask=None):
        # input_ids: (batch, seq_len)
        batch_size, seq_len = input_ids.shape
        
        # Embed tokens and positions
        positions = torch.arange(seq_len, device=input_ids.device)
        x = self.embedding(input_ids) + self.pos_embedding(positions)
        
        # Pass through transformer blocks
        for block in self.blocks:
            x = block(x, attention_mask)
        
        # Project to vocabulary
        logits = self.lm_head(x)
        return logits

Failure mode: Position overflow occurs when sequence length exceeds pos_embedding capacity. If you train on sequences of 512 tokens and later need 1024, the model crashes with an index error. Always set pos_embedding to your maximum expected sequence length.

Failure mode: Memory scales quadratically with sequence length. A 4096-length sequence requires 16x more memory than 1024 tokens in naive attention implementations. This constrains batch sizes and limits context windows.

Local verification checkpoint

Run the smallest example from this chapter in a local workspace and record the package version, runtime, data path, and observed output. If the result depends on model size, vector count, CPU/GPU backend, or available memory, note that constraint beside the exercise so the lesson remains reproducible.