Optimization for Machine Learning

Fitting a model comes down to making one number smaller. That number is the loss, and the settings you are allowed to change are the parameters. Optimization is the study of how to search for settings that make it small.

With many parameters, we cannot inspect the entire objective as a landscape. Instead, algorithms use information such as function values, gradients, curvature, and constraints. We begin with small quadratics to see what this information tells us. The derivative and gradient notation comes from Calculus for Deep Learning; the matrix background is in Linear Algebra for Machine Learning. Run the Python blocks in order; later examples use NumPy and SciPy.

Rolling downhill, one step at a time

Gradient descent subtracts a positive multiple of the current gradient from the parameters. The gradient determines the local descent direction, while the learning rate scales the update. Whether the new point actually has a lower objective depends on how large that update is.

Below, the landscape is a single bowl whose lowest point sits at \(x = 3\), and we start at \(x = 10\).

def f(x):     return (x - 3.0) ** 2      # a bowl whose lowest point is at x = 3
def slope(x): return 2.0 * (x - 3.0)     # steepness at x

x = 10.0
for step in range(7):
    print(f"  step {step}: x = {x:7.4f}   height = {f(x):8.4f}   slope = {slope(x):8.4f}")
    x = x - 0.3 * slope(x)
# step 0: x = 10.0000   height =  49.0000   slope =  14.0000
# step 1: x =  5.8000   height =   7.8400   slope =   5.6000
# step 2: x =  4.1200   height =   1.2544   slope =   2.2400
# step 3: x =  3.4480   height =   0.2007   slope =   0.8960
# step 4: x =  3.1792   height =   0.0321   slope =   0.3584
# step 5: x =  3.0717   height =   0.0051   slope =   0.1434
# step 6: x =  3.0287   height =   0.0008   slope =   0.0573

Every line can be checked by hand. At \(x = 10\) the slope is 14, and \(0.3 \times 14 = 4.2\), so the next position is \(10-4.2 = 5.8\). At 5.8 the slope has dropped to 5.6, so the next step is shorter. The steps shrink on their own as the ground flattens. At the \(\eta = 0.3\) used here the iterate comes in from one side and never crosses the bottom; the next section shows a step size where the steps still shrink and the iterate overshoots on every one of them.

Written in symbols, that loop is one line:

\[x_{t+1}=x_t-\eta\,\nabla f(x_t)\]

\(\nabla f(x_t)\) is the slope at the current position, \(\eta\) is the number we set to 0.3, called the learning rate or step size, and the minus sign is what makes it downhill rather than up. This is gradient descent, and it is the basis of many iterative fitting methods, including those used to train neural networks.

How big a step?

The step size was set to 0.3 with no justification. Changing it can turn convergence into oscillation or divergence, and the same bowl shows the whole range of behavior.

for lr in (0.1, 0.5, 0.9, 1.0, 1.1):
    x = 10.0
    for _ in range(20):
        x = x - lr * slope(x)
    print(f"  step size {lr:4.1f}   after 20 steps x = {x:12.4f}")
# step size  0.1   after 20 steps x =       3.0807
# step size  0.5   after 20 steps x =       3.0000
# step size  0.9   after 20 steps x =       3.0807
# step size  1.0   after 20 steps x =      10.0000
# step size  1.1   after 20 steps x =     271.3632

Four things happen across those five rows. A step size of 0.1 creeps toward the bottom and is still 0.08 away after twenty steps. A step size of 0.5 lands on 3.0000 exactly, and it gets there on the very first step. A step size of 0.9 ends up the same distance out as 0.1 did, but by jumping past the bottom and back each time. At 1.0 the iterate returns to exactly where it started, having bounced between 10 and \(-4\) forever. At 1.1 it alternates sides with increasing distance from 3; the displayed value is positive because the run ends after an even number of updates.

