The Elements of Differentiable Programming

Chapter 16First-order optimization

16.1 Gradient descent

Gradient descent is one of the simplest algorithms in our toolbox to minimize a function. At each iteration, it moves along the negative gradient direction, scaled by a stepsize γ\gamma: wt+1=wt−γ∇L(wt).(16.1) \w \pow {t+1} = \w \pow t - \gamma \nabla L(\w\pow t). \tag{16.1} The path taken by a gradient descent on a simple quadratic is illustrated in Fig. 16.1 for different choices of the stepsize.

Figure 16.1: Trajectory taken by gradient descent on the objective f(w)≔0.05w12+0.5w22f(w) \coloneqq 0.05w_1^2 + 0.5 w_2^2 with a small (left) or large (right) stepsize. In each case, the iterates follow the normal vectors to the contour lines (dashed lines): the negative gradients. A small stepsize leads to slow convergence but a larger stepsize induces oscillations.

16.1.1 Variational perspective

Consider the linear approximation of L(w)L(\w) around wt\w \pow t, L(w)≈L(wt)+⟨∇L(wt),w−wt⟩. L(\w) \approx L(\w \pow t) + \langle \nabla L(\w \pow t), \w - \w \pow t \rangle. One can easily check that the gradient descent update in Eq. (16.1) can be rewritten as the solution of a minimization problem, namely, wt+1=arg min⁡w∈WL(wt)+⟨∇L(wt),w−wt⟩+12γ∥w−wt∥22.(16.2) \w \pow {t+1} = \argmin_{\w \in \cW} L(\w \pow t) + \langle \nabla L(\w \pow t), \w - \w \pow t \rangle + \frac{1}{2\gamma} \|\w - \w \pow t\|_2^2. \tag{16.2} In words, a gradient descent update optimizes a trade-off between staying close to the current wt\w \pow t, thanks to the proximity term 12γ∥w−wt∥22\frac{1}{2\gamma} \|\w - \w \pow t\|_2^2, and minimizing the linearization of LL around wt\w \pow t. Intuitively, by choosing γ\gamma sufficiently small, we ensure that the minimizer of the regularized linear approximation stays in a neighborhood where the linear approximation is valid. This viewpoint is useful to motivate gradient descent extensions.

16.1.2 Convergence for smooth functions

As long as ∇L(wt)≠0\nabla L(\w\pow t) \neq \zeros, the function Lt(γ)≔L(wt−γ∇L(wt))L_t(\gamma) \coloneqq L(\w \pow t - \gamma \nabla L(\w\pow t)) has a negative derivative at 0, i.e., Lt′(0)=−∥∇L(wt)∥22L'_t(0)=-\|\nabla L(\w\pow t)\|_2^2. Hence, as long as ∇L(wt)≠0\nabla L(\w\pow t) \neq \zeros, there exists a stepsize ensuring a decrease in objective values at each iterate. However, without further assumptions, such a stepsize may depend on each iterate and may be infinitesimally small. To quantify the convergence of gradient descent with a constant stepsize, we restrict to the class of smooth functions. By applying Proposition 15.1 on the iterate of gradient descent, we obtain that L(wt+1)≤L(wt)−γ∥∇L(wt)∥22+βγ22∥∇L(wt)∥22. L(\w\pow {t+1}) \leq L(\w\pow t) - \gamma \|\nabla L(\w\pow t)\|_2^2 + \frac{\beta\gamma^2}{2} \|\nabla L(\w \pow t)\|_2^2. Therefore, for β\beta-smooth functions, by selecting γ≤1β\gamma \leq \frac{1}{\beta}, we get that L(wt+1)−L(wt)≤−γ2∥∇L(wt)∥22, L(\w\pow {t+1}) - L(\w\pow t) \leq - \frac{\gamma}{2}\|\nabla L(\w\pow t)\|_2^2, which illustrates the main mechanism behind gradient descent: each iteration decreases the objective by a constant times the norm of the gradient of the current iterate. This equation can further be summed over all iterates up to TT. This telescopes the objective values, leading to min⁡t∈{0,…,T−1}∥∇L(wt)∥22≤1T∑t=0T−1∥∇L(wt)∥22≤2γT(L(w0)−L(wT))≤2γT(L(w0)−L⋆),\begin{align*} \min_{t\in \{0, \ldots, T-1\}} \|\nabla L(\w \pow t)\|_2^2 & \leq \frac{1}{T}\sum_{t=0}^{T-1} \|\nabla L(\w\pow t)\|_2^2 \\ & \leq \frac{2}{\gamma T}\left(L(\w \pow 0) - L(\w \pow {T})\right) \\ & \leq \frac{2}{\gamma T}\left(L(\w \pow 0) - L^\star\right), \end{align*} where we recall that L⋆L^\star is the infimum of LL. Therefore, after sufficiently many iterations, gradient descent finds a point whose gradient norm is arbitrarily small.

