Linear Algebra in Machine Learning

Feature Matrices and Data Representation

Before any learning algorithm can operate, raw data must be encoded as numbers arranged in a matrix. This matrix — called the feature matrix or design matrix — is the fundamental linear algebra object in supervised learning. Given a dataset of n samples, each described by p features, the design matrix X is an n × p matrix where row i is the feature vector of sample i and column j is the j-th feature across all samples.

Every column of X lives in ℝⁿ and can be thought of as a direction in n-dimensional sample space; every row lives in ℝᵖ and is a point in p-dimensional feature space. The rank of X — the dimension of its column space — determines how much independent information the features carry. If rank(X) = r < p, then only r linearly independent directions exist in feature space; the remaining p − r features are redundant combinations of the others.

The choice of features fundamentally shapes what the model can learn. Linear models can only learn boundaries that are linear in the feature space. By manually constructing nonlinear features (polynomial terms x₁², x₁x₂, ...) or using a kernel function to implicitly map to a high-dimensional feature space, we expand expressivity. This is the famous kernel trick: rather than computing the n × ∞ feature matrix explicitly, we compute the n × n kernel matrix K where Kᵢⱼ = φ(xᵢ)ᵀφ(xⱼ) using a kernel function K(xᵢ, xⱼ) = φ(xᵢ)ᵀφ(xⱼ) that implicitly encodes an inner product in feature space.

Standardizing features — subtracting the mean and dividing by the standard deviation column-wise — is a linear algebra operation that centers and scales the design matrix. It dramatically improves the conditioning of XᵀX (brings its eigenvalues closer together) and is essential for gradient-based optimization. Without standardization, gradient descent oscillates wildly along the directions of high-variance features while barely moving along low-variance directions.

Dimensionality reduction methods like PCA (Principal Component Analysis) operate directly on the feature matrix. PCA computes the SVD X = UΣVᵀ, then projects onto the top r right singular vectors (columns of V): X_reduced = XV_r where V_r is the n × r matrix of the top r right singular vectors. The reduced matrix XV_r lives in an r-dimensional subspace that captures the maximum variance of the original data. This is a pure linear algebra transformation — no learning, just geometry.

Linear Regression: Closed-Form Solution

Linear regression is the simplest supervised learning problem: given design matrix X ∈ ℝⁿˣᵖ and target vector y ∈ ℝⁿ, find weights w ∈ ℝᵖ that minimize the residual sum of squares ‖Xw − y‖². This is a quadratic objective with a unique global minimum (when X has full column rank), and linear algebra gives us its exact closed-form solution.

Setting the gradient to zero: ∇_w ‖Xw − y‖² = 2Xᵀ(Xw − y) = 0, which gives the normal equations: XᵀXw = Xᵀy. When X has full column rank (all features are linearly independent), XᵀX is positive definite and the normal equations have a unique solution:

The condition number of XᵀX equals the square of the condition number of X. If X is ill-conditioned (some singular values near zero), XᵀX becomes numerically singular and direct inversion is unstable. The practical remedy is to use the QR decomposition of X (from Module 7): writing X = QR, we have XᵀX = RᵀQ ᵀQR = RᵀR, and the normal equations become Rᵀ(Rw) = Rᵀ(Qᵀy), which reduces to the upper triangular system Rw = Qᵀy — solved stably by back-substitution in O(p²) operations.

When X does not have full column rank (multicollinearity), the normal equations are singular and no unique w* exists. The minimum-norm solution is given by the pseudoinverse: w* = X⁺y = VΣ⁺Uᵀy, where Σ⁺ replaces each nonzero singular value σᵢ with 1/σᵢ and leaves zero singular values as zero. This solution lies in the row space of X and has the smallest ‖w‖ among all minimizers.

Ridge Regression and Regularization