The pattern has a simple cause. For this bowl the slope is \(2(x-3)\), so one step multiplies the signed error \(x-3\) by \(1-2\eta\). At \(\eta = 0.5\) that factor is zero, which is why one step suffices. For \(0<\eta<0.5\) the factor is a positive number less than one and the approach is one-sided; above 0.5 it turns negative and the iterate alternates sides. At \(\eta = 1\) the factor is exactly \(-1\), and past that its magnitude exceeds one and every step makes things worse.

The Hessian is the matrix of second derivatives. For a quadratic it is constant, and its eigenvectors give directions with curvatures equal to its eigenvalues. If all those curvatures are positive, the Hessian is positive definite. On such a quadratic, along each eigenvector of the Hessian, with curvature \(\lambda\), one step multiplies that component of the error by \(1-\eta\lambda\). Convergence to the minimum needs \(|1-\eta\lambda| < 1\) in every direction at once, so convergence from any starting point requires \(0 < \eta < 2/\lambda_{\max}\) — the sharpest direction sets the limit. For a general objective the Hessian describes a local quadratic approximation, and the curvature changes as you move. Its value at the current point alone does not guarantee that a finite step will decrease the loss.

The code below builds a two-dimensional bowl that curves 10 times more sharply along one axis than the other. The function is \(f(x) = \tfrac{1}{2}(x_1^2 + 0.1\,x_2^2)\), so its gradient is \((x_1,\,0.1\,x_2)\) — which is what A * x computes, with A holding the two curvatures. The curvatures are 1 and 0.1, so the binding limit is \(\eta < 2/1 = 2\).

import numpy as np

A = np.array([1.0, 0.1])                       # curvature along each axis
for lr in (0.5, 1.9, 2.0, 2.1):
    x = np.ones(2)
    for _ in range(50):
        x = x - lr * (A * x)
    print(f"lr {lr:4.1f}  |x| after 50 steps {np.linalg.norm(x):.4e}")
# lr  0.5  |x| after 50 steps 7.6945e-02
# lr  1.9  |x| after 50 steps 5.1538e-03
# lr  2.0  |x| after 50 steps 1.0000e+00
# lr  2.1  |x| after 50 steps 1.1739e+02

At 2.0 the sharp direction sits exactly on the boundary and holds its magnitude at 1.0 forever. At 2.1 it grows by a factor of 1.1 per step. The flat direction is nowhere near its own limit in any of these runs, and it does not matter — one unstable direction is enough.

A rapidly growing loss is a reason to try a smaller learning rate while also checking the data, gradient computation, and numerical arithmetic. Norm clipping rescales a gradient whose norm exceeds a chosen cap; another variant clips individual components. Here we mean norm clipping. For plain gradient descent, clipping it to \(c\) bounds the update norm by \(\eta c\) — but a bounded step is not a stable one, and clipping does not enforce \(\eta < 2/\lambda_{\max}\): on \(f(x) = x^2/2\) with \(\eta = 3\) and the gradient clipped to the range \(-1\) to \(1\), the iterate goes 0.5, \(-1\), 2, \(-1\), 2 and cycles forever. With momentum or Adam, clipping the current gradient does not directly impose that bound on the final update either, since the update is built from more than the current gradient.

Why some bowls take far longer than others

A stable learning rate need not give rapid convergence. On a positive-definite quadratic, the largest curvature limits the fixed learning rate, while smaller curvatures can make some error components decay slowly. The actual iteration count also depends on the initial error, the learning rate, and the requested accuracy.

The curvature ratio \(\kappa=\lambda_{\max}/\lambda_{\min}\) is the condition number of the positive-definite Hessian. With the best fixed learning rate for worst-case contraction, the worst-case iteration bound for a specified relative error grows roughly in proportion to \(\kappa\) when \(\kappa\) is large. The linear algebra article instead measured the condition number of the design matrix \(X\). For the least-squares objective \(\tfrac12\|Xw-y\|_2^2\) with full column rank, the Hessian is \(X^{\top}X\), and \(\kappa_2(X^{\top}X)=\kappa_2(X)^2\).