Non-convex case

Without further assumptions, i.e., in the non-convex case, the above result (i.e., convergence to a stationary point, measured by the gradient norm) is the best we may get in theory. Denoting Ts(ε)T_s(\varepsilon) the number of iterations needed for a gradient descent to output a point that is ε\varepsilon-stationary, i.e., ∥∇L(w^)∥2≤ε\|\nabla L(\hat \w)\|_2\leq \varepsilon, we have Ts(ε)≤O(ε−2)T_s(\varepsilon) \leq O(\varepsilon^{-2}).

Convex case

By adding a convexity assumption on the objective, we can use the lower bound provided by the convexity assumption to ensure convergence to a minimum. Namely, for a β\beta-smooth and convex function ff, and with stepsize γ≤1/β\gamma \leq 1/\beta, we have that (Nesterov, 2018) L(wT)−L⋆≤1γT∥w0−w⋆∥22. L(\w \pow T) - L^\star \leq \frac{1}{\gamma T}\|\w \pow 0 - \w^\star\|_2^2. That is, we get a sublinear convergence rate, and the associated computational complexity to find a minimum is T(ε)=O(1/ε)T(\varepsilon)= O(1/\varepsilon).

Strongly convex case

If we further strengthen the assumptions by considering β\beta-smooth, μ\mu-strongly convex functions, the convergence rate of a gradient descent can be shown to be (Nesterov, 2018), for any stepsize γ≤1/β\gamma \leq 1/\beta, L(wT)−L⋆≤(1−γμ)T(L(w0)−L⋆)≤exp⁡(−γμT)(L(w0)−L⋆).\begin{align*} L(\w \pow T) - L^\star & \leq \left(1 - \gamma \mu\right)^T \left(L(\w\pow 0) - L^\star\right) \\ & \leq \exp\left(- \gamma \mu T\right) \left(L(\w\pow 0) - L^\star\right). \end{align*} That is, we obtain a linear convergence rate and the associated computational complexity is T(ε)=O(ln⁡ε−1)T(\varepsilon) = O(\ln\varepsilon^{-1}). The above convergence rates may be further refined (Nesterov, 2018); we focused above on the simplest result for clarity.

Strong convexity can also be replaced by a weaker assumption, gradient-dominating property (Polyak, 1963), i.e., ∥∇L(v)∥22≥c(L(v)−L⋆)\|\nabla L(\v)\|_2^2 \geq c(L(\v) - L^\star) for some constant cc and any v∈W\v \in \cW. A convex, gradient-dominating function can also be minimized at a linear rate.

16.1.3 Momentum and accelerated variants

