Adafactor

Author: Aolei Cao (ac3237), Ziyang Li (zl986), Junjia Liang (jl4439) (ChemE 6800 Fall 2024)

Stewards: Nathan Preuss, Wei-Han Chen, Tianqi Xiao, Guoqing Hu

Problem formulation

1. Objective

Minimize the loss function $f(x)$, where $x\in R^{n}$ and $x$ is the weight vector to be optimized.

2. Parameters

• Gradient:

$G_{t}=\nabla f(x_{t-1})$

• Second moment estimate:

${\hat {V}}_{t}={\hat {\beta }}_{2t}{\hat {V}}_{t-1}+(1-{\hat {\beta }}_{2t})(G_{t}^{2}+\epsilon _{1}1_{n})$

• Where:
  • ${\hat {V}}_{t}$ is the running average of the squared gradient.
  • ${\hat {\beta }}_{2t}$ is the corrected decay parameter.
  • $\epsilon _{1}$ is a regularization constant.

• Step size:

$\alpha _{t}=\max(\epsilon _{2},{\text{RMS}}(x_{t-1}))\rho _{t}$

• Where:
  • $\rho _{t}$ is the relative step size.
  • $\epsilon _{2}$ is a regularization constant.
  • ${\text{RMS}}$ is the root mean square, defined as:
    • $u_{xt}={\frac {-g_{xt}}{\sqrt {{\hat {v}}_{xt}}}}$
    • ${\text{RMS}}(U_{t})={\text{RMS}}_{x\in X}(u_{xt})={\sqrt {{\text{Mean}}_{x\in X}\left({\frac {(g_{xt})^{2}}{{\hat {v}}_{xt}}}\right)}}$

3. Algorithms

Adafactor for Weighted Vectors

Inputs:

• Initial point: $X_{0}\in \mathbb {R} ^{n}$
• Relative step sizes: $\rho _{t}$ for $t=1$ to $T$
• Second moment decay: ${\hat {\beta }}_{2t}$ for $t=1$ to $T$, with ${\hat {\beta }}_{21}=0$
• Regularization constants: $\epsilon _{1},\epsilon _{2}$
• Clipping threshold: $d$

Algorithm:

• For $t=1$ to $T$:
  • Compute adaptive step size: $\alpha _{t}=\max(\epsilon _{2},{\text{RMS}}(X_{t-1}))\rho _{t}$
  • Compute gradient: $G_{t}=\nabla f_{t}(X_{t-1})$
  • Update second moment estimate: ${\hat {V}}_{t}={\hat {\beta }}_{2t}{\hat {V}}_{t-1}+(1-{\hat {\beta }}_{2t})(G_{t}^{2}+\epsilon _{1}1_{n})$
  • Compute normalized gradient: $U_{t}={\frac {G_{t}}{\sqrt {{\hat {V}}_{t}}}}$
  • Apply clipping: ${\hat {U}}_{t}={\frac {U_{t}}{\max(1,{\text{RMS}}(U_{t})/d)}}$
  • Update parameter: $X_{t}=X_{t-1}-\alpha _{t}{\hat {U}}_{t}$
• End for

Adafactor for Weighted Matrices

Inputs:

• Initial point: $X_{0}\in \mathbb {R} ^{n\times m}$
• Relative step sizes: $\rho _{t}$ for $t=1$ to $T$
• Second moment decay: ${\hat {\beta }}_{2t}$ for $t=1$ to $T$, with ${\hat {\beta }}_{21}=0$
• Regularization constants: $\epsilon _{1},\epsilon _{2}$
• Clipping threshold: $d$

Algorithm:

• For $t=1$ to $T$:
  • Compute adaptive step size: $\alpha _{t}=\max(\epsilon _{2},{\text{RMS}}(X_{t-1}))\rho _{t}$
  • Compute gradient: $G_{t}=\nabla f_{t}(X_{t-1})$
  • Update row-wise second moment: $R_{t}={\hat {\beta }}_{2t}R_{t-1}+(1-{\hat {\beta }}_{2t})(G_{t}^{2}+\epsilon _{1}1_{n}1_{m}^{T})1_{m}$
  • Update column-wise second moment: $C_{t}={\hat {\beta }}_{2t}C_{t-1}+(1-{\hat {\beta }}_{2t})1_{n}^{T}(G_{t}^{2}+\epsilon _{1}1_{n}1_{m}^{T})$
  • Update overall second moment estimate: ${\hat {V}}_{t}={\frac {R_{t}C_{t}}{1_{n}^{T}R_{t}}}$
  • Compute normalized gradient: $U_{t}={\frac {G_{t}}{\sqrt {{\hat {V}}_{t}}}}$
  • Apply clipping: ${\hat {U}}_{t}={\frac {U_{t}}{\max(1,{\text{RMS}}(U_{t})/d)}}$
  • Update parameter: $X_{t}=X_{t-1}-\alpha _{t}{\hat {U}}_{t}$
