Optimization

Last lecture fixed the model class, linear in the weights however we featurize the inputs, and left us staring at a picture. Predictions \(\hat{\mathbf{y}} = \mathbf{X}\mathbf{w}\) live in the column space of the design matrix \(\mathbf{X} \in \mathbb{R}^{n \times d}\), mean squared error measures the distance from the labels \(\mathbf{y} \in \mathbb{R}^n\) down to that subspace, and the best fit is the projection of \(\mathbf{y}\) onto it. We know exactly what we are looking for; we have not said how to compute it. We compare a closed-form formula with gradient descent and count the cost of finding the minimizing weights. The numerical accuracy and iteration count depend on the condition number of the design matrix, \[ \kappa(\mathbf{X}) = \frac{\sigma_{\max}}{\sigma_{\min}}, \] the ratio of the largest to the smallest singular value of \(\mathbf{X}\), from the Decompositions lecture. It controls how many digits the exact formula keeps and how many steps the iteration takes.

Squared Error Is a Bowl

Write the loss in the design-matrix form we set up last lecture: \[ \mathcal{L}(\mathbf{w}) = \frac1n\|\mathbf{X}\mathbf{w} - \mathbf{y}\|_2^2 , \] where \(\mathbf{w} \in \mathbb{R}^d\) is the weight vector. Problem 3 derived its gradient: \[ \nabla_\mathbf{w} \mathcal{L}(\mathbf{w}) = \frac2n \mathbf{X}^\top(\mathbf{X}\mathbf{w} - \mathbf{y}). \] The vector \(\mathbf{X}\mathbf{w} - \mathbf{y} \in \mathbb{R}^n\) is the residual, one entry per data point, and multiplying by \(\mathbf{X}^\top\) measures how much each feature correlates with it: the gradient’s \(j\)th entry is large exactly when feature \(j\) still lines up with what the model is missing.

Problem 3 also showed that the matrix \(\mathbf{X}^\top\mathbf{X} \in \mathbb{R}^{d\times d}\) is positive semi-definite: for any direction \(\mathbf{v} \in \mathbb{R}^d\), \[ \mathbf{v}^\top(\mathbf{X}^\top\mathbf{X})\mathbf{v} = (\mathbf{X}\mathbf{v})^\top(\mathbf{X}\mathbf{v}) = \|\mathbf{X}\mathbf{v}\|_2^2 \geq 0 . \] Since \(\frac2n\mathbf{X}^\top\mathbf{X}\) is the matrix of second derivatives of \(\mathcal{L}\), the curvature of the loss is nonnegative in every direction: \(\mathcal{L}\) is convex, a bowl facing up, with no ridges and no saddles, so a point where the gradient vanishes is the global minimum.

The Normal Equations

Set the gradient to zero at the minimizer \(\mathbf{w}^\star\) and solve: \[ \begin{align*} \frac2n \mathbf{X}^\top(\mathbf{X}\mathbf{w}^\star - \mathbf{y}) &= \mathbf{0} \\ \mathbf{X}^\top(\mathbf{X}\mathbf{w}^\star - \mathbf{y}) &= \mathbf{0} \\ \mathbf{X}^\top\mathbf{X}\mathbf{w}^\star &= \mathbf{X}^\top\mathbf{y} \\ \mathbf{w}^\star &= (\mathbf{X}^\top\mathbf{X})^{-1}\mathbf{X}^\top\mathbf{y} . \end{align*} \] The first move multiplies both sides by \(n/2\); the second distributes \(\mathbf{X}^\top\); the third multiplies on the left by \((\mathbf{X}^\top\mathbf{X})^{-1}\), assuming for the moment that the inverse exists. The middle line has a name. The normal equations are the linear system \[ \mathbf{X}^\top\mathbf{X}\mathbf{w}^\star = \mathbf{X}^\top\mathbf{y} , \] \(d\) equations in the \(d\) unknowns of \(\mathbf{w}^\star\). The line above the last, \(\mathbf{X}^\top(\mathbf{X}\mathbf{w}^\star - \mathbf{y}) = \mathbf{0}\), says every column of \(\mathbf{X}\) has inner product zero with the residual. This is the perpendicular projection onto the column space from last lecture. Problem 7 has you check that the two agree line for line.

The inverse exists exactly when the \(d\) columns are linearly independent, which requires at least \(d\) data points. Indeed, \(\mathbf{v}^\top(\mathbf{X}^\top\mathbf{X})\mathbf{v} = \|\mathbf{X}\mathbf{v}\|_2^2\) is zero exactly when \(\mathbf{X}\mathbf{v} = \mathbf{0}\); we assume independent columns today. When \(d > n\), infinitely many weight vectors can fit the data exactly, and the formula does not select among them. The Generalization lecture introduces a selection criterion. For now we have a formula and no sense of what it costs to evaluate.

