Part IV Theory 5 Math

Linear Algebra and Matrix Calculus

Every model is a stack of matrix products. Know what those products do to space, and most of ML stops being magic.

Linear algebra is the language of ML. Embeddings are vectors. Layers are matrices. Training is calculus on matrices. PCA, least squares, recommenders and LoRA all reduce to a few core ideas. This page builds those ideas from the ground up. Each section gives the math, the picture, and the place it shows up at work.

Interviewers rarely ask you to prove a theorem. They ask you to reason with one. Why does cosine beat dot product for this search? Why did your solver blow up? What is the gradient of this loss, and what shape is it? A strong answer names the right object and its shape, then says what it means.

Notation. Bold-free lowercase letters like x are column vectors in ℝn. Uppercase letters like A are matrices. A is m × n unless stated. AT is the transpose. I is the identity. ⟨x, y⟩ and xTy both mean the dot product.

Contents

  1. Vectors, norms and cosine similarity
  2. Matrices as linear maps
  3. Projections and least squares
  4. Eigenvalues and the spectral theorem
  5. The singular value decomposition
  6. Decompositions in practice
  7. Condition number and stability
  8. Matrix calculus
  9. Tensors, broadcasting and einsum
  10. Random projections and JL
  11. Low-rank structure in deep learning

Foundations — spaces and maps

Vectors, norms and cosine similarity

Plain definition. A vector is a list of numbers that you can add and scale. A norm measures its length. A dot product measures how much two vectors point the same way.

Norms

A norm ‖x‖ is any function that is positive, scales with |c|, and obeys the triangle rule ‖x + y‖ ≤ ‖x‖ + ‖y‖. Three norms do almost all the work in ML.

‖x‖1 = ∑i |xi|                 (L1, Manhattan)
‖x‖2 = ( ∑i xi² )1/2          (L2, Euclidean)
‖x‖∞ = maxi |xi|               (L-infinity, max)
‖x‖p = ( ∑i |xi|p )1/p        (general p ≥ 1)

In finite dimensions all norms are equivalent. Each one bounds the others up to a constant. For example ‖x‖∞ ≤ ‖x‖2 ≤ ‖x‖1 ≤ √n ‖x‖2. The constants grow with n, so the choice still matters in high dimensions.

Matrices have norms too. The Frobenius norm ‖A‖F = (∑ij aij²)1/2 treats the matrix as one long vector. The spectral norm ‖A‖2 = max‖x‖=1 ‖Ax‖ is the most a matrix can stretch any vector. It equals the largest singular value σ1. The nuclear norm ‖A‖* = ∑ σi is the L1 norm of the singular values. It is the convex stand-in for rank.

Dot product and the angle

xTy = ∑i xi yi = ‖x‖2 ‖y‖2 cos θ

Cauchy-Schwarz:   |xTy| ≤ ‖x‖2 ‖y‖2     (equality iff x, y are parallel)

cos_sim(x, y) = xTy / (‖x‖2 ‖y‖2)  ∈  [-1, 1]

Cauchy-Schwarz is why cosine sits in [-1, 1]. Cosine drops the lengths and keeps only the direction. Two vectors at 90 degrees have a dot product of zero. We call them orthogonal.

For unit vectors, cosine and Euclidean distance carry the same order. Expand the square.

‖x - y‖2² = ‖x‖² + ‖y‖² - 2 xTy = 2 - 2 cos(x, y)      when ‖x‖ = ‖y‖ = 1

So after you L2-normalize, nearest neighbor by L2 and top-k by cosine return the same items. This is why vector databases let you pick either metric.

θ x y (xᵀy / ‖x‖) along x̂ orthogonal part
The dot product is the length of y's shadow on x, times ‖x‖. Cosine keeps only the angle θ.

Why embeddings use cosine

Learned embeddings mix two signals. The direction holds meaning. The length often holds something else, like frequency or popularity. Word2vec gives frequent words longer vectors. A two-tower recommender can give popular items larger norms. Raw dot product rewards length, so it favors those popular items. Cosine strips length out and compares meaning alone.