We started with gradient descent as a simple example of first-order optimization algorithm. However, different optimization algorithms can be designed from the access to first-order oracles and the knowledge of the class of functions considered. For example, consider quadratic convex functions w↦12w⊤Aw+b⊤w\w \mapsto \frac{1}{2} \w^\top \A \w + \b^\top \w, that are a basic example of smooth strongly convex functions if A\A is positive definite. An optimal method in this case is the heavy-ball method of Polyak (1964), that can be written as vt+1≔νvt−γ∇L(wt)wt+1≔wt+vt+1.\begin{align*} \v \pow {t+1} &\coloneqq \nu \v \pow t - \gamma \nabla L(\w \pow t) \\ \w \pow {t+1} &\coloneqq \w \pow t + \v \pow {t+1}. \end{align*} The heavy-ball method uses an additional variable vt\v\pow t, that can be interpreted as the velocity of a ball driven by the negative gradient to converge towards a minimum. Intuitively, this additional velocity circumvents the oscillations that a gradient descent may present as illustrated in Fig. 16.2 compared to Fig. 16.1. For ν=0\nu=0, we recover usual gradient descent. For ν>0\nu>0, the velocities accumulate a form of an inertia momentum, where ν\nu is interpreted as the “mass” of the ball. In terms of convergence rates, the heavy-ball method can be shown to converge linearly similarly to gradient descent, but with a rate O(exp⁡(−Tμ/β))O(\exp(- T \sqrt{\mu /\beta})) for appropriate choices of ν,γ\nu, \gamma. In comparison, by choosing an optimal stepsize for the gradient descent, its convergence rate is O(exp⁡(−Tμ/β))O(\exp(-T\mu/\beta)) which is provably worse, as we always have μ/β≤1\mu/\beta \leq 1.

Beyond the case of quadratic functions, accelerated variants of gradient descent for convex or strongly convex functions have been developed by Nesterov (2018). Such variants have inspired the design of optimization algorithms in stochastic settings presented below.

Figure 16.2: Trajectory taken by gradient descent with momentum. Compared to gradient descent without momentum, for the same stepsize, the oscillations previously observed in Fig. 16.1 are no longer present, and the algorithm converges faster to the minimum.

16.2 Stochastic gradient descent

In machine learning, we are usually interested in minimizing the expected loss of the model over the data distribution ρ\rho: min⁡w∈W L(w)≔ES∼ρ[L(w;S)]. \min_{\w \in \cW} \ L(\w) \coloneqq \mathbb{E}_{S \sim \rho} \left[L(\w; S)\right]. For example, LL is often set to L(w;S)≔ℓ(Y,f(X,w))L(\w; S) \coloneqq \ell(Y, f(X, \w)), where ℓ\ell is a loss function, ff is a neural network and S=(X,Y)S=(X, Y) is a random pair, composed of an input XX and an associated target YY, sampled from ρ\rho. In this setting, since the data distribution ρ\rho is generally unknown and may be infinite, we cannot exactly evaluate the expected loss L(w)L(\w) or its gradient ∇L(w)\nabla L(\w).

In practice, we are often given a fixed dataset of nn pairs si=(xi,yi)\s_i = (\x_i, \y_i). This is a special case of the expected loss setting, since this can be seen as an empirical distribution ρ=ρn\rho = \rho_n L(w)=ES∼ρn[L(w;S)]=1n∑i=1nL(w;(Xi,Yi)). L(\w) = \mathbb{E}_{S \sim \rho_n} \left[L(\w; S)\right] = \frac{1}{n} \sum_{i=1}^n L(\w; (X_i, Y_i)). The gradient of L(w)L(\w) is then ∇L(w)≔1n∑i=1n∇L(w;(xi,yi)). \nabla L(\w) \coloneqq \frac{1}{n} \sum_{i=1}^n \nabla L(\w; (\x_i, \y_i)). In this case, we see that the full gradient ∇L(w)\nabla L(\w), as needed by gradient descent, is the average of the individual gradients. That is, the cost of computing ∇L(w)\nabla L(\w) is proportional to the number of training points nn. For nn very large, that is a very large amount of samples, this computational cost can be prohibitive. Stochastic gradients circumvent this issue.

16.2.1 Stochastic gradients

