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

Introduction

Adafactor is an efficient, adaptive learning rate optimization algorithm proposed by Noam Shazeer and Mitchell Stern in 2018. ¹

Unlike traditional Adam optimizers, Adafactor does not store complete second-order moment matrices. Instead, it employs a factorization approach that only maintains gradient statistics for the rows and columns of parameter matrices, significantly reducing memory usage. Moreover, Adafactor uses an adaptive learning rate, allowing it to dynamically adjust step sizes without the need for manually setting a global learning rate or relying heavily on hyperparameter tuning. Its design also defaults to not performing bias correction, yet it remains stable in scenarios involving large-batch training data.¹ This efficiency makes it an ideal choice for training ultra-large-scale models such as T5.²

Adafactor’s efficient memory usage and outstanding performance make it widely applicable in scenarios such as Natural Language Processing (NLP).² Compared to the Adam optimizer, Adafactor significantly reduces memory and computational resource requirements while maintaining comparable performance when training large-scale language models and vision models. ^3,6

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

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

Numerical Examples

Step-by-step instructions for determining the result of the first iteration.

Problem setup

Initial weights ($X_{0}$​):

$X_{0}={\begin{bmatrix}0.7&-0.5&0.9\\-1.1&0.8&-1.6\\1.2&-0.7&0.4\end{bmatrix}}$

Gradient for first iteration (​$G_{1}$):

Gradient of the loss function with respect to X

$G_{1}={\begin{bmatrix}0.3&-0.2&0.4\\-0.5&0.6&-0.1\\0.2&-0.4&0.3\end{bmatrix}}$

Hyperparameters setup

$\epsilon _{1}=10^{-30}$ (Minimum learning rate scaling factor))

$\epsilon _{2}=10^{-3}$ (Regularization constant)

$d=1$ (Clipping threshold)

$\rho _{t}=\min(10^{-2},1/{\sqrt {t}})$ (Relative step size)

${\hat {\beta }}_{2t}=1-t^{-0.8}$ (Second moment decay)

Step 1: Learning Rate Scaling

Define the relative step size

$\rho _{1}=\min(10^{-2},1/{\sqrt {1}})=10^{-2}$

Step 1.1: Root Mean Square(RMS) calculation for $X_{0}$

Root Mean Square(RMS) calculation for $X_{0}$

RMS formula

$RMS(X_{0})={\sqrt {{\tfrac {1}{n}}\sum _{i=1}^{n}X_{0}[i]^{2}}}$

Substitute the initial weights

$RMS(X_{0})={\sqrt {{\tfrac {1}{9}}(0.72^{2}+(-0.5)^{2}+0.9^{2}+(-1.1)^{2}+0.8^{2}+(-0.6)^{2}+1.2^{2}+(-0.7)^{2}+0.4^{2})}}$

$RMS(X_{0})={\sqrt {\frac {6.85}{9}}}\approx 0.806$

Step 1.2: Find the Learning Rate Scaling ($\alpha _{t}$​):

Learning rate formula

$\alpha _{1}=max(\epsilon _{2},RMS(X_{0}))\cdot p_{1}$

Substitute the RMS

$\alpha _{1}=max(0.001,0.806)\cdot 0.01=0.00806$

Step 2: Compute $G_{t}^{2}$​ (Element-wise Square of Gradient)

Compute the squared value of each element in the gradient matrix $G_{t}$.

$G_{1}^{2}={\begin{bmatrix}0.3^{2}&(-0.2)^{2}&0.4^{2}\\(-0.5)^{2}&0.6^{2}&(-0.1)^{2}\\0.2^{2}&(-0.4)^{2}&0.3^{2}\end{bmatrix}}$

$G_{1}^{2}={\begin{bmatrix}0.09&0.04&0.16\\0.25&0.36&0.01\\0.04&0.16&0.09\end{bmatrix}}$

Step 3: Find the moment estimate

Compute the exponential moving average of squared gradients to capture the variance or scale of gradients.

Step 3.1: Compute row moments ($R_{t}$)

This equation computes the row-wise second moments ($R_{t}$ ​) as an exponential moving average of past moments ($R_{t-1}$) and the current row-wise mean of squared gradients ( $G_{t}^{2}$​ ), with a balance controlled by (${\hat {\beta }}_{2t}$).