import numpy as np

def gd(kappa, steps=200):
    A = np.array([1.0, 1.0 / kappa])
    lr = 2.0 / (A.max() + A.min())             # the optimal fixed step
    x = np.ones(2)
    for _ in range(steps):
        x = x - lr * (A * x)
    return 0.5 * np.sum(A * x * x)

for k in (1, 10, 100, 1000):
    print(f"kappa {k:5d}  f after 200 steps {gd(k):.3e}"
          f"   rate bound {((k - 1) / (k + 1)) ** 400:.3e}")
# kappa     1  f after 200 steps 0.000e+00   rate bound 0.000e+00
# kappa    10  f after 200 steps 7.591e-36   rate bound 1.380e-35
# kappa   100  f after 200 steps 1.694e-04   rate bound 3.354e-04
# kappa  1000  f after 200 steps 2.249e-01   rate bound 4.493e-01

After 200 steps, the nonzero reported objectives range from about \(7.6\times10^{-36}\) to 0.225. At \(\kappa = 10\) the objective is effectively zero; at \(\kappa = 1000\) it has fallen from 0.5005 to 0.2249, so about 45% of its starting objective remains. The learning rate \(2/(\lambda_{\max}+\lambda_{\min})\) balances the two extreme contraction factors. Their common magnitude is \(q=(\kappa-1)/(\kappa+1)\). The objective is quadratic in the error, so after 200 steps its relative bound is \(q^{400}\). The rate bound column is the theoretical bound on the ratio \(f_t/f_0\) rather than on \(f_t\) itself. Dividing the reported objective by the starting value \(f_0 = \tfrac{1}{2}(1 + 1/\kappa)\) reproduces the rate-bound column at every \(\kappa\) here: for this quadratic the bound holds with equality.

These are deterministic quadratics, with no dataset or sampling noise. All have their minimum at the origin. What changes is the shape of the equal-loss contours: circles at \(\kappa=1\), and elongated ellipses at \(\kappa=1000\). The experiment isolates how curvature affects this fixed-step algorithm.

Rescaling parameters or features can improve conditioning when incompatible units are the problem. Standardizing columns does not remove their linear dependence or guarantee a well-conditioned design. Normalization layers change a network’s parameterization and gradient behavior, but they do not guarantee a smaller Hessian condition number. Adam divides a running average of each gradient component by the square root of a running average of its square, with bias corrections and a small denominator offset. This rescales updates using gradient history; it is not generally an estimate of the inverse Hessian.

Second-order methods use curvature more directly. An exact Newton step solves a positive-definite quadratic in one iteration in exact arithmetic. For a dense Hessian with \(p\) parameters, storing it requires \(O(p^2)\) memory and factoring it typically costs \(O(p^3)\); the cost of constructing it also depends on the objective and data. Here \(O(p^2)\) means storage grows proportionally to the square of the parameter count, ignoring constant factors. Methods such as Newton-CG can use Hessian-vector products without forming the full matrix.

One valley, or many?

The quadratics above are convex. This property helps interpret a stationary point, although it does not require a bowl with a unique bottom.

On a convex domain, one that contains the segment between any two of its points, a function is convex if the straight line joining any two points on its graph never dips below the graph itself. Formally \(f(t a + (1-t)b) \leq t f(a) + (1-t)f(b)\) for every \(t\) between 0 and 1, where \(t\) is the mixing weight that slides along the line from \(b\) to \(a\) — a different quantity from the curvature \(\lambda\) used earlier in this article. For a twice-differentiable function the definition is equivalent to the curvature being non-negative in every direction everywhere.