Usually, even if we do not know ρ\rho, we can sample from it, i.e., we have access to samples S∼ρS \sim \rho. We can then use a stochastic gradient of the form ∇L(w;S) \nabla L(\w; S) as a random estimate of ∇L(w)\nabla L(\w). This may look like a rough estimate but, on average, this is a valid approximation since ES∼ρ[∇L(w;S)]=∇L(w). \mathbb{E}_{S \sim \rho}\left[ \nabla L(\w; S)\right] = \nabla L(\w). We say that ∇L(w;S)\nabla L(\w; S) is an unbiased estimator of ∇L(w)\nabla L(\w). To further improve the approximation, we may also consider mini-batch estimates by sampling m≪nm \ll n data points Si≔(Xi,Yi)S_i \coloneqq (X_i, Y_i) and using 1m∑i=1m∇L(w;Si)\frac{1}{m} \sum_{i=1}^m \nabla L(\w;S_i), whose expectation still matches ∇L(w)\nabla L(\w), while potentially reducing the approximation error by averaging multiple stochastic gradients. Computationally, the main advantage is that the cost is now proportional to mm instead of nn.

In whole generality, one can consider stochastic first-order oracles defined below.

Definition 16.1 (Stochastic first-order oracles)

A stochastic first-order oracle of an expected objective L(w)L(\w) is a random estimate g(w;S)g(\w; S) of ∇L(w)\nabla L(\w) with SS sampled according to some distribution qq. A stochastic gradient is said to be an unbiased estimator if ES∼q[g(w;S)]=∇L(w). \mathbb{E}_{S \sim q}\left[g(\w; S)\right] = \nabla L(\w). The variance of a stochastic gradient is ES∼q[∥g(w;S)−∇L(w)∥22]. \mathbb{E}_{S \sim q}\left[\left\| g(\w; S) - \nabla L(\w) \right\|_2^2\right].

When q=ρq = \rho, we recover stochastic gradients. When qq is the product of mm independent samples according to pp, we recover mini-batch stochastic gradients. First-order stochastic optimization algorithms build upon stochastic first-order oracles to approximately find the minimum of the expected objective. In such a setting, the iterates of the algorithm are by definition random. Convergence rates therefore need to be expressed in probabilistic terms by considering for example the expected objective value according to the randomness of the oracles.

16.2.2 Vanilla SGD

Equipped with a stochastic first-order oracle, such as (mini-batch) stochastic gradients, we can define stochastic gradient descent as wt+1=wt−γg(wt;St)where St∼q. \w \pow {t+1} = \w \pow t - \gamma g(\w \pow t; S \pow t) \quad \mbox{where} \ S \pow t \sim q. We assume that StS \pow t is independent of wt\w \pow t. Compared to the usual gradient descent, the main impediment of the stochastic setting is the additional noise induced by the stochastic estimates: their variance.

For example, consider applying a stochastic gradient descent on the expectation of β\beta-smooth convex functions L(w;s)L(\w; \s) with unbiased oracles. To harness the randomness of the iterates, consider after TT iterations outputting the average of the first TT iterates, that is wˉT≔1T∑t=1Twt\bar \w \pow T \coloneqq \frac{1}{T} \sum_{t=1}^T \w\pow t. Moreover, suppose that the variance of the stochastic first-order oracles is bounded by σ2\sigma^2 for all minimizers w⋆\w^\star of LL. Denoting by ES0,…,ST−1\mathbb{E}_{S_0, \ldots, S_{T-1}} the randomness associated to the stochastic oracles, we have then that for a stepsize γ≤1/(4β)\gamma \leq 1/(4\beta)  (Lan, 2012), ES0,…,ST−1[L(wˉT)]−L⋆≤1γT∥w0−w⋆∥22+2γσ2. \mathbb{E}_{S_0, \ldots, S_{T-1}}[ L(\bar \w \pow T)] - L^\star \\ \leq \frac{1}{\gamma T}\|\w_0 - \w^\star\|_2^2 + 2\gamma \sigma^2. The resulting convergence rate illustrates that a stochastic gradient descent converges to the minimum of the expected objective up to a constant term depending on the variance of the oracle and the stepsize. One can diminish the variance by considering mini-batches: if the variance of a single stochastic gradient is σ12\sigma_1^2, considering a mini-batch of mm gradients reduces the variance of the corresponding oracle to σm2=σ12/m\sigma^2_m = \sigma_1^2/m. To decrease the additional term, one may also decrease the stepsizes over the iterations. For example, by choosing a decreasing stepsize like γt=t−1/2\gamma^t = t^{-1/2}, the convergence rate is then of the order O((∥w0−w⋆∥22+σ2ln⁡t)/t)O((\|\w_0-\w^\star\|_2^2 + \sigma^2\ln t)/\sqrt{t}). The stepsize can also be selected as a constant γ0\gamma_0 that decreases the average objective for the first T0T_0 iterations and is reduced by a multiplicative factor at regular intervals like γj=ργj−1\gamma_j = \rho \gamma_{j-1} for ρ∈(0,1)\rho \in (0, 1) to handle iterations between Tj,Tj+1T_j, T_{j+1}. Alternative stepsize schedules such as a cosine decay (Loshchilov and Hutter, 2016) have recently become popular.