When the system is ill-conditioned or underdetermined, adding a regularization term λ‖w‖² penalizes large weights, producing the ridge regression (Tikhonov regularization) objective: minimize ‖Xw − y‖² + λ‖w‖². The closed-form solution is w* = (XᵀX + λI)⁻¹Xᵀy. The regularization term λI adds λ to every eigenvalue of XᵀX, lifting the smallest eigenvalues away from zero and dramatically improving the condition number. This trades a small increase in bias for a large reduction in variance — the classic bias–variance tradeoff, expressed entirely in terms of eigenvalues.

Logistic Regression: Optimization Perspective

For binary classification, logistic regression models P(y = 1 | x) = σ(wᵀx) where σ(z) = 1/(1 + e⁻ᶻ) is the sigmoid function. Training maximizes the log-likelihood, which is equivalent to minimizing the binary cross-entropy loss:

The gradient of the logistic loss is ∇L(w) = −Xᵀ(y − σ(Xw)) = Xᵀ(σ(Xw) − y), a remarkably clean matrix-vector expression. The Hessian is ∇²L(w) = XᵀDX where D = diag(σᵢ(1−σᵢ)) ∈ ℝⁿˣⁿ is a diagonal matrix of per-sample curvature weights. This Hessian structure enables efficient Newton's method via iteratively reweighted least squares (IRLS): each Newton step solves a weighted least-squares problem (XᵀDX)δ = −∇L, which has exactly the form of a regularized linear regression with sample weights.

IRLS iterates: (1) compute σ̂ = σ(Xw_k); (2) form diagonal weight matrix D = diag(σ̂ ⊙ (1 − σ̂)); (3) solve the weighted normal equations (XᵀDX)δ = Xᵀ(y − σ̂); (4) update w_{k+1} = w_k + δ. Each iteration is a linear algebra problem — solving a p × p positive definite system. The Hessian changes between iterations (D depends on w_k), but remains positive definite throughout, so Cholesky works reliably. IRLS converges quadratically and typically needs only 5–10 iterations regardless of problem size.

For large n (millions of samples), forming XᵀDX costs O(np²) and is prohibitive. Instead, stochastic gradient descent with the gradient ∇L(w) = Xᵀ(σ(Xw) − y) is used, evaluating the gradient on a random mini-batch of samples. Each gradient evaluation costs O(batch × p) — far cheaper than the full Hessian. Modern implementations in PyTorch/JAX compute this with automatic differentiation on the vectorized expression; the underlying operation is always a matrix-vector product.

Neural Network Weight Matrices

A fully connected (dense) neural network is literally a sequence of matrix multiplications interspersed with elementwise nonlinearities. For a network with L layers, the forward pass computes:

In batch processing, the input is a matrix A^(0) ∈ ℝⁿˣᵈ₀ (n samples, d₀ input features) rather than a single vector. The layer computation becomes Z^(ℓ) = A^(ℓ-1) W_ℓᵀ + 1_n bᵀ_ℓ — a matrix-matrix product of shape (n × d_{ℓ-1}) × (d_{ℓ-1} × d_ℓ) = n × d_ℓ. On a GPU, this GEMM (General Matrix Multiply) operation is the dominant computational primitive, executed in highly optimized CUDA kernels. The entire forward pass of a large language model is, at its core, a sequence of such matrix multiplications.

Backpropagation as Matrix Calculus

Backpropagation — the algorithm that computes gradients in neural networks — is the chain rule of matrix calculus applied layer by layer. Given the loss L and the upstream gradient δ^(ℓ) = ∂L/∂Z^(ℓ) (a matrix of the same shape as Z^(ℓ)), the local gradients are:

• Weight gradient: ∂L/∂W_ℓ = (δ^(ℓ))ᵀ A^(ℓ-1) — an outer product summed over the batch
• Bias gradient: ∂L/∂b_ℓ = (δ^(ℓ))ᵀ 1_n — the sum of δ^(ℓ) over the batch
• Downstream gradient: δ^(ℓ-1) = δ^(ℓ) W_ℓ ⊙ σ'(Z^(ℓ-1)) — matrix product followed by elementwise multiply

