Linear Algebra for Transformers

A transformer is mostly matrix multiplication: embeddings are vectors, every layer applies matrices to them, and attention is a matrix of dot products. This note covers the linear algebra you need to read and reason about LLM code and papers - vectors and similarity, matrix shapes through a transformer layer, counting parameters and FLOPs, rank and low-rank factorization (the idea behind LoRA), SVD, and norms.

Vectors, Dot Products and Similarity

Where this shows up:

• Attention scores are dot products between a query and every key (Attention Mechanisms).
• Retrieval ranks documents by cosine similarity (or dot product) between embeddings (Embeddings & Vector Search). If embeddings are normalized to length 1, dot product and cosine similarity are the same, and many vector databases rely on that.
• The output layer computes a dot product between the final hidden state and every token's output embedding; the largest dot products become the most likely next tokens.

Matrices and Shapes

A matrix W of shape (m, n) is a linear map from ℝⁿ to ℝᵐ - or, with the row-vector convention most ML code uses, y = x W maps a length-m row vector to length n. The one rule to internalize:

(a, b) @ (b, c) -> (a, c)        inner dimensions must match; they disappear

Batched tensors follow the same rule on their last two dimensions; leading dimensions (batch, heads) are carried along.

Shapes through one transformer layer (B = batch, T = sequence length, d = model width, h = heads, d_h = d / h):

flowchart TD
    X["🔢 X<br/>(B, T, d)"] --> QKV["Q, K, V = X Wq, X Wk, X Wv<br/>each (B, T, d)"]
    QKV --> R["✂️ split heads<br/>(B, h, T, d_h)"]
    R --> S["🎯 scores = Q Kᵀ / √d_h<br/>(B, h, T, T)"]
    S --> P["softmax over last axis<br/>(B, h, T, T)"]
    P --> O["P V<br/>(B, h, T, d_h)"]
    O --> M["🔗 merge heads, @ Wo<br/>(B, T, d)"]
    M --> F["🧮 FFN: (B, T, d) -> (B, T, 4d) -> (B, T, d)"]

    style X fill:#d8dfe8,stroke:#b0bac8
    style S fill:#e8e0d4,stroke:#c8b89a
    style F fill:#dde4dc,stroke:#b0c4b0

The (T, T) score matrix is why attention memory and compute grow with the square of sequence length - and why FlashAttention's trick of never materializing it matters.

Counting Parameters and FLOPs

Parameters per layer. Attention has four d x d matrices (Q, K, V, output): 4d². A classic FFN expands to 4d and back: 2 x d x 4d = 8d². So a layer has about 12d² parameters, ignoring biases and norms.

Worked example - GPT-2 small (d = 768, 12 layers, vocabulary 50,257, context 1,024):

That matches GPT-2 small's published 124M. Modern models change the constants - gated FFNs (SwiGLU) use three matrices instead of two, and grouped-query attention shrinks K and V - but the habit of counting d x d blocks carries over.

FLOPs. Multiplying (m, k) @ (k, n) takes m x n x k multiply-adds, which is 2mnk FLOPs. For a whole model, a forward pass costs about 2N FLOPs per token (N = parameters), and the backward pass about twice that, so training costs ≈ 6ND FLOPs for D training tokens - the estimate behind Scaling Laws & Compute Budgets.

Rank and Low-Rank Factorization

The rank of a matrix is the number of linearly independent directions it can produce - the dimension of its output space. A d x d matrix can have rank up to d. A matrix of rank r can be written as a product of two thin matrices:

W (d_out x d_in), rank r  =  B (d_out x r) @ A (r x d_in)
parameters: d_out x d_in   vs   r x (d_out + d_in)

This is the whole idea of LoRA: freeze the pretrained W and learn a low-rank update ΔW = B A. Research on the "intrinsic dimension" of fine-tuning found that adapting a pretrained model needs far fewer degrees of freedom than the model has, which is why a small r works.

Worked example. One 4096 x 4096 projection, LoRA rank 16:

• Full matrix: 4096² = 16,777,216 parameters
• LoRA: 16 x (4096 + 4096) = 131,072 parameters - 0.78% of the matrix

That ratio, applied across every adapted matrix, is why a LoRA fine-tune fits on one GPU (LoRA & QLoRA Hands-On).

Singular Value Decomposition

Every matrix factors as W = U Σ Vᵀ: U and V have orthonormal columns (rotations), and Σ is diagonal with non-negative singular values σ₁ ≥ σ₂ ≥ ... (stretches). Read it as: rotate, stretch along each axis, rotate.

Three facts you will use:

• The rank is the number of non-zero singular values. Many small singular values mean the matrix is "nearly" low-rank.
• Truncated SVD is the best low-rank approximation (Eckart-Young, 1936): keeping the top r singular values gives the rank-r matrix closest to W. This underpins low-rank compression, and LoRA variants such as PiSSA that initialize the adapter from the top singular directions.
• The largest singular value is the spectral norm - the most the matrix can stretch any vector. Large spectral norms across layers are one way activations and gradients blow up.

Norms

LayerNorm and RMSNorm rescale each token's activation vector: LayerNorm subtracts the mean and divides by the standard deviation; RMSNorm (used by most current LLMs) only divides by the root-mean-square √(mean(xᵢ²)), which is cheaper and works as well in practice. Both keep activation magnitudes stable from layer to layer.

Check Yourself

Exercises

Study Notes

Must-know:

• (a, b) @ (b, c) -> (a, c); attention scores are (B, h, T, T) - quadratic in sequence length
• Parameters per classic layer ≈ 12d² (4d² attention + 8d² FFN); GPT-2 small ≈ 124M
• Matmul FLOPs = 2mnk; forward ≈ 2N per token; training ≈ 6ND
• Rank-r matrix = thin B @ A; LoRA trains r(d_in + d_out) parameters per adapted matrix
• SVD = rotate, stretch, rotate; truncated SVD is the best low-rank approximation; top singular value = spectral norm
• Cosine = dot product of unit vectors; RMSNorm divides by the root-mean-square

References

• Deisenroth, Faisal and Ong, Mathematics for Machine Learning (Cambridge University Press, 2020) - chapters 2-4
• Strang, Linear Algebra and Learning from Data (2019)
• Vaswani et al., Attention Is All You Need (2017)
• Kaplan et al., Scaling Laws for Neural Language Models (2020) - the 6ND estimate
• Hu et al., LoRA: Low-Rank Adaptation of Large Language Models (2021); Aghajanyan et al., Intrinsic Dimensionality Explains the Effectiveness of Language Model Fine-Tuning (2020)
• Eckart and Young, The Approximation of One Matrix by Another of Lower Rank (Psychometrika, 1936); Meng et al., PiSSA (2024)
• Zhang and Sennrich, Root Mean Square Layer Normalization (2019)

Last reviewed: 2026-10