The literature on alternative optimization schemes for stochastic optimization is still rapidly evolving, with new heuristics regularly proposed. We present below two popular techniques.

16.2.3 Momentum variants

Accelerated optimization algorithms developed in the deterministic setting may be extended to the stochastic setting. For example, the heavy-ball method can be adapted to the stochastic setting, leading to stochastic gradient descent with momentum (Sutskever et al., 2013) generally implemented as vt+1≔νvt+g(wt;St)wt+1≔wt−γvt+1.\begin{align*} \v \pow {t+1} &\coloneqq \nu \v\pow t + \g(\w \pow t; S \pow t) \\ \w \pow {t+1} &\coloneqq \w \pow t - \gamma \v \pow {t+1}. \end{align*} As mentioned earlier the momentum method can be modified to handle non-quadratic smooth strongly convex functions. This leads to Nesterov’s accelerated method in the deterministic setting. This has been adapted to the stochastic setting with a so-called Nesterov momentum (Sutskever et al., 2013) vt+1≔νvt+g(wt−γνvt;St)wt+1≔wt−γvt+1.\begin{align*} \v \pow {t+1} &\coloneqq \nu \v\pow t + \g(\w \pow t - \gamma \nu \v\pow t; S \pow t) \\ \w \pow {t+1} &\coloneqq \w \pow t - \gamma \v \pow {t+1}. \end{align*}

16.2.4 Adaptive variants

In any gradient descent-like algorithm, selecting the stepsize is key for good performance. While a constant stepsize may be used if the function is smooth, we may not know in advance the smoothness constant of the objective, which means that additional procedures may be required to select appropriately the stepsize. In the deterministic case, line-searches such as the Armijo or Wolfe’s rules (Wright and Nocedal, 1999) can be used to check whether the selected stepsize sufficiently decreases the objective at each iteration. Such rules have been adapted in the stochastic setting (Vaswani et al., 2019).

Another way to decrease the sensitivity of the algorithm with respect to the stepsize has been to estimate first- and second-order moments of the gradients and use the latter as a form of preconditioning to smooth the trajectory of the iterates. This led to the popular Adam optimizer (Kingma and Ba, 2014). It takes the form, mt+1≔ν1mt+(1−ν1)gtvt+1≔ν2vt+(1−ν2)(gt)2m^t≔mt/(1−ν1t)v^t≔vt/(1−ν2t)wt+1≔wt−γ m^t+1/(v^t+1+ε),\begin{align*} \m\pow {t+1} &\coloneqq \nu_1 \m \pow t + (1-\nu_1) \g \pow t \\ \v\pow {t+1} &\coloneqq \nu_2 \v \pow t + (1-\nu_2) (\g\pow t)^2 \\ \hat{\m}\pow t &\coloneqq \m\pow t / {(1-\nu_1^t)} \\ \hat{\v}\pow t &\coloneqq \v\pow t / {(1-\nu_2^t)} \\ \w \pow {t+1} &\coloneqq \w \pow t - \gamma \ \hat{\m}\pow {t+1} / \left({\sqrt{\hat{\v}\pow {t+1}} + \varepsilon} \right), \end{align*} where gt≔g(wt;St)\g \pow t \coloneqq \g(\w\pow t; S \pow t), (gt)2(\g\pow t)^2 denotes the element-wise square of gt\g \pow t and ν1,ν2,γ,ε\nu_1, \nu_2, \gamma, \varepsilon are hyper-parameters of the algorithm. Numerous variants exist, such as varying the stepsize γ\gamma above along the iterations.

