Adafactor: Difference between revisions

From Cornell University Computational Optimization Open Textbook - Optimization Wiki
Jump to navigation Jump to search
 
(30 intermediate revisions by the same user not shown)
Line 4: Line 4:


== Introduction ==
== Introduction ==
Adafactor is an efficient, adaptive learning rate optimization algorithm proposed by Noam Shazeer and Mitchell Stern in 2018. <sup>1</sup>
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.<sup>1</sup> This efficiency makes it an ideal choice for training ultra-large-scale models such as T5.<sup>2</sup>
Adafactor’s efficient memory usage and outstanding performance make it widely applicable in scenarios such as Natural Language Processing (NLP).<sup>2</sup> 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. <sup>3,6</sup>
== Problem formulation ==
== Problem formulation ==
=== 1. Objective ===
=== 1. Objective ===
Line 29: Line 35:
*** <math>u_{xt} = \frac{-g_{xt}}{\sqrt{\hat{v}_{xt}}}</math>
*** <math>u_{xt} = \frac{-g_{xt}}{\sqrt{\hat{v}_{xt}}}</math>
*** <math>\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)}</math>
*** <math>\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)}</math>
=== 3. Algorithms ===
=== 3. Algorithms ===
==== Adafactor for Weighted Vectors ====
==== Adafactor for Weighted Vectors ====
Line 67: Line 74:
** Update parameter: <math>X_t = X_{t-1} - \alpha_t \hat{U}_t</math>
** Update parameter: <math>X_t = X_{t-1} - \alpha_t \hat{U}_t</math>
* End for
* End for
==== 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.


=== 4. Proposed Hyperparameters for Adafactor ===
=== 4. Proposed Hyperparameters for Adafactor ===
* '''Regularization constant 1''': <math>\epsilon_1 = 10^{-30}</math>
* Regularization constant 1: <math>\epsilon_1 = 10^{-30}</math>
* 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: <math>\epsilon_2 = 10^{-3}</math>
* '''Regularization constant 2''': <math>\epsilon_2 = 10^{-3}</math>
* Clipping threshold: <math>d = 1</math>
* Help to stabilize parameter updates by controlling the effect of second-moment scaling in low-magnitude scenarios. Compared to <math>\epsilon_2</math>, a relatively larger value ensures the stability of noise and low-magnitude scenarios.
* Relative step size: <math>\rho_t = \min(10^{-2}, 1/\sqrt{t})</math>
* '''Clipping threshold''': <math>d = 1</math>
* Second moment decay: <math>\hat{\beta}_{2t} = 1 - t^{-0.8}</math>
* 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''': <math>\rho_t = \min(10^{-2}, 1/\sqrt{t})</math>
** <math>min(10^-2, ...)</math> can caps the learning rate at 10^-2, which is a empirical found for upper bound
** <math>\frac{1}{\sqrt{t}}</math> 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''': <math>\hat{\beta}_{2t} = 1 - t^{-0.8}</math>
** 1-...: ensures the decay factor remains close to 1
** <math>t^{-0,8}</math> the power 0.8 ensures a balance between rapid adaptation in early training and later iterations


