TL;DR. AdaGrad's square root compensates for the quadratic scaling of gradient outer products. To remove it, we need to match the curvature scale: for an averaged batch loss, use $B\bar{\mathbf g}\bar{\mathbf g}^\top$.

The AdaGrad Puzzle

Let the parameter vector be $\boldsymbol{\theta}$ and the batch gradient at step $t$ be $\mathbf g_t$. Full-matrix AdaGrad accumulates

$$ \mathbf S_t=\mathbf S_{t-1}+\mathbf g_t\mathbf g_t^\top, $$

but uses the inverse square root:

$$ \boldsymbol{\theta}_{t+1}

\boldsymbol{\theta}_t

\eta(\mathbf S_t+\lambda\mathbf I)^{-1/2}\mathbf g_t. $$

If $\mathbf S_t$ estimates the Hessian, Newton's method suggests using $\mathbf S_t^{-1}$. A natural question arises: Why does AdaGrad use a square root, and when can we remove it?

Why Uses Square Root?

Suppose a loss is multiplied by a positive constant:

$$ \mathcal L'(\boldsymbol{\theta})=c\mathcal L(\boldsymbol{\theta}), \qquad c>0. $$

Then

$$ \mathbf g_t'=c\mathbf g_t, \qquad \mathbf S_t'=c^2\mathbf S_t. \\ (\mathbf S_t')^{-1/2}\mathbf g_t'

(c^2\mathbf S_t)^{-1/2}(c\mathbf g_t)

\mathbf S_t^{-1/2}\mathbf g_t. $$

Note that the AdaGrad direction is unchanged. However, without the square root,

$$ (\mathbf S_t')^{-1}\mathbf g_t'

c^{-1}\mathbf S_t^{-1}\mathbf g_t. $$

Now the direction depends on the loss scale.

<aside> 💡

Takeaway: Square-root fixes the loss scaling problem.

</aside>

The Scaling Mismatch

The square root fixes AdaGrad's loss-scale dependence, but it exposes a deeper mismatch. Under the rescaling from the previous section,

$$ \nabla_{\boldsymbol{\theta}}^2\mathcal L \mapsto c\nabla_{\boldsymbol{\theta}}^2\mathcal L, \qquad \mathbf g\mathbf g^\top \mapsto c^2\mathbf g\mathbf g^\top. $$

The Hessian scales linearly with the loss, while the gradient outer product scales quadratically. Yet the outer product is often used to estimate the Hessian. How can both statements be true?