Major recent updates: added gradient whitening preconditioners to all classes; wrapped affine preconditioners as a class (also support complex matrices); preconditioner fitting step size normalization with the 2nd derivative info.
PSGD (Preconditioned SGD) is a general purpose (mathematical and stochastic, convex and nonconvex) 2nd order optimizer. Unlike many 2nd order optimizers, PSGD does not rely on damping terms to stabilize it.
Notations:
Criterion | Solution | Notes |
---|---|---|
Reduces to secant equation |
||
Reduces to Newton's method when |
||
|
||
Reduces to AdaGrad family, e.g., Adam(W), RMSProp, Shampoo, |
Table II: Implemented Lie groups preconditioners with storage and computation numbers for $\theta={\rm vec}(\Theta)$ and $\Theta\in\mathbb{R}^{m\times m}$
Lie Group | Update of |
Storages | Computations | Notes |
---|---|---|---|---|
See class Newton; set keep_invQ=True to calculate |
||||
Triangular matrices | See class Newton; set keep_invQ=False to make |
|||
Diagonal matrices, |
See class either LRA with rank_of_approximation=0 or XMat for implementations. | |||
|
|
See class Affine for implementations (also support complex matrices and diagonal matrices by setting preconditioner_max_size=0). | ||
|
See class LRA for implementations; typically |
|||
Scaling-and-normalization | no class implementation for now | Hard to make it a universal black-box preconditioner. |
For AdaGrad like preconditioner, we simply replace pair
This script generates the following plot showing the typical behaviors of different preconditioner fitting methods.
- With a static and noise-free Hessian-vector product model, both BFGS and PSGD converge linearly to the optimal preconditioner while closed-form solution
$P=\left(E[hh^T]\right)^{-0.5}$ only converges sublinearly with rate$\mathcal{O}(1/t)$ . - With a static additive noisy Hessian-vector model
$h=Hv+\epsilon$ , BFGS diverges easily. With a constant step size$\mu$ , the steady-state fitting errors of PSGD are proportional to$\mu$ . - With a time-varying Hessian
$H_{t+1}=H_t + uu^T$ and$u\sim\mathcal{U}(0,1)$ , PSGD locks onto good preconditioner estimations quicker than BFGS, also no divergence before convergence. The closed-form solution$P=\left(E[hh^T]\right)^{-0.5}$ is not good at tracking due to its sublinear convergence.
Optimizers with the criteria in Table 1 and preconditioner forms in Table 2 are wrapped into classes XMat, LRA (or UVd), Newton and Affine for easy use. The affine family are not black box preconditioners, and the users need to either group the gradients into matrices or reform the models such that the parameters are a list of matrices, e.g., this demo on wrapping torch.nn.functional.conv2d as an affine Conv2d class, and this one on wrapping torch._VF.rnn_tanh as an affine RNN class. Three main differences from torch.optim.SGD:
- The loss to be minimized is passed through as a closure to the optimizer to support more complicated behaviors, notably, Hessian-vector product approximation with finite difference method when the 2nd order derivatives are not available.
- Momentum here is the moving average of gradient so that its setting is decoupled from the learning rate, which is always normalized in PSGD.
- As any other regularizations, (coupled) weight decay should be explicitly realized by adding
$L2$ regularization to the loss. Similarly, decoupled weight decay is not included inside the PSGD implementations.
Rosenbrock function minimization: see how simple to apply PSGD to convex and stochastic optimizations. Typically, we just set preconditioner_init_scale=None, and may need to tune lr_params and lr_preconditioner a little.
LeNet5: demonstration of PSGD on convolutional neural network training with the classic LeNet5 for MNIST digits recognition. Also see this for another implementation and comparison with Shampoo.
Delayed XOR: demonstration of PSGD on gated recurrent neural network (RNN) learning with the delayed XOR problem proposed in the LSTM paper. Note that neither LSTM nor the vanilla RNN can solve this "naive" problem with most optimizers, while PSGD can, with either the LSTM or the vanilla RNN (also see this and this with simple RNNs).
Tensor rank decomposition benchmark: demonstrate the usage of all preconditioners on the tensor rank decomposition problem. PSGD performs better than BFGS.
Logistic regression: a large-scale (
Preconditioner fitting on Lie groups: see how multiplicative updates work on Lie groups for different types of preconditioners:
Preconditioner estimation efficiency and numerical stability: a playground to compare PSGD with BFGS and closed-form solution
How PSGD generalizes so well: This one serves as a good toy example illustrating it in the view of information theory. Starting from the same initial guesses, PSGD tends to find minima with smaller train cross entropy and flatter Hessians than Adam, i.e., shorter description lengths (DL) for train image-label pairs and model parameters. Similarly, this example shows that PSGD also generalizes better than Shampoo.
PSGD vs approximated closed-form solutions: This example shows that most approximated closed-form solutions, e.g., KFAC, Shampoo, CASPR to name a few, are pretty rough in terms of accuracy, although they do work in certain circumstances.
- Preconditioned stochastic gradient descent, arXiv:1512.04202, 2015. (General ideas of PSGD, preconditioner fitting losses and Kronecker product preconditioners.)
- Preconditioner on matrix Lie group for SGD, arXiv:1809.10232, 2018. (Focus on preconditioners with the affine Lie group.)
- Black box Lie group preconditioners for SGD, arXiv:2211.04422, 2022. (Mainly about the LRA preconditioner. I also have prepared these supplementary materials for detailed math derivations.)
- Stochastic Hessian fittings on Lie groups, arXiv:2402.11858, 2024. (Some theoretical works on the efficiency of PSGD. The Hessian fitting problem is shown to be strongly convex on group
${\rm GL}(n, \mathbb{R})/R_{\rm polar}$ .) - Curvature-informed SGD via general purpose Lie-group preconditioners, arXiv:2402.04553, 2024. (Plenty of benchmark results and analyses for PSGD vs. other optimizers.)
- Other implementations: Tensorflow 1.x and TensorFlow 2.x. I no longer maintain them as I have not used Tensorflow for some time.