Gradient flows, sometimes stochastic

2.1 Classic stochastic gradient descent

The classic SGD optimizer updates the parameters by taking steps proportional to the negative of the gradient of the loss function with respect to the parameters. Here’s the discrete-time update rule:

$$\begin{array}{}\text{(1)}& {\mathbf{x}}_{k+1}={\mathbf{x}}_k-\eta \mathrm{\nabla }\hat{U}({\mathbf{x}}_k)\end{array}$$

where

• ${\mathbf{x}}_k$ is the parameter vector at iteration $k$,
• $\eta$ is the learning rate, controlling the step size, and
• $\mathrm{\nabla }\hat{U}({\mathbf{x}}_k)$ is the stochastic gradient at iteration $k$, typically estimated using a mini-batch of data.

$\mathbf{x}$ is usually some set of interesting parameters, such as the weights of a neural network.

The stochastic gradient $\mathrm{\nabla }\hat{U}({\mathbf{x}}_k)$ can be modelled, for want of a better option, as a noisy version of the true gradient $\mathrm{\nabla }U({\mathbf{x}}_k)$,

$$\mathrm{\nabla }\hat{U}({\mathbf{x}}_k)=\mathrm{\nabla }U({\mathbf{x}}_k)+{\mathit{\xi }}_k$$

where

• $\mathrm{\nabla }U({\mathbf{x}}_k)$ is the true gradient of the loss function, and
• ${\mathit{\xi }}_k$ is the gradient noise, assumed to follow a Gaussian distribution ${\mathit{\xi }}_k\sim \mathcal{N}(0,\mathrm{\Sigma })$.

To transition from the discrete-time SGD updates to a continuous-time framework, we introduce a small time step $\mathrm{\Delta }t$ such that $k\mathrm{\Delta }t=t$. The goal is to express the parameter updates as differential equations by letting the time step size shrink to zero. Starting from the discrete update,

$$\begin{array}{rl}{\mathbf{x}}_{k+1}& ={\mathbf{x}}_k-\eta \mathrm{\nabla }\hat{U}({\mathbf{x}}_k)\\ \frac{{\mathbf{x}}_{k+1}-{\mathbf{x}}_k}{\mathrm{\Delta }t}& =-\frac{\eta }{\mathrm{\Delta }t}\mathrm{\nabla }\hat{U}({\mathbf{x}}_k)\\ \frac{d\mathbf{x}(t)}{\mathrm{\Delta }t}& =-\frac{\eta }{\mathrm{\Delta }t}(\mathrm{\nabla }U(\mathbf{x}(t))+\mathit{\xi }(t))\end{array}$$

Taking the limit as $\mathrm{\Delta }t\to 0$,

$$d\mathbf{x}(t)=-\eta \mathrm{\nabla }\hat{U}(\mathbf{x}(t))+\sqrt{\eta \mathrm{\Sigma }}d\mathbf{W}(t)$$

We model the stochastic term $\mathit{\xi }(t)$ using a standard Wiener process $\mathbf{W}(t)$,

$$\mathit{\xi }(t)dt=\sigma d\mathbf{W}(t)$$

where

• $\sigma \sqrt{\mathrm{\Sigma }}$ defines the noise intensity,
• $d\mathbf{W}(t)$ is the Wiener increment, satisfying $d\mathbf{W}(t)\sim \mathcal{N}(0,dt)$.

Example: Consider a quadratic loss function defined via where $A$ is a positive definite matrix and $\mathbf{b}$ a vector,

$$\begin{array}{rl}U(\mathbf{x})& =\frac{1}{2}{\mathbf{x}}^{\mathrm{\top }}A\mathbf{x}+{\mathbf{b}}^{\mathrm{\top }}\mathbf{x}\\ \mathrm{\nabla }U(\mathbf{x})& =A\mathbf{x}+\mathbf{b}\end{array}$$

Substituting into the SDE,

$$d\mathbf{x}(t)=-\eta A\mathbf{x}(t)dt-\eta \mathbf{b}dt+\sqrt{\eta \mathrm{\Sigma }}d\mathbf{W}(t)$$

This SDE describes an Ornstein-Uhlenbeck process, whose stationary distribution is Gaussian with mean $\mathbf{b}$ and whose variance is… annoying to calculate but occasionally useful (search for Lyapunov equation). There is a whole mini-industry based on that realisation — see Mandt, Hoffman, and Blei (2017).