High-dimensional trap. Random vectors in high dimension are nearly orthogonal. For iid Gaussian vectors, cosine concentrates near 0 with spread about 1/√d. A cosine of 0.3 in 768 dimensions is a strong signal, not a weak one. Embeddings also suffer from anisotropy. Many models put all vectors in a narrow cone, so every pair scores 0.6 or more. Mean-centering or whitening often fixes this.

Why it matters in practice

Interview check

Matrices as linear maps

Plain definition. A matrix is a function that maps vectors to vectors. It keeps straight lines straight and keeps the origin fixed. Column j of A is where the j-th basis vector lands.

A map f is linear when f(ax + by) = a f(x) + b f(y). Every linear map from ℝn to ℝm is a matrix. Read Ax in two ways.

Row view:     (Ax)i = aiT x            each output is a dot product with a row
Column view:  Ax = x1 a:1 + x2 a:2 + ... + xn a:n   output is a mix of the columns

The column view is the one to keep. The output of A is always some mix of its columns. Matrix product AB is composition. Apply B first, then A. That is why shapes must chain: (m × k)(k × n) = (m × n).

The four fundamental subspaces

Rank

The rank of A is the dimension of its column space. It also equals the dimension of the row space. That fact is not obvious, and the SVD proves it cleanly. The rank-nullity theorem ties rank to what gets lost.

rank(A) + dim null(A) = n           (number of columns)

rank(AB) ≤ min(rank(A), rank(B))
rank(A + B) ≤ rank(A) + rank(B)
rank(ATA) = rank(A)

A square n × n matrix is invertible when its rank is n. The same condition has many faces. The null space is just {0}. The determinant is nonzero. No eigenvalue is zero. The columns are linearly independent.

Rank tells you how many independent directions survive the map. A 1000 × 1000 matrix of rank 10 squeezes everything onto a 10-dimensional slice. You can store it as UVT with U, V of shape 1000 × 10. That costs 20,000 numbers instead of 1,000,000.

Key idea. A matrix of shape m × n and rank r only has r directions of real action. Everything in the null space is thrown away. Everything outside the column space is out of reach.

Determinant and trace

The determinant det(A) is the signed volume scale of the map. A unit cube maps to a shape with volume |det A|. A zero determinant means some direction got crushed flat. The trace tr(A) = ∑ aii is the sum of eigenvalues. The determinant is their product. The trace is cyclic: tr(ABC) = tr(CAB) = tr(BCA). That rule powers most matrix calculus tricks.

Why it matters in practice

Interview check

Projections and least squares

Plain definition. Projection drops a vector onto a subspace at a right angle. It finds the closest point in that subspace. Least squares is projection of the target onto the column space of the features.

The projection matrix

Project b onto the line through a. The closest point is p = (aTb / aTa) a. The error b - p is orthogonal to a. Now project onto the column space of a full-rank A.

P = A (ATA)-1 AT

P² = P        (projecting twice changes nothing: idempotent)
PT = P       (symmetric, so the projection is orthogonal)
I - P          projects onto the orthogonal complement, null(AT)
eigenvalues of P are 0 or 1, and tr(P) = rank(A)

If the columns of Q are orthonormal, the formula collapses to P = QQT. That is the main reason orthonormal bases are so useful.

Least squares and the normal equations

We want w that minimizes ‖Xw - y‖2². Here X is n × d with n > d. Usually y is not in col(X), so no exact fit exists. The best Xw is the projection of y onto col(X).

Derivation by calculus. Expand and take the gradient.

L(w) = (Xw - y)T(Xw - y) = wTXTXw - 2 yTXw + yTy
∇w L = 2 XTXw - 2 XTy = 0

Normal equations:   XTX ŵ = XTy
Solution:           ŵ = (XTX)-1XTy = X+y
Fitted values:      ŷ = Xŵ = P y        (P is the "hat matrix")