Every one of these operations is a matrix multiplication or elementwise product. The gradient of the loss with respect to W_ℓ is a matrix of the same shape as W_ℓ; stacking it with the parameter update w_{k+1} = w_k − η ∇_w L gives gradient descent in the flattened parameter space. Modern frameworks (PyTorch, JAX) automate this via automatic differentiation (autograd), which records the sequence of matrix operations in the forward pass and automatically differentiates them in reverse order.

Weight Matrix Geometry and Initialization

The geometry of the weight matrices matters enormously for training stability. Consider the singular value decomposition of a weight matrix W_ℓ = UΣVᵀ. The singular values σᵢ determine how much each "mode" of the transformation amplifies or shrinks signals passing through it. If σ_max ≫ 1, gradients explode in the backward pass (each layer multiplies the upstream gradient by up to σ_max); if σ_max ≪ 1, gradients vanish (each layer shrinks the gradient by σ_max). The product of singular values across L layers can be astronomically large or tiny — the vanishing/exploding gradient problem.

Smart initialization schemes address this directly:

• Xavier/Glorot initialization: draw weights from U[−√(6/(d_{in}+d_{out})), √(6/(d_{in}+d_{out}))], designed so the variance of activations and gradients is preserved through layers with linear activations.
• He initialization: scale by √(2/d_{in}), appropriate for ReLU networks where half the units are zeroed.
• Orthogonal initialization: draw W from the uniform distribution over orthogonal matrices (via QR of a random Gaussian matrix). Orthogonal matrices have all singular values exactly 1 — they preserve norms perfectly and are optimal for preventing vanishing/exploding gradients at initialization.

Batch normalization implicitly addresses conditioning at every layer by re-centering and re-scaling the preactivations: Z̃ = (Z − μ)/σ · γ + β where μ and σ are the batch mean and standard deviation. This keeps the effective condition number of each layer's weight matrix near 1 during training. Layer normalization, used in transformers, normalizes across the feature dimension instead of the batch dimension — same linear algebra, different axis.

Attention Mechanism: Weighted Sums via Matrices

The self-attention mechanism in transformers is a linear algebra operation that computes a weighted combination of value vectors, where the weights depend on the compatibility between queries and keys. Given a sequence of n tokens, each represented as a d-dimensional vector, the attention computation is:

Multi-head attention runs h independent attention heads in parallel, each with its own W_Q, W_K, W_V projections of smaller dimension d_k = d/h, then concatenates the outputs and applies a final linear projection. This is equivalent to computing a block-structured matrix operation: the concatenation is a horizontal stack of matrices, and the final projection is a matrix multiplication. The "attention pattern" — the n × n softmax matrix — determines which tokens influence each output position and is the mechanism through which transformers process context.

Matrix Factorization in ML

Many machine learning problems reduce to finding a low-rank matrix factorization. In collaborative filtering for recommendation systems, we observe a sparse ratings matrix R ∈ ℝⁿˣᵐ (n users, m items) and seek factors U ∈ ℝⁿˣᵏ and V ∈ ℝᵐˣᵏ such that R ≈ UVᵀ. Each row of U is a k-dimensional user embedding and each row of V is a k-dimensional item embedding. The dot product Uᵢ · Vⱼ predicts user i's rating for item j.

Training minimizes ‖R − UVᵀ‖²_F over observed entries (a non-convex problem due to the bilinear form UVᵀ), plus regularization terms λ_U‖U‖²_F + λ_V‖V‖²_F. Alternating least squares (ALS) fixes one factor and solves for the other in closed form: with V fixed, each user embedding Uᵢ is the solution to a small ridge regression problem; with U fixed, each item embedding Vⱼ is similarly solved. ALS converts the non-convex joint problem into a sequence of convex subproblems, each solved via normal equations.