Beyond Gradient Descent

If our objective function $f$ is convex, we are guaranteed that the Hessian of $f$, which we denote as $H$, is positive semi-definite. That is,

$$\begin{equation*} H_{t} = \frac{\partial^{2}}{\partial x^{2}} f(x_{t}) \succeq 0. \end{equation*}$$

When we substitute $H_{t}$ for $G_{t}$ in Eq. (5) (i.e., when we choose $D_{t}$ based on the Hessian of $f$), then the resulting subproblem (Eq. (2)) is simply a minimization over a 2nd-order Taylor approximation of $f$ at $x_{t}$. The resulting update is known as Newton’s Method.

When used for solving least-squares problems (for residuals $r_{i}(x)$), the Hessian may be calculated as

$$\begin{align*} f(x) &= \sum_{i} r_{i}(x) r_{i}(x) \\ \frac{\partial}{\partial x} f(x) &= 2 \sum_{i} r_{i}(x) \frac{\partial r_{i}(x)}{\partial x} \\ H_{t} = \frac{\partial^{2}}{\partial x^{2}} f(x) &= 2 \bigg( \sum_{i} \frac{\partial r_{i}(x)}{\partial x} \frac{\partial r_{i}(x)}{\partial x} + \sum_{i} r_{i}(x) \frac{\partial^{2} r_{i}(x)}{\partial x^{2}} \bigg) \end{align*}$$

Ignoring the second term on the last line and retaining only the first as an approximation for $H_{t}$ yields the Gauss-Newton algorithm. For linear models, the second term is identically zero.

We may also choose values of $G$ that derive from spaces other than $\mathcal{X}$. For example, when $x$ parameterizes a probability distribution $\rho(y ; x)$, from which the objective function $f$ derives, we may measure the magnitude of an update from $x_{t}$ to $x$ by the relative entropy from $\rho_{t}$ to $\rho$.

$$\begin{equation*} D_{t}(x) = \frac{1}{\eta} D_{\rm KL}\Big(\rho(y ; x) \parallel \rho(y ; x_{t})\Big) = \frac{1}{\eta} \int_{\mathcal{Y}} \rho(y ; x) \log \frac{\rho(y; x)}{\rho(y ; x_{t})} {\rm d}y. \end{equation*}$$

This choice yields the update rule for Fisher-Rao natural gradient descent:

$$\begin{equation} \label{eqn-7}\tag{7} x_{t+1} = x_{t} - \eta F^{\dagger}_{t} \nabla f(x_{t}), \end{equation}$$

where $F_{t}$ is the Fisher metric tensor

$$\begin{equation*} F_{t} = \int_{\mathcal{Y}} \rho(y ; x_{t}) \frac{\partial^{2}}{\partial x^{2}} \log \rho(y ; x) {\rm d}y. \end{equation*}$$

One way to understand the effect of the metric $G = F$ is that it (un)warps the space of marginal updates ${\rm d}x$, such that updates of the same induced magnitude contain the same amount of marginal information about $\rho(x_{t} + {\rm d}x)$.

In continuous time, Fisher-Rao natural gradient descent optimally approximates replicator dynamics (used to model evolution by natural selection) and continuous Bayesian inference (Raab et. al. (2022)). We may also choose different divergences in $\rho$ space, yielding other forms of “natural gradient descent” (Nurbekyan et. al. (2022)).

Vanilla gradient descent does not require a Euclidean metric.

As a counterexample, let $$G = \frac{\nabla f ~ \nabla^{\top} f}{\nabla^{\top} f \cdot \nabla f},$$ where the numerator is an outer product and the denominator an inner product of $\nabla f$ with itself.

It follows that $$G^{\dagger} = G \quad \text{and} \quad G^{\dagger} \nabla f = \nabla f.$$ Therefore, the update rule that solves $\nabla(\hat{f}_{t} + D_{t}) = 0$ (Eq. (2)) is vanilla gradient descent: $$x_{t+1} - x_{t} = -\eta G^{\dagger} \nabla f(x_{t}) = -\eta \nabla f(x_{t}).$$

An example of a multi-iterate approach is given by the classical “Momentum” variant of gradient descent. Consider, for example, a penalty term $D_{t}$ that quadratically penalizes step-sizes with a linear term that encourages steps in same continued direction, as in

$$\begin{equation*} D_{t}(x) = \frac{1}{\eta} \bigg( \frac{1}{2} \big\| x - x_{t} \big\|_{2}^{2} - \xi \big\langle x - x_{t}, x_{t} - x_{t-1} \big\rangle \bigg), \end{equation*}$$

where $\xi \in (0, 1)$. Alternatively, consider a quadratic penalty for the difference between the proposed update and the previous update with decay $\xi$, as in

