LLM Architecture Overview & Next-Token Prediction

The Big Idea

Large Language Models (LLMs) like GPT-4, Llama 3, and Claude often feel magical or intimidating, but beneath the massive scale lies an astonishingly simple game: guessing the next character or word. Imagine sitting with an apprentice storyteller who is listening to you tell a tale. Every time you pause mid-sentence, their job is to guess the exact next sound or letter that should follow. If you say 'Once upon a t...', they should bet heavily on 'i' followed by 'm' and 'e'.

Why does this simple guessing game produce human-like reasoning, code generation, and poetry? Because in order to accurately guess what letter comes next in a medical textbook, a legal contract, or a Python script, the neural network cannot just memorize rhymes. It is forced to understand grammar, syntax, facts about the world, and mathematical logic. Intelligence is the byproduct of optimizing next-token prediction at scale.

Modern generative models are 'decoder-only' Transformers, as pioneered in Attention Is All You Need (Vaswani et al. 2017) and popularized in minimal implementations like Andrej Karpathy's nanoGPT. Unlike older translation systems that read a full sentence in French before translating it into English, a decoder-only model processes text causally from left to right. It reads the past, processes it through stacked layers of attention and neural networks, and outputs scores for what token should arrive next.

Architecture Blueprint & Tensor Mechanics

The entire GPT architecture boils down to four distinct stages: embedding tokens and positions into continuous space using PyTorch Tensor & Broadcasting rules, passing them through stacked Transformer blocks, normalizing activations, and projecting to vocabulary logits. In Lesson 1.2, we start by engineering the character tokenizer.

import torch
import torch.nn as nn
from torch.nn import functional as F

class GPTLanguageModel(nn.Module):
    def __init__(self, vocab_size, n_embd, block_size, n_layer, n_head):
        super().__init__()
        # 1. Token embedding table: maps token IDs to vectors of size n_embd
        self.token_embedding_table = nn.Embedding(vocab_size, n_embd)
        # 2. Position embedding table: gives each sequence index a unique vector
        self.position_embedding_table = nn.Embedding(block_size, n_embd)
        # 3. Stack of Transformer blocks (attention + feed-forward)
        self.blocks = nn.Sequential(*[Block(n_embd, n_head, block_size) for _ in range(n_layer)])
        # 4. Final layer normalization
        self.ln_f = nn.LayerNorm(n_embd)
        # 5. Language model head: projects n_embd back to vocabulary scores
        self.lm_head = nn.Linear(n_embd, vocab_size)

    def forward(self, idx, targets=None):
        B, T = idx.shape
        # idx and targets are both (B, T) tensors of integers
        tok_emb = self.token_embedding_table(idx) # (B, T, C)
        pos_emb = self.position_embedding_table(torch.arange(T, device=idx.device)) # (T, C)
        x = tok_emb + pos_emb # (B, T, C) - via PyTorch tensor broadcasting
        x = self.blocks(x)    # (B, T, C) - through all transformer blocks
        x = self.ln_f(x)      # (B, T, C) - final layer norm
        logits = self.lm_head(x) # (B, T, vocab_size)

        if targets is None:
            loss = None
        else:
            B, T, C = logits.shape
            logits = logits.view(B * T, C)
            targets = targets.view(B * T)
            loss = F.cross_entropy(logits, targets)

        return logits, loss