AskAiml
Sign in Get started

Math for LLMs (only the parts you really need)

Vectors, matrices, dot products, softmax, gradients, and cross-entropy — the small toolkit that unlocks reading papers.

🔢 Module 0 9 min read Not started

Why this matters

You will hear "you need a PhD to understand LLMs" and "you need zero math to use them." Both are wrong. There is a small, focused math kit — linear algebra primitives, basic probability, the calculus of gradients, and a sprinkle of information theory — that lets you read papers, debug training, and have honest conversations with researchers. That is the goal of this lesson.

If you already know matrices, dot products, softmax, gradients, and cross-entropy, skim and move on. Otherwise: 4-6 hours here saves you weeks later.

Learning objectives

  1. Manipulate vectors and matrices with NumPy / PyTorch.
  2. Compute dot products, matmuls, and softmax intuitively.
  3. Reason about gradients and backpropagation at a high level.
  4. Use cross-entropy loss and log-probabilities correctly.
  5. Read shapes in transformer code without confusion.

1. Linear algebra — the dialect of deep learning

Vectors and matrices

A vector is a 1-D array of numbers (e.g., a token embedding [0.1, -0.3, 0.7, ...] of length d). A matrix is a 2-D array (e.g., a batch of B embeddings is shape (B, d)). A tensor is just an n-D generalisation; PyTorch and NumPy both use the term.

python
import torch
x = torch.tensor([1.0, 2.0, 3.0])           # shape (3,)
A = torch.randn(4, 3)                       # shape (4, 3)
y = A @ x                                   # matrix-vector product, shape (4,)
B = torch.randn(3, 5)
M = A @ B                                   # (4,3) @ (3,5) = (4,5)

Dot product (the operation)

a·b = sum(a_i * b_i). It measures alignment — how similar two vectors are. The attention mechanism is essentially: take dot products between queries and keys to score relevance.

python
a = torch.tensor([1.0, 0.0])
b = torch.tensor([0.7, 0.7])
sim = (a * b).sum()                         # 0.7 — partial alignment

Cosine similarity normalises it: cos(a,b) = (a·b) / (|a||b|). Used everywhere in retrieval.

Matmul shapes (read this twice)

The single most common bug: shape mismatch.

  • (B, T, d) @ (d, d) → (B, T, d) — apply a per-token linear transform.
  • (B, T, d) @ (B, d, T) → (B, T, T) — pairwise scores between T tokens.
  • (B, h, T, k) @ (B, h, k, T) → (B, h, T, T) — multi-head attention scores.

PyTorch broadcasts over leading dimensions: (B, T, d) @ (d, k) → (B, T, k) works because the (d, k) is broadcast across (B, T, *).

Softmax (turning scores into probabilities)

softmax(x_i) = exp(x_i) / sum_j(exp(x_j))

It outputs values in [0,1] that sum to 1 — a probability distribution. Numerically unstable in raw form (exp(1000) overflows). The trick: subtract the max first.

python
def softmax(x):
    x = x - x.max(dim=-1, keepdim=True).values
    e = x.exp()
    return e / e.sum(dim=-1, keepdim=True)

PyTorch ships torch.softmax(x, dim=-1) and torch.log_softmax(x, dim=-1) (the latter is more stable when you immediately take a log).

Layer norm vs batch norm vs RMSNorm

You will see all three; LLMs almost exclusively use LayerNorm or RMSNorm.

  • LayerNorm: normalise across the feature dim per token. y = (x - mean)/std * γ + β.
  • RMSNorm (used by Llama, Mistral, Qwen): y = x / sqrt(mean(x^2)) * γ. Drops the mean step → ~10% faster, equally good empirically.

Both stabilise training. You do not need to derive them by hand — just know what they do.

2. Probability — for sampling and loss

Random variables and distributions

A discrete distribution assigns probabilities to outcomes (e.g., a 50k-token vocabulary distribution from your LLM).

Useful properties:

  • sum p_i = 1
  • E[X] = sum x_i p_i (expectation)
  • Entropy: H(p) = -sum p_i log p_i — measures uncertainty.

Cross-entropy loss (the LLM loss)

For training a next-token predictor, the loss for a single token is:

loss = -log p(true_token | context)

Across a batch of N tokens:

loss = (1/N) * sum_i -log p_i(true_i)

In PyTorch:

python
import torch.nn.functional as F
logits = model(x)                                         # (B, T, V)
loss = F.cross_entropy(logits.view(-1, V), y.view(-1))    # scalar

F.cross_entropy takes raw logits (NOT probabilities) and a target index per token.

Why log? Multiplying probabilities is unstable; summing log-probs is stable.

KL divergence

KL(p || q) = sum p_i * log(p_i / q_i) — "how much does q surprise me when the truth is p?" Used in RLHF and DPO to keep a fine-tuned model close to a base model.

Sampling vs argmax

Greedy (argmax over the distribution) is deterministic. Sampling uses temperature T:

p_i ∝ exp(logits_i / T)

T → 0 becomes argmax; T → ∞ becomes uniform random. You will see this again in inference (Phase 5).

