How to set AdamW's weight decay as you scale model and dataset size
Authors: Xi Wang
Laurence Aitchison
Venue: ICML 2025 poster (2025)
Keywords: Hyperparameter transfer
AdamW
Optimization
Primary area: deep_learning

Abstract: The scaling of the optimal AdamW weight decay hyperparameter with model and dataset size is critical as we seek to build larger models, but is poorly understood. We show that weights learned by AdamW can be understood as an exponential moving average (EMA) of recent updates.  This gives critical insights for how to set the weight decay in AdamW, and how the weight decay should scale with model and dataset size.  In particular, the key hyperparameter for an exponential moving average is the EMA timescale. 
Intuitively, the EMA timescale can be understood as the number of recent iterations the EMA averages over. We find that the optimal timescale, measured in epochs, is roughly constant as we change model and dataset size. Moreover, given a learning rate, there is a one-to-one mapping from the EMA timescale to the weight decay hyperparameter. Thus, if the optimal EMA timescale is constant, that implies that as the dataset size increases, the optimal weight decay should fall and as the model size increases, the optimal weight decay should increase (if we follow the muP recommendation for scaling the learning rate). We validate these scaling rules on ResNet-18 and Vision Transformers trained on CIFAR-10 and ImageNet, and on NanoGPT pre-training on OpenWebText. Finally, we found that as training progresses, muP's learning rate scaling breaks down for AdamW unless weight decay is scaled appropriately.

How to set AdamW's weight decay as you scale model and dataset size
Xi Wang' Laurence Aitchison
Abstract
from small-scale student projects to the largest LLM train- ing runs. However, for this workflow to be effective, need to understand how to transfer optimal hyperparame- ters from smaller-scale prototyping runs to the final, largest model and dataset. This is very well studied for the optimal learning rate as you scale up width and depth (c.g. Yang et al., 2022: Bordelon et al., 2024; Noci et al., 2024; Everett et al., 2024). However, the largest LLM pretraining runs (e-g. Zhang et al., 2022; Touvron al., 2023a;b; Tow et al., 2023) use optimizer with a weight decay term such as AdamW (Loshchilov & Hutter, 2018), and it is not understood how to scale the Adam W weight decay hyperparameter with model sizes, and how to scale hyperparameters as we scale dataset size. Here, we this gap, showing how the optimal weight decay transfers across model width and dataset size.
The scaling of the optimal AdamW weight decay hyperparameter with model and dataset size is critical as we seek to build larger models, but poorly understood. We show that weights learned by AdamW can be understood moving average (EMA) of recent updates. This gives critical insights for how to set the weight de- cay in Adam W. and how the weight decay should scale with model and dataset size. In particular, the key hyperparameter for an exponential mov- ing average is the EMA timescale. Intuitively, the EMA timescale can be understood the num- ber of recent iterations the EMA averages over. We find that the optimal timescale, measured in epochs, is roughly constant as we change model and dataset size. Morcover, given a learning rate, there is a one-to-one mapping from the EMA timescale to the weight decay hyperparameter. Thus, if the optimal EMA timescale is constant, that implies that as the dataset size increases, the optimal weight decay should fall and as the model size increases, the optimal weight decay should increase (if we follow the muP recommendation for scaling the learning rate). We validate these scaling rules on R8 and Vision Transform- ers trained on CIFAR-IO and ImageNet, and on NanoGPT pre-training on WeenWeeTexxt. Finally, we found that as training progresses, muP's learn- ing rate scaling breaks down for AdamW W unless weight decay is scaled appropriately.
exponential
To understand how to transfer weight decay across model and dataset sizes, we argue that Adam should be under- stood as an Exponential Moving Average (EMA). Of course, both Adam and AdamW use EMAs to estimate the average gradient, my, and average squared gradient, iy That is not what we are talking about. Instead, we observe that in algo- rithms with decoupled weight decay (i.e. AdamW but not Adam), the weights themselves are EMA of recent up- dates (see Sec. 3.1). The key hyperparameter in an EMA is the EMA timescale, which intuitively describes the number of previous iterations that the EMA averages over, which is given by Tax = 1/(n)> (Sec. 3.1). In fact, we prove and empirically verify that under a 17-dependent initialization and a fully-scale invariant network, AdamW's optimization trajectory controlled solely by the timescale Tuer for all combinations of learning rate 17 and weight decay coefficient À (Sec. 3.2).
Perhaps more importantly, the EMA perspective tells us how to set the weight decay parameter. Specifically, we can take into account the dataset and batch sizes to mea- sure the timescale in epochs, i.e. Tepoch T where M = N/B denotes the number of iterations per epoch in training dataset of N samples under batch size of B, which measures the fraction of the dataset that contributes to the current weights through the EMA. We empirically observed that the optimal Tepach is stable across model and dataset sizes, which, consequently provides a guideline for seal- ing weight decay. In particular, under fixed learning rate schedule, the dataset size N increases, we find that the
1. Introduction
A common machine learning workflow is to prototype by training many smaller models, then the end do one final training run, with the largest possible model on the largest possible dataset. This workflow is used in many settings,
'Johns Hopkins University, US University Bristol, UK. Correspondence to: Xi Wang <xidulu @ comi.
Proceedings of the nnnd International Conference Machine Learning. Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).

---