Convexity buys one large thing: every local minimum is also a global minimum. For a differentiable, unconstrained convex objective, a zero gradient certifies a global minimum, and there is a duality theory that can certify optimality more generally. It does not guarantee that a minimum exists, that it is unique, or that a particular algorithm will reach it: unregularized logistic regression on separable data is convex with no finite minimizer, and where several minimizers exist, both the algorithm and its starting point can affect which one it reaches. The zero-gradient test also needs care about which problems it applies to. Linear regression, ridge and logistic regression have the differentiable unconstrained form; lasso’s penalty and the hinge loss of an unconstrained SVM are nonsmooth, while a constrained SVM formulation has feasibility conditions. These require subgradient optimality conditions or Karush–Kuhn–Tucker (KKT) conditions, which combine stationarity with constraints. Training a multilayer neural network, jointly fitting both factors in matrix factorization, and optimizing k-means centers and assignments generally give non-convex problems.

Convexity alone does not determine convergence speed — the \(\kappa = 1000\) problem above is convex, yet 200 steps leave the objective at about 45% of its initial value. It says nothing about whether the solution generalizes to new data, and it does not guarantee good numerical conditioning. A convex problem can be impractical at scale, and non-convex problems are solved well enough to be useful every day.

Not looking at all the data

Computing a full gradient directly from the per-row losses requires processing the whole dataset. With ten million rows that is expensive for a single step, and the step is only one of many thousands. The standard response is to estimate the slope from a small random sample of rows — a minibatch — whose average gradient is an unbiased estimate of the full gradient when rows are sampled uniformly and averaged, conditional on the current parameters. Nonuniform sampling needs appropriate weights to retain that property.

With a fixed learning rate and persistent gradient noise, stochastic gradient descent can keep fluctuating near a minimum even when the corresponding full-gradient method converges. The size of those fluctuations depends on the learning rate, gradient noise, and curvature; a large learning rate can instead cause divergence. If every row’s gradient vanishes at the same optimum, persistent noise is absent there and a suitable fixed learning rate can permit convergence.

The next example minimizes \(f(w)=\|Xw-y\|_2^2/(2n)\) on one fixed generated dataset. Its full gradient is \(X^{\top}(Xw-y)/n\); replacing all rows with a sampled batch gives the update in the code. exact is a numerical least-squares reference for this dataset, not the population coefficient that generated the observations.

import numpy as np

rng = np.random.default_rng(0)
n, d = 1000, 10
X = rng.normal(size=(n, d))
y = X @ rng.normal(size=d) + 0.5 * rng.normal(size=n)
exact = np.linalg.lstsq(X, y, rcond=None)[0]

def sgd(batch, lr=0.02, steps=3000):
    w = np.zeros(d)
    for _ in range(steps):
        idx = rng.integers(0, n, batch)
        w = w - lr * (X[idx].T @ (X[idx] @ w - y[idx]) / batch)
    return np.linalg.norm(w - exact)

for b in (1, 10, 100, 1000):
    print(f"batch {b:5d}  distance to the exact solution {sgd(b):.4f}")
# batch     1  distance to the exact solution 0.1531
# batch    10  distance to the exact solution 0.0367
# batch   100  distance to the exact solution 0.0168
# batch  1000  distance to the exact solution 0.0047

Three thousand steps at every batch size, and the leftover distance falls from 0.1531 to 0.0047. rng.integers samples with replacement, so even the last row is a random draw of 1000 rows rather than a full pass over the 1000 rows of the dataset. It therefore still has sampling noise; a remaining error can also include incomplete convergence after a finite number of updates. Each cell is one distance measured at the end of one run, so this is a single run consistent with the \(1/\sqrt{\text{batch}}\) shape of a standard error, not a measurement of the steady-state noise level.

A tenfold larger batch processes ten times as many sampled rows per update, but elapsed time need not grow tenfold: vectorized computation and hardware utilization matter. Compare loss or parameter error against elapsed time or processed examples, as well as against update count. Small batches can offer inexpensive early updates, while larger batches reduce gradient noise. This experiment holds the number of updates fixed and does not establish the fastest batch size.