The Cost of an Exact Answer

For \(n\) data points and \(d\) features, the closed form has the following arithmetic cost.

Forming \(\mathbf{X}^\top\mathbf{X} \in \mathbb{R}^{d\times d}\) means filling in \(d^2\) entries, and entry \((j,k)\) is the inner product of columns \(j\) and \(k\) of \(\mathbf{X}\), two vectors of length \(n\), costing \(n\) multiplications and \(n-1\) additions. That is \(O(n)\) per entry and \(O(nd^2)\) in total. Forming \(\mathbf{X}^\top\mathbf{y} \in \mathbb{R}^d\) is \(d\) inner products of length \(n\), so \(O(nd)\). Solving the \(d\times d\) system by Gaussian elimination sweeps \(d\) columns, and each sweep updates a block of up to \(d^2\) entries, so \(O(d^3)\). Adding the three and dropping the term that is dominated by the others: \[ O(nd^2 + d^3) . \] Doubling \(n\) doubles the work, but doubling \(d\) multiplies it by eight: features, not data points, are what make this formula expensive, which is worth remembering the next time a feature map tempts us into a few thousand polynomial terms. At \(n = 10^6\) and \(d = 10^4\), the term \(nd^2\) alone is \(10^{14}\) operations, a fit that finishes overnight instead of during lecture. (In practice nobody computes \((\mathbf{X}^\top\mathbf{X})^{-1}\) and then multiplies; solving the system \(\mathbf{X}^\top\mathbf{X}\mathbf{w}^\star = \mathbf{X}^\top\mathbf{y}\) directly is the same \(O(d^3)\) with better constants and accuracy.)

Forming \(\mathbf{X}^\top\mathbf{X}\) also affects numerical accuracy.

Numerical Accuracy: Normal Equations and SVD

Computers store real numbers with about \(16\) significant decimal digits, and every arithmetic operation rounds. Solving along the smallest singular direction divides by \(\sigma_{\min}\), amplifying whatever rounding error rode in on the data; the rule of thumb is that solving a system with condition number \(\kappa\) costs roughly \(\log_{10}\kappa\) significant digits. A design matrix with \(\kappa(\mathbf{X}) = 10^4\), unremarkable when one feature is in dollars and another in millions of dollars, costs four digits out of sixteen: affordable.

The normal equations solve a system involving \(\mathbf{X}^\top\mathbf{X}\) rather than \(\mathbf{X}\). Substitute the SVD \(\mathbf{X} = \sum_i \sigma_i \mathbf{u}_i \mathbf{v}_i^\top\) from the Decompositions lecture and use orthonormality of the left singular vectors, \(\mathbf{u}_i^\top\mathbf{u}_j = 1\) when \(i = j\) and \(0\) otherwise: \[ \mathbf{X}^\top\mathbf{X} = \Big(\sum_i \sigma_i \mathbf{v}_i\mathbf{u}_i^\top\Big)\Big(\sum_j \sigma_j \mathbf{u}_j\mathbf{v}_j^\top\Big) = \sum_{i,j} \sigma_i\sigma_j\, \mathbf{v}_i (\mathbf{u}_i^\top\mathbf{u}_j) \mathbf{v}_j^\top = \sum_i \sigma_i^2\, \mathbf{v}_i\mathbf{v}_i^\top . \] The double sum collapses because only the \(i = j\) terms survive. What is left is an eigendecomposition: the eigenvectors of \(\mathbf{X}^\top\mathbf{X}\) are the right singular vectors of \(\mathbf{X}\) and its eigenvalues are the squared singular values, exactly the link we verified by hand in the Decompositions lecture. Its condition number is therefore \[ \kappa(\mathbf{X}^\top\mathbf{X}) = \frac{\sigma_{\max}^2}{\sigma_{\min}^2} = \kappa(\mathbf{X})^2 . \] That squaring doubles the digits we lose: the comfortable \(\kappa(\mathbf{X}) = 10^4\) becomes \(10^8\), and eight digits out of sixteen is half our precision spent before the solve begins.

The SVD avoids forming \(\mathbf{X}^\top\mathbf{X}\). Define the pseudoinverse of \(\mathbf{X}\) from its SVD by inverting each singular value in place: \[ \mathbf{X}^+ = \sum_i \frac{1}{\sigma_i}\mathbf{v}_i\mathbf{u}_i^\top \in \mathbb{R}^{d \times n} . \]

Claim: When \(\mathbf{X}\) has linearly independent columns, the pseudoinverse reproduces the normal-equations formula exactly, so \(\mathbf{w}^\star = \mathbf{X}^+\mathbf{y}\) is the least-squares solution.