Derivation by geometry. The residual r = y - Xŵ must be orthogonal to every column of X. If it were not, you could move along that column and shrink it. So XTr = 0, which is XT(y - Xŵ) = 0. That is the normal equations again. The name comes from this: the residual is normal to the column space.

col(X) 0 y ŷ = Xŵ = Py r = y − ŷ Xᵀr = 0
Least squares is a right-angle drop. The fitted vector is the shadow of y in col(X). The residual stands straight up from the plane.

How to actually solve it

Do not form (XTX)-1. Forming XTX squares the condition number. More on that below. Use QR or the SVD instead. With X = QR, the normal equations become Rw = QTy. That is one triangular solve.

import numpy as np

rng = np.random.default_rng(0)
X = rng.normal(size=(100, 3))
w_true = np.array([1.0, -2.0, 0.5])
y = X @ w_true + 0.1 * rng.normal(size=100)

w, *_ = np.linalg.lstsq(X, y, rcond=None)   # SVD-based solver
r = y - X @ w                                # residual
print(np.round(w, 3))                        # close to w_true
print(np.abs(X.T @ r).max() < 1e-9)          # residual is orthogonal to col(X)
# [ 0.999 -2.013  0.484]
# True

Ridge and the pseudoinverse

Ridge adds λ‖w‖² to the loss. The normal equations become (XTX + λI)w = XTy. The matrix is now positive definite for any λ > 0, so a unique answer always exists. In SVD terms, ridge shrinks each direction by σi² / (σi² + λ). Weak directions get shrunk the most.

When X is rank deficient, many w fit equally well. The Moore-Penrose pseudoinverse X+ = VΣ+UT picks the one with the smallest norm. Gradient descent from w = 0 lands on that same minimum-norm answer. This is a simple case of implicit regularization.

Why it matters in practice

Interview check

Spectral structure — eigenvalues and the SVD

Eigenvalues and the spectral theorem

Plain definition. An eigenvector is a direction that a matrix only stretches. It does not turn. The eigenvalue is the stretch factor.
A v = λ v,     v ≠ 0

Characteristic polynomial:  det(A - λI) = 0      (n roots, counted with multiplicity)
tr(A)  = ∑ λi
det(A) = ∏ λi
Ak v = λk v            (powers act simply on eigenvectors)

If A has n independent eigenvectors, it is diagonalizable: A = VΛV-1. In the eigenvector basis, the map is just scaling along each axis. Not every matrix qualifies. The shear [[1, 1], [0, 1]] has only one eigenvector direction. A general real matrix can also have complex eigenvalues. A rotation is the classic case.

The spectral theorem

Symmetric matrices are the well-behaved case. Covariance matrices, Hessians, kernel matrices and graph Laplacians are all symmetric.

If A = AT (real), then
    A = Q Λ QT = ∑i λi qi qiT

    every λi is real
    Q is orthogonal: QTQ = I
    eigenvectors for distinct eigenvalues are orthogonal

Read the sum form slowly. A symmetric matrix is a weighted sum of rank-one projections onto orthogonal axes. It stretches along those axes by λi. It never shears. This is the cleanest picture of a matrix you will get.

Why distinct eigenvectors are orthogonal. Take Au = λu and Av = μv with λ ≠ μ. Then λ uTv = (Au)Tv = uTAv = μ uTv. The middle step uses A = AT. Since λ ≠ μ, we get uTv = 0.

The Rayleigh quotient

R(x) = xTAx / xTx

λmin ≤ R(x) ≤ λmax,   with equality at the matching eigenvectors
max‖x‖=1 xTAx = λmax

This turns eigenvalues into an optimization problem. PCA is exactly this: find the unit direction that maximizes the variance xTΣx. The answer is the top eigenvector of the covariance.

Positive definite matrices

PD matters because of the second-order test. At a point where the gradient is zero, a PD Hessian means a strict local min. An indefinite Hessian, with eigenvalues of both signs, means a saddle. Deep nets have far more saddles than bad local minima.

