The AdEMAMix Optimizer: Better, Faster, Older β€” part 5/5

Cleaned-text code points [60460, 67549). Bibliography omitted here; source.md includes it.

AdMeta. Chen et al. (2023b) take inspiration over DEMA and use nested EMAs in their AdMeta-S optimizer: { π’Ž 1 ( t ) = Ξ² 1 ​ π’Ž 1 ( t βˆ’ 1 ) + π’ˆ ( t ) 𝒉 ( t ) = ΞΊ ​ π’ˆ ( t ) + ΞΌ ​ π’Ž 1 ( t ) π’Ž 2 ( t ) = Ξ² 2 ​ π’Ž 2 ( t βˆ’ 1 ) + ( 1 βˆ’ Ξ² 2 ) ​ 𝒉 ( t ) 𝜽 ( t ) = 𝜽 ( t βˆ’ 1 ) βˆ’ Ξ· ​ π’Ž 2 ( t ) . \displaystyle\begin{cases}{\bm{m}}_{1}^{(t)}&=\beta_{1}{\bm{m}}_{1}^{(t-1)}+{\bm{g}}^{(t)}\\ {\bm{h}}^{(t)}&=\kappa{\bm{g}}^{(t)}+\mu{\bm{m}}_{1}^{(t)}\\ {\bm{m}}_{2}^{(t)}&=\beta_{2}{\bm{m}}_{2}^{(t-1)}+(1-\beta_{2}){\bm{h}}^{(t)}\\ {\bm{\theta}}^{(t)}&={\bm{\theta}}^{(t-1)}-\eta{\bm{m}}_{2}^{(t)}.\end{cases} (AdMeta-S) With ΞΌ \mu and ΞΊ \kappa parameterized by Ξ² 1 ∈ [ 0 , 1 [ \beta_{1}\in[0,1[ as such: ΞΌ \displaystyle\mu = 25 βˆ’ 10 ​ ( Ξ² 1 + 1 Ξ² 1 ) \displaystyle=25-10(\beta_{1}+\frac{1}{\beta_{1}}) ΞΊ \displaystyle\kappa = 10 Ξ² 1 βˆ’ 9 . \displaystyle=\frac{10}{\beta_{1}}-9. In their AdMeta-S experiments, they use Ξ² 1 = 0.9 \beta_{1}=0.9 , corresponding to ( ΞΌ , ΞΊ ) = ( 4.88 , 2.11 ) (\mu,\kappa)=(4.88,2.11) . Ξ² 2 \beta_{2} takes values ranging from 0.1 0.1 to 0.4 0.4 . As a results, unlike AdEMAMix, AdMeta is not leveraging very old gradients. Analysing the AdMeta algorithm, we see that it consists in two nested EMAs. We show the shape of nested EMAs in Fig. 31 . In sharp contrast with our approach we observe that (i) it reduces the weights given to recent gradients and (ii) it gives a small weight to old gradients.

#### C.3.2 Comparison with Lion

The Lion optimizer. Chen et al. (2023a) derived the following optimizer. We change notations to facilitate the comparison with AdEMAMix: { 𝜽 ( t ) = 𝜽 ( t βˆ’ 1 ) βˆ’ Ξ· β‹… ( sign ​ ( Ξ± ​ π’Ž ( t βˆ’ 1 ) + ( 1 βˆ’ Ξ± ) ​ π’ˆ ( t ) ) + Ξ» ​ 𝜽 ( t βˆ’ 1 ) ) π’Ž ( t ) = Ξ² ​ π’Ž ( t βˆ’ 1 ) + ( 1 βˆ’ Ξ² ) ​ π’ˆ ( t ) . \displaystyle\begin{cases}{\bm{\theta}}^{(t)}&={\bm{\theta}}^{(t-1)}-\eta\cdot\big(\text{sign}(\alpha{\bm{m}}^{(t-1)}+(1-\alpha){\bm{g}}^{(t)})+\lambda{\bm{\theta}}^{(t-1)}\big)\\ {\bm{m}}^{(t)}&=\beta{\bm{m}}^{(t-1)}+(1-\beta){\bm{g}}^{(t)}.\end{cases} (Lion) The Lion optimizer uses a sign function and updates its EMA after updating the parameters. Moreover, it does not normalize the updates. While Lion and AdEMAMix are quite different from each others, we can draw one similarity. Indeed, the interpolation Ξ± ​ π’Ž ( t βˆ’ 1 ) + ( 1 βˆ’ Ξ± ) ​ π’ˆ ( t ) \alpha{\bm{m}}^{(t-1)}+(1-\alpha){\bm{g}}^{(t)} is similar to combining two EMAs, one of them using Ξ² = 0 \beta=0 . We also explore the possibility of setting AdEMAMix’s Ξ² 1 \beta_{1} to 0 0 in Β§ 3 , we show in App. C.1.7 and App. C.1.4 (Fig. 15(b) ) that this often works despite a more unstable training. Interestingly, Chen et al. (2023a) find that larger Ξ² = 0.99 \beta=0.99 values work best. Beside this similarity, the two optimizers behave very differently, the biggest difference being the use of the sign function, which Chen et al. (2023a) claim can help regularize the training. We test the Lion optimizer on language modeling. To tune the hyperparameters, we took values from Chen et al. (2023a) as a starting point, as well as recipes provided by the Optax Jax library. A summary of our hyperparameter tuning is in Table 8 :

The training curve associated to the best hyperparameters is in Fig. 32 . We observe that Lion is not outperforming our carefully tuned AdamW baseline. Moreover, no attempt to increase Ξ² \beta beyond 0.99 0.99 was successful, as those models mostly diverged. This emphasizes one of our main contributions: the introduction of schedulers enabling the use of large momentum values.

#### C.3.3 Adding a third momentum term ( ∼ \sim AggMo)

Lucas et al. (2019, AggMo) propose to add an arbitrary number ( K K ) of momentum terms to the gradient descent algorithm: { π’Ž i ( t ) = Ξ² i π’Ž i ( t βˆ’ 1 ) + π’ˆ ( t ) for 1 ≀ i ≀ K 𝜽 ( t ) = 𝜽 ( t βˆ’ 1 ) βˆ’ Ξ· K ​ βˆ‘ 1 ≀ i ≀ K π’Ž i ( t ) . \displaystyle\begin{cases}{\bm{m}}_{i}^{(t)}&=\beta_{i}{\bm{m}}_{i}^{(t-1)}+{\bm{g}}^{(t)}\quad\text{for $1\leq i\leq K$}\\ {\bm{\theta}}^{(t)}&={\bm{\theta}}^{(t-1)}-\frac{\eta}{K}\sum_{1\leq i\leq K}{\bm{m}}_{i}^{(t)}.\end{cases} (AggMo)

In order to see whether more momentum terms can be beneficial, we modify our AdEMAMix optimizer to include a third momentum π’Ž 3 {\bm{m}}_{3} . We name the resulting optimizer Ad 3 EMAMix: { π’Ž 1 ( t ) = Ξ² 1 ​ π’Ž 1 ( t βˆ’ 1 ) + ( 1 βˆ’ Ξ² 1 ) ​ π’ˆ ( t ) , π’Ž ^ 1 ( t ) = π’Ž 1 ( t ) 1 βˆ’ Ξ² 1 t π’Ž 2 ( t ) = Ξ² 3 ​ π’Ž 2 ( t βˆ’ 1 ) + ( 1 βˆ’ Ξ² 3 ) ​ π’ˆ ( t ) π’Ž 3 ( t ) = Ξ² 4 ​ π’Ž 3 ( t βˆ’ 1 ) + ( 1 βˆ’ Ξ² 4 ) ​ π’ˆ ( t ) 𝝂 ( t ) = Ξ² 2 ​ 𝝂 ( t βˆ’ 1 ) + ( 1 βˆ’ Ξ² 2 ) ​ π’ˆ ( t ) 2 , 𝝂 ^ ( t ) = 𝝂 ( t ) 1 βˆ’ Ξ² 2 t 𝜽 ( t ) = 𝜽 ( t βˆ’ 1 ) βˆ’ Ξ· ⁑ ( π’Ž ^ 1 ( t ) + Ξ± ⁑ ( π’Ž 2 ( t ) + π’Ž 3 ( t ) ) 𝝂 ^ ( t ) + Ο΅ + Ξ» ​ 𝜽 ( t βˆ’ 1 ) ) . \displaystyle\begin{cases}{\bm{m}}_{1}^{(t)}=\beta_{1}{\bm{m}}_{1}^{(t-1)}+(1-\beta_{1}){\bm{g}}^{(t)},\quad\quad\hat{{\bm{m}}}_{1}^{(t)}=\frac{{\bm{m}}_{1}^{(t)}}{1-\beta_{1}^{t}}\\ {\color[rgb]{0,0,1}{\bm{m}}_{2}^{(t)}}{\color[rgb]{0,0,1}=\beta_{3}{\bm{m}}_{2}^{(t-1)}+(1-\beta_{3}){\bm{g}}^{(t)}}\\ {\color[rgb]{0,0,1}{\bm{m}}_{3}^{(t)}}{\color[rgb]{0,0,1}=\beta_{4}{\bm{m}}_{3}^{(t-1)}+(1-\beta_{4}){\bm{g}}^{(t)}}\\ {\bm{\nu}}^{(t)}=\beta_{2}{\bm{\nu}}^{(t-1)}+(1-\beta_{2}){{\bm{g}}^{(t)}}^{2},\quad\quad\hat{{\bm{\nu}}}^{(t)}=\frac{{\bm{\nu}}^{(t)}}{1-\beta_{2}^{t}}\\ {\bm{\theta}}^{(t)}={\bm{\theta}}^{(t-1)}-\eta\big(\frac{\hat{{\bm{m}}}_{1}^{(t)}{\color[rgb]{0,0,1}+\alpha({\bm{m}}_{2}^{(t)}+{\bm{m}}_{3}^{(t)})}}{\sqrt{\hat{{\bm{\nu}}}^{(t)}}+{\epsilon}}+\lambda{\bm{\theta}}^{(t-1)}\big).\end{cases} (Ad3EMAMix)

We Train a 110 110 M parameter model with same hyperparameters as in Table 2 , but we use Ξ± = 4 \alpha=4 instead of Ξ± = 8 \alpha=8 . We apply the same scheduler to Ξ² 4 \beta_{4} as for Ξ² 3 \beta_{3} . In Fig. 33 we show the resulting validation loss curves for various ( Ξ² 1 , Ξ² 3 , Ξ² 4 ) (\beta_{1},\beta_{3},\beta_{4}) triplets. Our experiments show no advantage of adding an extra Ξ² 4 \beta_{4} -EMA. Rather than carrying the message that we should use more momentum terms, our work shows howβ€”in order to use a large momentumβ€”a term should be added to stay sensitive to local fluctuations of the loss landscape. In that sense, we believe two terms should be enough.

## Instructions for reporting errors

We are continuing to improve HTML versions of papers, and your feedback helps enhance accessibility and mobile support. To report errors in the HTML that will help us improve conversion and rendering, choose any of the methods listed below:

Click the "Report Issue" ( ) button, located in the page header.

Tip: You can select the relevant text first, to include it in your report.

Our team has already identified the following issues . We appreciate your time reviewing and reporting rendering errors we may not have found yet. Your efforts will help us improve the HTML versions for all readers, because disability should not be a barrier to accessing research. Thank you for your continued support in championing open access for all.

Have a free development cycle? Help support accessibility at arXiv! Our collaborators at LaTeXML maintain a list of packages that need conversion , and welcome developer contributions .