$$\begin{equation*} D_{t}(x) = \frac{1}{2 \eta} \big\| (x - x_{t}) - \xi (x_{t} - x_{t-1}) \big\|_{2}^{2}. \end{equation*}$$

In either case,

$$\begin{equation*} \nabla D_{t}(x) = \frac{1}{\eta} \bigg((x - x_{t}) - \xi (x_{t} - x_{t-1})\bigg). \end{equation*}$$

If we solve Eq. (2) by setting $\nabla (\hat{f}_{t} + D_{t}) = 0$ for either choice of $D_{t}$, the resulting update rule is

$$\begin{equation*} x_{t+1} = x_{t} + \xi(x_{t} - x_{t-1}) - \eta \nabla \hat{f}_{t}, \end{equation*}$$

which can be decomposed, with an additional state variable $v$, to

$$\begin{equation} \label{eqn-8}\tag{8} \begin{aligned} v_{t+1} &= \xi v_{t} - \eta \nabla \hat{f}_{t}, \\ x_{t+1} &= x_{t} + v_{t+1}. \end{aligned} \end{equation}$$

This update rule is the classical “Momentum” algorithm (Botev et. al. (2016)) with decay parameterized by $\xi$.

There is yet another update penalty $D_{t}$ that may be used to derive the Momentum update rule. Specifically, let

$$\begin{equation*} D_{t}(x) = \frac{\xi_{1}}{2\eta} \big\| x - x_{t} \big\|^{2} + \frac{\xi_{2}}{2\eta} \big\| \underbrace{(x - x_{t})}_{v} - \underbrace{(x_{t} - x_{t-1})}_{v_{t}} \big\|^{2}. \end{equation*}$$

We can naturally extend this distance penalty to account for higher-order derivatives as

$$\begin{align*} D_{t}(x) &= \frac{\xi_{1}}{2\eta} \big\|x - x_{t} \big\|^{2} + \frac{\xi_{2}}{2\eta} \big\| v - v_{t} \big\|^{2} + ... + \frac{\xi_{n}}{2\eta} \big\| x^{(n-1)} - x^{(n-1)}_{t} \big\|^{2} \\ &= \frac{\xi_{1}}{2\eta} \big\|x^{(1)} \big\|^{2} + \frac{\xi_{2}}{2\eta} \big\| x^{(2)} \big\|^{2} + ... + \frac{\xi_{n}}{2\eta} \big\| x^{(n)} \big\|^{2}, \end{align*}$$

where we use the state variable $x^{(n)}$ to represent the $n$th time-derivative of $x$, defined empirically by the recursive formula

$$\begin{equation*} x^{(n)} = x^{(n-1)} - x_{t}^{(n-1)}. \end{equation*}$$

By this choice of $D_{t}$, if we solve Eq. (2) by setting $\nabla (\hat{f}_{t} + D_{t}) = 0$ at $x_{t+1}$, we obtain the equation

$$\begin{equation*} \eta \nabla \hat{f}_{t} = \sum_{k=1}^{n-2} \xi_{k} x^{(k)}_{t+1} + \xi_{n-1} x^{(n-1)}_{t+1} + \xi_{n} \big(x^{(n-1)}_{t+1} - x^{(n-1)}_{t}\big) \end{equation*}$$

and therefore, the update rules

$$\begin{equation} \label{eqn-9}\tag{9} \begin{aligned} x^{(n-1)}_{t+1} &= \frac{1}{\xi_{n} + \xi_{n-1}} \bigg( \xi_{n} x^{(n-1)}_{t} - \sum_{k=1}^{n-2} \xi_{k} x^{(k)}_{t+1} - \eta \nabla \hat{f}_{t} \bigg), \\ x^{(n-2)}_{t+1} &= x^{(n-2)}_{t} + x^{(n-1)}_{t+1}, \\ &\vdots \\ x_{t+1} &= x_{t} + x^{(1)}_{t+1}. \end{aligned} \end{equation}$$

For $n = 2$, Letting $\xi_{1} = (1 - \xi)$ and $\xi_{2} = \xi$, we recover the standard momentum update.

When using Eq. (8), one “zero-cost” (i.e., if we decouple candidate solutions $x_{t}$ from the points at which we query $\hat{f}$) choice of $\tilde{x}_{t}$ is given by

$$\begin{equation*} \tilde{x}_{t} = x_{t} + \xi v_{t}, \end{equation*}$$

such that

$$\begin{equation} \label{eqn-10}\tag{10} \hat{f}_{t}(x) = f(\tilde{x}_{t}) + \big\langle \nabla f(\tilde{x}_{t}), x - \tilde{x}_{t} \big\rangle. \end{equation}$$