For $G_{t}^{2}=\mathbb {R} ^{m\times n}$

$R_{t}={\hat {\beta _{2t}}}\cdot R_{t-1}+(1-{\hat {\beta }})\cdot ({\tfrac {1}{m}}\sum _{j=1}^{m}G_{t}^{2}[i,j]+\epsilon _{1})$

Since ${\hat {\beta }}_{2t}=1-t^{-0.8}$, for first iteration: ${\hat {\beta }}_{21}=0$. And because $\epsilon _{1}$ is too small, we can ignore it. The update of $R_{t}$ is:

$R_{1}={\tfrac {1}{m}}\textstyle \sum _{j=1}^{m}\displaystyle G_{1}^{2}[i,j]$

Row-wise mean ($R_{t}$):

$R_{1}={\begin{bmatrix}{\tfrac {0.09+0.04+0.16}{3}}\\{\tfrac {0.25+0.36+0.01}{3}}\\{\tfrac {0.04+0.16+0.09}{3}}\end{bmatrix}}={\begin{bmatrix}0.0967\\0.2067\\0.0967\end{bmatrix}}$

Step 3.2: Compute column moments ($C_{t}$)

The process is same as row moments.

$C_{t}={\hat {\beta }}\cdot C_{t-1}+(1-{\hat {\beta }})\cdot ({\tfrac {1}{n}}\sum _{j=1}^{n}G_{t}^{2}[i,j]+\epsilon _{1})$

Column-wise mean ($C_{t}$):

$C_{1}={\begin{bmatrix}{\tfrac {0.09+025+0.04}{3}}\\{\tfrac {0.04+0.36+0.16}{3}}\\{\tfrac {0.16+0.01+0.09}{3}}\end{bmatrix}}={\begin{bmatrix}0.1267\\0.1867\\0.0867\end{bmatrix}}$

Step 3.3: Second Moment Estimate (${\hat {V_{t}}}$​)

The Second Moment Estimate is calculated as the outer product of the row moments ($R_{t}$​) and column moments ($C_{t}$​).

${\hat {V}}_{t}=R_{t}\otimes C_{t}$

${\hat {V}}_{1}={\begin{bmatrix}0.0967\\0.2067\\0.0967\end{bmatrix}}\otimes {\begin{bmatrix}0.1267&0.1867&0.0867\\\end{bmatrix}}$

${\hat {V}}_{1}={\begin{bmatrix}0.0122&0.0180&0.0084\\0.0262&0.0386&0.0179\\0.0122&0.0180&0.0084\end{bmatrix}}$

Step 4: Update the vector ($U_{t}$)

Computed by scaling the gradient matrix $G_{t}$​ element-wise with the inverse square root of the second moment estimate (${\hat {V_{t}}}$​)

step 4.1: Find the vector value of $U_{t}$

Formula of $U_{t}$

$U_{t}={\frac {G_{t}}{\sqrt {{\hat {V_{t}}}+\epsilon _{1}}}}$

Substitute $C_{t}$ and $V_{t}$

$U_{1}={\frac {\begin{bmatrix}0.3&-0.2&0.4\\-0.5&0.6&-0.1\\0.2&-0.4&0.3\end{bmatrix}}{\sqrt {\begin{bmatrix}0.0122&0.0180&0.0084\\0.0262&0.0386&0.0179\\0.0122&0.0180&0.0084\end{bmatrix}}}}$

$U_{1}={\begin{bmatrix}2.711&-1.489&4.370\\-3.090&3.055&-0.747\\1.807&-2.978&3.278\end{bmatrix}}$

step 4.2: Clipped Update Vector ${\hat {U_{t}}}$

Scale the update vector ( $U_{t}$​ ) to ensure its RMS value does not exceed a predefined clipping threshold ($d$), maintaining stability in updates.

Formula of ${\hat {U_{t}}}$

${\hat {U_{t}}}={\frac {U_{t}}{max(1,{\tfrac {RMS(U_{t})}{d}})}}$

Compute RMS of $U_{t}$

$RMS(U_{1})={\sqrt {{\tfrac {1}{9}}\sum _{i=1}^{9}U_{t}[i]^{2}}}\approx 3.303$

