15. Optimization & Machine Learning: Advanced Topics

Chapter PDF · Complete book

15.1 Regularization¶

When the sample size $n$ is small relative to the model dimension $p$, regularization helps control the complexity of empirical risk minimization.

15.1.1 Penalized Least Squares¶

For simplicity, consider regularized least-squares regression:

where $R(x)$ is the regularizer and $\lambda\geqslant 0$ the regularization parameter.

The regularizer expresses prior structure of the weight vector $x$, such as small magnitude or sparsity. The strength $\lambda$ can be selected by cross-validation.

We assume that $R$ is finite, nonnegative, lower semicontinuous, and coercive, i.e. $R(x) \rightarrow +\infty$ as $|\!| x |\!| \rightarrow +\infty$.

The following proposition shows that as $\lambda\downarrow0$, regularization selects among the exact solutions when $Ax=y$ is feasible. This is particularly useful when $\ker(A)\neq\{0\}$ and there are infinitely many exact solutions.

15.1.2 Ridge Regression¶

Ridge regression uses the quadratic penalty $R(x)=\frac12|\!| x |\!|_{\mathbb{R}^p}^2$. For $\lambda>0$, the objective in (15.1) has a unique minimizer:

One has

Thus $x_\lambda$ depends linearly on $y$ and can be computed by solving a linear system. Two equivalent expressions are available.

The systems can be solved directly, for example by Cholesky factorization, or iteratively, for example by conjugate gradients. The latter exploits the quadratic structure more effectively than ordinary gradient descent.

For $n>p$, use the smaller system (15.4); for $n<p$, use (15.5).

Pseudo-inverse.

As $\lambda\downarrow0$, $x_\lambda$ converges to the minimum-norm least-squares solution $A^+y$. If $Ax=y$ is feasible, this is, by (15.2)

If $A$ has full column rank, so $\ker(A)=\{0\}$ and $n\geqslant p$, $A^\top A \in \mathbb{R}^{p \times p}$ is an invertible matrix, and $(A^\top A + \lambda\mathrm{Id}_p)^{-1} \rightarrow (A^\top A)^{-1}$, so that

If $A$ has full row rank, so $\ker(A^\top)=\{0\}$ and $n\leqslant p$, the alternative formula is

When $n=p$ and $A$ is invertible, both formulas reduce to $A^+=A^{-1}$. For any rank, the singular value decomposition (SVD) gives the Moore–Penrose pseudoinverse by inverting the nonzero singular values and leaving the zero ones unchanged.

15.1.3 Lasso¶

The Lasso uses an $\ell^1$ penalty:

The penalty promotes sparsity in the solutions $x_\lambda$ of

The solutions $x_\lambda$ may be nonunique, and typically have many zero coefficients.

Figure 15.1 illustrates the geometry of $\ell^q$ balls. They concentrate around the coordinate axes as $q \rightarrow 0$, favoring sparse solutions, but become nonconvex for $q<1$.

Sparsity can encode a prior on the unknown, as in imaging, or select a small subset of predictive features. Feature selection can make a model easier to interpret and cheaper to evaluate.

Whether Lasso or ridge regression performs better depends on sparsity, correlations, noise, and the prediction task.

The objective $f_\lambda$ is convex, but the $\ell^1$ penalty is nonsmooth, so ordinary gradient descent does not apply at all points.

Section 15.1.4 develops an appropriate modification of gradient descent.

A closed-form solution $x_\lambda$ is generally unavailable, but the case of an orthogonal design matrix $A$ reduces to coordinatewise thresholding.

The nonlinear map $S_\lambda$ is soft thresholding.

15.1.4 Iterative Soft Thresholding¶

We derive an iterative algorithm by minimizing a quadratic surrogate.

We aim at minimizing

For a fixed reference point $x'$, introduce the surrogate with parameter $\tau>0$:

The surrogate touches the objective, $f_\tau(x,x)=f(x)$. Its additional quadratic term is

