ai_ml / math foundations / 01_linear_algebra.md

Linear algebra for ML

6 interview angles 5 min read source

Linear algebra for ML

You will not be asked to prove theorems. You will be asked what a dot product means, why attention is a matrix multiply, and why your embeddings need normalising. Aim for operational understanding.

The objects

Object Shape In ML it usually is
scalar () a loss value, a learning rate
vector (n,) one embedding, one sample’s features
matrix (m, n) a batch of embeddings, a weight layer
tensor (b, s, d) batch x sequence x hidden — the shape of everything in a transformer

Most ML bugs are shape bugs. Get in the habit of writing the expected shape next to every line.

import numpy as np

X = np.random.randn(32, 768)     # (batch, dim) - 32 embeddings
W = np.random.randn(768, 256)    # (dim_in, dim_out)
H = X @ W                        # (32, 256)

Dot product: the one to actually understand

a @ b == sum(a_i * b_i) == |a| * |b| * cos(theta)

It is a similarity score. Large positive means “pointing the same way”, zero means orthogonal, negative means opposing. Every retrieval system in this repo rests on that sentence.

def cosine_similarity(a: np.ndarray, b: np.ndarray) -> float:
    return float(a @ b / (np.linalg.norm(a) * np.linalg.norm(b)))

Why normalise embeddings. Raw dot product mixes direction (semantic similarity) with magnitude (which often just reflects token count or how confident the encoder was). Normalise to unit length and dot product becomes exactly cosine similarity:

X_norm = X / np.linalg.norm(X, axis=1, keepdims=True)   # keepdims! or you broadcast wrong
similarities = X_norm @ query_norm                      # (32,) - one score per row

That’s also why vector databases let you pick a metric: on normalised vectors, cosine, dot product and Euclidean distance rank results identically, so you choose based on what the index is optimised for. See ../09_rag_embeddings/04_vector_databases.md.

Matrix multiplication is the whole game

(m, k) @ (k, n) -> (m, n). The inner dimensions must match; the outer ones survive.

A neural network layer is a matrix multiply plus a bias plus a nonlinearity:

h = np.maximum(0, X @ W + b)     # ReLU layer, one line

Attention is three matrix multiplies:

scores = Q @ K.T / np.sqrt(d_k)  # (seq, seq) - every token vs every token
weights = softmax(scores)
output = weights @ V             # (seq, d_v)

The Q @ K.T term is why transformer cost is quadratic in sequence length — you’re computing a similarity between every pair of positions. That single fact explains context-window pricing, KV caching, and most of the research on efficient attention.

Norms

Norm Formula Used for
L1 sum(abs(x)) Lasso; drives weights to exactly zero -> feature selection
L2 sqrt(sum(x**2)) Ridge; shrinks weights smoothly; the default distance
L-inf max(abs(x)) adversarial robustness bounds

The L1-vs-L2 distinction is a very common interview question — see ../01_ml_foundations/04_overfitting_regularization.md.

Eigenvectors, and what they buy you

For a square matrix A, an eigenvector v satisfies A @ v = lambda * v — the transform stretches it without rotating it. The eigenvalue lambda is the stretch factor.

Where it shows up: PCA finds the eigenvectors of the covariance matrix. The top ones are the directions of greatest variance, which is exactly “the directions that carry the most information”. See ../02_classical_ml/09_dimensionality_reduction.md.

SVD

Any matrix factorises as A = U @ S @ V.T. S is diagonal with singular values in decreasing order. Truncate to the top k and you get the best rank-k approximation of A.

This is the machinery behind:

  • PCA (SVD on centred data),
  • latent semantic analysis (the pre-neural document embedding),
  • LoRA — fine-tuning by learning a low-rank update dW = A @ B instead of the full weight matrix, which is why you can fine-tune a large model on one GPU. See ../07_training_finetuning/.
U, S, Vt = np.linalg.svd(A, full_matrices=False)
A_rank_k = U[:, :k] @ np.diag(S[:k]) @ Vt[:k, :]

Broadcasting — where the bugs live

NumPy aligns shapes from the right, stretching dimensions of size 1.

X = np.random.randn(32, 768)
mu = X.mean(axis=0)              # (768,)   -> broadcasts across rows. Correct.
X_centered = X - mu

row_norms = np.linalg.norm(X, axis=1)              # (32,)  WRONG for division
X / row_norms                                       # error, or silently wrong
X / row_norms[:, None]                              # (32,1) -> correct

keepdims=True exists precisely to avoid this. The failure mode is nasty because it often doesn’t raise — it broadcasts into a wrong-but-valid shape and your model quietly trains on garbage.

Interview angle

  • “What does a dot product mean geometrically?” — projection of one vector onto another; equals |a||b|cos(theta). It’s a similarity score, which is why it underpins retrieval and attention.
  • “Why normalise embeddings before comparing them?” — raw dot product conflates direction with magnitude, and magnitude often encodes irrelevant things like length. Unit-normalising makes dot product exactly cosine similarity.
  • “Why is transformer attention quadratic in sequence length?”Q @ K.T computes a score for every pair of positions, so it’s O(n^2 * d) in time and O(n^2) in memory for the score matrix. This drives context limits, pricing, and the whole efficient-attention literature.
  • “What is PCA doing, mathematically?” — eigendecomposition of the covariance matrix (equivalently, SVD of the centred data). Keeps the directions of greatest variance and drops the rest.
  • “How does LoRA reduce fine-tuning cost?” — instead of updating a d x d weight matrix, it learns a low-rank factorisation A (d x r) @ B (r x d) with r << d, so you train orders of magnitude fewer parameters and can merge the delta back at inference time.
  • “You divided a matrix by a vector and the result looks wrong. Why?” — broadcasting aligned from the right. (32, 768) / (32,) doesn’t do what you meant; you need (32, 1) via [:, None] or keepdims=True.