Since RMS($U_{t}$​)>d, scale $U_{t}$​ by ${\tfrac {1}{3.303}}$

${\hat {U_{1}}}={\begin{bmatrix}0.965&-0.53&1.556\\-1.1&1.088&-0.266\\0.664&-1.06&1.167\end{bmatrix}}$

Step 5: Weight Update ($X_{1}$)

Adjust the weights ($X_{t}$) by subtracting the product of the learning rate ($\alpha _{t}$) and the clipped update vector (${\hat {U_{t}}}$ ).

$X_{1}=X_{0}-\alpha \cdot {\hat {U_{t}}}$

The result for first iteration.

$X_{1}={\begin{bmatrix}0.7&-0.5&0.9\\-1.1&0.8&-1.6\\1.2&-0.7&0.4\end{bmatrix}}-0.00806\cdot {\begin{bmatrix}0.965&-0.53&1.556\\-1.1&1.088&-0.266\\0.664&-1.06&1.167\end{bmatrix}}$

$X_{1}={\begin{bmatrix}0.692&-0.496&0.887\\-1.091&0.791&-0.596\\1.195&-0.691&0.391\end{bmatrix}}$

Applications

Adafactor is an efficient adaptive optimizer designed specifically for large-scale deep learning tasks. Its unique memory-saving properties have made it widely used for training large-scale language models, image recognition models, and reinforcement learning policy networks. Compared to other optimizers (e.g., Adam), Adafactor delivers exceptional performance in large-scale computations while significantly reducing memory requirements. Below are several specific application scenarios of Adafactor:

1. Natural Language Processing (NLP)

In NLP tasks, Adafactor has been successfully applied to training ultra-large-scale language models, such as Google’s Transformer and T5 (Text-To-Text Transfer Transformer). By significantly reducing memory usage during the gradient update process, Adafactor enables efficient model training in resource-constrained environments. For example, the T5 model in Google’s research employed Adafactor to effectively train on large datasets through text-to-text conversion tasks.²

2. Training Large-Scale Language Models

Adafactor has been used to train large-scale language models like LLaMA, combining it with novel preconditioned diagonalization methods to significantly enhance training efficiency. Experiments showed that Adafactor achieved performance comparable to the Adam optimizer while consuming substantially less memory and computational resources.³

3. Humor Detection Tasks

Adafactor has been utilized to optimize ALBERT-based models for humor detection tasks. Configured as an adaptive learning rate optimizer and paired with a cross-entropy loss function, Adafactor was used to train models that achieved 99% accuracy and F1 scores. Moreover, training time was faster than with Adam, completing in approximately 43 minutes. Comparisons with Adam and AdaBound optimizers demonstrated that Adafactor excelled in terms of both time efficiency and performance, especially in accuracy, recall, and F1 scores for humor detection tasks .⁴

4. Multilingual Model Training

In training multilingual models, Adafactor improved scalability and efficiency, particularly by significantly reducing memory consumption when handling large-scale parameters.⁵

5. Pretraining Vision Models

When training ResNet50 and ViT on the ImageNet1k dataset, Adafactor successfully optimized these deep networks with its low memory requirements. Additionally, with new algorithms combining preconditioned diagonalization methods (e.g., AdafacDiag and AdafacDiag++), it outperformed the standard Adam optimizer in both convergence speed and final accuracy.⁶

Software Tools and Platforms

Adafactor has been integrated into the following mainstream deep learning frameworks, making it accessible to developers:

TensorFlow: Provides a built-in implementation of Adafactor, supporting T5 model optimization.⁷

PyTorch: PyTorch provides the Adafactor optimizer through the torch.optim.AdaFactor class.⁸

JAX/Flax: JAX provides an optimizer library called Optax, which includes the Adafactor optimizer.⁹

Future Prospects

As the scale of deep learning models continues to grow, Adafactor’s memory-saving and computational efficiency advantages will become increasingly important. In the training of ultra-large-scale models (e.g., GPT and Vision Transformers), Adafactor is expected to become an indispensable optimization tool. Furthermore, by combining with other optimization strategies, such as mixed precision training, Adafactor may further enhance its applicability in both industrial and research settings.