== Numerical Examples ==
== Numerical Examples ==
Step-by-step instructions for determining the result of the first iteration.
'''<big>Problem setup</big>'''
'''Initial weights ('''<math>X_0</math>​'''):'''
<math>X_0 = \begin{bmatrix} 0.7 &-0.5& 0.9\\ -1.1 & 0.8& -1.6\\1.2&-0.7& 0.4 \end{bmatrix}</math>
'''Gradient for first iteration (​<math>G_1</math>):'''
Gradient of the loss function with respect to X
<math>G_1 = \begin{bmatrix} 0.3&-0.2&0.4\\ -0.5&0.6&-0.1\\0.2&-0.4 &0.3 \end{bmatrix}</math>
'''<big>Hyperparameters setup</big>'''
<math>\epsilon_1 = 10^{-30}</math> (Minimum learning rate scaling factor))
<math>\epsilon_2 = 10^{-3}</math> (Regularization constant)
<math>d = 1</math> (Clipping threshold)
<math>\rho_t = \min(10^{-2}, 1/\sqrt{t})</math> (Relative step size)
<math>\hat{\beta}_{2t} = 1 - t^{-0.8}</math> (Second moment decay)
'''<big>Step 1:  Learning Rate Scaling</big>'''
Define the relative step size
<math>\rho_1 = \min(10^{-2}, 1/\sqrt{1})= 10^{-2}</math>
'''Step 1.1: Root Mean Square(RMS) calculation for <math>X_0</math>'''
Root Mean Square(RMS) calculation for <math>X_0</math>
RMS formula
<math>RMS(X_0) = \sqrt{\tfrac{1}{n}\sum_{i=1}^n  X_0[i]^2}</math>
Substitute the initial weights
<math>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)}</math>
<math>RMS(X_0) = \sqrt{\frac{6.85}{9}}\approx 0.806</math>
'''Step 1.2: Find the Learning Rate Scaling ('''<math>\alpha_t</math>​'''):'''
Learning rate formula
<math>\alpha_1 = max(\epsilon_2,RMS(X_0))\cdot p_1</math>
Substitute the RMS
<math>\alpha_1 = max(0.001,0.806)\cdot 0.01=0.00806</math>
'''<big>Step 2: Compute <math>G^{2}_t</math>​ (Element-wise Square of Gradient)</big>'''
Compute the squared value of each element in the gradient matrix '''<math>G_t</math>'''.
<math>G^{2}_1 = \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}</math>
<math>G^{2}_1 = \begin{bmatrix} 0.09& 0.04&0.16\\ 0.25&0.36&0.01\\0.04&0.16&0.09\end{bmatrix}</math>
'''<big>Step 3: Find the moment estimate</big>'''
Compute the exponential moving average of squared gradients to capture the variance or scale of gradients.
'''Step 3.1: Compute row moments (<math>R_t</math>)'''
This equation computes the row-wise second moments ('''<math>R_t</math>''' ​) as an exponential moving average of past moments ('''<math>R_{t-1}</math>''') and the current row-wise mean of squared gradients ( <small><math>G^{2}_t</math></small>​ ), with a balance controlled by (<math>\hat{\beta}_{2t}</math>).
For <math>G^{2}_t=\mathbb{R}^{m\times n} </math>
<math>R_t = \hat{\beta_{2t}} \cdot R_{t-1} + (1-\hat{\beta})\cdot (\tfrac{1}{m}\sum_{j=1}^m G^{2}_t[i,j]+\epsilon_1) </math>
Since <math>\hat{\beta}_{2t} = 1 - t^{-0.8}</math>, for first iteration: <math>\hat{\beta}_{21} = 0</math>. And because <math>\epsilon_1 </math> is too small, we can ignore it. The update of '''<math>R_t</math>''' is:
<math>R_{1} = \tfrac{1}{m}\textstyle \sum_{j=1}^m \displaystyle G^{2}_1[i,j] </math>
Row-wise mean ('''<math>R_t</math>'''):
<math>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} </math>
'''Step 3.2: Compute column moments (<math>C_t</math>)'''
The process is same as row moments.
<math>C_t = \hat{\beta}\cdot C_{{t-1}} + (1-\hat{\beta})\cdot (\tfrac{1}{n}\sum_{j=1}^n G^{2}_t[i,j]+\epsilon_1) </math>
Column-wise mean (<math>C_t</math>):
<math>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} </math>
'''Step 3.3: Second Moment Estimate ('''<math>\hat{V_t}</math>​''')'''
The Second Moment Estimate is calculated as the outer product of the row moments ('''<math>R_t</math>'''​) and column moments ('''<math>C_t</math>'''​).
<math>\hat{V}_t = R_t \otimes C_t</math>
<math>\hat{V}_1  = \begin{bmatrix} 0.0967\\0.2067\\0.0967 \end{bmatrix} \otimes    \begin{bmatrix} 0.1267&0.1867&0.0867\\ \end{bmatrix} </math>
<math>\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} </math>
'''<big>Step 4: Update the vector (<math>U_t </math>)</big>'''
Computed by scaling the gradient matrix '''<math>G_t</math>'''​ element-wise with the inverse square root of the second moment estimate (<math>\hat{V_t}</math>​)
'''step 4.1: Find the vector value of <math>U_t </math>'''
Formula of '''<small><math>U_t </math></small>'''
<math>U_t = \frac{G_t}{\sqrt{\hat{V_t}+\epsilon_1}} </math>
Substitute '''<small><math>C_t</math></small>''' and <small><math>V_t</math></small>
<math>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}}} </math>
<math>U_1 = \begin{bmatrix} 2.711&-1.489&4.370\\-3.090&3.055&-0.747\\1.807&-2.978&3.278  \end{bmatrix} </math>
'''step 4.2: Clipped Update Vector <math>\hat{U_t} </math>'''
Scale the update vector ( '''<math>U_t </math>'''​ ) to ensure its RMS value does not exceed a predefined clipping threshold (<math>d </math>), maintaining stability in updates.
Formula of '''<small><math>\hat{U_t} </math></small>'''
'''<small><math>\hat{U_t} = \frac{U_t}{max(1,\tfrac{RMS(U_t)}{d})        } </math></small>'''
Compute RMS of '''<math>U_t </math>'''
'''<small><math>RMS(U_1)  = \sqrt{\tfrac{1}{9}  \sum_{i=1}^9 U_t[i]^2}  \approx 3.303 </math></small>'''
Since RMS('''<math>U_t </math>'''​)>d, scale '''<math>U_t </math>'''​ by <math>\tfrac{1}{3.303} </math>
'''<math>\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} </math>'''
'''<big>Step 5: Weight Update (</big>'''<math>X_1 </math>'''<big>)</big>'''
Adjust the weights (<math>X_t </math>) by subtracting the product of the learning rate (<math>\alpha_t </math>) and the clipped update vector (<math>\hat{U_t} </math> ).
<math>X_1 = X_0 - \alpha \cdot      \hat{U_t}</math>
The result for first iteration.
<math>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}      </math>
<math>X_1 =  \begin{bmatrix} 0.692&-0.496&0.887 \\-1.091&0.791&-0.596\\ 1.195&-0.691&0.391\end{bmatrix}      </math>
== Applications ==
== 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.<sup>2</sup>
'''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.<sup>3</sup>
'''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 .<sup>4</sup>
'''4. Multilingual Model Training'''
In training multilingual models, Adafactor improved scalability and efficiency, particularly by significantly reducing memory consumption when handling large-scale parameters.<sup>5</sup>
'''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.<sup>6</sup>
=== '''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.<sup>7</sup>
'''PyTorch:''' PyTorch provides the Adafactor optimizer through the torch.optim.AdaFactor class.<sup>8</sup>
'''JAX/Flax:''' JAX provides an optimizer library called Optax, which includes the Adafactor optimizer.<sup>9</sup>
=== '''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.
== Conclusion ==
== Conclusion ==
== Reference ==
== Reference ==

Latest revision as of 20:07, 11 December 2024

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

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.1 This efficiency makes it an ideal choice for training ultra-large-scale models such as T5.2

Adafactor’s efficient memory usage and outstanding performance make it widely applicable in scenarios such as Natural Language Processing (NLP).2 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 , where and is the weight vector to be optimized.

2. Parameters

  • Gradient:

  • Second moment estimate:

  • Where:
    • is the running average of the squared gradient.
    • is the corrected decay parameter.
    • is a regularization constant.
  • Step size:

  • Where:
    • is the relative step size.
    • is a regularization constant.
    • is the root mean square, defined as:

3. Algorithms

Adafactor for Weighted Vectors

Inputs:

  • Initial point:
  • Relative step sizes: for to
  • Second moment decay: for to , with
  • Regularization constants:
  • Clipping threshold:

Algorithm:

  • For to :
    • Compute adaptive step size:
    • Compute gradient:
    • Update second moment estimate:
    • Compute normalized gradient:
    • Apply clipping:
    • Update parameter:
  • End for

Adafactor for Weighted Matrices

Inputs:

  • Initial point:
  • Relative step sizes: for to
  • Second moment decay: for to , with
  • Regularization constants:
  • Clipping threshold:

Algorithm:

  • For to :
    • Compute adaptive step size:
    • Compute gradient:
    • Update row-wise second moment:
    • Update column-wise second moment:
    • Update overall second moment estimate:
    • Compute normalized gradient:
    • Apply clipping:
    • Update parameter:
  • End for

4. Proposed Hyperparameters for Adafactor

  • Regularization constant 1:
  • Regularization constant 2:
  • Clipping threshold:
  • Relative step size:
  • Second moment decay:

Numerical Examples

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

Problem setup

Initial weights ():

Gradient for first iteration (​):

Gradient of the loss function with respect to X

Hyperparameters setup

(Minimum learning rate scaling factor))

(Regularization constant)

(Clipping threshold)

(Relative step size)

(Second moment decay)

Step 1: Learning Rate Scaling

Define the relative step size

Step 1.1: Root Mean Square(RMS) calculation for

Root Mean Square(RMS) calculation for

RMS formula

Substitute the initial weights

Step 1.2: Find the Learning Rate Scaling ():

Learning rate formula

Substitute the RMS


Step 2: Compute ​ (Element-wise Square of Gradient)

Compute the squared value of each element in the gradient matrix .



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 ()

This equation computes the row-wise second moments ( ​) as an exponential moving average of past moments () and the current row-wise mean of squared gradients ( ​ ), with a balance controlled by ().

For

Since , for first iteration: . And because is too small, we can ignore it. The update of is:

Row-wise mean ():

Step 3.2: Compute column moments ()

The process is same as row moments.

Column-wise mean ():

Step 3.3: Second Moment Estimate ()

The Second Moment Estimate is calculated as the outer product of the row moments (​) and column moments (​).


Step 4: Update the vector ()

Computed by scaling the gradient matrix ​ element-wise with the inverse square root of the second moment estimate (​)

step 4.1: Find the vector value of

Formula of

Substitute and


step 4.2: Clipped Update Vector

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

Formula of

Compute RMS of

Since RMS(​)>d, scale ​ by


Step 5: Weight Update ()

Adjust the weights () by subtracting the product of the learning rate () and the clipped update vector ( ).

The result for first iteration.




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

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

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

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

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.6

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

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

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

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.

Conclusion

Reference