Power iteration

Start with a random x. Repeat x ← Ax / ‖Ax‖. The part along the top eigenvector grows like λ1k. Every other part grows slower. So x converges to v1 at a rate set by |λ2 / λ1|. PageRank is power iteration on a Markov matrix. Spectral normalization in GANs uses one power step per update to estimate σ1.

Why it matters in practice

Interview check

The singular value decomposition

Plain definition. The SVD splits any matrix into three simple steps. First a rotation, then a stretch along the axes, then another rotation. It works for every matrix, square or not.
A = U Σ VT = ∑i=1r σi ui viT

A is m × n,  U is m × m orthogonal,  V is n × n orthogonal
Σ is m × n diagonal with σ1 ≥ σ2 ≥ ... ≥ σr > 0,   r = rank(A)

A vi = σi ui          AT ui = σi vi
ATA = V (ΣTΣ) VT     AAT = U (ΣΣT) UT

So the right singular vectors are eigenvectors of ATA. The left singular vectors are eigenvectors of AAT. The singular values are the square roots of their shared nonzero eigenvalues. The SVD also hands you all four subspaces. The first r columns of U span col(A). The last n - r columns of V span null(A).

unit circle, v₁ v₂ Vᵀ rotate to axes Σ stretch by σ₁, σ₂ U σ₁u₁, σ₂u₂
Every matrix maps the unit sphere to an ellipsoid. The axes of the ellipsoid are σiui. The inputs that land on them are vi.

Low-rank approximation: Eckart-Young

Keep only the top k terms of the sum. Call the result Ak = ∑i≤k σiuiviT. The Eckart-Young-Mirsky theorem says no rank-k matrix does better.

minrank(B) ≤ k ‖A - B‖F = ‖A - Ak‖F = ( σk+1² + ... + σr² )1/2
minrank(B) ≤ k ‖A - B‖2 = ‖A - Ak‖2 = σk+1

The error is exactly the energy in the singular values you dropped. So the spectrum tells you how compressible a matrix is. A fast decay means a small k captures most of it.

import numpy as np

rng = np.random.default_rng(1)
# a rank-5 signal plus small noise
A = rng.normal(size=(50, 5)) @ rng.normal(size=(5, 40))
A += 0.01 * rng.normal(size=(50, 40))

U, s, Vt = np.linalg.svd(A, full_matrices=False)
print(np.round(s[:7], 2))                 # sharp drop after 5 values

k = 5
A_k = (U[:, :k] * s[:k]) @ Vt[:k]         # best rank-k approximation
err = np.linalg.norm(A - A_k)             # Frobenius norm
tail = np.sqrt(np.sum(s[k:] ** 2))
print(round(err, 4), round(tail, 4))      # Eckart-Young: they match
# [53.54 43.29 33.89 30.78 20.09  0.12  0.11]
# 0.4045 0.4045

The link to PCA

Center the data matrix X (n × d) so each column has mean zero. The sample covariance is C = XTX / (n - 1). Take the SVD X = UΣVT.

C = V (Σ² / (n - 1)) VT

principal directions  = columns of V             (right singular vectors)
variance along PC i   = σi² / (n - 1)
PC scores             = XV = UΣ
explained variance    = σi² / ∑j σj²

So PCA is an SVD of the centered data. Use the SVD, not an eigen-solve of C. Forming XTX squares the condition number and loses small components to round-off. PCA also inherits Eckart-Young. Projecting onto the top k PCs gives the best rank-k reconstruction in squared error. Forgetting to center is the most common bug. Without centering, the first component just points at the mean.

Matrix factorization for recommenders

Put ratings in a matrix R of users by items. The low-rank assumption says taste has a few hidden factors. So R ≈ PQT with P ∈ ℝm×k and Q ∈ ℝn×k. The predicted rating is the dot product puTqi.

