LLM Architecture Deep Dive: From Tokens to Transformers

The Story That Explains LLMs

A Large Language Model (LLM) is a neural network — specifically a Transformer — trained on massive text corpora to model the probability distribution over sequences of tokens. Given a sequence of tokens as context, the model outputs a probability distribution over its vocabulary for the next token. Sampling from this distribution, repeatedly, produces fluent, coherent text.

From Words to Tokens — Tokenisation

Embeddings — Turning Tokens into Vectors

After tokenisation, each token ID is mapped to a dense vector called an embedding. This is a lookup table: token 4821 → a vector of 4096 numbers (for a 7B parameter model). These vectors live in a high-dimensional space where semantic similarity maps to geometric proximity.

import torch
import torch.nn as nn

# Simple embedding + sinusoidal position encoding
class TokenEmbedding(nn.Module):
    def __init__(self, vocab_size: int, d_model: int):
        super().__init__()
        self.embed = nn.Embedding(vocab_size, d_model)
        self.d_model = d_model

    def forward(self, x):  # x: [batch, seq_len] token IDs
        return self.embed(x) * (self.d_model ** 0.5)  # scale by √d_model

# Example usage
vocab_size, d_model, seq_len = 50257, 768, 10
embedder = TokenEmbedding(vocab_size, d_model)
tokens = torch.randint(0, vocab_size, (1, seq_len))  # 1 batch, 10 tokens
embedded = embedder(tokens)
print(embedded.shape)  # torch.Size([1, 10, 768])

The Transformer — Architecture Overview

The Transformer (Vaswani et al., 2017, "Attention Is All You Need") is the backbone of every modern LLM. It replaced RNNs and LSTMs by discarding recurrence entirely and relying solely on attention mechanisms. The result: massively parallelisable training and far superior long-range dependency modelling.

Self-Attention — The Heart of the Transformer

Self-attention is the mechanism that lets every token "look at" every other token in the sequence and decide how much information to gather from each. It computes three vectors for every token — Query (Q), Key (K), and Value (V) — and uses dot products to compute relevance scores.

import torch
import torch.nn.functional as F
import math

def scaled_dot_product_attention(Q, K, V, mask=None):
    """
    Q: [batch, heads, seq, d_k]
    K: [batch, heads, seq, d_k]
    V: [batch, heads, seq, d_v]
    """
    d_k = Q.shape[-1]

    # Step 1: Compute raw scores
    scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(d_k)
    # scores: [batch, heads, seq, seq]

    # Step 2: Apply causal mask (upper triangle = -inf)
    if mask is not None:
        scores = scores.masked_fill(mask == 0, -1e9)

    # Step 3: Softmax → attention weights
    attn_weights = F.softmax(scores, dim=-1)

    # Step 4: Weighted sum of values
    output = torch.matmul(attn_weights, V)
    return output, attn_weights

# --- Demo ---
batch, heads, seq_len, d_k = 1, 8, 5, 64
Q = torch.randn(batch, heads, seq_len, d_k)
K = torch.randn(batch, heads, seq_len, d_k)
V = torch.randn(batch, heads, seq_len, d_k)

# Causal mask: lower triangle = 1 (allowed), upper = 0 (masked)
mask = torch.tril(torch.ones(seq_len, seq_len)).unsqueeze(0).unsqueeze(0)
out, weights = scaled_dot_product_attention(Q, K, V, mask)
print(f"Output shape:  {out.shape}")
print(f"Weights shape: {weights.shape}")
print(f"Weight row sum (should be 1.0): {weights[0,0,2].sum():.4f}")

Multi-Head Attention — Many Perspectives at Once

A single attention head might learn to track grammatical subject-verb agreement. Another might track pronoun references. Another might group semantically related words. Multi-Head Attention runs H parallel attention heads, each with its own Q/K/V projection weights, then concatenates and projects their outputs.

class MultiHeadAttention(nn.Module):
    def __init__(self, d_model: int, n_heads: int):
        super().__init__()
        assert d_model % n_heads == 0
        self.n_heads = n_heads
        self.d_head  = d_model // n_heads

        # Projection matrices for Q, K, V and output
        self.W_q = nn.Linear(d_model, d_model, bias=False)
        self.W_k = nn.Linear(d_model, d_model, bias=False)
        self.W_v = nn.Linear(d_model, d_model, bias=False)
        self.W_o = nn.Linear(d_model, d_model, bias=False)

    def forward(self, x, mask=None):
        B, T, D = x.shape

        # Project and split into heads: [B, T, D] → [B, H, T, d_head]
        def split_heads(t):
            return t.view(B, T, self.n_heads, self.d_head).transpose(1, 2)

        Q = split_heads(self.W_q(x))
        K = split_heads(self.W_k(x))
        V = split_heads(self.W_v(x))

        # Attention for all heads in parallel
        attn_out, _ = scaled_dot_product_attention(Q, K, V, mask)

        # Recombine heads: [B, H, T, d_head] → [B, T, D]
        attn_out = attn_out.transpose(1, 2).contiguous().view(B, T, D)

        return self.W_o(attn_out)  # final output projection

