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:

1. For $t=1$ to $T$:
   1. Compute adaptive step size:

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

1. 1. Compute gradient:

  $G_{t}=\nabla f_{t}(X_{t-1})$

1. 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})$

1. 1. Compute normalized gradient:

  $U_{t}={\frac {G_{t}}{\sqrt {{\hat {V}}_{t}}}}$

1. 1. Apply clipping:

  ${\hat {U}}_{t}={\frac {U_{t}}{\max(1,{\text{RMS}}(U_{t})/d)}}$

1. 1. Update parameter:

  $X_{t}=X_{t-1}-\alpha _{t}{\hat {U}}_{t}$

1. 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:

1. For $t=1$ to $T$:
   1. Compute adaptive step size:

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

1. 1. Compute gradient:

  $G_{t}=\nabla f_{t}(X_{t-1})$

1. 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}$

1. 1. 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})$

1. 1. Update overall second moment estimate:

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

1. 1. Compute normalized gradient:

  $U_{t}={\frac {G_{t}}{\sqrt {{\hat {V}}_{t}}}}$

1. 1. Apply clipping:

  ${\hat {U}}_{t}={\frac {U_{t}}{\max(1,{\text{RMS}}(U_{t})/d)}}$

1. 1. Update parameter:

  $X_{t}=X_{t-1}-\alpha _{t}{\hat {U}}_{t}$

1. End for

4. Proposed Hyperparameters for Adafactor

• Regularization constant 1: $\epsilon _{1}=10^{-30}$
• Regularization constant 2: $\epsilon _{2}=10^{-3}$
• Clipping threshold: $d=1$
• Relative step size: $\rho _{t}=\min(10^{-2},1/{\sqrt {t}})$
• Second moment decay: ${\hat {\beta }}_{2t}=1-t^{-0.8}$