Adaptive stochastic gradient descent

1.1 As a Bayes procedure 1

Least squares is only ever one Bayesian saying well actually… away from being a Gaussian posterior updates.

Suppose the parameter vector at time $t$ to be ${\varphi }_t$ is a random variable and we want to see what kind of Bayesian interpretation we can give to the Adam update rule. Assume a Gaussian prior over the parameters $\varphi$,

$$\begin{array}{r}p(\varphi )=\mathcal{N}(\varphi ;{\overline{\varphi }}_t,{\mathbf{H}}_t^{-1}),\end{array}$$

where

• ${\overline{\varphi }}_t$ is the mean of the prior distribution.
• ${\mathbf{H}}_t^{-1}$ is the covariance matrix of the prior, reflecting our confidence in the current estimate ${\overline{\varphi }}_t$.

The likelihood function is based on observing a gradient $m_t$, which provides information about the parameters,

$$\begin{array}{r}p(m_t|\varphi )=\mathcal{N}(m_t;0,{\lambda }_t^{-1}{\mathbf{H}}_t),\end{array}$$

where the likelihood encourages $\varphi$ to move in the direction that reduces the loss function. Applying Bayes’ theorem,

$$\begin{array}{rl}p(\varphi |m_t)& \propto p(m_t|\varphi )p(\varphi )\\ & =\mathcal{N}(\varphi ;{\overline{\varphi }}_{t+1},{\mathbf{H}}_t^{-1}),\end{array}$$

where ${\overline{\varphi }}_{t+1}$ is the mean of the posterior distribution, given by

$$\begin{array}{r}{\overline{\varphi }}_{t+1}={\overline{\varphi }}_t-{\lambda }_t{\mathbf{H}}_t^{-1}{m}_t.\end{array}$$

This matches the update rule derived from the regularized least-squares problem and the Adam optimizer. Let’s consider when the likelihood function

$$\begin{array}{r}p(m_t|\varphi )=\mathcal{N}(m_t;0,{\lambda }_t^{-1}{\mathbf{H}}_t)\end{array}$$

might be a meaningful choice. In those terms, what likelihood function have we just assumed? We need to find a likelihood function $p(\text{data}|\varphi )$ such that its negative log-likelihood corresponds to the linear term ${\lambda }_t(\mathrm{\nabla }L({\varphi }_t){)}^{\mathrm{\top }}(\varphi -{\varphi }_t)$ in our optimization problem.

However, a standard Gaussian likelihood leads to a quadratic term in $\varphi$, not a linear one. What are we even doing, then?

Let’s define a synthetic observation $y_t$ such that $$\begin{array}{r}{y}_t=A(\varphi -{\overline{\varphi }}_t)+{ϵ}_t\end{array}$$

where

• $A$ is a matrix we will define.
• ${ϵ}_t$ is Gaussian noise: ${ϵ}_t\sim \mathcal{N}(0,{\mathbf{V}}_t)$. Our goal is to choose $A$ and ${\mathbf{V}}_t$ such that the negative log-likelihood corresponds to the linear term in our optimization problem.

The likelihood function is

$$\begin{array}{r}p(y_t|\varphi )=\mathcal{N}(y_t;A(\varphi -{\overline{\varphi }}_t),{\mathbf{V}}_t)\end{array}$$

The negative log-likelihood (up to a constant) is:

$$\begin{array}{r}-\log p(y_t|\varphi )=\frac{1}{2}(y_t-A(\varphi -{\overline{\varphi }}_t){)}^{\mathrm{\top }}{\mathbf{V}}_t^{-1}(y_t-A(\varphi -{\overline{\varphi }}_t))\end{array}$$

If we fix $A=I$ (identity matrix) the negative log-likelihood becomes