Decreasing the learning rate can reduce persistent fluctuations. Classical stochastic-approximation results often use \(\sum_t\eta_t=\infty\) and \(\sum_t\eta_t^2<\infty\), as with \(\eta_t=c/(t+1)\), \(c>0\). The first condition prevents the total step-size allowance from being finite; the second limits the accumulated effect of noise in the convergence analysis. Neither condition alone guarantees travel to the optimum. Theorems also require assumptions on the objective and gradient estimates. Finite training schedules such as cosine decay are instead evaluated for a chosen training budget; they do not automatically inherit these asymptotic guarantees.

When the answer has to obey rules

Some problems come with constraints attached: require a separating margin for linearly separable data, or minimize error subject to \(\|\beta\|_1 \leq t\). Write each constraint as \(g_i(x) \leq 0\). The Lagrangian attaches a multiplier \(\alpha_i \geq 0\) to each one and folds it into the objective.

\[L(x,\alpha)=f(x)+\sum_i \alpha_i g_i(x)\]

For example, minimize \(f(x)=(x-3)^2\) subject to \(x\leq1\). The unconstrained minimum at 3 is outside the allowed region, and the best allowed point is \(x=1\). Here \(g(x)=x-1\), so \(L=(x-3)^2+\alpha(x-1)\). At the solution, stationarity requires \(2(x-3)+\alpha=0\), giving \(\alpha=4\). The multiplier balances the objective’s slope at the boundary, and \(\alpha g(1)=4\times0=0\).

For a fixed \(\alpha\geq0\), the dual function is \(q(\alpha)=\inf_x L(x,\alpha)\), the lowest value the Lagrangian can approach over \(x\). It is a lower bound on the constrained minimum: at any feasible point, \(\alpha g(x)\leq0\), so \(q(\alpha)\leq L(x,\alpha)\leq f(x)\). The dual problem maximizes this bound over the allowed multipliers. In this example, minimizing over \(x\) gives \(x=3-\alpha/2\) and \(q(\alpha)=2\alpha-\alpha^2/4\). Its maximum occurs at \(\alpha=4\), where \(q=4=f(1)\). The bound matches the constrained optimum here. This equality, called strong duality, holds for convex problems under suitable conditions, such as a strictly feasible point for the usual convex inequality constraints; it is not automatic for every problem.

For least-squares lasso, the constrained problem \(\min_\beta \tfrac12\|X\beta-y\|_2^2\) subject to \(\|\beta\|_1\leq t\) and the penalized problem with \(\lambda\|\beta\|_1\) have corresponding solutions. For \(t>0\), strict feasibility gives a suitable nonnegative multiplier; at \(t=0\), a sufficiently large finite penalty makes zero optimal for this quadratic loss. The correspondence depends on the data and loss scaling and need not be one-to-one.

Complementary slackness, \(\alpha_i g_i(x)=0\), is one of the KKT conditions. A constraint is active when \(g_i(x)=0\); if it has slack, \(g_i(x)<0\), its multiplier must be zero at a KKT solution. An active constraint can still have a zero multiplier. In an SVM solution, points with nonzero dual coefficients are the support vectors. Their number depends on the data; soft-margin SVMs also include slack variables to allow margin violations.

When diagnosing an optimization run, check what the objective’s geometry permits, whether the learning rate is stable, how much sampling noise the updates contain, and whether constraints are satisfied. The experiments above separate these issues using small problems whose solutions can be checked directly. Their numerical results do not by themselves establish what will happen in a different model or dataset.

Exercises

1. Momentum against conditioning. Heavy-ball momentum carries a running vector \(v\) that accumulates the gradients seen so far, and \(\beta\) decides how much of it survives into the next step: \(v \leftarrow \beta v + \nabla f(x)\), then \(x \leftarrow x – \eta v\). At \(\beta = 0\) this is plain gradient descent. Compare the two on quadratics with condition number 10, 100, and 1000, counting the steps to reach an objective value of \(10^{-12}\) within a 500-step budget.

Expected observation: a plain method that does not finish at all on the hardest problem and a momentum method that does.

import numpy as np

