Let’s start with a second-order Taylor approximation of the loss, where $g$ is the gradient vector, $H$ is the Hessian (matrix of second-order partial derivative), and $\Delta$ is a vector of small parameter changes:

\begin{equation} L(\theta + \Delta) \approx L(\theta) + g^T\Delta + \frac{1}{2}\Delta^T H\Delta \tag{1} \end{equation}

Using only the first order, moving along the negative gradient decreases the loss for a sufficiently small step.

Including the second order is like fitting a parabola in 1D. If we assume we are near a minimum and that the $H$ is positive definite (all eigenvalues are positive, like a smiling parabola), the quadratic approximation has a minimum which we can find analytically.

Let’s set the loss derivative, in the proximity of the current parameters, to zero, in order to find the desired minimum we want to reach to:

\begin{equation} \nabla_\Delta\left(L(\theta) + g^T\Delta + \frac{1}{2}\Delta^T H\Delta\right) = g + H\Delta = 0 \quad\Rightarrow\quad \Delta = -H^{-1}g \tag{2} \end{equation}

And we found where we need to move to, in order to minimize the loss. This is the Newton step: the inverse Hessian basically rescales (=preconditions) the gradient. Positive definiteness guarantees a descent direction, but the full step can still overshoot. An indefinite Hessian can be damped by adding a sufficiently large positive multiple of the identity.

Now let’s move to Adam: Taylor expand the gradient:

\[g(\theta) = g(\theta^{\ast}) + H(\theta^{\ast})(\theta-\theta^{\ast}) \tag{3}\]

And we know that $g(\theta^{\ast}) = 0$, thus:

\[g(\theta) = H(\theta^{\ast})(\theta-\theta^{\ast}) = H\Delta\theta \tag{4}\]

where $\Delta\theta$ is the distance vector of the parameters from their minimum.

Define $H = H(\theta^{\ast})$ to be the Hessian at the minimum. It is constant. Now let’s expand the calculation for one coordinate:

\[g_i = \sum_j H_{i,j}\Delta\theta_j \tag{5}\]

And square it:

\[g_i^2 = \left(\sum_j H_{i,j}\Delta\theta_j\right)^2 = \sum_j\sum_k H_{i,j}H_{i,k}\Delta\theta_j\Delta\theta_k \tag{6}\]

Now take expectation over trajectories or noise in $\theta$ ($H$ does not change):

\[\mathbb{E}[g_i^2] = \sum_{j,k} H_{i,j}H_{i,k}\mathbb{E}[\Delta\theta_j\Delta\theta_k] \tag{7}\]

If we assume weak correlation, we have for $j \ne k$:

\[\mathbb{E}[\Delta\theta_j\Delta\theta_k] = 0 \tag{8}\]

So:

\[\mathbb{E}[g_i^2] = \sum_j H_{i,j}^2\mathbb{E}[(\Delta\theta_j)^2] \tag{9}\]

If coordinates are roughly decoupled (dominant diagonal in $H$ only), then:

\[\begin{aligned} \mathbb{E}[g_i^2] &= H_{i,i}^2\mathbb{E}[\Delta\theta_i^2] \\ \Rightarrow\quad (H^{-1})_{i,i} &= \frac{1}{H_{i,i}} = \frac{\overbrace{\sqrt{\mathbb{E}[\Delta\theta_i^2]}}^{\sim\,\text{constant}}}{\sqrt{\mathbb{E}[g_i^2]}} \end{aligned} \tag{10}\]

And this is why Adam divides by the square root of $v$, which keeps an EMA of the second moment.