$$\begin{array}{rl}-\log p(y_t|\varphi )& =\frac{1}{2}(y_t-(\varphi -{\overline{\varphi }}_t){)}^{\mathrm{\top }}{\mathbf{V}}_t^{-1}(y_t-(\varphi -{\overline{\varphi }}_t))\\ & =\frac{1}{2}(\varphi -{\overline{\varphi }}_t-y_t{)}^{\mathrm{\top }}{\mathbf{V}}_t^{-1}(\varphi -{\overline{\varphi }}_t-y_t)\end{array}$$

If we assume that $\varphi -{\overline{\varphi }}_t$ is small but $y_t$ is not, we can approximate the negative log-likelihood as

$$\begin{array}{r}-\log p(y_t|\varphi )\approx -(\varphi -{\overline{\varphi }}_t{)}^{\mathrm{\top }}{\mathbf{V}}_t^{-1}{y}_t\end{array}$$

This is an assumption that gives us a linear term in $\varphi$. (It’s not very satisfying though). The negative log of the posterior is (up to a constant)

$$\begin{array}{rl}-\log p(\varphi |y_t)& =-\log p(y_t|\varphi )-\log p(\varphi )\\ & \approx \frac{1}{2}(\varphi -{\overline{\varphi }}_t{)}^{\mathrm{\top }}{\mathbf{H}}_t(\varphi -{\overline{\varphi }}_t)-(\varphi -{\overline{\varphi }}_t{)}^{\mathrm{\top }}{\mathbf{V}}_t^{-1}{y}_t\end{array}$$

To find the posterior mean which is also the posterior maximum ${\overline{\varphi }}_{t+1}={\mathrm{argmax}}_{\varphi }-\log p(\varphi |y_t)$, we set the gradient of the negative log-posterior to zero,

$$\begin{array}{rl}0& ={\mathrm{\nabla }}_{\varphi }[\frac{1}{2}(\varphi -{\overline{\varphi }}_t{)}^{\mathrm{\top }}{\mathbf{H}}_t(\varphi -{\overline{\varphi }}_t)-(\varphi -{\overline{\varphi }}_t{)}^{\mathrm{\top }}{\mathbf{V}}_t^{-1}{y}_t]\\ & ={\mathbf{H}}_t(\varphi -{\overline{\varphi }}_t)-{\mathbf{V}}_t^{-1}{y}_t\\ \Rightarrow {\overline{\varphi }}_{t+1}& ={\overline{\varphi }}_t+{\mathbf{H}}_t^{-1}{\mathbf{V}}_t^{-1}{y}_t\end{array}$$

Recall that $y_t=\mathrm{\nabla }L({\overline{\varphi }}_t)+ϵ$, so

$$\begin{array}{r}{\overline{\varphi }}_{t+1}={\overline{\varphi }}_t-{\lambda }_t{\mathbf{H}}_t^{-1}{\mathbf{V}}_t^{-1}\mathrm{\nabla }L({\overline{\varphi }}_t)\end{array}$$ To recover the Adam update rule, we need to choose ${\mathbf{H}}_t^{-1}{\mathbf{V}}_t^{-1}=\mathrm{diag}\left(\frac{1}{\sqrt{v_t}+ϵ}\right)$.

While that was coherent, I am not sure I learned anything. This still seems arbitrary; for one thing, the Adam update uses a square root of a variance matrix, which does not naturally arise here.

1.2 As a Bayes procedure 2

Let us try something different and see if we can get a more plausible rationale. We propose that the uncertainty in the parameters is proportional to the square root of the gradient variance,

$$\begin{array}{r}\mathrm{Var}[{\varphi }_t]\propto \sqrt{v_t}.\end{array}$$

This means that the standard deviation (uncertainty) of each parameter ${\varphi }_t^i$ is proportional to $\sqrt{v_t^i}$, where $v_t^i$ is the estimated variance of the $i$-th parameter’s gradient. Assume a prior distribution over the parameters $\varphi$

$$\begin{array}{r}p(\varphi )=\mathcal{N}(\varphi ;{\overline{\varphi }}_t,{\mathrm{\Sigma }}_t),\end{array}$$

where