Proof of Claim

Inverting the sum of outer products we just computed inverts each eigenvalue and leaves the eigenvectors alone, so \((\mathbf{X}^\top\mathbf{X})^{-1} = \sum_i \sigma_i^{-2}\mathbf{v}_i\mathbf{v}_i^\top\). Multiply it by \(\mathbf{X}^\top = \sum_j \sigma_j\mathbf{v}_j\mathbf{u}_j^\top\) and collapse the double sum with orthonormality of the right singular vectors: \[ \begin{align*} (\mathbf{X}^\top\mathbf{X})^{-1}\mathbf{X}^\top &= \Big(\sum_i \frac{1}{\sigma_i^2}\mathbf{v}_i\mathbf{v}_i^\top\Big)\Big(\sum_j \sigma_j \mathbf{v}_j\mathbf{u}_j^\top\Big) \\&= \sum_{i,j} \frac{\sigma_j}{\sigma_i^2}\, \mathbf{v}_i (\mathbf{v}_i^\top\mathbf{v}_j) \mathbf{u}_j^\top \\&= \sum_i \frac{\sigma_i}{\sigma_i^2}\, \mathbf{v}_i\mathbf{u}_i^\top \\&= \sum_i \frac{1}{\sigma_i}\mathbf{v}_i\mathbf{u}_i^\top = \mathbf{X}^+ . \end{align*} \] The first equality substitutes both outer-product forms, the second regroups so the inner product \(\mathbf{v}_i^\top\mathbf{v}_j\) sits in the middle, and the third keeps only the \(i = j\) terms.

The two formulas agree in exact arithmetic but have different numerical behavior: the pseudoinverse route divides by \(\sigma_i\) once and has condition number \(\kappa(\mathbf{X})\), while the normal-equations route squares every singular value first. This is why numpy.linalg.lstsq and numpy.linalg.pinv are built on the SVD.

The class demo makes the gap visible with the polynomial feature map from last lecture, pushing \(\kappa(\mathbf{X})\) high enough that the explicit inverse and the pseudoinverse, run on the same data, stop agreeing.

Gradient Descent: An Iterative Alternative

Both routes so far cost \(O(nd^2)\), because both build a \(d\times d\) matrix out of the data before they solve anything. But we are standing on the inside of a bowl and we want the bottom; we do not have to solve for it if we are willing to walk toward it. The gradient points in the direction that increases \(\mathcal{L}\) fastest, by its defining property from the Linear Algebra lecture, so stepping the opposite way decreases the loss: \[ \mathbf{w}^{(t+1)} = \mathbf{w}^{(t)} - \alpha \nabla_\mathbf{w}\mathcal{L}(\mathbf{w}^{(t)}), \] where the superscript \((t)\) indexes iterations and \(\alpha > 0\) is a learning rate, the length of one step. This is gradient descent. For mean squared error we know the gradient exactly, so the update is \[ \mathbf{w}^{(t+1)} = \mathbf{w}^{(t)} - \frac{2\alpha}{n}\mathbf{X}^\top\big(\mathbf{X}\mathbf{w}^{(t)} - \mathbf{y}\big) . \] Each iteration predicts, measures the residual, and updates the weights. When the residual is orthogonal to every feature column, the gradient is zero and the iteration stops moving.

Now count what one step costs. The product \(\mathbf{X}\mathbf{w}^{(t)}\) is \(n\) inner products of length \(d\), so \(O(nd)\); subtracting \(\mathbf{y}\) is \(O(n)\); multiplying the result by \(\mathbf{X}^\top\) is another \(O(nd)\). A step is \(O(nd)\) and \(T\) steps are \[ O(Tnd) , \] against \(O(nd^2 + d^3)\) for the closed form. Iterating wins whenever \(T \ll d\), and it never forms \(\mathbf{X}^\top\mathbf{X}\), so it never squares the condition number either. Because \(\mathcal{L}\) is convex, repeated updates reach the same \(\mathbf{w}^\star\) the closed form computes, provided \(\alpha\) is small enough that a step does not fly past the bottom. We take convergence on faith today; the Gradient Descent lecture proves it and reuses this exact update for logistic regression and neural networks, where no closed form exists.

The number of iterations \(T\) can make gradient descent expensive despite its cheaper steps.

The Shape of the Bowl