The term $K(x,x')$ is nonnegative when $\lambda_{\max}(A^\top A) \leqslant 1/\tau$, equivalently $\tau \leqslant 1/|\!| A |\!|_{\text{op}}^2$, where $|\!| A |\!|_{\text{op}} = \sigma_{\max}(A)$ is the operator norm.

Thus $f_\tau(x,x')$ is a valid surrogate:

The surrogate is strongly convex in its first argument, so its minimizer is unique.

Define the iteration

which ensures by construction that

Equation (15.7) defines iterative soft thresholding. For $A\neq0$, the surrogate argument applies when $0<\tau\leqslant 1/|\!| A |\!|_{\mathrm{op}}^2$. Convergence to a Lasso solution holds for a fixed step in the larger range $0<\tau<2/|\!| A |\!|_{\mathrm{op}}^2$. If $A=0$, every positive step is admissible.

15.2 Stochastic Optimization¶

We study stochastic gradient methods for large datasets of size $n$ and for objectives defined as expectations.

15.2.1 Minimizing Sums and Expectations¶

Many learning objectives are finite averages:

or expectations:

A finite average is an expectation under the uniform measure $\pi=\frac1n\sum_{i=1}^n\delta_i$, with $f(x,i)=f_i(x)$. Conversely, drawing $(z_i)_i$ i.i.d. from $\pi$ and setting $f_i(x)=f(x,z_i)$ gives an empirical approximation of (15.9). Under integrability, the empirical objective converges pointwise by the law of large numbers. Convergence of minima or minimizers requires additional assumptions.

For example, empirical risk minimization with linear predictors uses

for $z=(a,y) \in \mathcal{Z}= (\mathcal{A}=\mathbb{R}^p) \times \mathcal{Y}$ (typically $\mathcal{Y}=\mathbb{R}$ or $\mathcal{Y}=\{-1,+1\}$ for regression and classification), where $\ell$ is some loss function.

We illustrate the methods with binary logistic classification, using

The same computational approach extends to general parametric models, including deep neural networks.

The stochastic methods below apply both to finite sums (15.8) and to expectations (15.9), under the corresponding unbiasedness and moment assumptions. We present them for finite sums. When $n$ is large, the aim is to make each iteration inexpensive without evaluating every summand.

When the summands $f_i(x)$ are similar, a single gradient can approximate the average well: $\nabla f_i \approx \nabla f$. Evaluating $\nabla f_i$ is typically $n$ times cheaper than forming the full gradient.

Stochastic optimization is nevertheless not always faster than batch gradient descent.

For moderate $n$, a deterministic method may use the available computation more effectively.

When $n$ is too large for a full pass to be practical, stochastic methods divide the work into small updates and begin improving the model immediately. Mini-batches also permit parallel computation.

15.2.2 Batch Gradient Descent (BGD)¶

Batch gradient descent (BGD), analyzed in Section 14.4, iterates

For a convex objective with an $L$-Lipschitz gradient and a minimizer, the condition $0<\tau_{\min}\leqslant\tau_k\leqslant\tau_{\max}<2/L$ ensures convergence. Strong convexity strengthens this to a linear rate.

For a finite-sum objective, the gradient is

Its cost is typically $O(np)$ when each $\nabla f_i$ can be evaluated in time linear in $p$.

For an ERM summand of the form (15.10), Taylor expansion gives

Here $\ell'(s,y)$ denotes differentiation with respect to the score $s$, with the response $y$ fixed. Consequently,

For the logistic loss, one has

15.2.3 Stochastic Gradient Descent (SGD)¶

For very large $n$, computing the full gradient $\nabla f$ as in (15.12) is prohibitive.

SGD replaces the full gradient by the gradient of one summand $f_{i}$, with $i$ sampled uniformly. This gives an unbiased estimate:

where $\mathbf{i}$ is a random variable distributed uniformly in $\{1,\ldots,n\}$.

Starting from $x_0$, stochastic gradient descent (SGD) iterates

where, for each iteration index $k$, $i(k)$ is drawn independently of previous draws, uniformly from $\{1,\ldots,n\}$.

The iterates $x_{k+1}$ are random. We therefore study convergence to a minimizer of $f$ in a probabilistic sense, such as mean square or probability, together with the corresponding rates.

A batch-gradient step costs $O(np)$, whereas an SGD step costs $O(p)$. This makes SGD attractive when $n$ is large and repeated passes through the data are costly. Redundancy between observations can sometimes yield accurate predictions after only $k \ll n$ updates, before all observations have been used.

The step-size schedule $\tau_k$ must balance two requirements: it usually needs to decrease to suppress stochastic noise, but not so quickly that the method stops making progress before reaching a minimizer.

A common schedule has asymptotic behavior $\tau_k \sim k^{-1}$ as $k\rightarrow +\infty$. One such schedule is

where $\tau_0>0$ sets the initial step and $k_0>0$ controls how long the steps remain approximately constant.

Figure 15.5 illustrates minimization of $f_1(x)+f_2(x)$ for $x \in \mathbb{R}$, with $f_1(x)=(x-1)^2$ and $f_2(x)=(x+1)^2$. The distribution of $x_k$ concentrates around the minimizer $x^\star=0$. The initial iterate $x_0$ is uniform on $[-1/2,1/2]$.

The following theorem gives a mean-square convergence rate. The second-moment assumption is imposed along the iterates, not uniformly on all of $\mathbb{R}^p$: a globally strongly convex function cannot have a globally bounded gradient.

The preceding SGD bound, and the averaging scheme below, give sublinear rates even when $f$ is strongly convex. By contrast, batch gradient descent has a linear rate under the assumptions of Theorem 14.16.

Figure 15.6 compares the objective values $f(x_k)$. The black dashed curve shows batch gradient descent, with the iteration axis rescaled to reflect the $n$-fold cost of a batch update.

15.2.4 Stochastic Gradient Descent with Averaging (SGA)¶

Diminishing steps can slow SGD as $\tau_k$ approaches zero.

Averaging reduces fluctuations in the estimate. Run SGD on auxiliary iterates $( \tilde x_k)_k$,

and return their Cesàro average:

This gives stochastic gradient descent with averaging (SGA), also called averaged SGD.

The average can be maintained without storing the full trajectory, using the recursive update

A typical step-size schedule is

These steps decrease more slowly, at rate $k^{-1/2}$.

Averaging can reduce sensitivity to the tuning parameters $(k_0,\tau_0)$, although their choice still affects performance.

Bach proves that for logistic classification, averaging can adapt to local curvature under suitable regularity and step-size assumptions. This is a problem-dependent result, rather than a universal improvement of every constant over SGD.

15.2.5 Stochastic Averaged Gradient Descent (SAG)¶

For a finite dataset with $n$ samples of total size $n \times p$, one can trade additional memory for reduced gradient noise by storing past gradient evaluations. This yields the stochastic averaged gradient (SAG) method.

Store the latest gradient of each summand in a table $(G^i)_{i=1}^n$, requiring $O(np)$ memory. Maintain its average $g$ as an approximation to the full gradient.

Initialize $G^i=0$ for all $i$ and $g=0$, so that $g=n^{-1}\sum_iG^i$ is the stored average. At each iteration, update the table and then $x$:

Because most entries were evaluated at earlier iterates, the table average is generally biased conditional on the current history. Its convergence analysis therefore differs from the unbiased-SGD proof above. SAG uses a fixed step size $\tau$. As for batch gradient descent, choose $\tau$ on the scale of $1/L$, where $L$ bounds the Lipschitz constants of the individual gradients $\nabla f_i$; the admissible numerical constant depends on the convergence theorem.

For smooth convex finite sums with a minimizer and a suitable step size, SAG admits an $O(1/k)$ expected objective bound for averaged iterates. Furthermore, in the presence of strong convexity (for instance after adding a positive quadratic penalty to logistic regression), it has a linear convergence rate, i.e.

for some $0 < \rho < 1$.

These improvements exploit the finite-sum structure: SAG stores information about each of the $n$ summands. SGD and averaged SGD also extend beyond finite $n$ to general expectations (15.9).

Figure 15.7 shows a comparison of SGD, SGA and SAG.

15.3 Automatic Differentiation¶

Gradient evaluation $\nabla f(x)$ is the main operation in batch and stochastic gradient methods. For linear models and shallow networks, explicit formulas usually reduce it to matrix-vector products. For deeper or more elaborate computations, expanded derivative formulas become unwieldy and can duplicate work. Automatic differentiation organizes the chain rule recursively to reuse intermediate results.

15.3.1 Finite Differences and Symbolic Calculus¶

For $f : \mathbb{R}^p \rightarrow \mathbb{R}$, we want to evaluate $\nabla f : \mathbb{R}^p \mapsto \mathbb{R}^p$. A finite-difference approximation with a small parameter $\varepsilon>0$ is

This requires $p+1$ evaluations of $f$. Here $\delta_k=(0,\ldots,0,1,0,\ldots,0)$ is the coordinate vector with its 1 at position $k$.

For large $p$, repeated evaluations are expensive. Finite differences also balance truncation error against floating-point cancellation as $\varepsilon$ decreases. Reverse-mode automatic differentiation computes the full gradient at a cost typically proportional to one evaluation of $f$.

Like symbolic differentiation, automatic differentiation applies exact derivative rules, subject to floating-point error. It organizes those rules around the operations of the program that evaluates the function, preserving reuse of intermediate computations.

15.3.2 Computational Graphs¶

Let $x=(x_1,\ldots,x_s)$ be the inputs to a program evaluating $f$. Number its intermediate variables $x_{s+1},\ldots,x_t$, with output $x_t=f(x)$, and let $x_k\in\mathbb{R}^{n_k}$. We seek the Jacobian blocks $\frac{ \partial f(x)}{\partial x_k} \in\mathbb{R}^{n_t\times n_k}$ for $k=1,\ldots,s$. The scalar case $n_k=1$ is easiest to follow; the formulas also apply to vector variables if the matrix products retain the displayed order. The total input dimension is $p=\sum_{k=1}^s n_k$.

Represent the computation as a sequence of elementary maps:

Each $f_k$ depends only on earlier variables; see Figure 15.8. The directed acyclic graph (DAG) connects each argument used by $f_k$ to its output $x_k$. The indexing is a topological order: every edge points from a smaller to a larger index.

Evaluating $f(x)$ is a forward traversal of the graph.

Automatic differentiation differentiates the supplied computational graph. It does not by itself choose an efficient implementation of the original formula; the derivative computation benefits from the efficiency and reuse already present in the forward program.

15.3.3 Forward Mode of Automatic Differentiation¶

Forward mode propagates derivatives of each intermediate variable with respect to a chosen input block, here $x_1$. For scalar inputs, repeating the traversal for all $p$ coordinates gives the full derivative. The formulas below also propagate a block Jacobian when $n_1>1$.

Initialize the input derivatives by

These are one and zero for scalar inputs. Propagate the derivatives using

The notation “$\text{parent}(k)$” denotes the nodes $\ell<k$ of the graph that are connected to $k$, see Figure 15.9, left.

The stored quantities are the derivatives $\frac{ \partial x_\ell}{\partial x_1}$. The symbol $\times$ denotes matrix multiplication in the vector-valued case.

The brackets $[\ldots]$ indicate a local derivative: $\frac{ \partial x_k}{\partial x_\ell}$ is the derivative of the elementary map $f_k$ with respect to that argument. It can be evaluated when needed, assuming the local derivative is available in closed form.

Suppose the local derivative operations $\frac{ \partial f_k}{\partial x_\ell}$ have comparable costs, as when the variable dimensions $n_k$ are uniformly bounded, and each node has a bounded number of parents. Computing the full Jacobian coordinate by coordinate requires $p$ traversals, each with cost comparable to a forward evaluation. This becomes expensive for large $p$.

Simple example.

We consider the function

on the open domain $x>0$, $y\log x>0$, where all elementary derivatives exist. Its computational graph is displayed in Figure 15.10. The iterations of the forward mode to compute the derivative with respect to $x$ read

A second forward pass gives the derivative with respect to $y$:

Dual numbers.

Dual numbers provide a convenient implementation of the forward pass. They form a real algebra whose elements have the form $x + \varepsilon x'$ where $\varepsilon$ is a symbol obeying the rule that $\varepsilon^2=0$. Here $(x,x') \in \mathbb{R}^2$ and $x'$ is intended to store a derivative with respect to some input variable. Multiplication and, when $x\neq0$, inversion are given by

For polynomial or rational $f$, the arithmetic rules imply

For another differentiable elementary function $f$, define its action on dual numbers by

This definition gives

which is the chain rule. More generally, for $f(x_1,\ldots,x_s)$ built from these elementary operations,

the dual-number evaluation computes $\frac{ \partial f}{\partial x_1} (x_1,\ldots,x_s)$ by forward-mode differentiation. The other coordinates are treated identically.

15.3.4 Reverse Mode of Automatic Differentiation¶

Reverse mode computes the derivatives $\frac{ \partial x_t}{\partial x_k}$ of the output with respect to the intermediate variables. For a scalar output and many input coordinates, this avoids the repeated forward traversals required to compute each $\frac{ \partial x_k}{\partial x_i}$ separately.

Initialize the derivative of the output node by

and apply the following recursion from the last node to the first:

The set $\text{son}(k)$ contains the children $m>k$ that depend directly on $x_k$, as shown in Figure 15.9, right. Contributions from every child must be added.

Back-propagation.

For a scalar output $x_t \in \mathbb{R}$, we have $\frac{ \partial x_t}{\partial x_k} = [\nabla_{x_k} f(x)]^\top \in \mathbb{R}^{1 \times n_k}$. The reverse recursion can then be written for gradient vectors:

Here $\left( \frac{ \partial f_m(x_1,\ldots,x_{m-1})}{\partial x_k} \right)^\top \in \mathbb{R}^{n_k \times n_m}$ is the adjoint Jacobian of $f_m$. This adjoint recursion is called backpropagation and is the standard form used in machine learning.

For a scalar output ($n_t=1$), reverse mode is usually the preferred way to compute the full gradient. It requires access to the intermediate variables $(x_k)_{k=s+1}^t$, which can be costly to store for a large graph. Checkpointing trades additional computation for reduced storage.

Simple example.

For the two-input example $f(x,y)$ in (15.19), reverse mode gives

A single reverse traversal computes both derivatives with respect to $x$ and $y$, whereas coordinatewise forward mode requires two passes.

15.3.5 Feed-forward Compositions¶

For a feedforward chain, the computed function is the composition

for functions $f_k : \mathbb{R}^{n_{k-1}} \rightarrow \mathbb{R}^{n_{k}}$.

The forward evaluation starts from $x_0=x \in \mathbb{R}^{n_0}$ and computes

giving the output $f(x) = x_t$.

Let $A_k \mathrel{:=}\partial f_k(x_{k-1}) \in \mathbb{R}^{n_{k} \times n_{k-1}}$ be the Jacobian of the local map. The chain rule gives

Forward mode multiplies the Jacobians from right to left, while reverse mode multiplies them from left to right

We note that the computation of the product $A \times B$ of $A \in \mathbb{R}^{n \times p}$ with $B \in \mathbb{R}^{p \times q}$ requires $O(npq)$ arithmetic operations with the standard dense algorithm.

As shown on Figure 15.11, the arithmetic costs of the two parenthesizations are proportional to

For roughly comparable intermediate widths, these costs favor reverse mode when $n_t\ll n_0$. A scalar learning objective typically has $n_t=1$, making reverse mode a natural choice; the exact cost comparison depends on all the intermediate dimensions.

15.3.6 Feed-forward Architecture¶

A parameterized feedforward architecture, such as a neural network, has the form

where $\theta_{k-1}$ is the parameter vector for a layer and the input $x_0 \in \mathbb{R}^{n_0}$ is fixed. The objective is

Here $L : \mathbb{R}^{n_t} \rightarrow \mathbb{R}$ is a loss, such as squared prediction error or logistic loss, and $\theta=(\theta_k)_{k=0}^{t-1}$ collects the parameters. Figure 15.12 shows the computational graph.

Reverse mode computes the gradient of $f$ by propagating derivatives through the variables $(x_k,\theta_k)$. Initialize

and recurse from $k=t$ down to 1:

where $z_k \mathrel{:=}\nabla_{x_{k}} f(\theta)$ is the gradient with respect to $x_{k}$.

Multilayer perceptrons.

For a fully connected network without biases, use

where $\theta_{k-1}\in\mathbb{R}^{n_k\times n_{k-1}}$ contains the weights and $\rho$ acts coordinatewise; see Figure 15.13. The derivative formulas below apply where $\rho$ is differentiable. At a ReLU kink, a chosen derivative convention defines the computed update, but the classical derivative need not exist.

For an adjoint vector $z_k \in \mathbb{R}^{n_k}$, typically $\nabla_{x_{k}} f$,

Link with adjoint state method.

A residual form of (15.21) can be interpreted as a time discretization of an ordinary differential equation.

Keep the state dimension fixed, $n_k=n$, and interpret $x_k$ as an approximation of a continuous trajectory $x(t)$ at time $t=k\tau$. The formal continuum limit takes $\tau\downarrow0$ with $k\tau$ approaching a fixed time.

Impose the residual structure

with a parameterized vector field $u(x,\theta,t) \in \mathbb{R}^n$. Formally, as $\tau \rightarrow 0$, this approaches the nonlinear ODE

with $x(t=0)=x_0$.

Writing $z(t) = \nabla_{x(t)} f(\theta)$ for the adjoint, the discrete relations (15.23) become a linear backward ODE:

The adjoint has terminal condition $z(T)=\nabla L(x(T))$. The parameter derivative is interpreted as the $L^2$ functional gradient with respect to the time-dependent control. Its discrete normalization is $\frac{1}{\tau} \nabla_{\theta_{k-1}} f \rightarrow \nabla_{\theta(t)} f(\theta)$.

15.3.7 Recurrent Architectures¶

A recurrent architecture reuses the parameter $\theta=\theta_k$ and map $f_k=h$ at every step of (15.21):