From the SDE, we can re-discretize the update using the Euler-Maruyama method, which is a quick-and-dirty option but OK for exposition. This should get us back to something like our original SGD rule (Equation 1). For an SDE of the form

$$d\mathbf{x}(t)=f(\mathbf{x}(t),t)dt+G(\mathbf{x}(t),t)d\mathbf{W}(t)$$

the Euler-Maruyama discretization over a time step $\mathrm{\Delta }t$ is

$${\mathbf{x}}_{k+1}={\mathbf{x}}_k+f({\mathbf{x}}_k,t_k)\mathrm{\Delta }t+G({\mathbf{x}}_k,t_k)\mathrm{\Delta }{\mathbf{W}}_k$$

where

• ${\mathbf{x}}_k=\mathbf{x}(t_k)$
• $t_k=k\mathrm{\Delta }t$
• $\mathrm{\Delta }{\mathbf{W}}_k=\mathbf{W}(t_{k+1})-\mathbf{W}(t_k)$

$\mathrm{\Delta }{\mathbf{W}}_k$ is a Gaussian random variable with zero mean and covariance $\mathrm{\Delta }t$ (since increments of Wiener processes are normally distributed with variance proportional to $\mathrm{\Delta }t$).

Pattern matching to our case,

• Drift Function: $f(\mathbf{x}(t),t)=-\eta \mathrm{\nabla }U(\mathbf{x}(t))$
• Diffusion Coefficient: $G(\mathbf{x}(t),t)=\sqrt{\eta \mathrm{\Sigma }}$ (assumed constant)

so the equation becomes

$$\begin{array}{rl}{\mathbf{x}}_{k+1}& ={\mathbf{x}}_k-\eta \mathrm{\nabla }U({\mathbf{x}}_k)\mathrm{\Delta }t+\sqrt{\eta \mathrm{\Sigma }}\mathrm{\Delta }{\mathbf{W}}_k\\ & ={\mathbf{x}}_k-\eta \mathrm{\nabla }U({\mathbf{x}}_k)\mathrm{\Delta }t+\sqrt{\eta \mathrm{\Delta }t}{\mathrm{\Sigma }}^{1/2}{\mathit{ϵ}}_k\end{array}$$

where

• $\mathrm{\Delta }{\mathbf{W}}_k\sim \mathcal{N}(\mathbf{0},\mathrm{\Delta }t\mathbf{I})$
• ${\mathit{ϵ}}_k$ is a vector of standard normal random variables, ${\mathit{ϵ}}_k\sim \mathcal{N}(\mathbf{0},\mathbf{I})$
• ${\mathrm{\Sigma }}^{1/2}$ denotes a matrix such that ${\mathrm{\Sigma }}^{1/2}({\mathrm{\Sigma }}^{1/2}{)}^{\mathrm{\top }}=\mathrm{\Sigma }$.

That is indeed the classic SGD update rule, but now we have a natural way of re-scaling it in time, and taking continuous limits etc.

For illustrative purposes, consider the quadratic loss function from earlier. Substituting into the state-space model:

$$\begin{array}{rl}{\mathbf{x}}_{k+1}& ={\mathbf{x}}_k-\eta \mathrm{\Delta }tA{\mathbf{x}}_k+\sqrt{\eta \mathrm{\Sigma }}\mathrm{\Delta }{\mathbf{W}}_k\\ & =(\mathbf{I}-\eta \mathrm{\Delta }tA){\mathbf{x}}_k-\eta \mathrm{\Delta }t\mathbf{b}+\sqrt{\eta \mathrm{\Sigma }}\mathrm{\Delta }{\mathbf{W}}_k.\end{array}$$

2.2 Adam

Adam is a popular optimizer that combines the benefits of adaptive learning rates and a flavour of momentum. The Adam optimizer updates the parameters using adaptive estimates of the first and second moments of the gradients. An SDE for Adam is surely published somewhere, but I couldn’t find it for some reason, so I derive it here.

cf Adam as an optimisation method, which is what it is designed for.

The resulting SDEs resemble underdamped Langevin dynamics (as seen in Langevin MCMC), because there is momentum and a friction term. They are a little weird though, because the adaptive mean estimates of the gradient drift over time.