• ${\overline{\varphi }}_t$ is the mean of the prior distribution (current estimate).
• ${\mathrm{\Sigma }}_t$ is the prior covariance matrix, reflecting our uncertainty about ${\varphi }_t$. We set the prior covariance matrix to

$$\begin{array}{r}{\mathrm{\Sigma }}_t={\sigma }^2\mathrm{diag}(\sqrt{v_t}),\end{array}$$

where ${\sigma }^2$ is a scaling factor. The observed gradient $m_t$ provides information about the parameters,

$$\begin{array}{r}{m}_t=\mathrm{\nabla }L({\overline{\varphi }}_t)+{ϵ}_t,\end{array}$$

where $\mathrm{\nabla }L({\overline{\varphi }}_t)$ is the true gradient of the loss function at ${\overline{\varphi }}_t$, and ${ϵ}_t$ represents the noise in the gradient estimation due to stochastic sampling, with ${ϵ}_t\sim \mathcal{N}(0,V_t),$ and $V_t=\mathrm{diag}(v_t)$. Now, we want the posterior mean update for $\varphi$,

$$\begin{array}{r}{\overline{\varphi }}_{t+1}={\overline{\varphi }}_t-K_t{m}_t,\end{array}$$

where $K_t$ is a gain matrix that adjusts the update based on the uncertainties in the parameters and the observations (gradients). The gain matrix $K_t$ is the usual, $$\begin{array}{r}{K}_t={\mathrm{\Sigma }}_t(V_t+{\mathrm{\Sigma }}_t{)}^{-1}.\end{array}$$

Given that ${\mathrm{\Sigma }}_t$ and $V_t$ are diagonal matrices in Adam, the inversion and multiplication are straightforward.

Assuming ${\mathrm{\Sigma }}_t\ll V_t$ (i.e., parameter uncertainty is small compared to gradient noise), we can approximate

$$\begin{array}{rl}{V}_t+{\mathrm{\Sigma }}_t& \approx V_t\\ K_t& ={\mathrm{\Sigma }}_t{V}_t^{-1}.\end{array}$$

Substituting ${\mathrm{\Sigma }}_t={\sigma }^2\mathrm{diag}(\sqrt{v_t})$ and $V_t=\mathrm{diag}(v_t)$, then ${\lambda }_t={\sigma }^2$ we have

$$\begin{array}{r}{K}_t={\sigma }^2\mathrm{diag}(\sqrt{v_t})\mathrm{diag}(v_t{)}^{-1}\\ & ={\sigma }^2\mathrm{diag}\left(\frac{1}{\sqrt{v_t}}\right).\\ {\overline{\varphi }}_{t+1}& ={\overline{\varphi }}_t-K_t{m}_t={\overline{\varphi }}_t-{\sigma }^2\mathrm{diag}\left(\frac{1}{\sqrt{v_t}}\right)m_t\\ & ={\overline{\varphi }}_t-{\lambda }_t\frac{m_t}{\sqrt{v_t}}.\end{array}$$

This matches the Adam update rule (neglecting the small constant $ϵ$).

By setting the parameter uncertainty proportional to the square root of the gradient variance, we adjust the parameter updates based on how much we “trust” the gradient information:

• High Gradient Variance ($v_t^i$ large): Indicates high uncertainty in the gradient estimate for parameter ${\varphi }_t^i$, leading to a small update due to larger $\sqrt{v_t^i}$ in the denominator.
• Low Gradient Variance ($v_t^i$ small): Indicates more reliable gradient information, leading to a larger update for that parameter.

The coincidence of the prior noise happening to be proportional to the square root of the gradient variance still feels a little bit arbitrary. We are using those 2nd moment estimates twice, once to invent the prior and then again for likelihood, and in a strange way.

1.3 As a Bayes procedure 3

OK, that was weirder than I expected, especially because of that danged square root. It can be eliminated, but with a more involved method (Lin et al. 2024). Maybe we should follow M. E. Khan and Rue (2024) and interpret it as a natural gradient method. TBC

1.4 As a gradient flow

See gradient flows

1.5 AdamW

TBD.