notes
← back to notes

optimizer study

Optimizer order 1D / 2D Preconditioner Used Momentum EMA? Curvature aware Notes
SGD 1st order any no
$v_{t+1} = \gamma v_t + \alpha g_t$
$\theta_{t+1} = \theta_t - v_{t+1}$
not in vanilla, but otherwise ema of gradients (momentum) no
PSGD 2nd order any Learned P- affine kronecker factored matrix lie group P / low rank approx P yes
$m \leftarrow 0.9m + 0.1g$
ema of momentum
yes - estimates $H^{-1}$ online using lie groups
Adam 1st order Diagonal preconditioning yes yes EMA of gradients = Mt and EMA of element-wise squared gradients = Vt no problem - l2 regularization added directly to the loss - weight decay scaled down for large gradients, vice versa
AdamW 1st order any diagonal P yes yes similar to adam no decouples weight decay
Adagrad 2nd order any $P^{-1} = \mathrm{diag}(\epsilon I + G_t)^{-1/2}$ no no - uses cumulative sum inverse hessian approx
Adafactor 1st order 2D
adaptive scaling used as diagonal preconditioner -
$X_t = X_{t-1} - \alpha_t G_t / (\sqrt{\hat V_t} + \epsilon)$
sqrtVt is low rank - rank-1
no - they drop momentum
yes - keeps ema of only row and col sum and later reconstructs low rank approx with this
$R_t = \beta_2 R_{t-1} + (1-\beta_2)(G_t^2)\mathbf{1}_n$
$C_t = \beta_2 C_{t-1} + (1-\beta_2)\mathbf{1}_n^\top(G_t^2)$
no replaces Vt from adam with its best rank-1 approximate V't to reduce memory usage
Shampoo 2nd order any
  • layer-wise Kronecker-factored preconditioner
  • has left and right precond
$L_t = L_{t-1} + G_t G_t^\top$
$R_t = R_{t-1} + G_t^\top G_t$
no
they do experiment with momentum tho
no yes - gradient outer products
SOAP 2nd order 2d
diagonal scaling like adam & shampoo's preconditioner in its eigenspace
$Q_L \leftarrow \mathrm{Eigenvectors}(L)$
$Q_R \leftarrow \mathrm{Eigenvectors}(R)$
yes
yes
$M \leftarrow \beta_1 M + (1-\beta_1)G$
$V \leftarrow \beta_2 V + (1-\beta_2)(G' \odot G')$
$L \leftarrow \beta_2 L + (1-\beta_2)GG^\top$
$R \leftarrow \beta_2 R + (1-\beta_2)G^\top G$
gradient outer products like shampoo uses shampoo with power 1/2 instead of 1/4
Muon 1st order 2D no - but orthogonalizes instead yes - nesterov momentum first moment EMA (momentum) no - muon uses geometry of the param space to make update step. steps into tangential plane and retracts to lie on the stiefel manifold. so not aware of loss curvature (??) muon(a 1st order opt) can act as an accumulation free version of shampoo(a 2nd order opt)
KFAC 2nd order 2D Kronecker-factored inverse Fisher yes? ~ approximates inverse fisher (kroneker)
EKFAC 2nd order 2D Kronecker eigenbasis + diagonal yes? ~ approx inverse fisher (eigenval)