minP,Q  ∑(u,i) observed ( rui - puTqi )²  +  λ ( ‖P‖F² + ‖Q‖F² )

Why it matters in practice

Interview check

Computation — solving systems that actually work

Decompositions in practice

Plain definition. A decomposition rewrites a matrix as a product of simple pieces. Triangular and orthogonal pieces are cheap and safe to solve with. You factor once, then solve many times.

LU

PA = LU      P permutation, L unit lower triangular, U upper triangular
Cost: ~(2/3) n³ flops to factor, ~2n² per solve

LU is Gaussian elimination written as a product. The permutation does partial pivoting. It swaps rows so you never divide by a tiny number. This is the default for a general square system. np.linalg.solve calls LAPACK's LU routine.

Cholesky

A = L LT     for symmetric positive definite A,   L lower triangular, positive diagonal
Cost: ~(1/3) n³ flops, half of LU, and no pivoting needed
log det A = 2 ∑i log Lii

Cholesky is the fastest and most stable choice when it applies. It also doubles as a PD test. If it fails, the matrix is not PD. Gaussian processes use it for the kernel solve and the log-determinant. Sampling from N(μ, Σ) is μ + Lz with z ~ N(0, I).

import numpy as np
from numpy.linalg import cholesky, solve

rng = np.random.default_rng(6)
M = rng.normal(size=(4, 4))
S = M @ M.T + 4 * np.eye(4)        # symmetric positive definite
b = rng.normal(size=4)

L = cholesky(S)                    # S = L L^T
z = solve(L, b)                    # L z = b   (use scipy solve_triangular in real code)
x = solve(L.T, z)                  # L^T x = z
print(np.allclose(S @ x, b))       # True
logdet = 2 * np.log(np.diag(L)).sum()
print(np.isclose(logdet, np.linalg.slogdet(S)[1]))   # True

QR

A = Q R      Q has orthonormal columns, R upper triangular
Least squares:  min ‖Ax - b‖  ⇒  R x = QTb
Cost: ~2mn² flops (Householder)

QR is the standard tool for least squares. Orthogonal Q preserves lengths, so it does not amplify error. You solve with R, whose condition number equals that of A. The normal equations would square it. Gram-Schmidt builds Q one column at a time but loses orthogonality in floating point. Householder reflections are the stable way.

Eigen and SVD solvers

Use eigh for symmetric matrices, not eig. It is faster, returns real values in sorted order, and gives orthonormal vectors. The SVD costs more, about O(mn²), but it is the most robust tool. It handles rank deficiency and gives the pseudoinverse.

Why you never invert a matrix

To solve Ax = b, you almost never want A-1 itself.

Say it like this. “I would not call inv. The matrix is symmetric positive definite, so I would Cholesky-factor it once. Then each solve is two triangular solves. The log-determinant comes from the diagonal for free.”

Why it matters in practice

Interview check

Condition number and numerical stability

Plain definition. The condition number says how much a problem amplifies small errors. A small number means a stable problem. A big number means tiny input noise can wreck the answer.
κ(A) = ‖A‖2 ‖A-1‖2 = σmax / σmin          (≥ 1)

Perturbation bound for Ax = b:
    ‖δx‖ / ‖x‖  ≤  κ(A) · ‖δb‖ / ‖b‖

κ(ATA) = κ(A)²
κ(Q) = 1 for any orthogonal Q

A rule of thumb: you lose about log10 κ digits of accuracy. Float64 carries about 16 digits. Float32 carries about 7. So κ = 108 leaves 8 good digits in float64 and none in float32. Condition is a property of the problem. Stability is a property of the algorithm. A backward-stable algorithm gives the exact answer to a nearby problem. With a bad κ, even that can be far from the truth.

import numpy as np

rng = np.random.default_rng(2)
t = np.linspace(0, 1, 50)
# two nearly identical columns: a near-collinear design
X = np.column_stack([np.ones(50), t, t + 1e-6 * rng.normal(size=50)])