def run(kappa, beta, steps=500):
    A = np.array([1.0, 1.0 / kappa])
    lr = (2.0 / (A.max() + A.min()) if beta == 0
          else 4.0 / (np.sqrt(A.max()) + np.sqrt(A.min())) ** 2)
    x, v = np.ones(2), np.zeros(2)
    for t in range(steps):
        v = beta * v + A * x
        x = x - lr * v
        if 0.5 * np.sum(A * x * x) < 1e-12:
            return t + 1
    return None                                # did not reach the target

for k in (10, 100, 1000):
    beta = ((np.sqrt(k) - 1) / (np.sqrt(k) + 1)) ** 2
    print(f"kappa {k:5d}  plain {run(k, 0.0)}"
          f"  momentum(beta={beta:.4f}) {run(k, beta)}")
# kappa    10  plain 68  momentum(beta=0.2699) 27
# kappa   100  plain None  momentum(beta=0.6694) 93
# kappa  1000  plain None  momentum(beta=0.8811) 315

2. Prove non-convexity by exhibit. Check the Hessian of the mean binary logistic negative log-likelihood at three parameter settings, explain algebraically why it is positive semidefinite for every setting, then find two distinct low-loss points of a two-parameter network loss and evaluate the loss at their midpoint.

Expected observation: a midpoint whose loss is far higher than the average of the two endpoints, which no convex function permits.

import numpy as np
from scipy.optimize import minimize

rng = np.random.default_rng(0)
X = rng.normal(size=(200, 3))

def logistic_hessian(w):
    p = 1 / (1 + np.exp(-X @ w))
    return X.T @ ((p * (1 - p))[:, None] * X) / len(X)

for _ in range(3):
    ev = np.linalg.eigvalsh(logistic_hessian(rng.normal(size=3) * 3))
    print(f"logistic Hessian min eigenvalue {ev.min():+.6f}")
# logistic Hessian min eigenvalue +0.006642
# logistic Hessian min eigenvalue +0.003064
# logistic Hessian min eigenvalue +0.009851

x = rng.normal(size=300)
y = np.tanh(2.0 * x) * 3.0 + 0.01 * rng.normal(size=300)
loss = lambda p: float(np.mean((np.tanh(p[0] * x) * p[1] - y) ** 2))

for start in ([1.0, 1.0], [-1.0, -1.0]):
    r = minimize(loss, start, method="Nelder-Mead",
                 options=dict(xatol=1e-10, fatol=1e-14))
    print("fitted point", np.round(r.x, 3), " loss", round(r.fun, 8))
# fitted point [2.    3.001]  loss 9.847e-05
# fitted point [-2.    -3.001]  loss 9.847e-05
print("loss at the midpoint", round(loss(np.zeros(2)), 6))   # 5.459496

3. Why iterate when a closed form exists. Compare storage for a dense Gram matrix with storage for one parameter-sized float64 vector, for \(p = 10^2\) to \(10^6\) features. Explain what this comparison says about scalability and which training allocations it omits.

Expected observation: quadratic growth in Gram-matrix storage and linear growth in vector storage; the vector alone is not the full training state.

for p in (100, 10_000, 100_000, 1_000_000):
    gram_gib = p * p * 8 / 2 ** 30
    state_mib = p * 8 / 2 ** 20
    print(f"p={p:9,d}  X.T @ X is {p * p:.2e} entries = {gram_gib:12,.1f} GiB"
          f"   one parameter vector = {state_mib:8.3f} MiB")
# p=      100  X.T @ X is 1.00e+04 entries =          0.0 GiB   one parameter vector =    0.001 MiB
# p=   10,000  X.T @ X is 1.00e+08 entries =          0.7 GiB   one parameter vector =    0.076 MiB
# p=  100,000  X.T @ X is 1.00e+10 entries =         74.5 GiB   one parameter vector =    0.763 MiB
# p=1,000,000  X.T @ X is 1.00e+12 entries =      7,450.6 GiB   one parameter vector =    7.629 MiB