Will Ahmad

Linear algebra for LLMs

A flashcard refresher on the linear algebra inside transformers, one idea per card, with skippable checks. Vectors, matmuls, shapes, attention, low rank, gradients.

  • Intermediate
  • 70 steps
  • ~454 min
Start →

  1. Part 0: Start here

    How the cards work, the notation, and a quick diagnostic.

    3 steps · ~15 min

    1. Welcome — What this course covers, how the flashcards work, and the map of all ten parts.
    2. Notation — Scalars, vectors and matrices; shapes as (rows, columns); row vectors in code vs column vectors in papers; how nn.Linear stores its weight.
    3. Where are you? — An optional 12-question diagnostic, about one question per part, to find the parts you can skim and the ones to take slowly.

  2. Part 1: Vectors and embedding space

    Embeddings, dot products, norms and cosine similarity, and what high dimensions do to them.

    6 steps · ~37 min

    1. Vectors and embeddings — A vector as a list, an arrow and a point; dimension d; the embedding table, and how a sentence becomes a T × d matrix.
    2. Add, scale, combine — Vector addition and scaling as numbers and as arrows, linear combinations, the residual stream as a running sum, and mean pooling.
    3. The dot product — Multiply matching entries and add; the same number as lengths times the cosine of the angle; its sign; projections; and query–key scores.
    4. Norms and cosine similarity — Length (L2 and L1 norms), unit vectors, cosine similarity, and why a dot product only equals the cosine for unit vectors.
    5. High dimensions — Random vectors in high dimensions are nearly orthogonal, and a random dot product grows like √d. This is why attention divides by √dₖ.
    6. Check: vectors — Six questions with new numbers each try. Dot products and norms by hand, dot vs cosine, and the √dₖ scaling. Skippable.

  3. Part 2: Matrices and matrix multiplication

    Shapes, transposes, the four ways to read a matrix product, and what a matmul costs.

    8 steps · ~54 min

    1. Matrices and shapes — Rows, columns and entries; a matrix as a stack of rows (tokens) or of columns (features); the shapes you meet in a transformer.
    2. Transpose — Flip a matrix over its diagonal cell by cell and watch its shape flip. Then symmetric matrices, and the kᵀ inside attention.
    3. Matrix × vector — Two ways to read Ax, as one dot product per row and as a mix of columns, then the row-vector version xW that code uses.
    4. Matrix × matrix — Each entry of AB is a row of A dotted with a column of B. The shape rule, AB ≠ BA, regrouping, and (AB)ᵀ = BᵀAᵀ.
    5. Four ways to see AB — The same product read as dot products, column mixes, row mixes and a sum of outer products. Gram matrices, the token view, and rank-1 pieces.
    6. Special matrices — Identity, diagonal (the same as elementwise scaling), the elementwise product ⊙ vs @, triangular matrices (the causal mask), and block matrices (heads).
    7. What a matmul costs — Count a matrix product as 2mkn FLOPs, see why a weight costs 2 FLOPs per parameter per token, and why the order of a chain of products matters.
    8. Check: matrices — Six questions drawn from seven, with new numbers each try. Shapes, a product entry by hand, transposes, * vs @, FLOPs and outer products. Skippable.

  4. Part 3: Matrices as transformations

    Linear maps, rotations, rank and projections, the geometry behind RoPE, LoRA and steering.

    7 steps · ~47 min

    1. Linear maps — A matrix as a function that moves every point of the plane. Its columns are where the basis vectors land; scale, shear, reflect and project.
    2. Composition — Doing one map after another is a matrix product, read right to left. Order matters, and stacked linear layers collapse into one matrix.
    3. Rotations — The 2×2 rotation matrix; rotations keep lengths and dot products, add their angles, and undo with a transpose. The identity behind RoPE.
    4. Span, basis and rank — Span, independence and basis; the column space; rank as the number of independent columns; why rank(AB) ≤ min(rank A, rank B) limits an attention head.
    5. Inverse, determinant, null space — The determinant as an area factor, when a matrix can be undone, the null space it destroys, and why orthogonal matrices invert for free. Brief, as you rarely invert in ML.
    6. Projections — Project a vector onto a direction, remove a direction, and build the projection matrix uuᵀ. How ablation, steering, probes and LayerNorm's centering use it.
    7. Check: transformations — Six questions drawn from seven, with new numbers each try. Columns as images, determinants, composition order, rank, projections, rotations. Skippable.

  5. Part 4: Tensors, shapes and einsum

    Broadcasting, reshape vs transpose, splitting heads, reductions and einsum.

    6 steps · ~40 min

    1. Tensors — Scalars, vectors and matrices are tensors of rank 0, 1 and 2. Read (B, T, d) axis by axis, index into it, and count its memory.
    2. Broadcasting — Align shapes from the right and stretch the size-1 axes. Biases, masks and per-token scales, plus the (T,) vs (T, 1) bug that runs without an error.
    3. Reshape vs transpose — Memory is one flat row-major list. Reshape re-cuts it; transpose changes how you walk it. Split tokens into heads correctly, and see the version that silently scrambles them.
    4. Reductions and axes — Sum, mean and max along an axis remove it; keepdim keeps it for broadcasting back. Softmax runs over the keys, LayerNorm over the features.
    5. Batched matmul and einsum — How @ works on 3- and 4-axis tensors, and einsum, which names every axis. The strings you'll meet in attention code, and how to count their FLOPs.
    6. Check: tensors and shapes — Six questions drawn from seven, with new shapes each try. Indexing, broadcasting, reshape vs transpose, splitting heads, keepdim, einsum. Skippable.

  6. Part 5: Attention, one operation at a time

    Softmax, Q, K and V, the √dₖ scaling, the causal mask, multi-head attention and the KV cache.

    9 steps · ~58 min

    1. Tokens to vectors — An embedding lookup is a one-hot row vector times the embedding table. A whole sequence is a one-hot matrix times E. Where position information comes in.
    2. Softmax — Exponentiate, then normalize. Shifting all scores changes nothing (so subtract the max), temperature sharpens or flattens, and −∞ gives exactly zero weight.
    3. Q, K and V — Three projections of the same tokens. Compute queries, keys and values for a three-token example by hand; these numbers carry through the rest of Part 5.
    4. Scores and the √dₖ scaling — Every query dotted with every key gives the T × T score matrix QKᵀ. Compute it for the example, then scale by 1/√dₖ and see why.
    5. The causal mask — Add −∞ above the diagonal, then softmax each row. Get the attention weights for the example, and see why zeroing weights after the softmax is wrong.
    6. Mixing values — The output is AV, so each token's output is a weighted average of value rows. Finish the example, then write the whole of attention in one line and count its cost.
    7. Multi-head attention — h heads attend in parallel, each in its own d/h-wide slice; their outputs are concatenated and mixed by W_O. Still 4d² parameters, and each head is low-rank.
    8. GQA and the KV cache — Generation reuses past keys and values from a cache. Size it with one formula, and see how grouped-query attention (GQA) and MQA shrink it.
    9. Check: attention — Six questions drawn from seven, with new numbers each try. A scaled score, softmax facts, the mask, a weighted average, parameter counts and KV cache size. Skippable.

  7. Part 6: The rest of the transformer

    Linear layers, the residual stream, normalization, SwiGLU, RoPE, logits, and counting a model.

    8 steps · ~51 min

    1. Linear layers in code — What nn.Linear computes, why its weight is stored (out, in), how to count its parameters, and why each output is a dot product with one row of the weight.
    2. The residual stream — Every sublayer reads from and adds back into one d-dimensional vector per token. The block equations, the unrolled sum, and why the identity path helps gradients.
    3. LayerNorm and RMSNorm — LayerNorm centers and rescales each token's features, then applies a per-feature gain; RMSNorm skips the centering. Both by hand, and what they do geometrically.
    4. The MLP and SwiGLU — Up-project, apply a nonlinearity, down-project. SwiGLU adds a gate multiplied elementwise. Shapes, parameter counts, and why most of a model's parameters live here.
    5. RoPE — Rotary position embeddings rotate each pair of query and key dimensions by an angle proportional to the position, so scores depend only on relative distance. Frequencies, and how code implements it.
    6. Logits — The final hidden state times the unembedding gives one score per vocabulary token. Weight tying, a toy example by hand, softmax to probabilities, and what it costs.
    7. Count a model — Add up every matrix. 12d² per classic block plus embeddings, checked against GPT-2 small; an exact count for Llama 3 8B; and 2N FLOPs per token.
    8. Check: the transformer — Six questions drawn from eight, with new numbers each try. Linear and SwiGLU parameter counts, LayerNorm geometry, residuals, RoPE, tying, sizing a model and its FLOPs. Skippable.

  8. Part 7: Decompositions and low rank

    Eigenvectors, the SVD, low-rank approximation, LoRA, conditioning and PCA.

    7 steps · ~45 min

    1. Eigenvectors — Directions a matrix only stretches. Find them for a 2×2, see why symmetric matrices have orthogonal ones, and why repeated multiplication makes the largest eigenvalue win.
    2. The SVD — Every matrix is a rotation, a stretch along the axes, and another rotation. Find the singular values of a 2×2 from AᵀA, and check the factorization multiplies back.
    3. Low-rank approximation — The SVD as a sum of rank-1 pieces, largest first. Keep the top k for the best rank-k approximation (Eckart–Young), measure the error, and count the storage saved.
    4. LoRA — Freeze W and train a low-rank update BA. Shapes, why B starts at zero, the α/r scale, parameter counts for a real model, and merging the update for free inference.
    5. Matrix norms and conditioning — How big is a matrix? The Frobenius norm, the spectral norm (its largest stretch), and the condition number that says how much it amplifies errors.
    6. PCA on embeddings — Center a cloud of embeddings, take its SVD, and project onto the top directions to plot it in 2D. Explained variance, and what to watch out for.
    7. Check: low rank — Six questions drawn from seven, with new numbers each try. Eigenvalues, singular values, truncation error, LoRA counts and init, the spectral norm. Skippable.

  9. Part 8: Gradients through matrices

    Gradients with shapes, the chain rule as matrix products, a linear layer backwards, and training cost.

    6 steps · ~36 min

    1. Gradients have shapes — The gradient of a scalar loss with respect to a parameter has that parameter's shape. Gradients of w·x and ½|w|², and one step of gradient descent by hand.
    2. The chain rule as matrix products — Jacobians of vector functions, why composing functions multiplies Jacobians, and why backpropagation multiplies from the loss end as vector–Jacobian products.
    3. A linear layer backwards — For Y = XW, derive ∂L/∂W = XᵀG and ∂L/∂X = GWᵀ by matching shapes, check them on numbers, and see why the backward pass costs twice the forward.
    4. Softmax and cross-entropy — The gradient of cross-entropy with respect to the logits is simply p − y. Derive it in two lines, compute it by hand, and meet the softmax Jacobian that attention backpropagates through.
    5. Training cost — Forward 2N plus backward 4N gives about 6N FLOPs per training token. Size a training run in FLOPs and GPU-hours, and count the memory per parameter.
    6. Check: gradients — Six questions drawn from seven, with new numbers each try. Gradient shapes, a descent step, a vector–Jacobian product, linear-layer gradients, p − y, and training compute. Skippable.

  10. Part 9: Numbers on real hardware

    Number formats, arithmetic intensity, quantization, stable and online softmax, tiling.

    6 steps · ~39 min

    1. Number formats — Sign, exponent and mantissa. FP32, BF16, FP16, FP8 and integer formats compared on range and precision, and what each does to a model's memory.
    2. Arithmetic intensity — FLOPs per byte moved decides whether a GPU waits on compute or on memory. Why decoding one sequence is memory-bound, how batching fixes it, and a speed limit you can compute.
    3. Quantization — Scale, round, and scale back. Quantize a vector to INT8 and INT4 by hand, see how one outlier ruins a shared scale, and why per-group scales fix it.
    4. Stable and online softmax — Why exp overflows in low precision, log-sum-exp, and the running-max trick that lets FlashAttention compute softmax one block at a time without ever holding a whole row.
    5. Tiled matmul — Block matrix multiplication by hand, and why loading tiles into fast memory raises a matmul's arithmetic intensity from about 1 to about the tile size.
    6. Check: hardware — Six questions drawn from seven, with new numbers each try. Weight memory, decode speed limits, intensity, INT8 rounding, online softmax and tiling. Skippable.

  11. Part 10: Capstone and reference

    Size a real model end to end, then keep the cheat sheet and glossary.

    4 steps · ~32 min

    1. Llama 3 8B, end to end — Follow a prompt through every matrix in Llama 3 8B, writing each shape, then size the model the way you would before deploying it: memory, decode speed, prefill and attention FLOPs, training compute, and a LoRA adapter.
    2. Final check — Eight questions drawn from ten, across the whole course, with new numbers each try. Shapes, heads, attention, the KV cache, LoRA, gradients, decode and prefill. Skippable.
    3. Cheat sheet — Every formula and rule of thumb from the course on one page, grouped by part, to keep open while you read model code.
    4. Glossary — Short definitions of every term used in the course, alphabetical, each with the deck that explains it.