How to set AdamW's weight decay you scale model and dataset size
optimal weight decay falls (Sec. 4.1). Moreover, if we fol- low the usual μP recommendation decrease the learning rate 1) as the model width increases, we find that the optimal weight decay should increase as the model size increases (Sec. 4.2). We validate both of these predictions in ResNet, vision transformer trained on CIFAR-I0 and ImageNet, and NanoGPT pre-training on WeenWeeText.
reason why is complex; see Yang et al., 2022, for details),
Ayf -i= V
(2)
thus, Ay; fan_ x AW,j- To ensure that Ay: 1, we therefore need AW;; 1/fan_in in. As the Adam updates, AWy scale with the learning rate, 77, this implies 1) must also scale with 1/fan_in.
In summary, our paper's contributions are:
To see why Adam updates scale with the learning rate, con- sider an Adam update,
The weights generated by AdamW can be written as EMA of recent weight updates.
AW = nmng / Dyj
(3)
The optimal EMA timescale changes little we scale the model and dataset size.
where we have neglected the small e. Here, myj is an EMA estimate of the gradient Pij. while ě:g is an EMA estimate of the expected squared gradient g Thus, mug ~ Pi and igg ~ o.3- That implies mhn / isg ~ 1, so looking back at Eq. 3, we have AW13 ~ 17. Hence, to get AW,j 1/fan_ in, we need 1 1/fan in, which is the origin of the 1/fan in scaling of the learning rates in uP with Adam.
The optimal weight decay scales with //ataseet size un- der a fixed learning rate schedule.
The optimal weight decay increases with model width under μP scaling of learning rate.
When using A with fixed weight decay, μP learning rate scaling breaks down, but proper weight decay scaling restores its effectiveness
Importantly, as Yang et al. (2022) considers the size of the first updates relative to the random initial weights, it does not give guidance about how to change the weight decay with the model size, as the weight decay only becomes relevant after many learning steps.
2. Background
Yang et al. (2022) considered how to transfer the learning rates in SGD and Adam across model sizes. They considered two key desiderata. First, the initial random weights, W E Dxfan_in, times an input, x E R should not grow or shrink as the model changes size. Here, fan in is the input dimension, and D the output dimension. We could write to indicate
3. Methods
3.1. Adam as an EMA
An Adam update for a single parameter at the /th iteration can be written as,
this requirement& as, po 1, where "scales with". This gives the usual WO ~ 1/Vfan in scaling of the standard deviation of the initial random weights. At the same time, they required that the change in the outputs, caused by the first weight update, AWr, should not grow or shrink as the model changes size; i.e. AWx x ~ 1. Yang et al. (2022) show that this requirement implies that the learning rate should scale as 1/fan in.
m
Wy = (1 - --)---- 1)t V0, +
(4)
where W, is a neural network weight, 7) is the learning rate (which potentially varies over time due to scheduling). À is the weight decay, e is small constant, me is bias-corrected EMA estimate of the expected gradient, and ê is bias- corrected EMA estimate of the expected squared gradient. Notice that here we adopt the parameterization used by ma- jor optimization libraries such as torch. opt im!, where the decay for w:-1 is controlled by the product n/A. This parameterization is different from the form suggested in the original AdamW paper (Loshchilov & Hutter, 2018), which we will discuss in detail in the related work section. Further, we follow the usual convention, established by Loshchilov & Hutter (2018) of scheduling the learning rate but not the weight decay.
The distinction between the 1/vfan in scaling of the initial weights vs. the 1 //fan_in scaling of the learning rate might be a bit puzzling. The intuitive reason for this distinction is given in Yang et al. (2022) Appendix J. In short, for the initial weights,
31 = Σ
fan W W TJ
we know that yi ~ Vfan_in x WD so to ensure that y ~ 1, 1y.
need WOO 3) ~ 1/ 1/Vfan_in in. However, this square root scaling only arises in very specific circumstances, when each term in the sum, W&J,, is zero-mean and uncorrelated. These requirements hold for the initial weights, as they are sampled L.LDD from a zero-mean distribution. However, these conditions do not hold for the update (the precise
Now we will show that these weight updates (Eq. 4) can be approximately understood as an EMA? Recall that a
'nttps://pytone
opt im. AdamW, FFor introduction to EMAs, see en, en,wik wiikipedia.oom .org/ wi wiki//Xxooentii.o aala__ssss___n____
2

---

How to set AdamW's weight decay you scale model and dataset size
generic exponential moving average estimate, ema, can be written as,
parameters learned by the first optimizer after the lth opti- mization step, and w/, to be the parameters learned by the second optimizer. These optimizers are initialized by scaling the same random init,
ema (1 1/meer) - +2/rune qr.
Here, ema,, forms an exponential moving average estimate of 9:. where the potentially time-varying EMA timescale is Theng We take this EMA (Eq. 5) and set:
Wp = =
w/b = o'E,
(8)
where & is random noise (e.g. IID Gaussian). Consider a scale-invariant network, in the sense that multiplying the weights, w, by an arbitrary positive constant, 1/c 0 gives the same output for all inpuls,
me
1/ 1/mmm== = T
ema = Wyy
q:
X V
then we recover the AdamW updates (Eq. 4). Of course, AdamW uses EMAs to compute rit, and 6z Our key insight was that additionally, the overall Adam updates for wi can themselves be understood as EMAs Intuitively, the AdamW EMA computes the average of qs, i.e the noisy minibatch update scaled by the weight decay.
net(æ; w) = net( t.w
(9)
We use the same trajectory of EMA timescales, 1/Twer = n = =n,X,, the same ratio of the first step, relative to the initial parameter scale, Yo The Now, the entire trajectory of weights for each optimizer is the same (up to scalar multiplier.) wi = _ where c o"
To build intuition, we consider a fixed timescale, Thers and write the EMA a weighted average of all minibatch up- dates (derivation in Appendix B),
We
Thus, the network outputs are the same at all points along the trajectory, net(;; wx) = net(a;; For the proof, see Appendix A.I, and for empirical confirmation, see Ap- pendix A.3. Thus, for scale-invarian networks, only two parameters are relevant to the trajectory: the EMA timescale (and its schedule), Ther; and the scale of the initial learning rate relative to the weight init. We would usually expect the influence of the init to diminish over time, leaving Ther: as the key hyperparameter.
Wt-
1/TRE
/f
(-x -
Vēr Vēr
1 my (
(7)
Thus, the unnormalized weights are e e-(t-()//m These weights highlight that recent updates (i.e. 1' close i) contribute more to the weighted average. Specifically, for
recent updates where t - t' is far smaller than Titers the weights, (t-rtr)/ are around 1. At the same time, up- dates from further into the past contribute less: when - L' is far larger than Ther, then the weights, e (t-xt)// decay to zero. Thus, intuitively, the EMA averages over roughly the last Teer minibatch updates.
4. Scaling weight decay across model and dataset sizes
Since timescale is the key hyperparameter in an EMA, we hypothesize that the optimal timescale should be trans- ferrable across model and dataset sizes. In this section, we empirically verify this hypothesis and consider the im- plications for the optimal weight decay as we increase the training dataset size (Eq. 11) and model size (Eq. 16).
Notice that AdamW's optimization process differs slightly from standard EMA process. Standard EMA typically assumes the updates qi 's are independent, while the 's (Eq. 6) in AdamW's EMA view are correlated since later updates depend on the weight ema,, which is affected by the current update q:. However, we can still use EMA as an approximation, and provides useful insight and intuition for understanding the hyperparameters.
For all experiments, we implemented the model using Py- Torch and performed optimization using its Adam W imple- mentation. For all tasks, we did not use weight decay on the normalization layers. Experiments were conducted on an internal cluster of Nvidia Titan (itan///0000///000iiill and only one GPU card used at time for each run. The results for image classification tasks are the averages from three distinct random seeds, for LLM pre-training tasks we only use one seed due to resource constraints. Comprehen- sive details on the hyperparameter search range and model specifications can be found in Appendix G.
GPUs,
3.2. Formalizing intuition from the EMA
In EMA, perhaps the key parameter is the timescale, A. Additionally, the initial value is also important; in Adam W, this corresponds to the size of the first update relative to the scale of the initialization, /./a. Indeed, Theorem Ishows that, under a scale-invariant network, if these two quantities are fixed, the whole optimization trajectory is determined.
4.1. Transferring weight decay across dataset sizes
Theorem 1. Consider two AdamW optimizers with differ- ent learning rates, weight decays, initialization scales and epsilons, (Mt, X, a. e vs. The X, o', d'))). Take we to be the
First, we study the relationship between the optimal À and the training dataset size. Remember that the timescale, Titer
3

---

How to set AdamW's weight decay you scale model and dataset size
Dataset size
80,000
160,000
320,000
640,000
1,280,000
Train Acc (top-1)
Train loss
Test Acc (top-1)
Test loss
1.0
0.6
0.4
Tune À
2.5
0.2
(A) ResNet
0.0
0.0
0.0
10-3
10
10A
10a
10-3
10
10
10a
10
10
10
10a
10-2
10
10
10a
A
1.0
0.6
0.4
Tune
Tepoch
0.2
0.0
0.0
0.0
10"4
10-?
10a
10
10-1
10-?
10
10'
10
10
10
102
10-1
10-2
10"
10?
Initial Tepoch = 1/nAM
Initial Topach = 1//AM
Initial Tapach = 1//AM
Initial Tapach = 1/rAMM
Train Acc (top-1)
Train loss
Test Acc (top-1)
Test loss
1.0
7
0.4
0.5
4
Tune À
(B) VIT
0.0
1
0.0
10-3
10
10
10a
10-3
10
10a
10a
10-3
10
10a
10a
10-3
10a
A
A
1.0
0.4
Tune
0.5
Tepoch
10-4
10
102
10?
10-4
10-2 10º
10
10"
10-1
10 10"
10?
Initial Tepoch 1/nAM
Initial Tepach 1/nAM
Initial Tepach 1/nAM
Initial Tapach 1/nAM
Figure 1: The optimal Tepoch transfers across dataset sizes. We trained ResNet-18 (A) and ViT (B) on subsets of downsampled ImageNet of various sizes (lines of different colors) under different weight decay (dots on the lines) under a fixed batch size of 100. An initial learning rate 10- is used with cosine decay scheduling. The performance metrics after 100 epochs are plotted against the weight decay À and the corresponding timescale Tepoch computed with the initial learning rate. The dashed lines show the optimal Tepoch subset size of 320, 000: In both models, the optimal Tepuch fairly stable across dataset sizes whereas the optimal À decreases dramatically as dataset size grows.
intuitively measures the number of previous weight updates that we average over to get the weights. To understand just what proportion of the whole dataset we are averaging over, we propose to work with the timescale in epochs,
size, Titer
M Fepuch And rearranging Titer L/(nA),
x= 1/ (y/Tier) = 1/ (m)
(11)
tells that if the optimal Tepach changes little the dataset size increases, we would predict that the optimal À should decrease as the dataset size increases. We confirmed this hypothesis on a ResNet-18 and a ViT trained on ImageNet (Fig. I), along with NanoGPT (124M parameters) trained on Open WebText (Fig. 2).
Tepoch Tiee/MM 1/ (nAN/B),
(10)
where is the number of training samples, B is the batch size, and M = N denotes the number of iterations in an epoch. Therefore, Tepoch tells us how many epochs of past updates AdamW's EMA averages over, where one epoch passes over all the data in the dataset.
For the ImageNet experiments, used the 32 x 32 down- scaled version of ImageNet provided by (Chrabaszcz et al., 2017) to ensure that we were able to perform the large number of training runs required to assess performance for different hyperparameter settings and dataset sizes. We trained the models on different subsets of ImageNet, where we randomly drew 80, 160, 320, and 640 samples from each
Our prediction is that as datasel size increases, the optimal Tepoch should remain approximately fixed. If we increase the dataset size with minibatch size fixed, then the number of minibatches the dataset/an epuch, M, increases. Rear- ranging Eq. (10), that implies that Tue increases with dataset

---

How to set AdamW's weight decay as you scale model and dataset size
Fig. 12 A, B in appendix), indicating that scaling v may not be necessary if we keep Tepoch constant through adjusting À when N scales.
Subset ratio
0.25
0.5
Validation loss
3.4
4.2. Transferring weight decay across model sizes
3.3
Of course, it is also critical to understand how to modify the AdamW hyperparameters as we increase the model size, The most obvious approach is uP (Yang et al., 2022), which predicts that the optimal learning rate should decrease as 1/fan in, by considering the behavior of the first few learn- ing steps relative to the random initial weights. However, the theory behind μuP did not consider how weight decay affects the learned weights in the later phase of the optimization.
3.2
1091
1100
10-2
10
10
10
10
10
A
Initial Tepach = 1/nAM
Figure 2: For LLM pre-training, optimal weight decay shifts with dataset sizes but Tepoch transfers. We trained a 124M NanoGPT on subsets of Open with different sizes (line colors) for 4 epochs under a fixed batch size, an initial learning rate of 6 x 10-4 and various À (dots lines, varied in powers of 2). As dataset size increases, weight decay that gives optimal validation loss (red crosses) decreases, whereas Topach (Eq. 10) is stable across scales.
Formally, we take Thase and Ahase as the learning rate and weight decay tuned on a smaller base model, with fan-in of fan uan_inbse- Then, we scale the model width by a factor of 8,
s = fan ifa inhase.
(12)
The most direct approach using μP scaling with Adam (e.g. used by Lingle, 2024) scales the learning rate as rec- ommended by uP, while keeping the weight decay fixed,
of the 1, 000 classes except for the largest subset of 1.28M samples in total, where we randomly drew a subset from
the whole dataset. We used a fixed batch size of 100 for all runs, an initial learning rate of 10-3, swept À from 10-3 to 10S, and then plotted the performance metries after 100 epochs. The top row of Fig. 1A, IB, shows performance vs. X. Note that the optimal À decreases dramatically as the dataset size increases. The bottom row of Fig. 1A, 1B, shows performance metrics vs. Tepoch. Critically, the optimal Topoch is far more stable than the optimal X. Here, we used a cosine schedule which decays to 0.1 of the initial learning rate. In the Appendix, we see similar patterns for a constant schedule, (Fig. 13) and cosine decay zero (Fig. 14).
17 Musse /
A Abuse.
(13)
Ther = 1/ = s/ (Doase =
Titerbase
(14)
However, the implied EMA timescale, now changes with model size, while we hypothesize that the optimal timescale should not vary with the model size. Therefore, we conjec- ture that the scaling in Eq. (14) would break the transferabil- ity of optimal learning rate across model sizes if we run the optimization for long enough (see Figure 5AB, top rows).
How do we resolve this issue? We have two desiderata. First, from μP, we should scale 17 with 1/fan_in in. Second, we have the timescale Titer 1/(n>) fixed. It might at first seem difficult to reconcile these two desiderata. Indeed, it is impossible in the usual setting where À is fixed. However, it is possible to reconcile them if we allow the weight decay, X, to strengthen for larger models.
Next, we consider training 12-layer NanoGPT with 124M parameters on subsets of Open OpenWeb Wex (Gokaslan et al., 2019) of various sizes. We trained the model using the first 1/4, 1/2, and the whole dataset for approximately epochs under a fixed batch size of 122880 tokens with a cosine learning rate decay schedule from 6 10-1 to 6 x 10-5, We swiped À between 2-8 to 2, and plotted the final validation loss against optimal À and Tepoch (Fig. 2), where we again observe that as dataset size grows, optimal À decreases but Tepach is more stable
1] = Thuse /
A == S Xbases
(15)
Titer = 1/(12) = 1/ (V/huse Abuse) =
Therpbase-
(16)
We tested our conjecture that the optimal EMA timescale is constant across model sizes, while the optimal weight decay is not, by training a ResNet-18 18 of varying widths using JPP's s codebase on 320, 000 samples subset of down- scaled ImageNet, where we randomly drew 320 samples from each class. We scaled the models' widths by factors of 2, giving E {0.5, 1, 2}. We fixed These = 10-3 and used 17 = / /s to change the learning rate with network size. We swept Apase from 10-2 2 to 10'. We then consid- ered two ways of modifying the weight decay as we scaled
Note that in the experiments considered, we used a fixed initial learning rate under different dataset sizes. Some recent works (Bjorck et al., 2024; Schaipp et al., 2025) suggest that, under a fixed weight decay and batch size, the optimal learning rate decreases as training run gets longer. We empirically verified this observation in our experiment setting (top rows in Fig. 12 A, B in appendix). However, we also notice that if we scale Abase x 1/N, the optimal n becomes consistent across dataset sizes (bottom rows in
5

---

How to set AdamW's weight decay you scale model and dataset size
0=0.5
$=1.0
$=2.0
Train Acc (top-1)
Train loss
Test Acc (top-1)
Test loss
1.0
0.4
5.5
2
0.3
4.0
n - 10 3/s.
A = Anuse
0.0
0
0.2
2.5
10
10
10
2
10
10
10
2
10
1
10
10
1
20
1.0
4
0.4
5.5
0.5
2
0.3
4.0
n=10 10 2/5. À SAbase
0.2
2.5
10
10
10
10
10
100
10
10
10
10
10:
10
Ausse
Agsse
Abase
Absse
Figure 3: The optimal X increases with model size whereas the optimal timescale is more stable. We trained ResNet- 18 on subset of ImageNet 32x32, with varying width factor (lines of different colors) under a fixed base learning rate 10-3 with varying weight decay (dots on the lines) and plotted the metrics after 50 epochs vs. weight decay strength. The top row scales the hyperparameters using the direct μP approach (Eq, 13; i.e. fixed A), while the bottom row scales the hyperparameters to ensure Tier is fixed (Eq. 15; À increases with model size). Note that as Taase 10-3 3 is fixed, there is a direct relationship between the optimal Absse and the optimal Tieerbbase.
Doese 3 x 10"4
Ahase =3 x 10
dramatically across network sizes. In contrast, Fig. 3 (bot- tom row) increases the weight decay as the network size increases (Eq. 15), as that leaves The unchanged (Eq. 16). It is evident that the optimal Aamee and hence Taer is far more stable in this setting. Notably, while e.g. the optimal Apase for test loss on the bottom row does appear different from the others if anything, this would indicate that the relation- ship between s and À is even stronger than expected. In particular, this point indicates that even as we scale À with s, the optimal Auuse may still be increasing with s. That would seem to indicate that the relationship between s and Agase might be superlinear (e.g. Auuse x s", where a > 1), though we would need more, larger-scale experiments to definitively establish any such super-linearity, which is out- of scope for the present work.
Validation loss (shiited)
256 (3.84)
0.0
256 (3.94)
512 (3.50)
(3.60)
1024 (3.23)
1024 (3.36)
(A)
-0.05
Dsssels,
Zowce
-0.10
(3.84)
256
Validatinn loss
0.000
512 (3.50)
0.0
512 (3.60)
(shitted)
1024 (3:23)
1024 (3.36)
(8)
-0.025
n = Ds
s
-0.050-
-0.1
10-2
10-2
10-1
100
Abase
Apuse
Figure 4: For LLM pre-training, optimal weight decay increases with model size whereas the optimal timescales transfer. We trained 8-layer GPT's on OpenWebText, with various widths (line colors), under two Trase with actual learning rates for each width scaled following μP with s 1024 wilth We plot the final validation loss against various Abase (dots on lines, varied by powers of 2). We align all lines to the rightmost point for better visibility, where the numbers in the brackets denote the actual optimal validation loss each width. We considered training under two parameterizations, if we keep À decoupled from then the optimal Aruse shows clear shift with model size (top rows), if X increases with s. optimal Abase becomes much stable across widths.
We conducted similar experiments for LLM pre-training, we considered 8-layer GPTs, with hidden states of widths in [256, 512, 1024), achieved by changing the number of attention heads and fixing each head's dimension. We con- sidered two base learning rates, 3 x 10-4 and 3 10-5 trained for 250K iterations, and swiped lambda in the range 2 to 1. The results shown in Fig. 4, broadly, if we fix À irrespective of (top row, Eq. 13), optimal À increases with model sizes. However, as we scale À x s (bottom row, Eq. 14), the optimal Abuse becomes consistent across widths, i.c. confirming that optimal timescale transfers.
Importantly, scaling the weight decay correctly has im- portant implications even for the original μP predictions about how the optimal learning rate transfers across model sizes. Specifically, we trained a ResNet-18 on CIFAR-10
the network and plotted the metrics after 50 epochs against weight decay. Fig. 3 (top row) leaves À fixed as the network
size increases (Eq. 13). We can see the optimal Anose varies
6

---

How to set AdamW's weight decay you scale model and dataset size
5-0.25
5-1.0
5m4.0
Train Acc
Train loss
Test Acc
Test loss
1.0
10a
0.95
0.75
(A) CIFAR-10
1 Thasels
0.8
10-1
0.85
0.50
==1.0
0.6
0.75
0.25
1.0
0.95
0.75
109
n = Mhasals,
0.8
10-1
0.85
0.50
A
0.75
0.25
10
10-4
10-2
10-4
10
10-2
10-
10
10-2
10-*
10 Abass
10
Ooase
Obose
Abase
s=0.5
s=1.0
$=2.0
Train Acc (top-1)
Train loss
Test Acc (top-1)
Test loss
0.70
4.5
0.4
4.5
(B) ImageNet
3.0
3.5
==1.0
n= Nhasol,
0.20
1.5
0.2
2.5
0.70
4.5
0.4
4.5
n= Mbase
0.45
3.0
0.3
A=
1.5
10
10
10-3
10
10-+
10-3
10->
10
10-3
10->
10
10-3
Oasse
Phase
Toase
Ahase
Figure 5: Adam W breaks the learning rate scaling of μP. Following the experiment setting in (Yang et al., 2022), we trained a ResNet-18 with varying width factor (lines of different colors) under various base learning rates Thase (x-axis) on CIFAR-10 10 (A) and a 320, 000 samples subset of ImageNet 32x32 (B). We then plotted the metrics after 200 (for CIFAR-10) and 50 (for ImageNet) epochs against Thase. The top row scales the hyperparameters using the direct μPP approach (Eq. 13; i.e. fixed A), while the bottom row scales the hyperparameters to ensure Tier fixed (Eq. 15; À increases with model size). In both datasets, the direct approach breaks the stability of optimal Thase in terms of test metries due changing the timescale whereas our scaling allows for consistent Masse across model sizes.
and on 320K images downscaled from ImageNet, with fixed Abase = 1.0 and swept Thase from 10-5 to 10-2. In the top row of Fig. 5A and B, we use the standard juP scaling (Eq. 13), with a fixed weight decay. Remarkably, we found that the optimal base learning rate varied dramatically across model sizes. This indicates that yP scaling of the optimal learning rate breaks down in AdamW. Indeed, the original uP paper (Yang et al., 2022) did not examine AdamW, and this finding is confirmed by recent work on hyperparameter transfer in large-scale transformers (Lingle, 2024; Blake et al., 2024). We hypothesized that this breakdown was due the timescale changing with model size (Eq. 14), and thus could restore the usual μP scaling of the learning rate by increasing the weight decay with model size (Eq. 15) to fix the timescale when sweeping Thase. To confirm this finding, in the bottom row of Fig. 5A and B, we increase the weight decay with model size (Eq. 15). We found that the optimal base learning rate was now far more stable when varying model sizes, confirming our hypothesis. We observed simi- lar behavior when training ViTs of different widths on the
ImageNet subset, the results are presented in Appendix E, and follow-up work, citing us, observed similar behavior in large-scale training of transformers (Blake et al., 2024).
Lastly, we additionally considered constant and cosine decay to zero learning rate schedule for experiments Fig. 3 and Fig. 5, the results are shown in Fig. 16, 17 and Fig. 18, 19, respectively in the appendix.
Finally, the size of the weight decay has implications for the magnitude of the learned weights, and hence whether the outputs Wa r remain O (1) which was one of the key desideratum used to derive μP (Yang al., 2022). partic- ular, in Adam W the magnitude of the learned parameters, Wi scales with 1/X (or Wij 1/A). In framework, this scaling arises because the quantity that AdamW takes the EMA of (i.e. qi in Eq.6) scales with 1/A (we provide empirical justifications in Appendix. D), but this also arises in Kosson et al. (2023). As we propose gests that the learned weights (not the initial weights) scale as, Wa 1/2 1/fan in. Thankfully, this is precisely the
fan in, this sug-

---

How to set AdamW's weight decay as you scale model and dataset size
scaling we need ensure that, yi = سز blow up or shrink the model size increases (specifically, y: ~ 1). In particular, the learned weights, Wa are a sum of many updates, AW Thus, following the same reasoning as that for Eq. (2), we have yi ~ Wij x fan_in. Thus, to ensure that y/ ~ 1, we need Wij ~ 1/fan_in, in, which is precisely what we get using our proposed scaling.
fan in Wy does not f=1
the potential risk of optimization diverging due to learning rate set too large. However, there exists an alternative ap- proach to keep Tepoch unchanged as we scale B, that is to have À x B. Indeed Bergsma et al. (2025a) demonstrates the feasibility of such an approach, which shows optimal À scales linearly with respect to B, and allows for much more stability and flexibility compared with adjusting v when seal- ing B. Additionally, our manuscript only considered high epoch training where the same dataset is repeated multiple times, and suggests that optimal Tepach should stay constant and lies between one and the total number of epochs as dataset sizes increase. However, this may not be the case, LLM pre- -raiiinnnss 's one epoch setting, where each training data point is only presented once. Particularly, Bergsma et al. (2025a) show that optimal Tepuch should not stay con- stant in one epoch setting, but rather follows a power law optimal s (TPP) -0.527 where the optimal Tepoch decreases as the poken-per-paramete (TPP) increases, and is always smaller than 1 when TPP > 1. Intuitively, this suggests that when the training run is short, it would be better to av- crage over all past updates (high Tepoeh), whereas as training gets longer, it becomes more optimal to forget noisy updates seen in the beginning of the training (low Tapach).
5. Related work
5.1. Work after us
There have been number of follow-up works that follow up on the concepts presented in this paper. First, Blake et al. (2024) confirmed that increasing weight decay as we increase model width fixes learning rate transfer. Second, Bergsma et al. (2025b) has built on our EMA interpreta- tion of AdamW to understand and propose better learning rate schedules. Third, Bjorck al. (2024) looked at the scaling of optimal learning rates with dataset size. They found that the optimal learning rate decreases with dataset size (N), with a fixed batch size, and a fixed weight decay. This broadly what you would expect from our findings, that you need to increase the EMA timescale as the dataset size increases, and you can do that by reducing the weight decay, or reducing the learning rate. However, they found that the dependence of learning rate on dataset size was sublin- ear. We suspect this is because you modify the learning rate with weight decay fixed, you simultancously change two things: the EMA timescale (which you probably do want to change) and the behaviour of the optimizer close to initialization, which you probably do not want to change (Appendix A). Additionally, our experiments verified that Tepoch transfers across model widths on small-scale prob- lems while recent works (Dey et al., 2025; Narayan et al., 2025) further conclude that Tepoch is also transferrable in billion-parameter I.LM pre-training.
(IPP))-.527,
5.2. Work before us
Further, our EMA view of Adam W provides an explanation for several striking observations made in the literature.
First, the original AdamW paper (Loshchilov & Hutter, 2018) 8) proposed parameterization in terms of the learning rate, 1 and 7 = y7 describing the weight decay. They, as well as a recent blog post (Schaipp, 2024), made an empirical observation that v) and y were more decoupled than 7) and X (in the sense that the optimal v) depended strongly on À and less strongly on 7). We are able to provide theoretical understanding of their observation, by noting that y is related to our timescale, 7 = 1/Tur We are therefore able to identify two separate processes one associated with 1 and another associated with 7 = 1/7 First, describes the size of the initial updates relative to the random initialization, which is relevant close to initialization (see Appendix A for details). In contrast, 7 1/rwe describes the EMA timescale, which is relevant as the model moves away from the initialization. Finally, Loshchilov & Hutter (2018), did not use these insights to propose how to scale the weight decay with model and dataset size (our key contribution).
Lastly, very recent work by Bergsma et al. (2025a) con- sidered two critical problems unstudied in our current manuscript: The effect of batch size as well as the opti- mal value and transferability of Tepoch when you only train the model for one epoch, as in LLM pre-training settings. In our experiments, we kept the batch size B fixed when tuning other hyperparameter, however in real world pre- training, batch size B is often tuned extensively to ensure maximum utilization of parallelism, Therefore it is impor- tant understand the how to adjust other hyperparameter as we scale B. Fortunately, the heuristic of keeping Tepoch (Eq. (10)) constant provides guidance on how to achieve this: In order to keep Tepach unchanged as B scales, one can have x B as suggested by McCandlish et al. (2018), however the picture is often more complicated in practice
Second, our viewpoint (Sec. 3.2) suggests that the effects of the learning rate, 17, on the magnitude of the updates relative to the random initialization should become less relevant than the EMA timescale as training proceeds (Appendix A). While this seems strange (of course the learning rate mat- ters a lot), there are some suggestions in the literature that indeed the EMA timescale may be more important in some
due to e.g. AdamW's effective learning rate induced by the gradient normalizer (Malladi et al., 2022), as well as
8

---

How to set AdamW's weight decay as you scale model and dataset size
settings. Specifically, Wortsman et al. (2024) ran proxies for large-scale L.MM pretraining. Fig. 6 in their paper showed that if y/λ = 1/rue is fixed, the final validation loss is relatively insensitive to the learning rate, 1). In contrast, if X is fixed, the validation loss shows much more pronounced sensitivity to the learning rate, 7). This makes sense if the key hyperparameter is Titer as modifying the learning rate, n, with fixed implies Tieer = (/(r)) also changes, How- ever, different from work, Wortsman et al. (2024) do not explicitly study how Ther transfers across model sizes or if keeping À fixed would break pP. Indeed, the 7 = 7) parameterization is often used in SGD, where it is identi- fied as the effective learning rate (Zhang et al., 2019b; Li & Arora, 2019; Li et al., 2020; Wan et al., 2021; Li et al., 2022). Recent works extend the effective learning rate per- spectíves AdamW by viewing AdamW as sign gradient descent (D' Angelo et al., 2024) or rotation speed on fixed radius sphere (Kosson et al., 2023)². However, these literatures do not study how optimal y7 transfers across model and dataset sizes.
Concurrent work (D' A'geel et al., 2024) suggests that weight decay does not serve as regularization in deep learn- ing, but rather controls minibatch noise, which grows with yy. Our work provides a similar viewpoint, in the sense that, if you are averaging over more data points (higher Titer = I/ma), then the iteration-to-iteration noise will of course be smaller. But the EMA framework also provides you with guidelines for how to set or equivalently the timescale as you scale model and dataset size.
Of course, μP itself is related work (Yang et al., 2022), and is discussed extensively in the Background section. Van Laarhoven (2017) argue that in Adam, the weight decay hyperparameter changes the scale of the learned weights, and thereby has an influence on the effective learning rate. This effect pops up in our framework in the 1/A factor in qe (Eq. 6), the quantity over which take the EMA. Addi- tionally, recent work (Busbridge et al., 2023) studies the sce- narios where an EMA is used for averaging model weights, such as in self-supervised learning, and proposes hyperpa- rameter scaling rules to achieve the same performance under different batch sizes. Our work differs in that we are inter- preting AdamW as EMA over past updates rather than using explicitly using EMA for averaging weights. Finally, hints that the optimal weight decay decreases with dataset size have appeared previously in empirical literature when the authors have conducted thorough hyperparameter sweeps, though they did not propose a quantitative, theoretically motivated scaling rule (Zhai et al., 2019).
Third, Lingle (2024) assessed the accuracy of μP in large- scale LLM pretraining. Remarkably, they found that the usual μP scaling of the optimal learning rate with model size broke down for AdamW. We provide an explanation and fix for this by noting that pP theory considers only the behaviour of the first few optimization steps, relative to the random initialization. The issue is that if the weight decay, A, is fixed, then Ther changes with model size (as Ther depends on 17 and 17 changes with model size following the usual yP recommendations). We propose fix for the issue by proposing to strengthen weight decay as the model size increases, and validate the fix experimentally.
6. Conclusions
We showed that AdamW's weight updates can be understood as an EMA, showed that the optimal EMA timescale is fixed as we scale model and dataset size, and explored the implications for hyperparameter scaling of the weight decay.
Fourth, many papers have considered scaling with batch size (which we hold fixed) (Zhang et al., 201 9a; Malladi et al., 2022; Wang & Aitchison, 2024). These papers find that when the batch size small, the optimal learning rate is also small. Then, the batch size increases, the optimal learning rate also increases. However, there's limit, the critical batch size, at which point increases in the batch size no longer allow you to increase the learning rate. Once you hit the critical batch size, you no longer get any benefit from increasing the batch size, and you instead must increase the number of iterations (McCandlish et al., 2018; Shallue et al., 2019). Note that while very recent work suggests that the optimal batch size does increase with dataset size, they con- firm that this increase is sublinear, implying that the number of training iterations will increase (Zhang et al., 2025). As soon as you need to increase the number of iterations, our recommendation to hold Tepach constant comes into play.
Impact Statement
This paper presents work whose goal is to advance the field of Machine Learning. There are many potential societal consequences of our work, none which we feel must be specifically highlighted here.
"As pointed out by one of the anonymous reviewers, the timescale parameter in EMA process can been seen mea- surement of the relative update size studied by see Appendix. C for detailed discussion.
Kosson et al. (2023),
9

---

How to set AdamW's weight decay you scale model and dataset size
References
Everett, K. E., Xiao, L., Wortsman, M., Alemi, A. A., Novak, R., Liu, P. J., Gur, I,, Sohl-Dickstein,, J., Kaelbling, L. P., Lee, J., and Pennington, J. Scaling exponents across parameterizations and optimizers. In ICML, 2024.
Arora, S., Li, Z.., and Lyu, K. Theoretical analysis of auto rate-tuning by batch normalization. In ICLR, 2019.
Balles, L. and Hennig, P. Dissecting adam: The sign, mag- nitude and variance of stochastic gradients. In ICML, 2018.
Gokaslan, A., Cohen, V., Pavlick, E., and Tellex, S. Open- webtext corpus. http://Skyl on . thub. Lo/OpenneeTexxtCopppo 2019.
Bergsma, S., Dey, N., Gosal, G., Gray, G., Soboleva, D., and Hestness, J. Power lines: Scaling laws for weight decay and batch size in IIm pre-training. arXiv preprint arXiv:2505. 13738, 2025a.
Jia, X., Song, S., He, W., Wang, Y., Rong, H., Zhou, F,, Xie, L., Guo, Z., Yang. Y. Yu, L., et al. Highly scal- able deep learning training system with mixed-precision: Training imagenet in four minutes. arXiv preprint arXiv:/801) 0800.110555 2018
Bergsma, S., Dey, N. S.. Gosal, G., Gray, G., Soboleva, D., and Hestness, J. Straight zero: Why linearly decaying the learning rate to zero works best for LLMs. In ICLR, 2025b.
Kosson, A., Messmer, B., and Jaggi, M. Rotational equilib- rium: How weight decay balances learning across neural networks. In NeurlPS 2023 Workshop on Mathematics of Modern Machine Learning, 2023.
Bjorck, J., Benhaim, A., Chaudhary, V., Wei, F., and Song, X. Scaling optimal Ir across token horizons. arXiv preprint arXiv:2409. 19913, 2024.
Li, Z. and Arora, S. An exponential learning rate schedule for deep learning. In ICLR, 2019.
Blake, C., Eichenberg, C., Dean, J., Balles, L., Prince, L. Y., Deiseroth, B., Cruz-Salinas, A. F., Luschi, C., Weinbach, S., and Orr, D. u-mu p: The unit-scaled maximal update parametrization. arXiv preprint arXiv:2407. 1ot.22477..7..... 17465, 2024.
Li, Z., Lyu, K., and Arora, S. Reconciling modern deep learning with traditional optimization analyses: The in- trinsic learning rate. In NeurlPS, 2020.
Li, Z., Bhojanapalli, S., Zaheer, M., Reddi, S., and Kumar, Robust training of neural networks using scale invariant architectures. In ICML, 2022.
Bordelon, B., Noci, L., Li. M. B., Hanin, B., and Pehle- van, C. Depthwise hyperparameter transfer in residual networks: Dynamics and scaling limit. In ICLR, 2024.
Lingle, L. A large-scale exploration of y-transfer. 2024.
Busbridge, D., Ramapuram, J., Ablin, P., Likhomanenko, T.,
Loshchilov, and Hutter, Decoupled weight decay regu- larization. JCLR, 2018.
Dhekane, F. G., Suau Cuadros, X., and Webb, R. How to scale your ema. In NewrIPS, 2023.
Malladi, S., Lyu, K., Panigrahi, A., and Arora, S. On the sdes and scaling rules for adaptive gradient algorithms. In NeurlPS, 2022.
Chrabaszez, P., Loshchilov, I., and Hutter, F. A downsam- pled variant of imagenet as an alternative to the cifar datasets. arXiv preprint arXiv:/707.000 (9881,,00 2017.
McCandlish, S., Kaplan, J., Amodei, D., and Team, O. D.
Cubuk, E. D., Zoph, B., Mane, D., Vasudevan, V., and Le, Q. V. Autoaugment: Learning augmentation strategies from data. CVPR, 2019.
An empirical model of large- batch training. arXiv preprint 1rXi:
2018.
Müller, R., Korbblith, S., and Hinton, G. E. When does label smoothing help? In NeurIPS,2019.
D' Angelo, F., Andriushchenko, M., Varre, A., and Flammar- ion, N. Why do need weight decay in modern deep learning? In NeurlPS, 2024.
Narayan, S., Gupta, A., Paul, M., and Blalock, D. junit scaling: Simple and scalable fp8 Ilm training. arXiv preprint arXiv:2502.05967, 2025.
Dehghani, M., Djolonga, J., Mustafa, B., Padlewski, P., Heek, J., Gilmer, J., Steiner, A. P., Caron, M., Geirhos, R., Alabdulmohsin, I., et al. Scaling vision transformers to 22 billion parameters. In ICML, 2023.
Noci, L., Meterez, A., Hofmann, T., and Orvieto, A. Super consistency of neural network landscapes and learning rate transfer. In NeurlPS, 2024.
Dey, N., Zhang, B. C., Noci, L., Li, M., Bordelon, B., Bergsma, S., Pehlevan, C., Hanin, B., and Hestness, Don't be lazy: Completep enables compute-efficien deep transformers. arXiv preprint arXiv:2505 00668 2025.
Schaipp. How to jointly tune learning rate and weight decay for AdamW, ht https: www.aanaaannnnaaaae github. io/posts/2022 / 2024. 2024.
10

---

How to set AdamW's weight decay you scale model and dataset size
Schaipp, F., Hägele, A., Taylor, A., Simsekli, U., and Bach, F. The surprising agreement between convex optimiza- tion theory and learning-rate scheduling for large model training, arXiv preprint (tx.xv:2501.88900 2025.
Zhang, H., Morwani, D., Vyas, N., Wu, J., Zou, D., Ghai, U., Foster, D., and Kakade, S. M. How does critical batch size scale in pre-training? In ICLR, 2025.
Zhang, S., Roller, S., Goyal, N., Artetxe, M., Chen, M., Chen, S., Dewan, C., Diab, M., Li, X., Lin, X. V., al. Opt: Open pre-trained transformer language models. arXiv preprint arXiv:2205.0 01068, 2022.
Shallue, C. J., Lee, J., Antognini, J., Sohl-Dickstein, J., Frostig, R., and Dahl, G. E. Measuring the effects of data parallelism on neural network training. Journal of Machine Learning Research, 20(112):1-44 2019.
Touvron, H., Lavril, T., Izacard, G., Martinet, X., Lachaux, M.-A,, Lacroix, T., Rozière, B., Goyal, N., Hambro, E., Azhar, F., et al. LLaMA: Open and efficient founda- tion language models. arXiv preprint 1r311 2023a.
Touvron, H., Martin, L., Stone, K., Albert, P., Almahairi, A., Babaci, Y., Bashlykov, N., Batra, S., Bhargava, P., Bhosale, S., et al. LLaMA 2: Open foundation and fine- tuned chat models. arXiv preprint arXiv:2307.09288, 2023b.
Tow, J., Bellagente, M., Mahan, D., and Ruiz, C. Techni- cal report for stablelm-3b-4el Technical Report, 2023.
Van Laarhoven, T. 1.2 regularization versus batch and weight normalization. arXiv preprint (tt1:1706.053500 2017.
Wan, R., Zhu, Z.,, Zhang, X., and Sun, Spherical mo- tion dynamics: Learning dynamies of normalized neural network using sgd and weight decay. 2021.
Wang, X. and Aitchison, L. Batch size invariant adam. In OPT 2024: Optimization for Machine Learning. 2024.
Wortsman, M., Liu, P. J., Xiao, L., Everett, K. E., Alemi, A. A.. Adlam, B., Co-Reyes, J. D., Gur, I., Kumar, A., Novak, R., Pennington, J., Sohl-Dickstein, J., Xu, K., Lee, J., Gilmer, J., and Kornblith, S. Small-scale proxies for scale transformer training instabilities. In ICLR, 2024.
Yang. G.. Hu, E. J.. Babuschkin, I., Sidor, S., Liu, X., Farhi, D., Ryder, N., Pachocki, J., Chen, W., and Gao, J. Tensor programs v: Tuning large neural networks via zero-shot hyperparameter transfer. arXiv preprint arXiv:2203.03-466 2022.
Zhai, x., Oliver, A., Kolesnikov, A., and Beyer, L. S41: Self- supervised semi-supervised learning. In ICCV, 2019.
Zhang. G., Li, L., Nado, Z., Martens, J., Sachdeva, S., Dahl, G., Shallue, C., and Grosse, R. B. Which algorithmic choices matter at which batch sizes? insights from a noisy quadratic model. Advances in neural information processing systems, 32, 2019a.
Zhang, G., Wang. C., Xu, B., and Grosse, R. Three mecha- nisms of weight decay regularization. ICLR, 2019b.

---

How to set AdamW's weight decay you scale model and dataset size
A. Confirming the connection between Adam W and weight decay
A.1. Proof of Theorem
Here, we consider two Adam training trajectories with hyperparameters Vj, X, and m/. x', We prove that for a specific way of initializing the optimizer state, and if the optimizer hyperparameters are related by,
T/ A
(17)
X' = cA
(18)
= cc.
(19)
then the full trajectories are the same. Notice that 1/Teet 17, '//a remains unchanged with this choice of hyperparame-
ters.
We consider AdamW updates of the form,
(20a)
my 3 + (1-31)99
(20b)
0t 8204-1 + +(1 - -B2)
(20c)
m2
m, - B
(20d)
Ut
0 1 - $2
(20e)
m Vve +
Wt (1 Am W-t11 + Tt
(20f)
(And the analogous equations for the primed optimizer). The optimizer state is the variables that persist across timesteps, namely m, v. and w. Specifically, the first optimizer has state my, Vt. and w/ while the second optimizer has state m/. d/. and wy. Our goal to prove that for a scale-invariant network, as À or the trajectory of 7/: changes while the trajectory of 1/Tunee A fixed, the states for the different optimizers are related by scale parameters,
m/_1 cm4-1,
(21a)
1/_1
CiU-11
(21b)
1 w c W
(21c)
where e X/AA nt/1/ The argument proceeds inductively, by assuming that the scaling relations in Eq. (21) hold at timestep t - 1, then proving that as consequence they hold at timestep t.
We start by considering the updates for m and v, which require the scaling of the gradients. In particular, assuming, that Eq. (21c) holds at timestep I, Appendix A.2 tells us that the relationship between gradients is,
g. v eV cV cgt.
(22)
Thus, the updates for m' are,
m = 3/m4-1 (1 1--1)gg.
(23)
substituting Eq. (21a) and Eq. (22)
m = cmcmm-1 + (1 - 1-)BB1)cg = (1
-01--000 = cm,
(24)
So the relationship continues to hold for at timestep For updates,
of = 320 + (1 (g)
(25)
substituting Eq. (21b) and Eq. (22),
of = B202 11 + (1- 32)09/ 2 (B20 v + (1- (1-3))97) 2v
(26)
12

---

How to set AdamW's weight decay you scale model and dataset size
So the relationship continues to hold for timestep
Next, consider rn and é, which are not in themselves state variables, as they can be computed directly from m that timestep.
m4 1 8 my
1 1- 81 CM cm
(27a)
6 1 8 u/
1 81 é = c' Ur.
(27b)
So the scaling relationships for in and i are analogous those for m and v.
Finally, we consider the weight updates themselves. The weight updates for the two optimizers are,
m
Wy = (1 1/Teas 1 + 1)t
(28a)
02 +
m
w/- = 1/meare ) + T/(
(28b)
Vil +
Substituting 7/- = 3/Tw/// and r// = 1/Twese/
1 me
Wi (1 1/Tume W + 1/Turs
(29a)
X
v
0z +
1
m
w = (1 1/Tee Jui + T/Til
(29b)
X
/
8 +
Now, substituting Eq. (19) and Eq. (18) into the form for w/,
1 cmly
w (1 / + 1/Tees
(30)
cd
2 + ce
Substituting Eq. (21c) and Eq. (27),
1 cmi
w/ (1 1-//mmme w + 1/Tuni
(31)
c² vc + Ce
Cancelling c,
1 m
w/ (1 1/rane) W1 + 1/Terit
(32)
C³ v Uz
And pulling out 1/c,
w/ = (a WW + 1/ront
! wr.
(33)
Vor
As required.
Thus, we have shown that the scaling relations in Eq. (21) hold at timestep L- - 1, they also hold at timestep f. It only remains to complete the inductive proof by showing that the scaling relations hold at initialization. The m and v state variables are usually initialized zero, mo = do 0, and the scaling relations of course hold zero. However, networks are usually initialized with a fixed parameter scale which does not depend on X. To ensure full ccale-invarian training, we therefore need to rescale the parameter initializations, ensure
w Unn
WD-
(34)
13

---

How to set AdamW's weight decay you scale model and dataset size
Loss
Gradient
Normalized Gradient
2.5
2N
1.5
W2
W2
0.5
-0.5
w1
W1
w1
Figure 6: An image of a 2D scale-invariant loss, clearly showing that the gradients get larger closer to the origin.
A.2. Gradients increase as weights decrease in scale-invariant networks
One important component of our proof is the notion that as weights decrease, the gradient of scale-invariant loss increases (e.g. see Arora et al., 2019). The geometry is straightforward if you look at an image (Fig. 6). But to spell it out formally, consider a scale-invariant loss,
((w) f(u)
(35)
BuL 1/0. where w = au. Now, take the partial derivative and use
The;
Ol(w) dC(u) ou, dd(())
(36)
Ow:
Ou: Ow: a
Ou;
Thus, bigger w (larger a) implies smaller gradients wrt w.
A.3. Empirical confirmation of Theorem
In this section, we empirically confirm Theorem I using a ResNet-18 and VIT (with QK-layernorm) trained on CIFAR- 10. We used AdamW to perform the optimization for 200 epochs with a cosine learning rate schedule to decay the initial learning rate by factor of 10. We used a fixed batch size of 100 and tested initial learning rates in range (10-6, 6 10-5
10-1). For each 7), we considered À in range (10-7/1 /n, 10 10-6 po-.//,..,,. 10-2/1) 2/n). giving us Tepoch in range x 10 2 1 .... 10a} for all rys.
We begin by looking the performance of ResNet under different values of Tepoch and under the standard configuration. Fig. 7B and Fig. 8B show performance metrics vs. timescale for different learning rates after 200 epochs. We can see from Fig. 7B that the behavior for different learning rates, is very different. In addition, Fig. 8B shows that, while for most rys, the optimal Tepoch in terms of test metrics is the same, the exact performance at the optimal Tepoch differs.
This suggests that standard network setups require considerable modification before training trajectories actually become invariant 17 for fixed 1/The (or Tepoca). In particular, recall that Theorem has two major assumptions about the model: 1. needs be scale-invariant (Eq. 9): 2. The initialization needs be 7--aaaee (Eq. 34). Therefore, need to make three key modifications in order for the assumptions to hold:
1. The output weights are not scale-invariant in standard setups, as they are not usually followed by normalization layer; we therefore introduced a normalization layer after these weights (Appendix A.3.1; row C in Fig. 7, 8)
2. Normalization layers such as batchnorm have scale and bias parameters, which are usually not scale-invariant, as such we adopted a decoupled learning rate for these weights (Appendix A.3.3; row D in Fig. 7)
3. We introduced an initialization that depended on the learning rate, r7 (Appendix A.3.2; row E in Fig. 7, 8).
Combining all these modifications together, we found that networks with the same Topoch but different values of 1 demon- strated exactly the same learning trajectories (Fig. 7F) and performance metries (Fig. 8F).
14

---

How to set AdamW's weight decay you scale model and dataset size
We additionally considered ViT under a similar setting. Notice that ViT needs more extensive modifications ensure full icale-invariance; full details are presented in Appendix. A.3.4. The performance iteration and timescale for various versions of the model are presented in Fig. 9 and 10. Mirroring the ResNet results, we again found that, under the suggested modifications, the performance becomes invariant to 7 under fixed Topoch-
See below for more details of the network modifications required to ensure that experimental results can mirror Theorem I in Sec. A. 1. OCTPUT-BATCHNOKE FOR SCALE-INVARIAN OUTPUT WEIGHTS
Most neural network weights/parameters are scale-invariant. This means that can multiply the weights/ parameters by an arbitrary constant without changing the network output (Eq. 9). It turns out that most neural network weights/parameterss are scale-invariant because they are followed by normalization such as batchnorm or layernomm, which will get rid of any change in scale. However, this is not true for the output layer, as the output layer is not usually followed by a normalization layer.
Therefore, we applicd a "global batchnorm" " layer after the output layer, i.c., upon the logits. In particular, denote the logits by Z e BBBxC where B the batchsize and C' is the number of classes, we first flatten (and reshape) Z into a matrix of shape BC x and then feed the vector into a standard ID batchnorm, which computes the mean and variance across all BC elements. Finally, we reshape the output of batchnorm back to B x C and send it the final softmax layer. The effect of output-batchnorm is shown in row C in Fig. 7 and 8.
A.3.2. LEARNING AAT-AI---RAT--EEEEENNNN INITIALIZATION
Note that Theorem 1 requires the ratio between the initial learning rate and the initial parameter scale to be the same in order for two AdamW optimizers to show the same trajectory. To satisfy the condition in implementation, we fix p = n/ar = 10 10-3 so the initial variance depends on the learning rate, = 1/10-3
Confirming this requirement, initializing the network weights in this way seemed to reduce the dependency on 1y for larger values of n under fixed Tepoch (Fig, 8E).
However, it is worth emphasizing that, p itself is another hyperparameter, and different p will change the performance (indeed, if we fix the initial variance, then P changes as we modify 7)).
A.3.3. DECOUPLING THE LEARNING RATES FOR THE BATCH/L PARAMETERS
In most networks, there are some parameters that are fundamentally non-scale-invariant These usually include the scale and bias parameters in normalization layers. All of the reasoning Theorem I rely on network output being invariant to scale, So how do we optimize these non-scale-invarian parameters?
It turns out that we already do something different with these parameters. In particular, it is common practice* to optimize them using the same learning rate, as the other parameters, but to drop weight decay. This standard setting is depicted in Fig. 7E and Fig. 8E, and we can see that 7 does still change the behavior of the optimizer.
The issue is that modifying 7), changes the training trajectory of these non-scale-invarian parameters. The solution is therefore to "decouple" the learning rate for the non-scale-invariantt parameters. In particular, we fix the initial learning rate for these non-scale-invariant parameters to 10-3 regardless of the learning rate 17 for other scale-invariant parameters. Doing so gives Fig. 7F and Fig. 8F, which exhibit almost perfect decoupling, with performance almost entirely independent of 1 for fixed Tepoch.
A.3.4. SCALE- SCALE-INVARIANT INVARIANT VIT
In order to make VIT iaveeiivvriiant, we need more extensive modifications beyond simply adding output normalization layer. In particular, we added layernorm after all linear layers in the network, which includes
The embedding layer.
"We cannot find literature explicitly studying the effect of dropping weight decay for normalization layers' parameters, however, its importance has been noticed in vision tasks, e.g. by Jia et al. (2018) and https: / //gittubbbbu wtwpahuuu//ssmmmmmmms networks/ ssues/ / as well language model training: https: gi thub, com/ /ka ka herppttty/minne///lloo/// /tt.s//ter//mmm///////////
m ngpt/mode1 wttp://mmmlll........... 412222
15

---

How to set AdamW's weight decay you scale model and dataset size
The query, key and value projection layers.
The output projection layer after the attention operation.
The linear layer following each attention block.
B. The weights in an EMA
Using Eq. (5), and assuming emag = 0.
ema1 1/Tu=Q1
(37)
emag 1/THe ((1 - ///mmm/ + 92)
(38)
emag = 1/Tue: ((1 - 1/Taaa ) + (1 - 1/nw)99 + 93)
(39)
Thus,
1. Σ (1m)//m)) 1-1
ema& = 1/nee
qr.
(40)
In deep learning settings, Tie is much larger than one (implying that we average over many iterations). In that setting, the first-order Taylor expansion is very accurate,
1 1/71a Ri e 1/m
(41)
Thus,
ema, 2 1/Tuer
Σ (0-1/ qr = 1/mm E ( ge-
(42)
Mm
with exact equality as Tier approaches infinity.
C. Connection between EMA timescale and relative update size
Consider an EMA given by 0r+1 = (1 - 7)a; + 7b, Assume that the b, are independently random and identically distributed over time with mean zero, and focus on the relative updates in the steady state (large t). The relative update size for the scalar case defined as
E (a - 3
(43)
E[a?]
To compute need to determine B[a?). We start by expressing a, in terms of by using the EMA formula:
[-
ax
M0- 7 7000
(44)
K=U
For simplicity, we assume dg 0 in the steady state, leading to:
t-1
at = Σ (1 - ") 1-
(45)
k=0
Squaring both sides gives:
1-1
2
of =
Σ
- 7) b
(46)
Expanding the square, we have:
af = 7²
Σ k=0
(1- 2k 62
11 -
1- +2 Σ (1 (1- (1-
(47)
ogk<jet
16

---

How to set AdamW's weight decay you scale model and dataset size
Due to the independence and zero mean assumption of by, the cross terms b for + j have zero expectation.
Thus, have:
B[a 2
7)
(48)
k=0
Assuming E[b] = o² and using the geometric series sum formula for large t, we get:
7²0² 1 (1 7)2
E Elay
(49)
This simplifies to:
3202
B|a 7(2 - 7)
(50)
Next will simplify B[(a++1 a,) we begin by expanding the inner term
(a++1 - az)² = 72 2b,a + az).
(51)
Taking expectations, we have:
E[ - az) = - 2E [b,a]] + E[aa)).
(52)
Assuming by and are independent and by has mean zero, E [b:a:] = 0. Thus:
El(a++1 at 7 (E] R(a)))
(53)
Using B[b]] = o² and substituting B[a]] from the previous derivation:
B[(a++1-0 - -0)²]= (0² 0²+ 720²
(54)
7(2 -7)
This simplifies to:
aa(a = 7202 1+ ?
(55)
2 Y
Finally, simplifying further gives:
[[ - a1)²]= = 2220²
(56)
2 79
Thus, the expression for r becomes:
2
2
(57)
o²
which indicates that the timescale parameter in the EMA is directly linked with the relative update size.
D. Weight norm
Notice that the EMA perspective of Adam (Eq. (6)) suggests that AdamW averages over recent 4: = Vo.+e* and if we is of scale (11) (Balles & Hennig, 2018), the weights acquired from AdamW should be of magnitude O voete
assume
To verify this, we trained a ResNet on CIFAR-10 for 1000 epochs under a batch size of 100, with an initial timescale Titer = 1/(1/)) = 10 We used a cosine learning rate schedule, starting from rio = 10-3/> and decaying to 0. We then swiped between {10-3 3 3 x 10-3 10 ..... 0.3, 1.0} and recorded the average absolute value of all dimensions in the convolutional and linear layer weights (i.e. weights where adopt weight decay) under different X.
Fig. 11 left plots the final weight magnitudes against X, with both X and Y axes log-scaled. By comparing it with the reference line 0.01/z (gray, dashed), we can see the magnitude against À line shares the same slope, indicating that the weight magnitude approximately follows E[ || = c/\, where e is a constant. Fig. 11 right plots the weight magnitudes throughout optimization, where we can see that all configurations start with the same magnitude at initialization, but gradually converge to our predicted magnitudes as optimization proceeds.
17

---

How to set AdamW's weight decay you scale model and dataset size
E. Additional experiment results
E.1. ViT scaling experiments on ImageNet
We additionally considered training a VIT on the 320K subset of ImageNet 32x32. In particular, we used a ViT with OK-layernorm similar to the one used in the main text, and we varied the width/size of the model by multiplying the number of hidden dimensions and internal linear layers' width with a factor s. Similar to the setting in the main text, we tested the
width factor in {0.5, 1.0, 2.0], used fixed Auuse 1.0, and swept Mheas in (2.5 x 10-6) x 2 with i E (0. 11).
The results are presented in Fig. 15, where we plotted the performance after 50 epochs against Thase. We again considered the direct uP scaling (Eq. 13; top row) and our proposed scaling (Eq. 15; bottom row), and we can see that the standard scaling breaks the stability of optimal learning rate in terms of test metrics whereas our proposed scaling is much more stable.
F. Note on muP library
We used the mup library² from the authors of Yang al. (2022) for varying-model-size experiments. In particular, for ResNet experiments in Fig. 3 and 5, we directly use the ResNet codebase under the examples folder together with the provided AuAdamii optimizer to run our experiments, as this library takes care of details such as the scaling of the learning rate and initialization for the input/output layers. For ViT experiments in Fig. 15, we manually constructed the required model shape file using the provided make _asee__bae______eeee base_shapes function and used the provided MuReadout module as the classification head.
Importantly, this library does come with a decoupled keyword argument. Given the connections between Theer and the pa- rameterization for AdamW originally proposed in the original "Decoupled Weight Decay Regularization" paper (Loshchilov & Hutter, 2018), you would have thought that you could implement our proposed scalings using decoup 1 ed=True. However, it turns out that as of writing, to get our proposed scaling for X (Eq. 15), you need to set decoupled=Fod eesee This may be fixed in the future, but in any case, if using the mup library, it is critical to check the mup source to see precisely what scalings you are getting for A.
G. Extended experiment setups
G.1. Model specification
For ResNet-18 experiments, we utilized the implementation from hitps://githut cww com/kuangiiu Pytorch-ci mrrch-cifan far/. For both CIFAR-10 and ImageNet, we used random cropping and horizontal flip as data aug- mentation and we used cross-entropy as the loss function. For μP experiments with ResNet, we use the codebase' 7 providedd by Yang et al. (2022), which is an adaptation of the codebase provided by the GitHub user kuangl iu.
For VIT, we adopted the implementation from https: / gig.......... com/ om hwt.hbb7vvi--CIPAR htte///A//ree//aaan We also incorporated QK layernorm suggested by Dehghann et al. (2023) and Wortsman et al. (2024) in order to stabilize the training when sweeping learning rates. The loss function is again chosen as cross-entropy loss. Additionally, when training on CIFAR-10, we follow the suggestions in the original codebase to use auto augmentation (Cubuk et al., 2019) and label smoothing (Müller et al., 2019) with a = 0.1 to alleviate overfitting. We indeed found these techniques crucial for reaching the level of test accuracy reported by the repo. For ImageNet experiments, we used label smoothing with standard data augmentation: random cropping and horizontal flip.
For dataset size scaling experiments with NanoGPT in Fig. 2, we used the configuration code provided by (D' Angelo et al., 2024).. The model structure is identical to the standard 124M version of NanoGPT, which contains 12 layers, each layer has 12 attention heads with head dimension of 64, i.e. an embedding size of 768. Identical to the original NanoGPT, (D' Angelo et al., 2024) uses micro-batch size of 12 sequences and 40 gradient accumulation steps at each iteration. However, to reduce memory consumption, it uses a fixed context length of 256 tokens instead of the original 1024-token context length.
optim. py optim.py4173 httpg:/..ht....ooococoooa/rrr/rr
httpps://githan http:://hh..tcttceoo/////////e/a
train_gpte small block256.py
18

---

How to set AdamW's weight decay you scale model and dataset size
For width transfer experiments in Fig. 4, used the codebase from the recent project nanoGPT-mup', In particular, trained transformers with 8 layers, with fixed attention head dimension of 64, and varied the number of attention heads between (4, 8, 16). We additional included QK layernorm for better stability when sweeping hyperparameters, The model is trained with micro-batch size of 8 sequences, 16 gradient accumulation steps, and a fixed context length of 256 tokens.
G.2. Hyperparameter range
In Sec. 4.1, for the experiments in Fig. 1, we used À (2.5 10 2 with i E {0. 1...., 15}.
In Sec. 4.2, for the experiments in Fig. 3, we used Abase 10-3 x 2' with i E {0, (0..,...,111 11).
For the CIFAR-10 experiments in Fig. 5A we used Thase = (2.5 x 10 x 4' with E {0, .,)
For the ImageNet experiments in Fig. we used Messe (2.5 10-9) 2' with i E {0, I......... 11).
For the NanoGPT experiments in Fig. 2, used 2' with {-8, -T,...:11 1).
For the NanoGPT experiments in Fig. 4. We used Abase
2' with i E {-8, -7, ...... -13 for Jhuse 3 x 10 and
Abase = 2' with i E {-6, -5.... 1) for Toase = 3 10
H. Licenses
ResNet-18 from httpgg/g./.tttttoooo///e///////r MIT licensed.
ViT from htangp/p.//.ttta..ooooaoo///////
is MIT licensed.
Ctt.K///tttseo///o.cccooooottttt (No license evident).
CIFAR-10 htt.l:0tttttsssw...sttt
ImageNet license is available at ottpspsswww..innge......ooo/oood
The mup library MIT licensed.
NanoGPT pre-training code from http://tt////ttttaooooo//aoaaaaa -eeaac is MIT licensed.
NanoGPT-muup code from http://p///gggtttoooooo////ttttt MIT licensed.
Phtpp:://gt http://hht......oe/eeeeeeeeaaaaa
19

---

How to set AdamW's weight decay you scale model and dataset size
n= 10
1=
n= 10
n= 10
==10-3
n= 10-2
10
10
=0
1.0
10
10a
Train loss
(A)
Standard output
10-2
Standard init.
Standard BN LR
10
200
100 Epoch
200
100 Epoch
200
100 Epoch
200
Epoch
Tepoch = 200
Tepoch = 20
Tepach = 2
Tepoch = 0.2
10
Train loss
(B)
Standard output Standard init.
10-2
Standard
10-5
10a
Train loss
(C) Normed output,
10 2
Standard init. Standard LR
10-5
10
Train loss
(D) Normed output. Standard init.
10-2
Decoupled BN LR.
5
10
Train loss
(E) Normed output. Rescaled init.
10-2
Standard BN LR
10.5
10
Train loss
(F)
Normed output Rescaled init.
10.2
Decoupled BN LR.
10-5
100 Epoch
200
100 Epoch
200
100 Epoch
200
0
100
200
Epoch
Figure 7: Proposed modifications decouple the training trajectory from 17 under fixed Tepoch We trained ResNet under various learning rate 7) (lines) and timescale Tepoch (columns) and plotted the training loss against epochs. We considered standard vs. normalized output layer (row B vs. row C, Sec. A.3.1) using 1/ for batchnorm parameters vs. a separate decoupled learning rate (row C vs. row D, Sec. A.3.3): using standard initialization scale vs. 17dependent initialization (row vs. row D. Sec. A.3.2; and lastly, combining all modifications (row F), which allows the trajectory to be independent of under a fixed timescale.
20

---

How to set AdamW's weight decay you scale model and dataset size
n= 10
== 10-5
n= 10
1
n= 10-2
10
Train Acc
Train loss
acc
Test loss
1.0
0.95
0.75
(A)
Standard output Standard init.
0.90
0.50
0.9
Standard BN LR
10
0.85
0.25
10"5
10
101
10-5
10
10
10
10
2
10
10-5
10
10a
À
À
À
À
1.0
0.95
0.75
(B)
Standard output
10
0.90
0.50
Standard init.
0.9
Standard BN LR
0.85
0.25
1.0
0.95
0.75
(C) Normed output.
10
0.90
0.50
Standard init.
0.9
Standard BN LR
10-5
0.85
0.25
1.0
0.95
0.75
(D) Normed output. Standard init. Decoupled BN LR.
0.90
0.50
0.9
10-5
0.85
0.25
1.0
0.95
0.75
(E) Normed output. Rescaled Init.
1
0.90
0.50
Standard BN LR
10-5
0.85
0.25
1.0
0.95
0.75
(F)
10
Normed output Rescaled init,
0.90
0.50
0.9
Decoupled BN LR.
10-5
0.85
0.25
10
100
102
10
10
10P
102
10
10
10a
102
10a
10
10a
102
10f
Tepoch
1/nAM
Tepoch
1/nAM
Tepoch
1/nAM
Tepoch
1/nAM
Figure 8: Under proposed modifications, the performance is only controlled by the timescale. We trained ResNet- 18 CIFAR-10 and plot the metries the end of the optimization under different initial learning rates (lines) against the timescale (x-axis). Similar to Fig. 7, from row B to F, we considered different levels of modification to the network. From row B to E, while the optimal Tepach (marked by red crosses) lies in close range across different 7, the exact performance still varies by 7J. Whereas in row F, after adopting all three modifications, the performance metrics become invariant to under the same Topoch-
21

---

How to set AdamW's weight decay you scale model and dataset size
n= 10-6
==
n= 10
1=
77= 10-2
n= 10
A= 10-2
== 10-1
À = =100
À = 10
(A)
Train loss
Standard output. Standard init.
1.5
Standard LN LR.
Standard structure.
0.5
100 Epoch
Epoch
Epoch
Epoch
Tepoch 0.2
Tepoch = 2
Tepoch 20
Tepach = 200
2.5
(B)
Train loss
Standard output,
Standard init.
1.5
Standard LN LR.
Standard structure.
0.5
2.5
(C)
Train loss
Normed output.
Rescaled init.
Decoupled LN LR. Standard structure.
0.5
2.5
(D) Normed output. Rescaled init. Decoupled LN LR. SI structure.
Train loss
1.5
0.5
100
200
100
200
100
200
100 Epoch
200
Epoch
Epoch
Epoch
Figure 9: Proposed scale-invariant structure decouples the training trajectory from under fixed Tepoch for VIT. We trained a ViT under various learning rates 7) (lines) and timescale Tepoch (columns) and plotted the training loss against epochs. We first considered standard model vs. model with the three modifications used in ResNet experiments (row vs. row C, Sec. 1,3.1, A.3.3, A.3.2), where the training loss traces still show discrepancy for different v). We then considered further modification (Sec. A.3.4, row D) to ensure invarinnne (SI), which successfully decoupled the loss trajectory from under the same Tepoch
22

---

How to set AdamW's weight decay you scale model and dataset size
10
10-5
10
7 10
1 10-2
10-1
Train Acc
Train loss
Test acc
Test loss
(A)
1.00
2.0
0.9
1.25
Standard output. Standard init.
0.75
1.5
1.00
Standard LN LR. Standard structure.
0.50
1.0
0.5
0.7
0.75
10-2
10a
10
10a
10-2
10a
10-2
10a
A
1.00
2.0
(B)
1.25
Standard output.
0.75
1.5
Standard init.
0.8
Standard LN LR.
0.50
1.0
Standard structure.
0.75
0.5
0.7
1.00
2.0
V
(C)
1.25
Normed output.
0.75
1.5
Rescaled init.
0.8
1.00
Decoupled LN LR. Standard structure.
1.0
0.5
0.7
0.75
1.00
2.0
(D)
1.25
Normed output. Rescaled init. Decoupled LN LR. SI structure.
0.75
1.5
1.00
1.0
0.5
0.7
10
2
10a
102
10a
10
2
10a
10?
10*
10
10"
102
10a
10
10"
102
10*
Tepoch
1/nAM
Tepoch
1/nAM
Tepach 1/nAM
Tepoch = 1/MAM
Figure 10: Under proposed modifications, VIT shows performance controlled only by the timescale and irrelevant to the learning rate. Similar to Fig. 9 but now plotted the final performance metrics against Tepoch (x-axis) under different (lines). In rows B, and C, the model is not fully scale-invariant and the initialization is not r7dependent, as such the performance varies between y7s under a fixed Tepach. In row D. when all modifications are incorporated, i.e. assumptions in Theorem. are satisfied, the performance becomes only dependent on Topoch-
23

---

How to set AdamW's weight decay you scale model and dataset size
End of optimization value
Trace throughout optimization
10a
À
y = 0.01/x
10
10-1
10
10
Averaged absolute value of weights
100
10
10-3
10
10
2
10
100
0
200
400
600
800
1000
Epoch
Weight magnitudes scale inversely with À AdamW optimization. We trained a ResNet-18 on CIFAR-10 using AdamW with cosine learning rate decay over 1,000 epochs. The weight decay parameter À varies from 10-3 to 1, with initial learning rate = 10-5/A. Left: Final average absolute weight values E[[W,]] in convolutional and linear layers versus X, showing approximate 1/A scaling (same slope as the reference dashed gray line). Right: Evolution of weight magnitudes throughout training for different À values, different runs start with the same magnitude and gradually converge (11A).
Figure 11:
24

---

How to set AdamW's weight decay you scale model and dataset size
Dataset size
N- N=0000000
N=320,000
N=640,000
N=1,280,000
Train Acc
Train loss
Test acc
Test loss
6
100
0=
4
(A) ResNet
0.5
10-4
0.2
2
1.0
10a
0.4
A
0.1 x 1x8,0,000 000
0.5
10-1
0.2
2
10-3
10-3
10-1
10-3
10-3
10-1
10
10
10-1
10-5
10
101
0
n
0
n
1.0
10
4 x 10
0.4
3 x 10
== 0.1
0.5
2 100
0.2
4
0.0
(B) VIT
1.0
x
10
0.4
0.1 x 128, 000 dataset
3 x 10
0.5
10
0.2
4
0.0
10-5
10-3
10-1
10-1
10-3
10-1
10-1 10-1
10-3 10-1
0
n
Figure 12: Under the suggested weight decay scaling, the optimal learning rate is stable across training length. Mirroring Fig. in the main text, trained the model for 100 epochs with different dataset sizes under fixed batch size. Using fixed weight decay (top in subfig. A, B), the optimal learning rate decreases with the dataset size. Under suggested weight decay scaling (bottom rows in subfig. A, B), where À Jataset the optimal learning rate becomes more stable across dataset sizes. Note that we select the values for À as 0.1 as they were close-to-optimg for the experiments in Fig. 1.
25

---

How to set AdamW's weight decay you scale model and dataset size
Dataset size
80,000
160,000
320,000
640,000
1,280,000
Train Acc (top-1)
Train loss
Test Acc (top-1)
Test loss
1.0
5.0
0.4
Tune
2.5
(A) ResNet
0.0
0.0
10
10
10
10a
10
10-1
10a
10-1
10-1
101
10a
10
10a
À
1.0
5.0
Tune
0.5
0.2
Tepoch
10
10
10
10a
10
10-1
10a
10a
10
10-1
10
10a
10
10
10a
10a
Initial Tepoch 1/nAM
Initial Tepoch 1/nAM
Initial Tepoch 1/nAM
Initial Tepoch
1/nAM
Train Acc (top-1)
Train loss
Test Acc (top-1)
Test loss
0.4
0.5
Tune
(B) VIT
10
10
10
10a
10
10
104
10a
10
10
10a
10a
10-
10
10a
10a
À
1.0
0.4
Tune Tepoch
0.2
0.0 4 10
0.0
10
10 10
102
10 4
10-2
10"
10a
10
10
102
Initial Tepoch 1/nAM
Initial Tepoch 1/nAM
Initial Tepoch 1/nAM
Initial Tepoch 1/nAM
Figure 13: Optimal and Tepoch V.S. dataset sizes under constant learning rate. Similar to Fig. but use a constant learning rate of 10-3 rather than cosine decay to 0.1 of the initial learning rate.
26

---

How to set AdamW's weight decay you scale model and dataset size
Dataset size
80,000
160,000
320,000
640,000
1,280,000
Train Acc (top-1)
Train loss
Test Acc (top-1)
Test loss
1.0
0.6
5.0
0.4
0.5
Tune À
2.5
0.2
(A) ResNet
0.0
0.0
10
10
10
10a
10-1
10
10a
10-1
10
10
10a
10-1
10-1
10a
A
À
A
À
1.0
0.6
5.0
0.4
Tune Tepoch
0.5
0.2
0.0
0.0
0.0
10
10
10
10
10
10
10
10
10
10
10
10
10
10
10
10
Initial Tepach 1/nAM
Initial Topoch 1/nAM
Initial Tepach 1/ΠAM
Initial Topach 1/nAM
Train Acc (top-1) 1)
Train loss
Test Acc (top-1) 1)
Test loss
1.0
7.5
0.4
0.5
Tune
0.2
(B) VIT
0.0
0.0
2.5
10
10
10
10a
10
10
10
10a
10
10
101
10a
10
10
10
10a
7.5
Tune
0.5
0.2
5.0
Tapoch
0.0
0.0
2.5
10
10 Initial Tepoch 1/nAM
10-2
10a
10
10 10"
10
10-4
10-2
10"
10²
10 10-1
10
10"
Initial Tepoch 1/nAM
Initial Tepoch 1/nAM
Initial Tepoch 1/nAM
Figure 14: Optimal À and Tepach v.S. dataset sizes under a cosine decay schedule to 0. Similar to Fig. but use a cosine learning rate decay schedule from 10-3 rather than to 10-4
S= 0.5
5=100
5=2.0
Train Acc (top-1)
Train loss
Test Acc (top-1)
Test loss
1.0
0.2
ViT ImageNet
Mbaa/s, à = 1.0
0.5
0.1
5
1.0
0.2
n= Abasals,
0.5
0.1
5
À ==
0.0
10-5 10-4
10-3
10-2
10-5
10-4
10-3
10-2
10-5
10-4
10-3
10-2
10-3
10-4
10-3
10-2
Moase
Aoose
neaas
Aoase
Figure 15: AdamW breaks the learning rate scaling μP on VIT. Similar to the setting Fig. 5, here we trained a VIT with different width factors the 320K subset of ImageNet 32x32 under the direct μP scaling (Eq. 13; top row) and our proposed scaling (Eq. 15; bottom row). The direct scaling breaks the transferability of optimal Mome due changing the timescale, whereas our scaling allows stable optimal Thase across model sizes.
27

---

How to set AdamW's weight decay you scale model and dataset size
s=0.5
s=1.0
$=2.0
5=4.0
Train Acc (top-1)
Train loss
Test Acc (top-1)
Test loss
1.0
0.5
10
0.3
4.0
n= 10 1/5.
A Apase
10-1
0.0
2.5
10
10
10
10-2
10-1
10
10
10
10
10
10-1
10
1.0
5.5
0.5
2
0.3
4.0
n= 10- 15,
À SAbase
0.0
2.5
10
10-1
10
-0-2
10-1
10
10
10-1
10"
10-2
10-1
10
Abase
Anase
Anase
Aase
Figure 16: Similar setting to Fig. 3, but use constant learning rate.
5=0.5
$=1.0
5=2.0
$=4.0
Train Acc (top-1)
Train loss
Test Acc (top-1)
Test loss
1.0
0.5
5.0
10c
0.5
10-1
0.4
3.5
n= 10 3/s,
À Abase
10 2
0.0
1911
2.0
10
10
100
10
10
100
10
10
100
10
10
10
1.0
0.5
10a
5.0
0.5
10 1
0.4
3.5
n. 7= 10- /s, À Sbbse
0.0
19/H
10-2
2.0
10
-2
10
100
10
10
100
10
10
100
10 -
10
10
Absse
Abase
Abase
Aase
Figure 17: Similar setting to Fig. 3, but use cosine decay to zero.
s= 0.5
s=1.0
$=2.0
$=4.0
Train Acc (top-1)
Train loss
Test Acc (top-1)
Test loss
0.75
0.4
4.5
(B) ImageNet
0.50
0.3
n) 1 Absse
À 1.0
0.25
0.2
2.5
0.75
0.4
4.5
n Absse
0.50
0.3
À
0.25
0.2
2.5
10-5
10
10 3
10-5
10
10
3
10-5
-5
10
10-5
-0
10 2
Masse
Aoss
Mhase
Abase
Figure 18: Similar setting to Fig. 5, but use constant learning rate.
28

---

How to set AdamW's weight decay you scale model and dataset size
$=0.5
$=1.0
5=2.0
$=4.0
Train Acc (top-1)
Train loss
Test Acc (top-1)
Test loss
0.75
0.4
4.5
(B) ImageNet
0.50
0.3
3.5
À 1.0
n Mbase/5
0.25
0.2
2.5
0.4
4.5
n Nbase
0.50
0.3
3.5
==
0.2
2.5
10
10
10-3
10-5
10
10-3
10-5
10
10-3
10-5
10
10
Moas
Obas
nhaas
Qbase
Figure 19; Similar setting to Fig. 5, but use cosine decay to zero.
29


Reviews:
[{"review_id": "5e5WVDGyme", "reviewer": "Reviewer_9CyY", "summary": "The authors study the scaling behavior of the optimal AdamW weight decay hyperparameter with respect to model and dataset sizes. They provide a theoretical insight by framing AdamW's learned weights as an exponential moving average (EMA) of recent updates, identifying the EMA timescale as the key underlying hyperparameter. Empirically, the authors demonstrate that the EMA timescale (in epochs) remains roughly constant across scales.\n\n## update after rebuttal\nI confirm that I have read the authors' response to my review and have revised my review accordingly, where appropriate.", "questions": "", "limitations": "", "rating": 5, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thanks for your extremely positive review!\n\nGreat catch with Chen et al. (Theorem B.6) and Liu et al. (Theorem A.1)! We have added a discussion of these papers to our working draft and adjusted our contributions section. In short, we aren't surprised that someone has used an EMA-like result/form as an intermediate result in other calculations.  Our contribution is in noticing that this simple connection between AdamW and EMA provides a powerful lens on **hyperparameter scaling** specifically for weight decay.\n\nWe have also added a discussion of Chen et al. (ICLR 2024) and Xie and Li (ICML 2024) to the Related Work.  As far as we can see, the \"implicit constraints of AdamW\" noticed in these works work, refers primarily to the relationship between *the scale of the weight* and *weight decay*, which we briefly discussed in Appendix C, Fig 11, but did not claim as a contribution.  To our understanding, this work did not study the transfer/scaling of weight decay across problem sizes, which is our primary contribution.  Though let us know if you're thinking of a different result here.\n\n> In theorem 1, is the \"scale-invariant network\" assumption too strong?\n\nWe aren't quite sure what you mean by \"too strong\" here.  Do you mean that we could prove the same result with a weaker assumption?  While this may be possible, we don't see how.  Intuitively, the need for the scale-invariant network assumption arises from the $1/\\lambda$ scaling of the weights in AdamW.  If the scale of the weight matters (i.e. the network is not scale invariant), then $\\lambda$ on its own affects the learning trajectory, in addition to the EMA timescale, and the ratio between the learning rate and the initialization scale $\\rho = \\eta / \\sigma$.\n\n> It would be great if the author could give a table to summarize the scaling \"rules\" of weight decay in terms of model size, data size, batch size, etc.\n\nWe have added this table to our working draft.\n\n> Have you explored extending these insights to other adaptive optimizers, such as AdaGrad or AdaFactor, to investigate if similar implicit bias or parameter norm constraints exist?\n\nNot as of yet, but we do expect the same considerations to transfer across all settings with constant magnitude updates and decoupled weight decay, including SignGD, Lion, Sophia, Muon, SOAP etc.  For instance, it was recently noted that Muon with decoupled weight decay has the same $1/\\lambda$ scaling of the max eigenvalue (see section Weight Decay from [1])\n\n[1] Why We Chose Muon: Our Chain of Thought? https://x.com/Kimi_Moonshot/status/1897929976948965870", "claims_and_evidence": "In theorem 1, is the \"scale-invariant network\" assumption too strong?", "methods_and_evaluation_criteria": "Most experiments are conducted on vision tasks, but they have multiply epochs, while for LLM pre-training, we only do single epoch given the large data size. \n\nIt would be great if the author could give a table to summarize the scaling \"rules\" of weight decay in terms of model size, data size, batch size, etc.", "theoretical_claims": "The authors' interpretation of AdamW's weight updates as an exponential moving average (EMA) provides valuable insights. However, similar EMA-based interpretations of adaptive optimizers and their implicit constraints have been recently explored in related literature. In particular, Chen et al. (ICLR 2024, Theorem B.6, \"Lion Secretly Solves Constrained Optimization: As Lyapunov Predicts\") and Liu et al. (NeurIPS 2024, Theorem A.1, \"Communication Efficient Distributed Training with Distributed Lion\") explicitly derived EMA forms for optimizer updates, highlighting their implicit constrained optimization behavior. Additionally, Xie and Li (ICML 2024, \"Implicit Bias of AdamW: ℓ∞-Norm Constrained Optimization\") similarly analyzed AdamW's implicit constraints.\n\nTo appropriately contextualize and acknowledge prior work, the authors should clearly cite and discuss these papers, clarifying how their EMA formulation aligns with or differs from these recent studies.", "experimental_designs_or_analyses": "This leads to a direct rule for scaling AdamW weight decay: the optimal weight decay decreases with increasing dataset size and increases with model size (under the muP learning rate scaling recommendations).", "relation_to_broader_scientific_literature": "EMA of update has been well studied, for example in Theorem B.10. of \"Lion Secretly Solves Constrained Optimization: As Lyapunov Predicts\".", "essential_references_not_discussed": "The primary contribution is identifying and characterizing the scaling rule for weight decay. While 'Implicit Bias of AdamW: ℓ∞ Norm Constrained Optimization' illustrates the implicit bias introduced by weight decay, earlier work such as 'Lion Secretly Solves Constrained Optimization: As Lyapunov Predicts' establishes the connection between the norm of trained neural network parameters and the weight decay coefficient λ, demonstrating this relationship for both Lion and AdamW.", "other_strengths_and_weaknesses": "Novelty and Impact:\n\nThe topic is relevant and timely, addressing an important practical challenge in the era of large-scale neural networks. The paper contributes substantially to our understanding of hyperparameter transferability across model and data scales.", "other_comments_or_suggestions": "EMA of update can be seem as kind of nice perspective to see weight decay. Have the authors use continuous time", "questions_for_authors": "Have you explored extending these insights to other adaptive optimizers, such as AdaGrad or AdaFactor, to investigate if similar implicit bias or parameter norm constraints exist?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "OJQShSnXrN", "reviewer": "Reviewer_YYQ3", "summary": "This paper studies the AdamW optimizer. the authors provide empirical studies on the hyperparameters of the AdamW under different settings. Specifically, the paper first reformulate the AdamW itself is an EMA. then provide experiments on resnet, vit and LLM training, showing some empirical weight decay parameters tuning rules when increasing the dataset and model size. Finally, the authors discussed the relationship between weight decay and \\muP learning.", "questions": "", "limitations": "", "rating": 2, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thanks for your careful review.\n\n# Value of the EMA perspective\n\nWe fully agree that the reformulation of AdamW to EMA is almost trivial.\nOur key insight was noticing that this almost trivial connection provides novel, powerful insights into hyperparameter transfer for weight decay, that we then validated extensively. \nIt is precisely the simplicity of our approach that makes it easy for people actually training neural networks to understand and apply.\nIndeed, we know that at least one \"unicorn\" is using training recipes inspired by the work here for training large scale foundation models.\n\nFinally, we believe our work is **the very first paper** studying the problem of transferring weight decay parameters across problem sizes.\n\n# Experiment design\n\nAs you noted, we took the rules for hyperparameter transfer suggested by the EMA view, and validated them through extensive experiments on the three major classes of machine learning models (ResNets, LLMs and VITs).\n\n# Theoretical understanding\n\nWhile our results are supported by Theorem 1 (proof in Appendix A1), we do agree that our results are mainly around validating the simple but powerful viewpoint on hyperparameter transfer provided by the connection between AdamW and EMA.", "claims_and_evidence": "yes", "methods_and_evaluation_criteria": "yes", "theoretical_claims": "yes", "experimental_designs_or_analyses": "yes", "relation_to_broader_scientific_literature": "None", "essential_references_not_discussed": "None", "other_strengths_and_weaknesses": "Strength\n- the authors conducted extensive numerical experiments to support their claims.\n- the reported results could be helpful in ML applications.\n\n\nWeaknesses:\n- the paper is mainly empirical results based. though the authors provided experiments on resnet, LLM, vit, it could be difficult to generalize to all other machine learning models.\n- the reformulation of AdamW to EMA looks trivial and incremental to me. the author didn't claim what are the benefits by this reformulation.\n- though there are extensive experimental evidence, the theory part of this paper is weak. the authors didn't provide solid theoretical understanding to the phenomenon they observed.", "other_comments_or_suggestions": "none", "questions_for_authors": "none", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "IBzrG6UGmt", "reviewer": "Reviewer_xMai", "summary": "This paper proposes a simple framework for understanding weight decay in AdamW through its similarity to an exponential-moving-average of weight updates. The hypothesize that the timescale of the EMA (analogous to the half life) is a useful quantifier of the training dynamics. In particular, the authors show that the timescale should be preserved when scaling either the model width or dataset size. This is empirically validated by experiments across several (smallish) deep learning tasks. This offers useful guidance on how to scale the weight decay hyperparameter across task sizes, which is of significant value to the community.\n\n**Post-rebuttal update:** The authors have promised changes that sufficiently address all my concerns. I have raised my score from 3 (weak accept) to 4 (accept).", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thanks for your extraordinarily extensive and thoughtful review! We agree with all your points and have extensively updated our working manuscript in response.\n\n## Claims And Evidence\n\n**Approximate EMA:** Fixed.\n\nWe do NOT **assume that the initial weights are zero in Eq 9**!  We assume that the initializations for the two networks use the same random noise, but different scalings (Eq 8).\n\nWe are assuming **learning rate rate does not change as we train on more data** (see 375r-379r, which mirrors a point from Kosson), and we assume **muP scaling for width**.\n\n We have (re) emphasized all these points in the manuscript.\n\n### Q1 Prior work\n\nThanks for highlighting **Kosson**: This is a great paper that should have had a bigger impact.\nWe cite Kosson in the related work (line 424l), stating that the key difference with our work is that Kosson does not discuss hyperparameter transfer.\nNonetheless, we agree that there are more connections and have extended the discussion of Kosson in the Related Work, along with writing an \"extended prior work\" section that fully elaborates on this connection.\nLet us know if you'd like specific wording as we frame this connection!\n\n**Wortsman** studies the relationship between the final performance and learning rate.\nIn their Fig 6, they find that performance is insensitive to learning rate, $\\eta$, when fixing $\\eta \\lambda$, while it is more sensitive if $\\lambda$ is fixed. \nWe see this as directly connecting to the **Kosson** result, that $\\eta$ matters early on and $\\eta \\lambda$ matters later on.\nThe connection to our contributions on EMA and hyperparameter transfer are *definitely there*, but are far more indirect (though let us know if you're actually thinking of a different result from Wortsman).\n\n### Q2\nWe believe the result is **more general than AdamW**, applying to any setting where the updates have roughly constant magnitude as the scale of the underlying parameters varies, eg Lion Sophia Muon etc but not SGD.\n\n### Q3: EMA timescale constant\n\nThe EMA viewpoint suggests a reason why the time horizon should be constant. We have edited the manuscript to include this fuller discussion.\n\nAn EMA can very approximately be understood as averaging over the last $\\tau_{iter}$ samples (Eq 7), as it downweights considerably any datapoints more than $\\tau_{iter}$ steps ago.\n\nWhat does this mean for training neural networks?\n\nWe assume that each datapoint provides useful information, so you can't drop or considerably downweight datapoints without harming performance. As discussed above, the EMA downweights considerably datapoints that it saw more than $\\tau_{iter}$ iterations ago. Thus, to avoid considerably downweighting datapoints, $\\tau_{iter}$ should be larger than the number of iterations required for all the datapoints (i.e. one epoch).\n\nThe above reasoning suggests that you can set $\\tau_{iter} = \\infty$, or $\\lambda = 0$, as that would would average over every update. But it would also reduce AdamW to Adam.\n\nOf course, we know that AdamW (i.e non-zero $\\lambda$) works better.\nTo resolve this conundrum, we conjecture that you don't want to average over all updates. We believe that updates from very early in training are detrimental, as they are based on early settings of the weights. To forget these initial updates we can use a $\\tau_{iter}$ which is smaller than the total number of iterations in the training run.\n\nThese two observations give a \"natural range\" for the optimal $\\tau_{iter}$: somewhere between the number of iterations for one epoch and the total number of iterations.\nThis natural range is simpler if you measure the timescale in terms of epochs: it's somewhere between $1$ and the total number of epochs.\nWe find experimentally that the optimal timescale does lie in this range (Fig 1,2).\nAs this natural range is fixed wrt model and dataset size, we conjectured that the optimal EMA timescale was fixed.\n\n\n### Q4\n\nWe transfer the sequence of timescales through training.\n\nWe believe it's optimal to have a scheduled timescale across training: scheduling makes it easier to forget detrimental initial updates while doing long-timescale averaging towards the end of training.\n\nAll existing \"decoupled weight decay\" implementations (e.g Composer from MosaicML) change the timescale by changing the learning rate. While you could also use $\\lambda$, we leave that for future work.\n\n### Q5\nWe agree that you could **modify timescale using either $\\eta$ or $\\lambda$**. Our motivation for modifying $\\lambda$ is explained in the first paragraph of the Related work (line 375r-379r).\n\nRelated to this, we ran an extra [experiment](https://anonymous.4open.science/r/3MDG98-66F3/response_pdf.pdf) showing that if you fix $\\lambda$, then the optimal $\\eta$ changes as dataset size increases.  In contrast, if you use our scaling for $\\lambda$, then the optimal $\\eta$ is fixed.  We will run a complete set of experiments on this for the camera-ready.", "claims_and_evidence": "In general the claims made in the paper are decently supported by evidence. Below are some (minor) comments about the exact phrasing.\n\n**The weights generated by AdamW can be written as an EMA of recent weight updates:** Typically EMA is viewed as a smoothing operation (e.g. in the wikipedia reference provided in the paper). In this case the input sequence typically doesn't depend on the EMA output. This is not the case in AdamW, later updates depend on the previous ones. With a typical EMA in signal processing, you would expect different values of the time constant to vary the amount of smoothing. In AdamW it can give you an entirely different input sequence as well (fundamentally different, more like a dynamical system than EMA). To derive Equation 9, there is also the key assumption of the initial weights being zero (almost never the case in deep learning). Finally the learning rate changing over time would skew the exact EMA form. Overall I feel it would be more fair to say that it can be **approximated as an EMA** where these core assumptions or simplifications are prominently stated.\n\n**The optimal weight decay scales with 1/dataset size:** I think this is conditioned on the learning rate being constant. There are some other works that claim the learning rate should be lower for longer runs (see Schaipp 2501.18965 for a recent example). If that modification is made I am not sure the optimal weight decay is lower. The paper focuses on showing the EMA time coefficient for longer runs is lower, but there are two avenues for lowering this (via either $\\eta$ or $\\lambda$).\n\n**The optimal weight decay increases with model width:** Similar to the previous point, I believe this assumes the learning rate is scaled via muP. If the learning rate is kept constant (no muP) the optimal weight decay may stay the same.", "methods_and_evaluation_criteria": "The proposed framework of viewing weight decay as an approximates EMA makes sense and the scaling rules are appropriate and useful for deep learning.", "theoretical_claims": "I read through them but did not verify the math step by step.", "experimental_designs_or_analyses": "I reviewed the experiments and they seem reasonable. It would be nice to see experiments on larger scale networks or datasets, but I understand this may not be feasible.", "relation_to_broader_scientific_literature": "The paper makes a decent attempt at discussing related concepts in literature. However there are several places where I feel the connections are not sufficiently detailed.\n\n**Using decoupled weight decay in muP**: This paper proposes scaling the weight decay inversely with the learning rate in PyTorch AdamW. This is equivalent to keeping the weight decay constant in the original “fully-decoupled” AdamW. Other works found this empirically before this paper (e.g. Wortsman 2309.14322). I feel this should be made more explicit in the paper, especially since one of the highlighted claimed contributions is showing “When using AdamW with fixed weight decay, μP learning rate scaling breaks down, but proper weight decay scaling restores its effectiveness”. I believe this exact conclusion could be drawn from prior work as well.\n\n**Effective learning rates:** Viewing AdamW as an update-EMA is **very closely related** to using “relative” (or “rotational”) effective learning rates (Wan 2006.08419, Kosson 2305.17212). The paper briefly mentions these works but does not explain the relation (which is crucial for understanding how this work relates to a large body of prior work on weight decay). The relative learning rate specifies how large the relative change in the weights is on average. Intuitively it shouldn't be too hard to see why this is conceptually related to an EMA, decaying the prior weights at each timestep by e.g. 10% may result in something like a 10% relative change in the weights on average. This can also be formally shown. Given some conditions on the input sequence which are also needed for most of the EMA interpretations to make sense, there is a direct mapping from the EMA decay coefficient to the relative learning rate in the steady-state (see brief sketch below). The other way also works but is less specific. EMA is one way of achieving a given relative learning rate, but there are others e.g. for scale-invariant systems you can exponentially increase your learning rate schedule instead of doing EMA (Li 1910.07454) or keep the weight norm constant and modify the update size accordingly (Kosson 2305.17212).\n\nEMA to Relative Updates: Consider an EMA given by $a_{t+1} = (1-\\gamma)a_t + \\gamma b_t$. Let's assume that the $b_t$ values are independently random and identically distributed over time with mean zero, and focus on the relative updates in the steady state (large $t$). The relative update size for this scalar case is best defined as $r = \\sqrt{\\mathbb{E}[(a_{t+1}-a_{t})^2] / \\mathbb{E}[a_t^2]}$. To compute $r$ we need to compute $\\mathbb{E}[(a_{t+1}-a_t)^2] = \\gamma^2 (\\mathbb{E}[b_t^2] + \\mathbb{E}[a_t^2])$, using independence and mean zero. We can relate $\\mathbb{E}[a_t^2]$ to $\\mathbb{E}[b_t^2]$ by expanding $a_t$ in terms of $b_t$, squaring, removing the cross terms that are zero due to independence, and finally approximating the sum as infinite (large $t$). This gives $\\mathbb{E}[a_t^2] = \\frac{\\gamma^2}{\\gamma^2 - 2\\gamma} \\mathbb{E}[b_t^2]$, which results in $r=\\sqrt{2\\gamma}$. Plugging this in **gives an exact match between the EMA timescale proposed in this work and the prior work on effective learning rates in AdamW** (Kosson 2305.17212). Overall I feel the EMA timescale, EMA decay rate, and the relative learning rate are essentially just slightly different characterizations of the same underlying phenomenon. The time coefficient may be a useful quantity as the authors argue, but it feels more like a slight variation of existing works rather than a fundamentally new perspective on weight decay. Like the time coefficient, the effective learning rate also has its strengths. For example it is better defined when the learning rate is changing, it can be measured directly, it allows you to obtain the same learning dynamics without weight decay and / or eliminate the arbitrariness in the weight norms, and is useful for transferring hyperparameters from one optimizer to another.\n\n**Similar or identical conclusions to prior work:** Some of the phenomena discussed when explaining how \"EMA view of AdamW provides an explanation for several striking observations made in the literature\" have already been explained by prior work, for example Kosson 2305.17212. These include identifying two separate processes for the learning rate and weight decay (lines 387L to 398L) and the scheduling effects where $\\eta$ matters at the beginning but $\\eta\\lambda$ later in training (line 402L). Note that as discussed by Kosson 2305.17212, changing the global learning rate while keeping the weight decay fixed also affects the learning dynamics of gains and biases (i.e. 3 effects not just the 2 mentioned), something that matters in practice as observed in Appendix A but is not discussed around lines 376R and 393L.", "essential_references_not_discussed": "Yes, I feel that key components of this work are closely related or overlap with prior literature on effective learning rates (see above), but this is not sufficiently discussed.", "other_strengths_and_weaknesses": "The main strength of the paper is providing a simple view on weight decay and practical guidance on how to set it backed by sufficient although small scale experiments. The main issues are undisclosed relations to prior work and sometimes a lack of clarity or specificity. I think these can be easily addressed which I have already reflected in my rating.\n\nThe discussion in the paper feels somewhat limited. In particular there are important questions that are not brought up or addressed:\n- The paper conjectures that the EMA time horizon should be kept constant but offers little to no explanation for why this should be the case.\n- With a learning rate schedule the time coefficient changes throughout training. Is this important?", "other_comments_or_suggestions": "I think characterizing (Blake 2024) and (Bjorck 2024) as \"follow-up works that follow up on the concepts presented in this paper and cite us\" is a bit of a stretch (they seem more like concurrent works).\n\nYou use \"swiped\" in several places where I think you might mean \"swept\".\n\nThe paper could use a thorough read through for spelling mistakes.", "questions_for_authors": "Q1: Do you agree with my characterization of the relationship between the proposed EMA view and relative learning rates or do you know of a better way to explain this relation?\n\nQ2: Which parts of this paper are specific to AdamW and would not apply to SGD with momentum, Lion or other optimizers?\n\nQ3: How would you justify the time coefficient remaining constant during scaling?\n\nQ4: Does the fact that the time coefficient varies throughout training matter and which value should be transferred?\n\nQ5: Is a lower weight decay strictly better for longer runs / larger datasets or is decreasing the learning rate as good?\n\nI would consider raising my score if the authors can convincingly answer these questions and address the other concerns I described before about clarity and relation to prior work.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: The paper notes a simple connection between EMA and weight decay to motivate a new hyperparameter transfer scheme. It performs thorough experiments to validate the approach. Given the empirical importance of weight decay, the paper's questions, investigation, and results can significantly aid practitioners. Two positive reviewers justify their scores thoroughly while one negative reviewer argues the paper doesn't have enough theoretical contribution, but the paper's main contribution is empirical. The paper can do a better job of contextualizing itself with prior work that implies related statements (e.g., Wortsman et al.), which the authors have promised to do.
Decision: Accept (poster)