We start from the original Adam equations, build the continuous-time version of the stochastic differential equations (SDEs), introduce the Wiener noise to model gradient noise, and then write the SDEs. Finally we can discretize using Euler-Maruyama.

The Adam optimizer updates the parameters using adaptive estimates of the first and second moments of the gradients. Here are the discrete-time update rules:

First moment update: $$\begin{array}{}\text{(2)}& {\mathbf{m}}_k={\beta }_1{\mathbf{m}}_{k-1}+(1-{\beta }_1)\mathrm{\nabla }\hat{U}({\mathbf{x}}_k)\end{array}$$

where

• ${\mathbf{m}}_k$ is the first moment estimate (momentum) at step $k$,
• $\mathrm{\nabla }\hat{U}({\mathbf{x}}_k)$ is the noisy gradient estimate at step $k$,
• ${\beta }_1$ is the decay rate for the momentum (typically close to 1).

Second moment update: $$\begin{array}{}\text{(3)}& {\mathbf{s}}_k={\beta }_2{\mathbf{s}}_{k-1}+(1-{\beta }_2)(\mathrm{\nabla }\hat{U}({\mathbf{x}}_k){)}^2\end{array}$$

where

• ${\mathbf{s}}_k$ is the second moment estimate (variance),
• ${\beta }_2$ is the decay rate for the second moment.

$$\begin{array}{}\text{(4)}& {\mathbf{x}}_{k+1}={\mathbf{x}}_k-\eta \frac{{\mathbf{m}}_k}{\sqrt{{\mathbf{s}}_k}+ϵ}\end{array}$$

where

• ${\mathbf{x}}_k$ is the parameter vector at step $k$,
• $\eta$ is the learning rate,
• $ϵ$ is a small constant to prevent division by zero.

Next, we approximate the Adam update equations as continuous-time SDEs by introducing a small time step $\mathrm{\Delta }t$ and converting the discrete update rules into differential equations. We introduce Wiener noise to model the stochasticity in the gradient estimates. And we are very loose! We do not bother to show the SDE is well-posed.

The first moment (momentum) evolves as a stochastic process influenced by the true gradient $\mathrm{\nabla }U(\mathbf{x}(t))$ and stochastic noise $\mathbf{W}(t)$.

First moment update SDE:

$$\begin{array}{}\text{(5)}& d\mathbf{m}(t)=-{\gamma }_m\mathbf{m}(t)dt+(1-{\beta }_1)\mathrm{\nabla }U(\mathbf{x}(t))dt+\sigma d\mathbf{W}(t)\end{array}$$

where

• ${\gamma }_m\approx -\frac{\log ({\beta }_1)}{\mathrm{\Delta }t}$ is the decay rate for momentum,
• $\mathrm{\nabla }U(\mathbf{x}(t))$ is the true gradient,
• $\sigma d\mathbf{W}(t)$ represents the Wiener noise capturing gradient stochasticity.

Second moment update SDE:

The second moment (variance) evolves based on the square of the gradient and also involves the stochastic noise term. The squared gradient noise term simplifies to a deterministic contribution due to Itô calculus.

$$\begin{array}{}\text{(6)}& d\mathbf{s}(t)=-{\gamma }_s\mathbf{s}(t)dt+(1-{\beta }_2)(\mathrm{\nabla }U(\mathbf{x}(t)){)}^{2}dt+{\sigma }^{2}dt\end{array}$$

where

• ${\gamma }_s\approx -\frac{\log ({\beta }_2)}{\mathrm{\Delta }t}$,
• The noise term ${\sigma }^{2}dt$ arises from the squared Wiener process $d\mathbf{W}(t{)}^2$ and simplifies to a deterministic term due to Itô calculus.

Position update SDE:

The position update depends on the momentum and variance estimates, with the learning rate $\eta$ controlling the step size.

$$\begin{array}{}\text{(7)}& d\mathbf{x}(t)=-\eta \frac{\mathbf{m}(t)}{\sqrt{\mathbf{s}(t)}+ϵ}dt\end{array}$$

The division here is element-wise, and $ϵ$ is a small constant (vector) to prevent division by zero.

In this setup, the Wiener noise $d\mathbf{W}(t)$ enters the momentum update directly, while its squared form appears in the variance update. The shared noise term reflects the fact that both the first and second moment estimates are based on the same underlying noisy gradient information.