16.3 Projected gradient descent

Oftentimes, we seek to find the solution of a minimization problem subject to constraints on the variables, of the form min⁡w∈CL(w),(16.3) \min_{\w \in \cC} L(\w), \tag{16.3} where C⊆W=RP\cC \subseteq \cW = \RR^P is a set of constraints. We say that an approximate solution w^\widehat \w to Eq. (16.3) is feasible if w^∈C\widehat \w \in \cC. Naturally, the design of algorithms for the constrained setting now depends, not only on information about LL, but also on information about C\cC.

Similarly to LL, different oracles can be considered about C\cC. One of the most commonly used oracles is the Euclidean projection

Definition 16.2 (Euclidean projection)

The Euclidean projection onto the set C\cC is defined by projC(w)≔arg min⁡v∈C∥w−v∥22. \mathrm{proj}_{\cC}(\w) \coloneqq \argmin_{\v \in \cC} \|\w-\v\|_2^2.

This projection, which is well-defined when C\cC is a convex set, can be used in projected gradient descent, that we briefly review below. Typically, the projection on a particular set C\cC requires a dedicated algorithm to compute it.

Other possible oracles are linear maximization oracles (LMO) used in Frank-Wolfe algorithms and Bregman projection oracles, used in mirror descent algorithms. The algorithm choice can be dictated by what oracle about C\cC is available.

16.3.1 Variational perspective

Projected gradient descent is a natural generalization of gradient descent, based on the Euclidean projection oracle. Its iterates read wt+1≔projC(wt−γ∇L(wt)). \w \pow {t+1} \coloneqq \mathrm{proj}_{\cC}(\w \pow t - \gamma \nabla L(\w \pow t)). At each iteration, we attempt to decrease the objective by moving along the negative gradient direction, while ensuring that the next iterate remains feasible, thanks to the projection step.

Similarly to the variational perspective of gradient descent in Eq. (16.2), the projected gradient descent update is equivalent to wt+1=arg min⁡w∈CL(wt)+⟨∇L(wt),w−wt⟩+12γ∥w−wt∥22. \w \pow {t+1} = \argmin_{\w \in \cC} L(\w \pow t) + \langle \nabla L(\w \pow t), \w - \w \pow t \rangle + \frac{1}{2\gamma} \|\w - \w \pow t\|_2^2. This shows that projected gradient descent minimizes a trade-off between staying close to wt\w \pow t and minimizing the linearization of LL around wt\w \pow t, while staying in C\cC.

In terms of convergence rates, they remain the same as gradient descent (Nesterov, 2018). For example, projected gradient descent on a smooth convex function still converges at a rate R(T)=O(1/T)R(T) = O(1/T).

There are numerous extensions of vanilla projected gradient descent. Similarly to gradient descent, the stepsize can be automatically adjusted using linesearch techniques and there exists accelerated variants. If we replace ∇L(w)\nabla L(\w) with a stochastic gradient ∇L(w;S)\nabla L(\w; S), we obtain a stochastic projected gradient descent.

16.3.2 Optimality conditions

In the unconstrained case, a minimum necessarily has a zero gradient. In the constrained setting, there may not be any feasible parameters with zero gradient. Instead, the optimality of a point is characterized by the fact that no better solution can be found by moving along the gradient at that point, while staying in the constraints. Formally, it means that for any γ>0\gamma > 0, a minimizer w⋆\w^\star of LL on C\cC satisfies w⋆=projC(w⋆−γ∇L(w⋆)). \w^\star = \mathrm{proj}_{\cC}(\w^\star - \gamma \nabla L(\w^\star)). It can be shown that this condition is equivalent (Nesterov, 2018) to ⟨∇L(w⋆),w−w⋆⟩≥0∀w∈C. \langle \nabla L(\w^\star), \w - \w^\star \rangle \geq 0 \quad \forall \w \in {\cC}.

16.3.3 Commonly used projections

We now briefly review a few useful Euclidean projections.

