The Elements of Differentiable Programming

Chapter 17Second-order optimization

We review in this chapter methods whose iterations take the form wt+1≔wt−γtBt∇L(wt), \w^{t+1} \coloneqq \w^t - \gamma^t B^t \nabla L(\w^t), where γt\gamma^t is a stepsize and BtB^t is a pre-conditioning matrix involving second-order derivatives.

17.1 Newton’s method

17.1.1 Variational perspective

We saw in Eq. (16.2) that gradient descent can be motivated from a variational perspective, in which we use a linear approximation of the objective around the current iterate, obtained from the current gradient. Similarly, if we have access not only to the gradient but also to the Hessian of the objective, we can use a quadratic approximation of the objective around the current iterate. More precisely, given a function L(w)L(\w), we may consider minimizing the second-order Taylor approximation of L(w)L(\w) around the current iterate wt\w^t, L(w)≈L(wt)+⟨∇L(wt),w−wt⟩+12⟨w−wt,∇2L(wt)(w−wt)⟩. L(\w) \approx L(\w^t) + \langle \nabla L(\w^t), \w-\w^t\rangle + \frac{1}{2}\langle \w-\w^t, \nabla^2 L(\w^t)(\w -\w^t) \rangle. Newton’s method simply iteratively minimizes this quadratic approximation around the current iteration wt\w^t, namely, wt+1=arg min⁡w∈WL(wt)+⟨∇L(wt),w−wt⟩+12⟨w−wt,∇2L(wt)(w−wt)⟩.\begin{align*} \w^{t+1} = \argmin_{\w \in \cW} L(\w^t) + \langle \nabla L(\w^t), \w{-}\w^t\rangle + \frac{1}{2}\langle \w{-}\w^t, \nabla^2 L(\w^t)(\w {-}\w^t) \rangle. \tag{17.1}\end{align*} If the Hessian is positive definite at wt\w^t, which we denote by ∇2L(wt)≻0\nabla^2 L(\w^t) \succ 0, then the minimum is well-defined and unique (this is for example the case if LL is strictly convex). The iterates can then be written analytically as wt+1=wt−∇2L(wt)−1∇L(wt). \w^{t+1} = \w^t - \nabla^2 L(\w^t)^{-1}\nabla L(\w^t). If the Hessian is not positive definite, the minimum may not be defined. Ignoring this issue and taking the analytical formulation could be dangerous, as it could amount to computing the maximum of the quadratic instead if, for example, the quadratic was strictly concave (i.e., ∇2L(w)≺0\nabla^2 L(\w) \prec 0).

17.1.2 Regularized Newton method

A simple technique to circumvent this issue consists in adding a regularization term to the Hessian. Namely, from a variational viewpoint, we can add a proximity term 12∥w−wt∥22\frac{1}{2} \|\w -\w^t\|_2^2, encouraging it to stay close to the current wt\w^t. The iterates of this regularized Newton method then take the form wt+1=arg min⁡w∈W L(wt)+⟨∇L(wt),w−wt⟩+12⟨w−wt,∇2L(wt)(w−wt)⟩+ηt2∥w−wt∥22,\begin{align*} \w^{t+1} = \argmin_{\w \in \cW} ~ &L(\w^t) + \langle \nabla L(\w^t), \w{-}\w^t\rangle + \frac{1}{2}\langle \w{-}\w^t, \nabla^2 L(\w^t)(\w {-}\w^t) \rangle \\ &+ \frac{\eta^t}{2} \|\w-\w^t\|_2^2, \end{align*} where ηt\eta^t controls the regularization strength. Assuming ηt>0\eta^t > 0 is strong enough to make ∇2L(wt)+ηtI⁡\nabla^2 L(\w^t) + \eta^t \idm positive-definite, we have wt+1=wt−dt, \w^{t+1} = \w^t - \d^t, where we defined the direction dt≔(∇2L(wt)+ηtI⁡)−1∇L(wt).(17.2) \d^t \coloneqq (\nabla^2 L(\w^t) + \eta^t \idm)^{-1} \nabla L(\w^t). \tag{17.2} Other techniques to circumvent this issue include using cubic regularization and modifying the spectral decomposition of the Hessian, by thresholding the eigenvalues or taking their absolute values. We refer the interested reader to, e.g., (Nesterov, 2018; Wright and Nocedal, 1999) for more details.