\(T\) can be large because a single learning rate has to serve every direction at once, and a bowl need not be equally steep in every direction. To see the shape we are walking on, measure position by the error \(\mathbf{e} = \mathbf{w} - \mathbf{w}^\star \in \mathbb{R}^d\) and expand the loss around its minimum: \[ \begin{align*} \mathcal{L}(\mathbf{w}^\star + \mathbf{e}) &= \frac1n\big\|(\mathbf{X}\mathbf{w}^\star - \mathbf{y}) + \mathbf{X}\mathbf{e}\big\|_2^2 \\&= \frac1n\|\mathbf{X}\mathbf{w}^\star - \mathbf{y}\|_2^2 + \frac2n\mathbf{e}^\top\mathbf{X}^\top(\mathbf{X}\mathbf{w}^\star - \mathbf{y}) + \frac1n\|\mathbf{X}\mathbf{e}\|_2^2 \\&= \mathcal{L}(\mathbf{w}^\star) + \frac1n\mathbf{e}^\top\mathbf{X}^\top\mathbf{X}\mathbf{e} . \end{align*} \] The first equality substitutes \(\mathbf{w} = \mathbf{w}^\star + \mathbf{e}\) and splits the vector inside the norm, the second expands the square of a sum, and the third drops the middle term because \(\mathbf{X}^\top(\mathbf{X}\mathbf{w}^\star - \mathbf{y}) = \mathbf{0}\) is exactly the normal equations. Now substitute the outer-product form \(\mathbf{X}^\top\mathbf{X} = \sum_i \sigma_i^2\mathbf{v}_i\mathbf{v}_i^\top\) we derived above: \[ \mathcal{L}(\mathbf{w}^\star + \mathbf{e}) - \mathcal{L}(\mathbf{w}^\star) = \frac1n\sum_i \sigma_i^2\,(\mathbf{v}_i^\top\mathbf{e})^2 = \frac12\sum_i \lambda_i\,(\mathbf{v}_i^\top\mathbf{e})^2 . \] The second equality just names the curvature of the loss along the right singular vector \(\mathbf{v}_i\): \[ \lambda_i = \frac{2\sigma_i^2}{n} . \] The excess loss is a sum of independent one-dimensional parabolas: moving along \(\mathbf{v}_i\) changes only the \(i\)th term. The bowl is steep along the top singular direction (\(\lambda_{\max}\)) and shallow along the bottom one (\(\lambda_{\min}\)), and its aspect ratio is the number we have been tracking all lecture: \[ \frac{\lambda_{\max}}{\lambda_{\min}} = \frac{\sigma_{\max}^2}{\sigma_{\min}^2} = \kappa(\mathbf{X})^2 . \] The same \(\kappa(\mathbf{X})^2\) limits the precision of the normal equations and gives the loss bowl its aspect ratio.

In the plot, both panels show the level sets of such a quadratic loss in two dimensions and the path gradient descent takes from the same starting point, with the minimum at the black cross. On the left the two curvatures are equal, the level sets are circles, and one step of the right size arrives. On the right the steep direction curves \(15\) times harder than the shallow one, the level sets stretch into a narrow valley, and the path zig-zags across it: a step size large enough to make progress down the shallow axis overshoots the steep axis and bounces off the opposite wall, so \(\alpha\) must be small enough for the steep direction, which leaves it far too small for the shallow one. The count of steps grows with the aspect ratio \(\lambda_{\max}/\lambda_{\min}\) and not with the number of parameters \(d\): a two-parameter problem with a stretched bowl is harder than a thousand-parameter problem with a round one. Problem 8 makes this precise by choosing the fixed learning rate that is best against every possible eigendirection, then checks the resulting rate against the demo.

Fixing the Geometry Instead of the Step Size

The class demo watches this stretched bowl slow gradient descent on a two-feature dataset where one feature sits on a scale a thousand times larger than the other: the large-scale weight converges quickly and the small-scale one barely moves after tens of thousands of steps.

The repair is to change the bowl rather than the step size. Dividing each feature column by its standard deviation only rescales the columns, so the column space and the achievable minimum are unchanged, but it shrinks \(\kappa(\mathbf{X})\) toward \(1\) and rounds the bowl; the demo shows the rescaled run reaching the same minimum in a handful of steps. Rescaling costs one pass over the data and helps more here than any amount of tuning \(\alpha\), and it worked because our two features differed only in their units. What would you do if instead two features were near-duplicates of each other, so that \(\sigma_{\min}\) is small for a reason no choice of units can repair?

Conditioning Later in the Course

The Gradient Descent lecture adds momentum, improving the dependence from \(\kappa\) to \(\sqrt{\kappa}\); the Depth-enablers lecture keeps a deep network’s Jacobians well conditioned; Muon builds an optimizer that flattens a matrix’s singular values to \(1\) before stepping; and the Generalization lecture traces the double-descent peak to a \(\sigma_{\min}\) near zero.

Both solvers minimize error on the data we already have. Next lecture asks whether training error is the right objective for selecting a model. A model that fits the training data perfectly may still perform poorly on held-out data. When the loss geometry has very different curvatures, rescaling or otherwise improving the conditioning can help more than tuning a single step size.