print(f"{np.linalg.cond(X):.1e}")          # cond(X)
print(f"{np.linalg.cond(X.T @ X):.1e}")    # about cond(X) squared
# 1.8e+06
# 3.3e+12

With κ(X) ≈ 106, QR keeps about 10 digits. The normal equations face 1012 and keep about 4. In float32 they would keep none.

Stability tricks you use every day

Why it matters in practice

Interview check

Calculus and notation — the language of backprop

Matrix calculus

Plain definition. Matrix calculus finds derivatives when the inputs or outputs are vectors or matrices. The gradient of a scalar loss has the same shape as the thing you differentiate by. That one rule catches most mistakes.

Layout and the shape rule

This page uses the gradient layout. For a scalar f and a matrix X, ∇Xf has the same shape as X. Entry (i, j) is ∂f / ∂Xij. This matches PyTorch, where W.grad has the shape of W. For a vector function f: ℝn → ℝm, the Jacobian J is m × n, with Jij = ∂fi/∂xj.

The differential trick

Do not grind through indices. Write the differential df, then massage it into the form df = tr(GT dX). Then ∇Xf = G. For vectors, df = gTdx gives ∇f = g. Two facts do most of the work. A scalar equals its own trace. The trace is cyclic.

The gradients to know cold

f(x)                       ∇x f                        Hessian
----------------------------------------------------------------------------
aTx                        a                            0
xTx = ‖x‖²                 2x                           2I
xTAx                       (A + AT) x                  A + AT
xTAx,  A symmetric         2Ax                          2A
‖Ax - b‖²                  2AT(Ax - b)                 2ATA
‖x‖2                       x / ‖x‖2                    (I - x̂x̂T) / ‖x‖2

f(X)                       ∇X f
----------------------------------------------------------------------------
tr(AX)                     AT
tr(XTAX)                   (A + AT) X
aTXb                       a bT
‖AX - B‖F²               2AT(AX - B)
log det X                  X-T    (= X-1 for symmetric PD X)
tr(X-1A)                   -X-TATX-T

Derive xTAx. Perturb x by dx. Then df = dxTAx + xTA dx. The first term is a scalar, so it equals its transpose xTATdx. So df = xT(A + AT) dx, and the gradient is (A + AT)x.

Derive ‖Ax - b‖². Let r = Ax - b. Then df = 2rTdr = 2rTA dx. So the gradient is 2ATr. Set it to zero and you get the normal equations again.

Derive log det X. Use det(I + E) ≈ 1 + tr(E) for small E.

det(X + dX) = det(X) det(I + X-1dX) ≈ det(X) (1 + tr(X-1dX))
d log det X = tr(X-1 dX) = tr((X-T)T dX)
⇒  ∇X log det X = X-T

This gradient appears in the Gaussian log-likelihood. With precision Λ = Σ-1, the MLE sets Λ-1 - S = 0. So the MLE of the covariance is the sample covariance S. Log-det is concave on PD matrices, so this is a clean convex problem. The graphical lasso adds an L1 penalty to the same objective.

Softmax Jacobian and cross-entropy

p = softmax(z),   pi = ezi / ∑k ezk

∂pi/∂zj = pi(δij - pj)       ⇒     J = diag(p) - p pT

Cross-entropy with one-hot y:  L = -∑i yi log pi
∇z L = JT(-y / p) = p - y

The Jacobian is symmetric and PSD. Its rows sum to zero, because the probabilities always sum to one. So J has 1 in its null space. Shifting all logits by a constant changes nothing. The gradient p - y is clean and bounded. That is why softmax and cross-entropy are always fused. The same p - y form shows up for logistic regression and for every GLM with its canonical link.

import numpy as np

def softmax(z):
    e = np.exp(z - z.max())          # shift for stability
    return e / e.sum()

rng = np.random.default_rng(3)
z = rng.normal(size=4)
p = softmax(z)
J = np.diag(p) - np.outer(p, p)      # analytic Jacobian