We transform the system into state-space form, stacking the position, momentum, and variance into a state vector. We express the system as a matrix operation involving deterministic and stochastic parts.

We stack the position ${\mathbf{x}}_k$, momentum ${\mathbf{m}}_k$, and variance ${\mathbf{s}}_k$ into a single state vector ${\mathbf{z}}_k$,

$${\mathbf{z}}_k=\left[\begin{array}{c}{\mathbf{x}}_k\\ {\mathbf{m}}_k\\ {\mathbf{s}}_k\end{array}\right]$$

Based on the SDEs, we now write the discrete-time process updates using the Euler-Maruyama discretisation.

Position Update:

$${\mathbf{x}}_{k+1}={\mathbf{x}}_k-\eta \frac{{\mathbf{m}}_k}{\sqrt{{\mathbf{s}}_k}+ϵ}\mathrm{\Delta }t$$

Momentum Update (with Wiener noise):

$${\mathbf{m}}_{k+1}={\mathbf{m}}_k-{\gamma }_m{\mathbf{m}}_k\mathrm{\Delta }t+(1-{\beta }_1)\mathrm{\nabla }U({\mathbf{x}}_k)\mathrm{\Delta }t+\sigma \sqrt{\mathrm{\Delta }t}\mathit{\xi }$$ Where $\mathit{\xi }\sim \mathcal{N}(0,\mathbf{I})$ is the Gaussian noise driving the stochastic gradient estimates.

Variance Update (with simplified Wiener noise contribution):

$${\mathbf{s}}_{k+1}={\mathbf{s}}_k-{\gamma }_s{\mathbf{s}}_k\mathrm{\Delta }t+(1-{\beta }_2)(\mathrm{\nabla }U({\mathbf{x}}_k){)}^2\mathrm{\Delta }t+{\sigma }^2\mathrm{\Delta }t$$

We can now express the process update equations in state-space matrix update form. Define the system matrix $\mathbf{A}$:

$$\mathbf{A}=\left[\begin{array}{ccc}\mathbf{I}& -\eta \mathrm{\Delta }t\frac{1}{\sqrt{{\mathbf{s}}_k}+ϵ}\mathbf{I}& \mathbf{0}\\ \mathbf{0}& (1-{\gamma }_m\mathrm{\Delta }t)\mathbf{I}& \mathbf{0}\\ \mathbf{0}& \mathbf{0}& (1-{\gamma }_s\mathrm{\Delta }t)\mathbf{I}\end{array}\right]$$

Define the gradient-dependent vector ${\mathbf{b}}_k$,

$${\mathbf{b}}_k=\left[\begin{array}{c}\mathbf{0}\\ (1-{\beta }_1)\mathrm{\nabla }U({\mathbf{x}}_k)\mathrm{\Delta }t\\ (1-{\beta }_2)(\mathrm{\nabla }U({\mathbf{x}}_k){)}^2\mathrm{\Delta }t+{\sigma }^2\mathrm{\Delta }t\end{array}\right]$$

Define the noise matrix $\mathbf{G}$,

$$\mathbf{G}=\left[\begin{array}{c}\mathbf{0}\\ \sigma \sqrt{\mathrm{\Delta }t}\mathbf{I}\\ \mathbf{0}\end{array}\right]$$

The final state-space form for the Adam-type optimizer, including the Wiener noise, is

$${\mathbf{z}}_{k+1}=\mathbf{A}{\mathbf{z}}_k+{\mathbf{b}}_k+\mathbf{G}{\mathit{\xi }}_k$$

where

• ${\mathbf{z}}_k=\left[\begin{array}{c}{\mathbf{x}}_k\\ {\mathbf{m}}_k\\ {\mathbf{s}}_k\end{array}\right]$ is the state vector,
• $\mathbf{A}$ is the system matrix governing the deterministic part of the update,
• ${\mathbf{b}}_k$ is a gradient-dependent vector,
• $\mathbf{G}$ maps the stochastic noise ${\mathit{\xi }}_k\sim \mathcal{N}(0,\mathbf{I})$ into the system.

This is a little misleading to my eye. It looks like a linear update, and in fact like two uncoupled systems. But we inject ${\mathbf{s}}_k$ back into the position update, so the systems are both nonlinear and coupled. This does make things more complicated, but how depends on what we were doing this for.

2.3 Nesterov momentum

See Q. Li, Tai, and Weinan (2019).