Substituting Eq. (10) into Eq. (8), we obtain Nesterov’s variant of gradient descent with momentum:

$$\begin{equation*} \begin{aligned} v_{t+1} &= \xi v_{t} - \eta \nabla f(x_{t} + \xi v_{t}), \\ x_{t+1} &= x_{t} + v_{t+1}. \end{aligned} \end{equation*}$$

We may iteratively approximate Eq. (12) with a sequence of subproblems in the form of Eq. (2):

$$\begin{equation*} x_{t+1} = \underset{x}{\rm argmin}\bigg[ \underset{\lambda \succ 0, \mu}{\rm max} ~ \bigg( \hat{\mathcal{L}}_{t}(x, \lambda, \mu) + D_{t}(x) - R_{t}(\lambda) - S_{t}(\mu) \bigg) \bigg], \end{equation*}$$

introducing update penalties for state variables $x$, $\lambda$, and $\mu$.

For the simple choices

$$\begin{align*} D_{t}(x) &= \frac{1}{2\eta}\big\|x - x_{t}\big\|_{2}^{2}, \quad \eta > 0, \\ R_{t}(\lambda) &= \frac{1}{2\alpha}\big\|\lambda - \lambda_{t}\big\|_{2}^{2}, \quad \alpha > 0, \\ S_{t}(\mu) &= \frac{1}{2\beta}\big\|\mu - \mu_{t}\big\|_{2}^{2}, \quad \beta > 0, \end{align*}$$

and local, linear approximations

$$\begin{align*} \hat{f}_{t}(x) &= f(x_{t}) + \big\langle \nabla f(x_{t}), x - x_{t} \big\rangle; \\ \hat{g}_{t}(x) &= g(x_{t}) + \big\langle \nabla g(x_{t}), x - x_{t} \big\rangle; \\ \hat{h}_{t}(x) &= h(x_{t}) + \big\langle \nabla h(x_{t}), x - x_{t} \big\rangle; \\ \hat{\mathcal{L}}_{t}(x, \lambda, \mu) &= \hat{f}_{t}(x) + \lambda^{\top} \hat{g}_{t}(x) + \mu^{\top} \hat{h}_{t}(x), \end{align*}$$

we have the update rules

$$\begin{equation} \label{eqn-14}\tag{14} \begin{aligned} x_{t+1} &= x_{t} - \eta \bigg( \nabla f(x_{t}) + \lambda_{t+1}^{\top} \nabla g(x_{t}) + \mu^{\top}_{t+1} \nabla h(x_{t})\bigg), \\ \lambda_{t+1} &= \max\bigg( \lambda_{t} + \alpha \hat{g}(x_{t+1}), 0\bigg), \\ \mu_{t+1} &= \mu_{t} + \beta \hat{h}(x_{t+1}), \end{aligned} \end{equation}$$

where $\max$ is taken element-wise and the gradients are taken with respect to $x$ (and these gradient components remain uncontracted with the dimensions of $\lambda$ or $\mu$).

In practice, it is common to eliminate the mutual dependence between variables in Eq. (14) by replacing $x_{t+1} \mapsto x_{t}$, $\lambda_{t+1} \mapsto \lambda_{t}$, and $\mu_{t+1} \mapsto \mu_{t}$ on the right-hand side of each equation, thus yielding the standard primal-dual algorithm, which maintains iterates for the dual variables $\lambda$ and $\mu$ in addition to the primal variable $x$. The error of this approximation has order $\mathcal{O}(\eta( \alpha + \beta))$, though this is not strictly necessary.

When $\eta \alpha < |\nabla g(x_{t}) |^{2}$ and $\eta \beta < |\nabla h(x_{t}) |^{2},$ the true solution to Eq. (14) may be found by iterating the recursive map

$$\begin{equation*} v^{k+1}_{t+1} \leftarrow -\eta \left( \begin{aligned} \nabla f(x_{t}) &~+ \\ \max\Big(0, \lambda_{t} + \alpha \big( g(x_{t}) + \langle \nabla g(x_{t}), v^{k}_{t+1} \rangle \big)\Big) \nabla g(x_{t}) &~+ \\ \Big( \mu_{t} + \beta \big( h(x_{t}) + \langle \nabla h(x_{t}), v^{k}_{t+1} \rangle \big)\Big) \nabla h(x_{t})& \end{aligned} \right). \end{equation*}$$

such that $x_{t+1} = x_{t} + v_{t+1}^{\infty}$.

For convex functions $f$, $g$, and $h$, and sufficiently small $\eta, \alpha, \beta$, iterating subproblem Eq. (14) will cause $x$, $\lambda$, and $\mu$ to converge to finite values and solve the target constrained optimization problem Eq. (12).