h = 1e-6
I = np.eye(4)
J_num = np.column_stack([
    (softmax(z + h * I[i]) - softmax(z - h * I[i])) / (2 * h)
    for i in range(4)
])
print(np.allclose(J, J_num, atol=1e-8))   # True
print(np.round(J.sum(axis=0), 12))        # columns sum to 0

The chain rule with shapes

Backprop is the chain rule applied right to left. It never builds a full Jacobian. Each layer takes the upstream gradient and returns a vector-Jacobian product (VJP). Take a batched linear layer.

Y = X W + b        X: (B, n)    W: (n, m)    b: (m,)    Y: (B, m)
upstream:  G = ∂L/∂Y              G: (B, m)

∂L/∂W = XT G               (n, B)(B, m) = (n, m)   matches W
∂L/∂X = G WT               (B, m)(m, n) = (B, n)   matches X
∂L/∂b = G.sum(axis=0)        (m,)                   matches b

You can rebuild these from shapes alone. The gradient for W must be (n, m). The only way to get that from X and G is XTG. The bias sums over the batch because it was broadcast across the batch. Broadcasting in the forward pass becomes summing in the backward pass.

Why it matters in practice

Interview check

Tensors, broadcasting and einsum

Plain definition. In ML code, a tensor is an array with any number of axes. Broadcasting stretches small arrays to match big ones without copying. Einsum names each axis with a letter and says which ones to sum.

Shapes as documentation

Real models carry many axes. A transformer activation is (B, T, D): batch, tokens, features. Attention heads split it into (B, H, T, Dh). Images are (B, C, H, W). The habit that saves hours is to write every shape in a comment. Then check it with an assert.

Broadcasting rules

  1. Line up the shapes from the right.
  2. Two axes match if they are equal or one of them is 1.
  3. A missing axis on the left counts as 1.
  4. Size-1 axes stretch to match. No memory is copied.
(B, T, D) + (D,)       → (B, T, D)       add a bias per feature
(B, T, D) * (B, T, 1)  → (B, T, D)       scale each token by a mask or gate
(n, 1)    - (1, m)     → (n, m)          all pairwise differences
(n,)      - (n, 1)     → (n, n)          the classic silent bug
The silent bug. y_pred of shape (n, 1) minus y of shape (n,) gives an (n, n) matrix. Then .mean() returns a number that looks fine. Training runs, but the loss is wrong. Assert shapes before any reduction.

Einsum as the working notation

Einsum writes a tensor product the way you write the math on paper. Each input gets a string of axis letters. Letters in the output are kept. Letters that vanish are summed over.

np.einsum("ij,jk->ik", A, B)          matrix product            Cik = ∑j AijBjk
np.einsum("bij,bjk->bik", A, B)       batched matmul
np.einsum("i,i->", x, y)              dot product
np.einsum("i,j->ij", x, y)            outer product
np.einsum("ii->", A)                  trace
np.einsum("ij->ji", A)                transpose
np.einsum("bi,ij,bj->b", X, A, X)     xTAx for every row in a batch
np.einsum("btd,de->bte", X, W)       apply a linear layer to every token

Here is multi-head attention scores and outputs in two lines of einsum.

import numpy as np

rng = np.random.default_rng(5)
B, H, T, D = 2, 4, 5, 8
q = rng.normal(size=(B, H, T, D))
k = rng.normal(size=(B, H, T, D))
v = rng.normal(size=(B, H, T, D))

scores = np.einsum("bhid,bhjd->bhij", q, k) / np.sqrt(D)
w = np.exp(scores - scores.max(-1, keepdims=True))
w /= w.sum(-1, keepdims=True)                 # softmax over j
out = np.einsum("bhij,bhjd->bhid", w, v)

print(scores.shape, out.shape)
print(np.allclose(scores, q @ k.swapaxes(-1, -2) / np.sqrt(D)))
# (2, 4, 5, 5) (2, 4, 5, 8)
# True

