← Back to the Applied Science guide
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.
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.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.
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.
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.
τ. The loss is a softmax over cosine scores. So cosine at serving time matches the training objective.[-1, 1]. That makes thresholds portable across queries and model versions.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.‖x - y‖² = 2 - 2cos, a decreasing function of cosine.‖A‖2 = σ1. Frobenius is the root sum of all σi².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).
col(A) ⊂ ℝm. Every output Ax can reach. Ax = b has a solution only if b is in here.null(A) ⊂ ℝn. All x with Ax = 0. These inputs are invisible to the map.col(AT) ⊂ ℝn. It is orthogonal to the null space. Together they fill ℝn.null(AT) ⊂ ℝm. It is orthogonal to the column space. Least-squares residuals live here.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.
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.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.
X rank deficient. Then XTX is singular and the weights are not unique.W ∈ ℝm×n with m < n has a null space of dimension at least n - m. Those input directions never reach the next layer.log |det J|. Architectures pick triangular Jacobians so the determinant is the product of the diagonal.500 × 20 feature matrix has rank 18. What does that tell you? A: Two features are exact linear mixes of others. The null space has dimension 2, so OLS weights are not unique. Drop or merge the redundant features, or add ridge.rank(ATA) = rank(A)? A: If ATAx = 0 then xTATAx = ‖Ax‖² = 0, so Ax = 0. The two null spaces match, and rank-nullity does the rest.uvT? A: One, if both are nonzero. Every column is a multiple of u.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.
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.
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 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.
Pii is the leverage of row i. High-leverage points can pull the fit alone. Cook's distance builds on this.2XT(Xw - y) to zero. Or require the residual to be orthogonal to col(X). Both give XTXw = XTy.XTX is positive semidefinite. Adding λI lifts every eigenvalue by λ > 0, so the matrix is invertible.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.
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.
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.
xTAx ≥ 0 for all x. Same as: all eigenvalues ≥ 0. Written A ≽ 0.xTAx > 0 for all x ≠ 0. Same as: all eigenvalues > 0. Same as: a Cholesky factor A = LLT exists.BTB is PSD because xTBTBx = ‖Bx‖² ≥ 0. Covariance matrices and kernel matrices are Gram 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.
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.
λmax / λmin of the Hessian. The stable step size is below 2 / λmax.W scales by λk. A spectral radius above 1 explodes. Below 1, the signal vanishes.εI to fix tiny negative eigenvalues from round-off.xTΣx is the variance of the projection onto x, and variance is never negative.|λ2/λ1|. A small gap means slow convergence.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).
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
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.
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² )
Q and solves a ridge problem for each user row, then swaps. Each step is a small closed-form solve, so it parallelizes well.min ½(‖P‖F² + ‖Q‖F²) over all factorizations of Z = PQT equals the nuclear norm ‖Z‖*. So weight decay on factors is a convex rank penalty in disguise.W with its truncated SVD UkΣkVkT. Parameters drop from mn to k(m + n).k factors in near-linear time. Scikit-learn uses it for TruncatedSVD.ATA = VΣ2VT. For a symmetric PSD matrix they coincide. For a symmetric matrix with negative eigenvalues, σi = |λi|.σi². No rank-10 matrix does better.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.
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
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.
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.
To solve Ax = b, you almost never want A-1 itself.
2n³ flops. An LU factor takes about (2/3)n³. Each extra right-hand side costs only O(n²) after that.inv(A) @ b has larger error than a backward-stable solve, especially when A is ill-conditioned.A often has a dense inverse. A sparse factor can stay sparse. The inverse can blow memory.xTA-1x is ‖L-1x‖² with Cholesky. log det A comes from the factor diagonal. Neither needs the inverse.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.”Hd = -g. Large models use conjugate gradient, which needs only Hessian-vector products.κ(A). The normal equations work with κ(A)², which can lose all digits.κ(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.
log ∑ ezi = m + log ∑ ezi - m with m = max z. Softmax uses the same shift. Without it, e1000 overflows.cross_entropy on logits, not log(softmax(z)). The fused form avoids log 0.λI lifts σmin and caps κ. It trades a little bias for a lot of stability.κ. Standardizing them is a cheap preconditioner.κ or VIF before you interpret any weight.log of zero, or an fp16 overflow.κ(X). Use ridge, drop or combine features, and solve with QR.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.
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.
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.
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
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.
h = σ(a), the Jacobian is diagonal. The VJP is just G ⊙ σ'(a), an elementwise product.Hv for about twice the cost of a gradient. Differentiate ∇f(x)Tv once more. You never form H.10-6.xTAx when A is not symmetric? A: (A + AT)x. Only the symmetric part of A affects the quadratic form.p - y. The softmax Jacobian diag(p) - ppT cancels the 1/p from the log.Y = XW with X of shape (B, n), what is ∂L/∂W? A: XTG, shape (n, m). It sums the outer products over the batch.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.
(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
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 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).
(B, 1, D) with a (1, N, D) makes a (B, N, D) tensor. For large N, use a matmul instead to keep memory flat.X (n, d) and Y (m, d) without loops. A: (X**2).sum(1)[:, None] + (Y**2).sum(1)[None, :] - 2 * X @ Y.T. Clip at zero for round-off."bhqd,bhkd->bhqk" compute? A: Attention logits. For each batch and head, it takes the dot product of every query with every key.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.
±1 and 0 work just as well. Very sparse versions keep only about 1/√d of the entries. They are much faster to apply.k buckets with a random sign. It is a sparse random projection. It preserves inner products in expectation.rTx for random r. Two vectors at angle θ agree on a bit with chance 1 - θ/π. Many bits give a fast cosine sketch.A by a thin random matrix to find its range. Sketch-and-solve least squares shrinks the rows first.d? A: No. It depends only on n and ε, as O(log n / ε²).1 - (π/3)/π = 2/3.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%)
W0 + (α/r)BA into one matrix. Latency is unchanged. Or keep adapters separate and hot-swap many tasks on one base.A and B. QLoRA goes further and stores W0 in 4-bit.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
dh matrix. Multi-head attention sums several such low-rank pieces.dh is too small for the context, a head cannot express some patterns. This is the low-rank bottleneck.V × d. For recommenders, V can be billions of user or item IDs. These tables dominate model memory.V × H with V × E times E × H, with E ≪ H. That is an explicit low-rank table.HET, with hidden states H of width d. So the log-probability matrix over contexts has rank at most about d + 1. It cannot match every true distribution. Mixture-of-softmaxes was proposed to lift this limit.B initialized to zero in LoRA? A: So ΔW = 0 at the start. The model begins exactly at the pretrained weights. A is random so gradients for B are nonzero.dh, the head dimension. It is the product of a T × dh and a dh × T matrix.4096 × 4096 layer with r = 16. How many trainable params? A: 16 × (4096 + 4096) = 131,072, about 0.78% of the full layer.y onto col(X). The residual is orthogonal to it: XT(y - Xw) = 0.k approximation. PCA is the SVD of centered data.κ = σmax/σmin sets the digits you lose. Normal equations square it.(A + AT)x, 2AT(Ax - b), X-T, and p - y. The gradient always has the shape of the variable.O(log n / ε²) dimensions, whatever d is.