In this section, we show how a rather interesting penalty choice for the dual variable $\lambda$, as used above, completely eliminates the need to track $\lambda$ as an independent variable all and provides a straight-forward method for constrained optimization based on local gradients of the objective and the constraint.

Consider the problem

$$\begin{equation} \begin{array}{lll} {\rm minimize} &f(x) \\ {\rm subject~to} &g(x) \leq 0, \end{array} \end{equation}$$

for vector $x$ and scalar-valued $f$ and $g$.

Given standard assumptions such as convexity, this problem may be solved the sequence of subproblems

$$\begin{align*} x_{t+1} =&~\underset{x}{\rm argmin}\quad \hat{f}(x) + \lambda \hat{g}(x) + \frac{1}{2\eta} \| x - x_{t} \|^{2}, \\ \lambda \equiv&~ \underset{v \geq 0}{\rm argmax}\quad f(x_{t}) + v g(x_{t}) - \frac{1}{2} \Big\| \nabla f(x_{t}) + v \nabla g(x_{t}) \Big\|^{2}, \end{align*}$$

which, as in Eq. (4), we express using the local, linear approximations $\hat{f}$ and $\hat{g}$. Note that we have no need to track the value of $\lambda$ between iterates of the above subproblems. This is because, instead of penalizing the update magnitude $\|\lambda - \lambda_{t}\|$, we penalize the distance of $\lambda$ away from the value that renders $x_{t}$ a critical point of the Lagrangian. We have omitted an explicit time-index on $\lambda$ despite the fact that its value may change with each iteration.

Solving the above subproblem, $x_{t+1}$ has the explicit solution given by

$$\begin{align*} x_{t+1} &= x_{t} - \eta \nabla \Lambda; \\ \nabla \Lambda &\equiv \nabla f(x_{t}) + \lambda \nabla g(x_{t}); \\ \lambda &\equiv \max\bigg[0, \frac{ g(x_{t}) - \big\langle \nabla f(x_{t}), \nabla g(x_{t}), \big\rangle}{\big\langle \nabla g(x_{t}), \nabla g(x_{t}) \big\rangle} \bigg], \end{align*}$$

where $\nabla \Lambda$ represents an estimate for the gradient of the Lagrangian at $x_{t}$.

Manipulating the above equations, we see that

$$\begin{align*} \big\langle \nabla \Lambda, \nabla g(x_{t}) \big\rangle &= - \frac{1}{\eta} \big\langle x_{t+1} - x_{t}, \nabla g(x_{t}) \big\rangle. \\ \big\langle \nabla \Lambda, \nabla g(x_{t}) \big\rangle &= \big\langle \nabla f(x_{t}), \nabla g(x_{t}) \big\rangle + \lambda \big\langle \nabla g(x_{t}), \nabla g(x_{t}) \big\rangle. \\ &= \big\langle \nabla f(x_{t}), \nabla g(x_{t}) \big\rangle + \max\bigg[ 0, g(x_{t}) - \big\langle \nabla f(x_{t}), \nabla g(x_{t}) \big\rangle \bigg] \\ &= \max\bigg[ \big\langle \nabla f(x_{t}), \nabla g(x_{t}) \big\rangle, g(x_{t}) \bigg]. \end{align*}$$

From which it follows

$$\begin{equation*} -\frac{1}{\eta} \big\langle x_{t+1} - x_{t}, \nabla g(x_{t}) \big\rangle \geq g(x_{t}). \end{equation*}$$

By similar logic,

$$\begin{equation*} -\frac{1}{\eta} \big\langle x_{t+1} - x_{t}, \nabla f(x_{t}) \big\rangle \geq \|\nabla f(x_{t}) \|^{2}. \end{equation*}$$

This suggests the following, alternative expression for our subproblem:

$$\begin{align*} x_{t+1} =~\underset{x}{\rm argmin}&\quad \big\langle x, \nabla f(x_{t}) \big\rangle + \frac{1}{2\eta} \| x - x_{t} \|^{2}, \\ ~{\rm subject~to}&\quad \big\langle x - x_{t}, -\nabla g(x_{t}) \big\rangle \geq \eta g(x_{t}). \end{align*}$$

Intuitively, when $g(x_{t})$ is positive, in order to make progress towards the constraint $g \leq 0$, we ensure that the update $x_{t+1} - x_{t}$ is aligned with the negative gradient of $g$. When $g$ is negative, and the constraint already satisfied, the update cannot be too aligned with increasing $g$. Subject to these constraints $x_{t+1}$ may be chosen to make progress towards decreasing the objective $f$.