The Feed-Forward Network — Where Knowledge Lives

After attention gathers context, a Feed-Forward Network (FFN) processes each token independently. It's a 2-layer MLP with a wide hidden layer (4×d_model), applying a non-linear activation. Researchers believe this is where factual knowledge is stored — attention routes information, the FFN transforms it.

class SwiGLU_FFN(nn.Module):
    """SwiGLU feed-forward — used in LLaMA, Mistral, Gemma."""
    def __init__(self, d_model: int, hidden: int):
        super().__init__()
        self.gate  = nn.Linear(d_model, hidden, bias=False)
        self.value = nn.Linear(d_model, hidden, bias=False)
        self.proj  = nn.Linear(hidden, d_model, bias=False)

    def forward(self, x):
        # Swish(gate(x)) acts as a learned, smooth gating signal
        return self.proj(
            F.silu(self.gate(x)) * self.value(x)
        )

# Sanity check
ffn = SwiGLU_FFN(d_model=4096, hidden=11008)  # LLaMA-7B dims
x   = torch.randn(1, 10, 4096)
print(ffn(x).shape)  # [1, 10, 4096] — same shape in, same shape out

# Count parameters in FFN alone (7B model has 32 layers)
total_ffn_params = sum(p.numel() for p in ffn.parameters())
print(f"FFN params per layer: {total_ffn_params:,}")
print(f"All 32 layers FFN total: {total_ffn_params * 32:,}")

Normalisation — Keeping Training Stable

Deep networks suffer from vanishing and exploding gradients. Normalisation layers stabilise the distribution of activations, enabling training of 100+ layer networks. Modern LLMs have largely moved from LayerNorm to RMSNorm.

class RMSNorm(nn.Module):
    """RMS Normalisation — used in LLaMA, Mistral, Gemma."""
    def __init__(self, d_model: int, eps: float = 1e-6):
        super().__init__()
        self.eps   = eps
        self.scale = nn.Parameter(torch.ones(d_model))  # learnable γ

    def forward(self, x):
        # RMS = sqrt(mean(x²) + eps)
        rms = x.pow(2).mean(-1, keepdim=True).add(self.eps).sqrt()
        return (x / rms) * self.scale

Positional Encoding — Giving Tokens a Sense of Place

Attention is permutation-invariant by design — shuffle the tokens and you get the same attention scores (just reordered). Positional encodings inject sequence order so the model knows "this token is at position 47".

Scaling Laws — Why Bigger Works Better

Grouped Query Attention & KV Cache

During inference, the model generates tokens one at a time. Without caching, it would recompute K and V for all previous tokens at every step — O(T²) computation. The KV Cache stores past K and V tensors, reducing generation to O(T) per step.

Pre-Training — The Foundation Phase

Pre-training is the most expensive phase: billions of parameters, trillions of tokens, thousands of GPUs, months of wall-clock time. The goal is to build a model with rich world knowledge and language understanding. Everything downstream — instruction following, RLHF — fine-tunes this foundation.

# Simplified pre-training loop (conceptual)
import torch
from torch.optim import AdamW
from torch.optim.lr_scheduler import CosineAnnealingLR

model    = LLM(vocab_size=128256, d_model=4096, n_layers=32, n_heads=32)
optimiser = AdamW(model.parameters(), lr=3e-4,
                   betas=(0.9, 0.95), weight_decay=0.1)
scheduler = CosineAnnealingLR(optimiser, T_max=1_000_000, eta_min=3e-5)

for step, batch in enumerate(dataloader):
    input_ids = batch["input_ids"]   # [B, T]
    labels    = input_ids[:, 1:]       # next-token labels: shift right
    logits    = model(input_ids[:, :-1])  # [B, T-1, vocab]

    # Cross-entropy loss over all positions
    loss = F.cross_entropy(
        logits.reshape(-1, logits.shape[-1]),
        labels.reshape(-1),
        ignore_index=-100   # padding tokens
    )

    optimiser.zero_grad()
    loss.backward()
    torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
    optimiser.step()
    scheduler.step()

    if step % 100 == 0:
        print(f"Step {step}: loss={loss.item():.4f}, perplexity={loss.item().exp():.2f}")

Post-Training — Instruction Tuning & Alignment

A pre-trained LLM is a raw text completer, not an assistant. If you prompt it "What is the capital of France?", it might continue "...was debated by historians...". Post-training transforms it into a helpful, safe, instruction-following model.

Mixture of Experts — Efficient Scaling