17.1.3 Approximate direction

We observe a main impediment for implementing such a second-order optimization algorithm: even if we had access to the Hessian of the objective for free and this Hessian was positive definite, computing the exact direction dt\d^t in Eq. (17.2) requires computing an inverse-Hessian vector product (IHVP) with the gradient ∇L(wt)\nabla L(\w^t). Doing so exactly requires solving a linear system (∇2L(wt)+ηtI⁡)dt=∇L(wt), (\nabla^2 L(\w^t) + \eta^t \idm) \d^t = \nabla L(\w^t), which a priori takes O(P3)O(P^3) time. In practice, however, we can compute IHVPs approximately, as explained in Section 9.4.

17.1.4 Convergence guarantees

While implementing Newton’s method comes at a higher computational cost, it can also benefit from faster convergence rates. Briefly, if Newton’s method is initialized at a point w0∈W\w^0 \in \cW close enough to the minimizer w⋆\w^\star of a μ\mu-strongly convex function with MM-Lipschitz continuous Hessian (namely ∥w0−w∗∥2≤2μ3M\|\w^0 -\w^*\|_2\leq \frac{2\mu}{3M}), then Newton’s method converges at a quadratic rate (Nesterov, 2018), that is, R(t)≤O(exp⁡(−2t))R(t) \leq O(\exp(-2^t)) (see Section 15.5 for a brief introduction to performance guarantees). This is far superior to gradient descent. Such an efficiency motivated the development of interior point methods, that have been a breakthrough in constrained optimization, thanks to the use of log-barrier penalties (Nesterov, 2018).

17.1.5 Linesearch

In practice, we may not have access to an initial point close enough to the minimizer. In that case, even for strictly convex functions for which Newton’s steps are well-defined, taking wt+1=wt−dt\w^{t+1} = \w^t - \d^t may not ensure a decrease of the objective values. Nevertheless, the direction dt\d^t may define a descent direction as defined below.

Definition 17.1 (Descent direction)

A point d∈W\d \in \cW defines a descent direction −d-\d for an objective LL at w\w, if there exists a positive stepsize γ>0\gamma >0 such that L(w−γd)≤L(w). L(\w - \gamma \d) \leq L(\w). If LL is differentiable, −d-\d is a descent direction if ⟨−d,∇L(w)⟩<0\langle -\d, \nabla L(\w) \rangle <0.

For Newton’s method without regularization, dt=∇2L(wt)−1∇L(wt)\d^t = \nabla^2 L(\w^t)^{-1}\nabla L(\w^t) is then a descent direction at wt\w^t, as long as ∇L(wt)≠0\nabla L(\w^t) \neq 0 and ∇2L(wt)≻0\nabla^2 L(\w^t)\succ0. If ∇2L(wt)⊁0\nabla^2 L(\w^t)\not \succ 0, choosing ηt>0\eta^t >0 such that ∇2L(wt)+ηtI⁡≻0\nabla^2 L(\w^t) + \eta^t \idm \succ 0, also ensures that dt=−(∇2L(wt)+ηtI⁡)−1∇L(wt)\d^t = - (\nabla^2 L(\w^t) +\eta^t\idm)^{-1}\nabla L(\w^t) is a descent direction (as long as ∇L(wt)≠0\nabla L(\w^t) \neq 0). Newton’s method is then generally equipped with a linesearch method that attempts to take steps of the form wt+1=wt−γtdt \w^{t+1} = \w^t - \gamma^t \d^t with γt\gamma^t chosen as the largest stepsize among {ρτ,τ∈N}\{\rho^\tau, \tau \in \mathbb{N}\} for ρ∈(0,1)\rho \in (0, 1) until a sufficient decrease of the objective is satisfied such as, for c∈(0,1)c \in (0, 1), L(wt−γtdt)≤L(wt)−cγt⟨∇L(wt),∇2L(wt)−1∇L(wt)⟩. L(\w^t - \gamma^t \d^t) \leq L(\w^t) - c \gamma^t \langle \nabla L(\w^t), \nabla^2 L(\w^t)^{-1}\nabla L(\w^t) \rangle. For strongly convex functions, such an implementation exhibits two phases: a first phase during which Newton’s steps are “damped” by using a stepsize γt<1\gamma^t <1 and a second phase of super-fast convergence during which stepsizes γt=1\gamma^t = 1 are taken, and the objective decreases very fast. Even far from the optimum, Newton directions can advantageously adapt to the local geometry of the objective to speed up convergence compared to a regular gradient descent as explained below.