3. Calculus — gradients without the agony

You almost never differentiate by hand. PyTorch (autograd) does it. But you must:

  1. Understand that gradient = direction of steepest increase.
  2. Believe that gradient descent moves opposite the gradient: θ ← θ − lr · ∇θ L.
  3. Recognise the chain rule: gradients of compositions are products of local gradients. That is what backpropagation computes.
python
import torch
x = torch.tensor([2.0], requires_grad=True)
y = x**3 + 2*x
y.backward()
print(x.grad)        # 3*x^2 + 2 = 14

That is autograd. Every layer in PyTorch knows its local Jacobian; .backward() walks the computational graph.

Optimisers

  • SGD: θ ← θ − lr · g. Vanilla.
  • Momentum: add a velocity buffer (helps with valleys).
  • Adam / AdamW (default for LLMs): maintains 1st and 2nd moment estimates of the gradient. AdamW decouples weight decay (now standard).
  • Adafactor / Lion / Sophia: memory-efficient or Hessian-aware variants used at scale.
  • 8-bit Adam (bitsandbytes): keeps optimiser states in 8-bit; saves a lot of VRAM.

You will use torch.optim.AdamW 95% of the time.

Learning rate schedules

LLMs use warmup → cosine decay almost universally. Warmup avoids early instability; cosine reaches a smooth small value at the end.

python
from torch.optim.lr_scheduler import OneCycleLR  # or use HF schedulers

4. Information theory snippets

You will encounter these terms; here is what they mean:

  • Entropy H(X) — average bits needed to encode samples from X. Higher = more random.
  • Perplexity — exp(H). The standard intrinsic metric for language models. PPL = exp(loss). Lower is better. Llama-3 8B base gets ~7-8 PPL on WikiText.
  • Mutual information I(X;Y) — how much knowing X reduces uncertainty about Y.
  • Negative log-likelihood (NLL) — same as cross-entropy with one true label.

Quick mental model: when somebody says "training reduced loss from 3.0 to 2.4," they mean the average NLL per token went down — i.e., the model assigns higher probability to the true next token.

5. Reading shapes in LLM code (the cheat sheet)

You will see these symbols repeatedly. Burn them in:

  • B — batch size (number of sequences).
  • T — sequence length (tokens per sequence).
  • V — vocabulary size (often 32k-200k).
  • d (d_model, n_embd, hidden_size) — embedding dimension (768, 4096, ...).
  • h — number of attention heads.
  • k (d_head) — per-head dimension. Usually d / h.
  • L — number of transformer layers.

A typical decoder-only forward pass works on:

input ids:  (B, T)            integers in [0, V)
embed:      (B, T, d)
attn out:   (B, T, d)
mlp out:    (B, T, d)
final:      (B, T, V)         logits

Multi-head attention internally reshapes:

Q: (B, T, d) -> (B, T, h, k) -> (B, h, T, k)
QK^T:                         -> (B, h, T, T)
attended:                     -> (B, h, T, k) -> (B, T, h, k) -> (B, T, d)

When code looks confusing, annotate the shape on every line. It is the senior trick.

Hands-on lab (3 hours)

math_drills.ipynb:

  1. Generate A: (4, 5) and B: (5, 3). Print A @ B and verify shape.
  2. Implement softmax(x, dim=-1) from scratch. Compare with torch.softmax.
  3. Implement log_softmax(x, dim=-1) numerically stably.
  4. Manually compute cross-entropy on a 5-class toy and compare to F.cross_entropy.
  5. Use autograd: define f(x,y) = (x*y + sin(x))^2. Compute the partials at (2, 3) automatically and by hand. Verify they match.
  6. Plot loss curves of SGD vs Adam on loss(w) = (w-3)^2 + 0.5*sin(5w) starting at w=0.
  7. Bonus: implement a tiny RMSNorm class and verify that scaling input by 10x leaves output approximately unchanged in direction.

Common pitfalls

  1. Forgetting dim=-1 in softmax / log_softmax — collapses on the wrong axis.
  2. Passing probabilities to cross_entropy — it expects raw logits.
  3. Re-using a graph after .backward() without retain_graph=True — autograd frees the buffers.
  4. Mixing float16 and float32 without autocast — silent precision bugs.
  5. A.T @ B when you wanted A @ B.T — always print shapes.

Self-check

  1. What is the difference between a dot product and cosine similarity?
  2. Why subtract the max before computing softmax?
  3. Express perplexity in terms of cross-entropy loss.
  4. What does AdamW add to plain SGD?
  5. Why does almost every LLM use LayerNorm or RMSNorm?

References

  • 3Blue1Brown — Essence of Linear Algebra (YouTube playlist).
  • 3Blue1Brown — Essence of Calculus + Neural Networks playlist.
  • Deep Learning (Goodfellow, Bengio, Courville) — Chapters 2-4.
  • Mathematics for Machine Learning (Deisenroth et al.) — free online.
  • PyTorch tutorials: Autograd basics.
  • The Matrix Cookbook — reference, not a textbook.

Sign in to save your progress and earn badges.