Wolfram Cloud Document

ORIGINAL ARTICLE: Jarek Duda, Improving SGD convergence by online linear regression of gradients in multiple statistically relevant directions, arXiv:1901.11457v11 [cs.LG]. https://doi.org/10.48550/arXiv.1901.11457
GITHUB: https://github.com/JarekDuda/SGD-OGR-Hessian-estimator

Abstract: The AI revolution is based on neural networks as extremely large parametrizations, often of billions and more parameters, optimizing some objective function. The optimization process, called neural network training, is generally done by SGD: stochastic gradient descent - going along local gradients, with stochasticity as noise in the calculated gradients (based on batches as small subsets of dataset). Optimization of this training process could translate into gigantic savings worldwide. It is tempting to include second order derivatives (Hessian), but it turns out very difficult due to especially: Hessian size, estimation from noisy gradients, saddle attraction, and badly conditioned Hessian. There is discussed novel approach overcoming these difficulties: OGR (online gradient regression) - estimating Hessian from linear regression of gradients, done in adaptive way: with exponentially weakening weights. For high dimensions there are discussed two approaches (finally to be combined): diagonal (dOGR) approximating Hessian as diagonal represented as vector, and subspace (sOGR) updating locally interesting linear subspace of chosen dimension with full Hessian model there.

This notebook briefly introduces to the methodology, containing Mathematica code and working demonstration at the end, further sources are available in GitHub, more details can be found in arXiv:1901.11457 and video lecture. Here is popular overview of standard methods: https://www.ruder.io/optimizing-gradient-descent/

The task is finding local minimum of function

R

, like FindMinimum[] Mathematica function, in potentially extremely large dimension

like millions/billions, being able to ask only about "noisy" gradients in a chosen sequence of points

≈

∇

)

, with noise as an unknown error - for simplicity here assumed to be from

(0,σ)

Gaussian distribution: noise:=RandomVariate[NormalDistribution[0,σ],D]. The simplest approach is often called "vanilla SGD": just subtract gradient multiplied by some learning rate

t+1



-η

. To handle noise, there is popular momentum approach replacing gradients with their exponential moving average:

t+1

β

+(1-β)

then

t+1



-η

. Further successful optimizations divide step in each direction by square root of mean squared gradient in this direction: in AdaGrad, RMSProp, and finally very popular ADAM method. While these are 1st order methods, it seems tempting to search for 2nd order based on Newton's method - trying to model the second derivative for smarter choice of step size (learning rate), ideally getting to minimum of parabola in a single step. However, there are many difficulties, like:

size of full Hessian, difficulty and cost of its estimation from noisy gradients, attraction to saddles of basic Newton method, inversion of usually badly conditioned Hessian. Therefore, while there were many attempts for 2nd order methods, like nonlinear conjugated gradients, Gauss-Newton, TONGA, quasi-Newton, K-FAC, Shampoo, so far none of them was good enough to get a wider popularity. We will discuss OGR (online gradient regression) novel approach for Hessian estimation from noisy gradients, also of its approximations (d,s,c variants), which overcome these difficulties, at least in low dimensions providing much better convergence - hopefully will lead to better optimizers. The basic idea is very simple: use linear regression for gradients vs positions, this way getting estimation of 2nd derivative. To get local estimators, this linear regression uses exponentially weakening weights, what can be inexpensively realized by updating a few exponential moving averages.

Let us start with 1D case: assuming

f(θ)=

λ(θ-p)

2

function, from a sequence of values

(

)

and gradients





, we would like to estimate its parameters: second derivative

and position of minimum

Here is sketch of derivation of such Hessian estimator in general multidimenstional case - directly we get non-symmetric, further there is found symmetric. It requires to update two vectors:

as mean position and gradient, also two covariance matrices

θθ

In the shown further implementations, to prevent saddle attraction there are used absolute values of Hessian eigenspectrum as in saddle-free Newton method. To handle badly conditioned Hessian to be inverted, for regularization its eigenspectrum is bound from below by “cut” parameter and divided by “div”. Here is simple implementation using the above Hessian estimator for 2D popular Beale function, example trajectories, and comparison with momentum and ADAM method - with evaluation as sorted values after 30 steps starting in points from {-2,...,2} x {-1,0,1} lattice:

Above full Hessian model, referred as fOGR, is practical rather only in low dimensions. For high dimensional situation we need approximations. There are considered two basic approaches: diagonal variant dOGR approximating Hessian as diagonal, allowing to model it as vector, this way modeling separate parabola for each canonical direction. Second approximation approach is subspace variant sOGR: modeling full Hessian in

d<D

dimensional subspace (e.g.

d10

dimensional), which is locally interesting - can be found e.g. with online PCA of sequence of gradients, here using a simpler inexpensive approach instead: regularly adding new gradient times some small

to all vectors of the current basis V, and Gram-Schmidt orthogonalizing this basis.

Finally dsOGR simultaneously updates both d and s Hessian model, in the current subspace V uses average of their predicted steps, and in the orthogonal directions use step from d variant. Below are such simple implementations and evaluation as sorted values after 50 steps for 3D Beale function and 343 starting points:

CITE THIS NOTEBOOK

SGD optimizers with online gradient regression Hessian estimation
by Jarek Duda
Wolfram Community, STAFF PICKS, May 10 2023
https://community.wolfram.com/groups/-/m/t/2917908