class MoELayer(nn.Module):
    """Sparse Mixture of Experts FFN layer."""
    def __init__(self, d_model: int, n_experts: int, top_k: int, hidden: int):
        super().__init__()
        self.n_experts = n_experts
        self.top_k     = top_k

        # Router: maps each token to expert logits
        self.router = nn.Linear(d_model, n_experts, bias=False)

        # E independent FFN experts
        self.experts = nn.ModuleList([
            SwiGLU_FFN(d_model, hidden) for _ in range(n_experts)
        ])

    def forward(self, x):
        B, T, D = x.shape
        x_flat = x.view(-1, D)  # [B*T, D]

        # Route: pick top-K experts per token
        logits  = self.router(x_flat)                   # [B*T, E]
        weights = F.softmax(logits, dim=-1)
        top_w, top_idx = torch.topk(weights, self.top_k, dim=-1)  # [B*T, K]
        top_w = top_w / top_w.sum(dim=-1, keepdim=True)  # renormalise

        # Aggregate expert outputs
        out = torch.zeros_like(x_flat)
        for k in range(self.top_k):
            expert_id = top_idx[:, k]   # [B*T]
            for e in range(self.n_experts):
                mask = (expert_id == e)
                if mask.any():
                    out[mask] += top_w[mask, k:k+1] * self.experts[e](x_flat[mask])

        return out.view(B, T, D)

Quantisation — Smaller, Faster, Cheaper

A 70B parameter model in FP32 requires ~280 GB VRAM — four A100s. Quantisation reduces the precision of weights and/or activations, dramatically shrinking memory and speeding up inference. The tradeoff: some accuracy loss.

# Load a 70B model in 4-bit quantisation using bitsandbytes
from transformers import AutoModelForCausalLM, BitsAndBytesConfig
import torch

quantisation_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_compute_dtype=torch.bfloat16,   # compute in bf16 for speed
    bnb_4bit_quant_type="nf4",              # NormalFloat4 — best quality at 4-bit
    bnb_4bit_use_double_quant=True,         # quantise the quantisation constants too
)

model = AutoModelForCausalLM.from_pretrained(
    "meta-llama/Meta-Llama-3-70B-Instruct",
    quantization_config=quantisation_config,
    device_map="auto",                       # auto-distribute across available GPUs
)

print(f"Model loaded. Memory: {model.get_memory_footprint() / 1e9:.1f} GB")

Fine-Tuning Efficiently — LoRA & PEFT

Full fine-tuning a 70B model requires 70B×4 bytes of weights + gradients + optimiser states ≈ 1–2 TB VRAM. Parameter-Efficient Fine-Tuning (PEFT) methods freeze most of the model and add a tiny number of trainable parameters.

from peft import LoraConfig, get_peft_model, TaskType
from transformers import AutoModelForCausalLM

# LoRA config: rank 16, applied to Q and V projection matrices
lora_config = LoraConfig(
    r=16,               # rank of the decomposition (A·B where A:[d,r], B:[r,d])
    lora_alpha=32,      # scaling factor: effective_lr = alpha/r
    target_modules=["q_proj", "v_proj"],   # which weight matrices to adapt
    lora_dropout=0.05,
    bias="none",
    task_type=TaskType.CAUSAL_LM
)

base_model = AutoModelForCausalLM.from_pretrained("meta-llama/Meta-Llama-3-8B")
model = get_peft_model(base_model, lora_config)
model.print_trainable_parameters()

Inference & Decoding Strategies

After training, the model outputs logits (unnormalised scores) over its vocabulary at each step. How you convert those logits to actual text is a decoding strategy — and the choice dramatically affects the character of the output.

import torch
import torch.nn.functional as F

def sample_next_token(logits, temperature=0.8, top_p=0.9, top_k=50):
    """
    logits: [vocab_size] — raw model output for next position
    Returns: sampled token id (int)
    """
    # 1. Temperature scaling
    logits = logits / temperature

    # 2. Top-K filtering: zero out all but top K
    if top_k > 0:
        top_k_vals = torch.topk(logits, top_k).values[-1]
        logits[logits < top_k_vals] = -float("inf")

    # 3. Top-P (nucleus) filtering
    probs = F.softmax(logits, dim=-1)
    sorted_probs, sorted_idx = torch.sort(probs, descending=True)
    cumulative = torch.cumsum(sorted_probs, dim=-1)
    # Remove tokens once cumulative prob exceeds top_p
    remove_mask = cumulative - sorted_probs > top_p
    sorted_probs[remove_mask] = 0
    probs = torch.zeros_like(logits).scatter_(0, sorted_idx, sorted_probs)
    probs = probs / probs.sum()   # renormalise

    # 4. Sample
    return torch.multinomial(probs, num_samples=1).item()

LLM Architecture Comparison — Modern Landscape

Evaluation — How We Measure LLM Quality

Golden Rules for LLM Practitioners