Contraction order matters for cost. (AB)x costs O(n³). A(Bx) costs O(n²). Pass optimize=True to np.einsum and it searches for a cheap order. Libraries like einops add readable reshapes, such as rearrange(x, "b t (h d) -> b h t d", h=H).

Why it matters in practice

Interview check

Structure in high dimensions — randomness and low rank

Random projections and Johnson-Lindenstrauss

Plain definition. You can multiply high-dimensional points by a random matrix to get far fewer dimensions. All pairwise distances stay nearly the same. The target size depends on the number of points, not the original dimension.

The lemma

For any 0 < ε < 1 and any n points in ℝd, if
    k ≥ 4 ln n / (ε²/2 - ε³/3)          (so k = O(ε-2 log n))
then a map f: ℝd → ℝk exists with, for every pair u, v,
    (1 - ε) ‖u - v‖²  ≤  ‖f(u) - f(v)‖²  ≤  (1 + ε) ‖u - v‖²

A random Gaussian map f(x) = R x / √k, with Rij ~ N(0, 1), works with high probability.

Why it works. Fix one vector u. Each coordinate of Ru is N(0, ‖u‖²). So ‖Ru‖² / ‖u‖² is a chi-squared with k degrees of freedom. Divided by k, it has mean 1. It concentrates fast. The chance it leaves [1 - ε, 1 + ε] falls like exp(-c k ε²). There are about n²/2 pairs. A union bound over them needs k of order log n / ε². The dimension d never appears.

import numpy as np

rng = np.random.default_rng(4)
n, d, k = 200, 10_000, 1_000
P = rng.normal(size=(n, d))                  # 200 points in 10,000 dims
R = rng.normal(size=(d, k)) / np.sqrt(k)     # random Gaussian map
Q = P @ R                                    # same points in 1,000 dims

def pdist(Z):
    sq = (Z ** 2).sum(axis=1)
    D2 = sq[:, None] + sq[None, :] - 2 * Z @ Z.T
    i, j = np.triu_indices(len(Z), 1)
    return np.sqrt(np.maximum(D2[i, j], 0))

ratio = pdist(Q) / pdist(P)
print(round(ratio.min(), 3), round(ratio.max(), 3))   # all near 1
# 0.921 1.086

All 19,900 distances stay within about 9%. The worst-case bound asks for over 4,000 dimensions at ε = 0.1. In practice the constants are loose, and far fewer dimensions do the job.

Variants

Why it matters in practice

Interview check

Low-rank structure in deep learning

Plain definition. Many big matrices in deep learning act as if they had far fewer directions than their size. Writing them as a product of two thin matrices saves memory and compute. It often loses little accuracy.

LoRA: low-rank fine-tuning

W' = W0 + ΔW = W0 + (α / r) B A

W0: d × k (frozen)     B: d × r (init 0)     A: r × k (random init)     r ≪ min(d, k)
trainable params: r (d + k)   instead of   d k

Example: d = k = 4096, r = 8  ⇒  65,536 vs 16,777,216 params   (about 0.39%)
x W₀ (d × k) frozen A (r × k) B (d × r) + h trainable, rank r
LoRA adds a thin trainable path BA beside the frozen weight. Since B starts at zero, training starts from the base model exactly.

Attention is low-rank by construction

scores = (X WQ)(X WK)T / √dh = X (WQ WKT) XT / √dh

WQ, WK: d × dh     ⇒     WQWKT is d × d with rank ≤ dh
the T × T logit matrix has rank ≤ dh          (before softmax)
WV WO per head is also d × d with rank ≤ dh

Embedding tables

Key idea. Low rank appears in three roles. It is a cost saver (LoRA, factorized tables). It is a hidden limit (softmax bottleneck, head size). It is an inductive bias (matrix factorization, PCA). Name which role you mean.

Why it matters in practice

Interview check

Recap

← T4 — Optimization T6 — The Practitioner's Toolkit →