• End for

Why Adafactor is more memory efficient, compared to Adam

Row-wise and Column-wise Second Moment Updates

• $R_{t}={\hat {\beta }}_{2t}R_{t-1}+(1-{\hat {\beta }}_{2t})(G_{t}^{2}+\epsilon _{1}1_{n}1_{m}^{T})1_{m}$
• $C_{t}={\hat {\beta }}_{2t}C_{t-1}+(1-{\hat {\beta }}_{2t})1_{n}^{T}(G_{t}^{2}+\epsilon _{1}1_{n}1_{m}^{T})$

Instead of storing the full $G_{t}^{2}$, Adafactor computes the row and column respectively, which reduces the memory requirements from $O(n\times m)$ to $O(n+m)$

Factored Representation of the Second Moment

• ${\hat {V}}_{t}={\frac {R_{t}C_{t}}{1_{n}^{T}R_{t}}}$

This updates the second momentum based on the outer product $R_{t}C_{t}$.

• However, this is not $O(n\times m)$ since
  • The operation is performed element-wise, so it actually never materializes ${\hat {V_{t}}}$ as a $n\times n$ matrix
  • It also only storing $R_{t}$and $C_{t}$ instead of storage the full second-moment matrix

4. Proposed Hyperparameters for Adafactor

• Regularization constant 1: $\epsilon _{1}=10^{-30}$
• Ensures numerical stability by preventing division by zero in the calculation of second-moment estimates, so the numerical value should be very close to zero
• Regularization constant 2: $\epsilon _{2}=10^{-3}$
• Help to stabilize parameter updates by controlling the effect of second-moment scaling in low-magnitude scenarios. Compared to $\epsilon _{2}$, a relatively larger value ensures the stability of noise and low-magnitude scenarios.
• Clipping threshold: $d=1$
• A threshold of 1 balances stability and learning efficiency. It avoids excessive suppression of large gradients, which could hinder learning, while still protecting against extreme updates that could destabilize the model.
• Relative step size: $\rho _{t}=\min(10^{-2},1/{\sqrt {t}})$
  • $min(10^{-}2,...)$ can caps the learning rate at 10^-2, which is a empirical found for upper bound
  • ${\frac {1}{\sqrt {t}}}$ This step size promote convergence of the model. This rate ensures a balance between sufficient exploration in early iteration and stability in later iterations
• Second moment decay: ${\hat {\beta }}_{2t}=1-t^{-0.8}$
  • 1-...: ensures the decay factor remains close to 1
  • $t^{-0,8}$ the power 0.8 ensures a balance between rapid adaptation in early training and later iterations

Why Clipping

Adafactor employs clipping to maintain numerical stability, especially since it is designed for use with very large models and often works with unscaled learning rates.

• Clipping prevents the update step from becoming very large, which would destabilize training
• Clipping mitigates the effects of very large gradients preventing numerical instability

Therefore, implementing clipping helps ensure stability and efficient training without requiring per-parameter scaling like Adam.

Why Adafactor is more memory efficient, compared to Adam

Row-wise and Column-wise Second Moment Updates

• $R_{t}={\hat {\beta }}_{2t}R_{t-1}+(1-{\hat {\beta }}_{2t})(G_{t}^{2}+\epsilon _{1}1_{n}1_{m}^{T})1_{m}$
• $C_{t}={\hat {\beta }}_{2t}C_{t-1}+(1-{\hat {\beta }}_{2t})1_{n}^{T}(G_{t}^{2}+\epsilon _{1}1_{n}1_{m}^{T})$

Instead of storing the full $G_{t}^{2}$, Adafactor computes the row and column respectively, which reduces the memory requirements from $O(n\times m)$ to $O(n+m)$

Factored Representation of the Second Moment

• ${\hat {V}}_{t}={\frac {R_{t}C_{t}}{1_{n}^{T}R_{t}}}$

This updates the second momentum based on the outer product $R_{t}C_{t}$.

• However, this is not $O(n\times m)$ since
  • The operation is performed element-wise, so it actually never materializes ${\hat {V_{t}}}$ as a $n\times n$ matrix
  • It also only storing $R_{t}$and $C_{t}$ instead of storage the full second-moment matrix