16.4 Proximal gradient method

The constrained setting (with C\cC a convex set) can be recast as unconstrained optimization, by extending our analysis to functions taking infinite values. Let us denote the indicator function of the set C\cC by ιC(w)≔{0if w∈C+∞otherwise. \indic_{\cC}(\w) \coloneqq \begin{cases} 0 &\mbox{if } \w \in \cC \\ +\infty &\mbox{otherwise} \end{cases}. Clearly, the constrained problem in Eq. (16.3) can then be rewritten as min⁡w∈WL(w)+ιC(w). \min_{\w \in \cW} L(\w) + \indic_{\cC}(\w). This suggests that constrained optimization is a special case of composite objectives of the form min⁡w∈WL(w)+Ω(w), \min_{\w \in \cW} L(\w) + \Omega(\w), where Ω\Omega is a convex but potentially non-differentiable function. We assume that we have access to an oracle associated with Ω\Omega called the proximal operator.

Definition 16.3 (Proximal operator)

The proximal operator associated with Ω ⁣:W→R\Omega \colon \cW \to \RR is proxΩ(w)≔arg min⁡v∈W12∥w−v∥22+Ω(v). \mathrm{prox}_{\Omega}(\w) \coloneqq \argmin_{\v \in \cW} \frac{1}{2} \|\w -\v\|_2^2 + \Omega(\v).

This leads to the proximal gradient method, reviewed below.

16.4.1 Variational perspective

With this method, the update reads wt+1=proxγΩ(wt−γ∇L(wt)). \w \pow {t+1} = \mathrm{prox}_{\gamma \Omega}(\w \pow t - \gamma \nabla L(\w \pow t)). This update again enjoys an intuitive variational perspective, namely, wt+1=arg min⁡w∈WL(wt)+⟨∇L(wt),w−wt⟩+12γ∥w−wt∥22+Ω(w). \w \pow {t+1} = \argmin_{\w \in \cW} L(\w \pow t) + \langle \nabla L(\w \pow t), \w - \w \pow t \rangle + \frac{1}{2\gamma} \|\w - \w \pow t\|_2^2 + \Omega(\w). That is, we linearize LL around wt\w \pow t, but keep Ω\Omega as is.

The proximal gradient method is popularly used when the objective function contains a sparsity-inducing regularizer Ω\Omega. For example, for the LASSO (Tibshirani, 1996), which aims at predicting targets y=(y1,…,yn)⊤∈RN\y = (y_1, \ldots, y_n)^\top \in \RR^N from observations X=(x1,…,xn)⊤∈RN×P\X = (\x_1, \ldots, \x_n)^\top\in \RR^{N\times P}, we set L(w)=12∥Xw−y∥22L(\w) = \frac{1}{2}\|\X \w - \y\|_2^2 and Ω(w)=λ∥w∥1\Omega(\w) = \lambda \|\w\|_1, where λ>0\lambda > 0 controls the regularization strength. In this case, proxΩ\mathrm{prox}_{\Omega} is the so-called soft-thresholding operator (see below).

Convergence guarantees of the proximal gradient method remain the same as for gradient descent, such as a O(1/T)O(1/T) rate for smooth convex functions.

16.4.2 Optimality conditions

An optimal solution of the problem is characterized by the fixed point equation w⋆=proxγΩ(w⋆−γ∇L(w⋆)), \w^\star = \mathrm{prox}_{\gamma \Omega}(\w^\star - \gamma \nabla L(\w^\star)), for all γ>0\gamma > 0 (Nesterov, 2018). In other words, the proximal gradient method (which includes gradient descent and projected gradient descent as special cases), can be seen as fixed point iteration schemes. Such a viewpoint suggests using acceleration methods from the fixed point literature such as Anderson acceleration (Pollock and Rebholz, 2021). It is also useful when designing implicit differentiation schemes as presented in Chapter 8.

16.4.3 Commonly-used proximal operators

We now briefly review a few useful proximal operators.

For a review of more proximal operators, see for instance (Bach et al., 2012; Parikh et al., 2014).

16.5 Summary