17.1.6 Geometric interpretation

To understand the efficiency of Newton’s method compared to gradient descent, consider the minimization of a simple quadratic L(w)=12aw12+12bw22 L(\w) = \frac{1}{2} a w_1^2 + \frac{1}{2} b w_2^2 for a≫b≥0a\gg b \geq 0, as illustrated in Fig. 17.1. A gradient descent moves along the directions ∇L(w)=(aw1,bw2)⊤\nabla L(\w) = (a w_1, b w_2)^\top and its stepsize is limited by the variations in the first coordinate leading to some oscillations. If we were simply rescaling the gradient by (a,b)(a, b), i.e., taking steps of the form wt+1=wt−γdiag⁡(a−1,b−1)∇L(wt), \w^{t+1} = \w^t - \gamma \diag(a^{-1}, b^{-1})\nabla L(\w^t), the variations in both coordinates would be normalized to one and the stepsize could simply be chosen to γ=1\gamma = 1 to directly get w⋆\w^\star. In other words, by adapting the geometry of the directions with the geometry induced by the objective, we can circumvent the oscillations.

That’s exactly what Newton’s method does by modifying the gradient direction using the inverse of the Hessian. Formally, at iteration tt, consider the modified objective L~(v)=L(Av) for A=∇2L(wt)−1/2, \tilde L(\v) = L(A \v) \ \mbox{for} \ A = \nabla^2 L(\w^t)^{-1/2}, with LL strictly convex and AA the inverse matrix square root of the Hessian. One easily verifies that a Newton step is equivalent to a gradient step on L~\tilde L, that is, vt+1=vt−∇L~(vt)   ⟺   wt+1=wt−(∇2L(wt))−1∇L(wt) \v^{t+1} = \v^t - \nabla \tilde L(\v^t) \ \iff \ \w^{t+1} = \w^t - (\nabla^2 L(\w^t))^{-1} \nabla L(\w^t) where wt=Avt=∇2L(wt)−1/2vt. \w^t = A \v^t = \nabla^2 L(\w^t)^{-1/2} \v^t. In the geometry induced by AA, the objective is generally better conditioned as illustrated in Fig. 17.1. This explains the efficiency of Newton’s method. In particular for any strongly convex quadratic, a Newton step reaches the optimum in one iteration, while a gradient step can take many more iterations.

Figure 17.1: Left: Minimization of a quadratic L(w)≔12aw12+12bw22L(\w) \coloneqq \frac{1}{2}a w_1^2 + \frac{1}{2} b w_2^2 by gradient descent. For a≫b≥0a\gg b \geq 0, gradient descent typically oscillates. Right: minimization by Newton’s method amounts to changing the geometry of the problem to avoid oscillations.

17.1.7 Stochastic Newton’s method

Consider now an expected loss 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]. In that case, an estimate of the Hessian can be constructed just like for the gradient using that ES∼ρ[∇2L(w;S)]=∇2L(w). \mathbb{E}_{S \sim \rho} \left[\nabla^2 L(\w; S)\right] = \nabla^2 L(\w). Les us then denote g(w;S)≈∇L(w),H(w;S′)≈∇2L(w) g(\w; S) \approx \nabla L(\w), \quad H(\w; S') \approx \nabla^2 L(\w) some stochastic estimates of the gradient and the Hessian, respectively, with S,S′S, S' independently drawn from ρ\rho or from mini-batch approximations with varying mini-batch sizes. One implementation of a stochastic Newton method can then be wt+1=wt−γt(H(wt;S′)+ηtI⁡)−1g(wt;S), \w^{t+1} = \w^t - \gamma^t (H(\w^t; S') + \eta^t \idm)^{-1}g(\w^t; S), for ηt≥0\eta^t \geq 0 such that (H(wt;S′)+ηtI⁡)−1≻0(H(\w^t; S') + \eta^t \idm)^{-1} \succ 0 and γt\gamma^t fixed or chosen to satisfy some sufficient decrease condition. We refer the interested reader to, e.g., (Xu et al., 2020), for more details and variants.

17.2 Gauss-Newton method

Newton’s method  (17.1) is usually not properly defined for non-convex objective functions, since the Hessian may not be positive definite at the current iterate. We saw in Section 9.2 that the Gauss-Newton matrix can be used to define a positive-semidefinite approximation of the Hessian. Here, we revisit the Gauss-Newton method from a variational and partial linearization perspective. While the original Gauss-Newton method originates from nonlinear least-squares, we will first describe an extension to arbitrary convex loss functions, since it is both more general and easier to explain.

17.2.1 With exact outer function

Consider a composite objective of the form L(w)≔ℓ(f(w)), L(\w) \coloneqq \ell(f(\w)), where ℓ:M→R\ell: \cM \rightarrow \RR is a convex function, such as a convex loss function applied on a given sample, and f:W→Mf: \cW \rightarrow \cM is a nonlinear function, such as a neural network with parameters w∈W\w \in \cW, evaluated on the same sample. We saw that gradient descent and Newton’s method amount to using linear and quadratic approximations of L(w)L(\w) around the current iterate wt\w^t, respectively. As a middle ground between the two, the Gauss-Newton method uses the linearization of ff around wt\w^t f(w)≈f(wt)+∂f(wt)(w−wt) f(\w) \approx f(\w^t) + \partial f(\w^t) (\w -\w^t) but keeps ℓ\ell as is to obtain the objective wt+1≔arg min⁡w∈Wℓ(f(wt)+∂f(wt)(w−wt))=arg min⁡w∈Wℓ(∂f(wt)w+f(wt)−∂f(wt)wt)=arg min⁡w∈Wℓ(Jtw+δt),\begin{align*} \w^{t+1} &\coloneqq \argmin_{\w \in \cW} \ell(f(\w^t) + \partial f(\w^t) (\w -\w^t)) \\ &= \argmin_{\w \in \cW} \ell(\partial f(\w^t) \w + f(\w^t) -\partial f(\w^t)\w^t) \\ &= \argmin_{\w \in \cW} \ell(J^t \w + \deltav^t), \end{align*} where we defined the shorthands Jt≔∂f(wt)J^t \coloneqq \partial f(\w^t) and δt≔f(wt)−∂f(wt)wt\deltav^t \coloneqq f(\w^t) - \partial f(\w^t)\w^t. We call ℓ(Jtw+δt)\ell(J^t \w + \deltav^t) the partial linearization of L=ℓ∘fL = \ell \circ f at wt\w^t, as opposed to the full linearization of LL used in gradient descent.

Since the composition of a convex function and of linear function is convex, this objective is convex even if L(w)L(\w) is nonconvex. In practice, we often add a proximity term as regularization to define wt+1≔arg min⁡w∈Wℓ(Jtw+δt)+ηt2∥w−wt∥22.(17.3) \w^{t+1} \coloneqq \argmin_{\w \in \cW} \ell(J^t \w + \deltav^t) + \frac{\eta^t}{2} \|\w - \w^t\|_2^2. \tag{17.3} We can see this update as an approximation of the proximal point update arg min⁡w∈WL(w)+ηt2∥w−wt∥22, \argmin_{\w \in \cW} L(\w) + \frac{\eta^t}{2} \|\w - \w^t\|_2^2, where L(w)L(\w) has been replaced by its partial linearization. Solving Eq. (17.3) using gradient-based solvers requires computing the gradient of w↦ℓ(Jtw+δt)\w \mapsto \ell(J^t \w + \deltav^t), which is w↦(Jt)∗∇ℓ(Jtw+δt)\w \mapsto (J^t)^* \nabla \ell(J^t \w + \deltav^t). Computing this gradient by autodiff therefore requires performing a forward pass to compute the JVP JtwJ^t \w and a backward pass to compute the VJP (Jt)∗∇ℓ(z)(J^t)^* \nabla \ell(\z). See Section 2.3 for an introduction to these operators and Chapter 8 for an introduction to autodiff.

The Gauss-Newton method with arbitrary convex outer loss is often called modified Gauss-Newton (Nesterov, 2007) or prox-linear (Drusvyatskiy and Paquette, 2019). The classical Gauss-Newton and Levenberg-Marquardt (Levenberg, 1944; Marquardt, 1963) methods originate from nonlinear least-squares and are recovered when ℓ(z)\ell(\z) is quadratic (Kelley, 1995), such as ℓ(z)≔12∥z−y∥22\ell(\z) \coloneqq \frac{1}{2} \|\z - \y\|^2_2, for y\y some reference target. The Gauss-Newton method corresponds classically to not using regularization (i.e., ηt=0\eta^t = 0) and the Levenberg-Marquardt method uses regularization (usually called damping, potentially changing ηt\eta^t across iterations). See e.g., (Messerer et al., 2021), for a survey of different variants.

17.2.2 With approximate outer function

Another variant of the Gauss-Newton method consists in replacing the convex loss ℓ\ell with its quadratic approximation around zt≔f(wt)\z^t \coloneqq f(\w^t), qt(z)≔ℓ(zt)+⟨∇ℓ(zt),z−zt⟩+12⟨z−zt,∇2ℓ(zt)(z−zt)⟩≈ℓ(z) q^t(\z) \coloneqq \ell(\z^t) + \langle \nabla \ell(\z^t), \z-\z^t\rangle + \frac{1}{2}\langle \z-\z^t, \nabla^2 \ell(\z^t) (\z -\z^t)\rangle \approx \ell(\z) to define the update wt+1≔arg min⁡w∈Wqt(Jtw+δt)+ηt2∥w−wt∥22. \w^{t+1} \coloneqq \argmin_{\w \in \cW} q^t(J^t \w + \deltav^t) + \frac{\eta^t}{2} \|\w - \w^t\|_2^2. Notice that ℓ\ell has been replaced by its quadratic approximation qtq^t. This objective is always a convex quadratic, unlike the objective of the Newton method in Eq. (17.1), which is a priori a nonconvex quadratic, if ff is nonlinear. Simple calculations show that wt+1=arg min⁡w∈W L(wt)+⟨qt,Jt(w−wt)⟩+12⟨w−wt,(Jt)∗QtJt(w−wt)⟩+ηt2∥w−wt∥22\begin{align*} \w^{t+1} = \argmin_{\w \in \cW} ~ &L(\w^t) + \langle \q^t, J^t(\w {-} \w^t) \rangle \\ & + \frac{1}{2} \langle \w {-} \w^t, (J^t)^* Q^t J^t (\w {-} \w^t) \rangle +\frac{\eta^t}{2} \|\w - \w^t\|^2_2 \end{align*} where qt≔∇ℓ(f(wt))∈M=RZ,Qt≔∇2ℓ(f(wt))∈RZ×Z. \q^t \coloneqq \nabla \ell(f(\w^t)) \in \cM=\RR^Z, Q^t \coloneqq \nabla^2 \ell(f(\w^t)) \in \RR^{Z \times Z}. The closed form solution is wt+1=wt−((Jt)∗QtJt+ηtI⁡)−1(Jt)∗qt=wt−(∇GN2(ℓ∘f)(wt)+ηtI⁡)−1∇L(wt),\begin{align*} \w^{t+1} &= \w^t - ((J^t)^* Q^t J^t + \eta^t \idm)^{-1} (J^t)^* \q^t \\ &= \w^t - (\gn (\ell \circ f) (\w^t) + \eta^t \idm)^{-1} \nabla L(\w^t), \end{align*} where we used the (generalized) Gauss-Newton matrix of L=ℓ∘fL = \ell \circ f, defined in Section 9.2.

17.2.3 Linesearch

Similarly to Newton’s method, the iterates of a Gauss-Newton method may diverge when used alone. However, the direction −(∇GN2(ℓ∘f)(wt)+ηtI⁡)−1∇L(wt)-(\gn(\ell\circ f)(\w^t) + \eta^t \idm)^{-1} \nabla L(\w^t) defines a descent direction for any ηt>0\eta^t >0 and can be combined with a stepsize γt\gamma^t (typically chosen using a linesearch) to obtain iterates of the form wt+1=wt−γt(∇GN2(ℓ∘f)(wt)+ηtI⁡)−1∇L(wt). \w^{t+1} = \w^t - \gamma^t(\gn(\ell\circ f)(\w^t) + \eta^t \idm)^{-1} \nabla L(\w^t).

17.2.4 Stochastic Gauss-Newton

In deep learning, the objective generally consists in an expectation over samples of the composition between a loss function and a network function: L(w)=ES∼ρ[L(w;S)]=E(X,Y)∼ρ[ℓ(f(w;X);Y)] L(\w) = \mathbb{E}_{S \sim \rho} [L(\w; S)] = \mathbb{E}_{(X, Y) \sim \rho}[\ell(f(\w; X); Y)] where S=(X,Y)S = (X, Y) denotes a sample pair of input XX with associated label YY. In that case, as already studied in Section 9.2, the Gauss-Newton matrix ∇GN2L\gn L is the expectation of the individual Gauss-Newton matrices ∇GN2L(w;x,y)≔∂f(w;x)⊤∇2ℓ(f(w;x))∂f(w;x),∇GN2L(w)≔E(X,Y)∼ρ[∇GN2L(w;x,y)].\begin{align*} \gn L(\w; \x, \y) & \coloneqq \partial f(\w; \x)^\top \nabla^2 \ell(f(\w; \x)) \partial f(\w; \x), \\ \gn L (\w) & \coloneqq \mathbb{E}_{(X, Y)\sim \rho}[\gn L(\w; \x, \y)]. \end{align*} We can estimate the gradient and the Gauss-Newton matrix by, respectively, g(w;S)≈∇L(w)\g(\w; S) \approx \nabla L(\w), and G(w;S′)≈∇GN2L(w)G(\w; S') \approx \gn L(\w) for S,S′∼ρS, S' \sim \rho or using mini-batch approximations. A stochastic Gauss-Newton method therefore performs iterates wt+1≔wt−γt(G(wt;S′)+ηtI⁡)−1g(wt;S), \w^{t+1} \coloneqq \w^t - \gamma^t(G(\w^t; S') + \eta^t \idm)^{-1} \g(\w^t; S), for ηt≥0\eta^t \geq 0 and γt\gamma^t fixed or selected to satisfy some criterion.

17.3 Natural gradient descent

Natural gradient descent (Amari, 1998) follows a similar principle as gradient descent: linearize the objective around the current iterate and minimize this approximation together with a proximity term. It differs from gradient descent in the choice of the proximity term: rather than using a squared Euclidean distance between the parameters, it uses a Kullback-Leibler divergence between the probability distributions these parameters define.

Negative log-likelihood

We consider objectives of the form min⁡w∈WL(w)=ES∼ρ[L(w;S)]=ES∼ρ[−log⁡qw(S)], \min_{\w \in \cW} L(\w) = \EE_{S \sim \rho} \left[L(\w; S)\right] = \EE_{S \sim \rho} \left[-\log q_\w(S)\right], where ρ\rho is an unknown data distribution (but from which we can sample) and where qwq_\w is a probability distribution parameterized by w\w. As reviewed in Chapter 3, the negative log-likelihood can be used as a loss function (many loss functions can be seen from this perspective, including the squared and logistic loss functions). In the unsupervised setting, where S=YS=Y, we simply use qw(Y)q_\w(Y) as is. In the supervised setting, where S=(X,Y)S=(X,Y), we use the product rule P(X,Y)=P(X)P(Y∣X)\PP(X, Y) = \PP(X) \PP(Y|X) to parameterize qw(S)q_\w(S) as qw(x,y)≔ρX(x)pθ(y), q_\w(\x, \y) \coloneqq \rho_X(\x) p_\thetav(\y), where ρX\rho_X is the marginal distribution for XX, pθ(y)p_\thetav(\y) is the PMF/PDF of a probability distribution and θ=f(w;x)\thetav = f(\w; \x) is for instance a neural network with parameters w∈W\w \in \cW and input x∈X\x \in \cX.

17.3.1 Variational perspective

Natural gradient descent is motivated by updates of the form wt+1=arg min⁡w∈WL(wt)+⟨∇L(wt),w−wt⟩+KL(qwt,qw), \w^{t+1} = \argmin_{\w \in \cW} L(\w^t) + \langle \nabla L(\w^t), \w - \w^t\rangle + \mathrm{KL}(q_{\w^t},q_\w), where KL(p,q)≔∫p(z)log⁡p(z)q(z)dz\mathrm{KL}(p,q) \coloneqq \int p(\z) \log \frac{p(\z)}{q(\z)} d\z is the Kullback-Leibler (KL) divergence. Unlike gradient descent, the proximity term is therefore between the current distribution qwtq_{\w^t} and a candidate probability distribution qwq_{\w}. The above problem is intractable in general, as the KL may not have a closed form. Nevertheless, its quadratic approximation can be shown (Amari, 1998) to admit a simple form, KL(qwt,qw)≈12⟨w−wt,∇F2L(wt)(w−wt)⟩ \mathrm{KL}(q_{\w^t}, q_\w) \approx \frac{1}{2}\langle \w - \w^t, \fisher L (\w^t)(\w -\w^t)\rangle where we used the Fisher information matrix ∇F2L(w)\fisher L (\w), studied in Section 9.3. Equipped with this quadratic approximation of the KL divergence, natural gradient descent amounts to computing iterates as wt+1≔arg min⁡w∈W L(wt)+⟨∇L(wt),w−wt⟩+12⟨w−wt,∇F2L(wt)(w−wt)⟩+ηt2∥w−wt∥22,\begin{align*} \w^{t+1} \coloneqq \argmin_{\w \in \cW} ~ &L(\w^t) + \langle \nabla L(\w^t), \w - \w^t\rangle \\ & + \frac{1}{2}\langle \w - \w^t, \fisher L (\w^t)(\w -\w^t)\rangle + \frac{\eta^t}{2}\|\w-\w^t\|^2_2, \end{align*} where a quadratic proximity term was added to ensure a unique solution. This is a strictly convex problem as ∇F2L(wt)\fisher L(\w^t) is positive semi-definite. The closed-form solution is wt+1=wt−(∇F2L(wt)+ηtI⁡)−1∇L(wt). \w^{t+1} = \w^t - (\fisher L (\w^t) + \eta^t \idm )^{-1}\nabla L(\w^t). Because the Gauss-Newton and Fisher information matrices are equivalent when pθp_\thetav is an exponential family distribution (Proposition 9.6), the Gauss-Newton and natural gradient methods coincide in this case.

17.3.2 Stochastic natural gradient descent

In practice, we may not have access to ∇L(wt)\nabla L(\w^t) in closed form as it is an expectation over ρ\rho. Moreover, ∇F2L(wt)\fisher L(\w^t) may not be computable in closed form either. To estimate the Fisher information matrix, we can use that (see Section 9.3) using the shorthand θ≔f(w,X)\thetav \coloneqq f(\w, X), ∇F2L(w)=EX∼ρXEY∼pθ[∇L(w;X,Y)⊗∇L(w;X,Y)]. \fisher L(\w) = \EE_{X \sim \rho_\cX} \EE_{Y \sim p_\thetav} [\nabla L(\w; X, Y) \otimes \nabla L(\w; X, Y)]. We can then build estimates g(wt;S)≈∇L(wt,S)\g(\w^t; S) \approx \nabla L(\w^t, S) and F(wt;S′)≈∇F2L(wt)F(\w^t; S') \approx \fisher L(\w^t) for SS sampled from ρ\rho and S′S' sampled from qwt(x,y)=pX(x)ρθ(y)q_{\w^t}(\x, \y) = p_\cX(\x) \rho_{\thetav}(\y). A stochastic natural gradient descent can then be implemented as wt+1=wt−γt(F(wt;S′)+ηtI⁡)−1g(wt;S), \w^{t+1} = \w^t - \gamma^t(F(\w^t; S') + \eta^t \idm )^{-1}\g(\w^t; S), where γt\gamma^t is a stepsize, possibly chosen by linesearch.

In deep learning, the product with the inverse Fisher or Gauss-Newton matrices can remain costly to compute. Several approximations have been proposed, such as KFAC (Botev et al., 2017; Martens and Grosse, 2015), which uses a computationally efficient structural approximation to these matrices.

17.4 Quasi-Newton methods

17.4.1 BFGS

A celebrated example of quasi-Newton method is the BFGS method (Broyden, 1970; Fletcher, 1970; Goldfarb, 1970; Shanno, 1970), whose acronym follows from its authors’ names. The rationale of the BFGS update stems once again from a variational viewpoint. We wish to build a simple quadratic model of the objective ht(w)=L(wt)+⟨∇L(wt),w−wt⟩+12⟨w−wt,Qt(w−wt)⟩h^t(\w) = L(\w^t) + \langle \nabla L(\w^t), \w - \w^t\rangle + \frac{1}{2} \langle \w - \w^t, Q^t (\w -\w^t) \rangle for some QtQ^t built along the iterations rather than taken as ∇2L(wt)\nabla^2 L(\w^t). One desirable property of such a quadratic model would be that its gradients at consecutive iterates match the gradients of the original function, i.e., ∇ht(wt)=∇L(wt)\nabla h^t(\w^t) = \nabla L(\w^t) and ∇ht(wt−1)=∇L(wt−1)\nabla h^t(\w^{t-1})= \nabla L(\w^{t-1}). A simpler condition, called the secant condition consists in considering the differences of these vectors, that is, ensuring that ∇ht(wt)−∇ht(wt−1)=∇L(wt)−∇L(wt−1)  ⟺  Qt(wt−wt−1)=∇L(wt)−∇L(wt−1)  ⟺  wt−wt−1=Bt(∇L(wt)−∇L(wt−1)),\begin{align*} &\nabla h^t(\w^t) - \nabla h^t(\w^{t-1}) = \nabla L(\w^t) - \nabla L(\w^{t-1})\\ \iff& Q^t (\w^t - \w^{t-1}) = \nabla L(\w^t) - \nabla L(\w^{t-1})\\ \iff& \w^t - \w^{t-1} = B^t(\nabla L(\w^t) - \nabla L(\w^{t-1})), \end{align*} for Bt=(Qt)−1B^t = (Q^t)^{-1}. Building BtB^t, a surrogate of the inverse of the Hessian satisfying the secant equation, can then be done as Bt+1≔(I⁡−ρtst(yt)⊤)Bt(I⁡−ρtst(yt)⊤)+ρtst(st)⊤ B^{t+1} \coloneqq \left(\idm - \rho^t \s^t (\y^t)^\top \right) B^t \left(\idm - \rho^t \s^t (\y^t)^\top\right) + \rho^t \s^t (\s^t)^\top where st≔wt+1−wtyt≔∇L(wt+1)−∇L(wt)ρt≔1⟨st,yt⟩.\begin{align*} \s^t &\coloneqq \w^{t+1} - \w^t \\ \y^t &\coloneqq \nabla L(\w^{t+1}) - \nabla L(\w^t) \\ \rho^t &\coloneqq \frac{1}{\langle \s^t, \y^t \rangle}. \end{align*} A typical implementation of BFGS stores Bt∈RP×PB^t \in \RR^{P \times P} in memory, which is prohibitive when PP is large.

17.4.2 Limited-memory BFGS

In practice, the limited-memory counterpart of BFGS, called LBFGS (Liu and Nocedal, 1989), is often preferred. The key observation of LBFGS is that we do not need to materialize BtB^t in memory: we only need to multiply it with the gradient ∇L(wt)\nabla L(\w^t). That is, we can see BtB^t as a linear map. Fortunately, the product between BtB^t and any vector v\v can be computed efficiently if we store (s1,y1,ρ1),…,(st,yt,ρt)(\s^1, \y^1, \rho^1), \dots, (\s^t, \y^t, \rho^t) in memory. In practice, a small history of past values is used to reduce memory and computational cost. Because LBFGS has the benefits of second-order-like methods with much reduced cost, it has become a de-facto algorithm, outperforming most other algorithms for medium-scale problems without particular structure (Liu and Nocedal, 1989).

17.5 Approximate Hessian diagonal inverse preconditioners

One application of the approximations of the Hessian diagonal developed in Section 9.7 is to obtain cheap approximations of the Hessian diagonal inverse, Bt≔diag⁡(∣H11t∣−1,…,∣HPPt∣−1). B^t \coloneqq \diag(|H_{11}^t|^{-1}, \ldots, |H_{PP}^t|^{-1}). Such a scaling would for instance be sufficient to make the quadratic example presented in Fig. 17.1 work. Many optimization algorithms, including the popular ADAM, can be viewed as using a preconditioner that approximates the inverse of the Hessian’s diagonal.

17.6 Summary