Thanks for addressing the comments!

Our overall philosophy has been to adopt the maximal update parameterization, where weights are initialized in order to ensure invariance to changes in model width.

I understand this part, I am suggesting having a weight-decay dependent scaling to the initialization together with the 1 / sqrt{fan_in} MuP scaling, which I believe would not break MuP. But yes, this would beyond the scope of the current submission.

Actually, in very recent work, we have found that while the LR does not scale perfectly linearly with batch size when using D2Z, the weight decay does do so (and, moreover, leads to lower loss!).

Interesting finding!

I have increased my score from 6 to 8. I recommend acceptance for this paper. Please improve the writing and presentation of Sec 3 accordingly, that will hugely improve the readability and easiness of understanding for the proposed idea.

However, I have to admit that I only know about academia-size deep learning optimization problems but I have zero expertise and experience in LLM scaling law-related topics, therefore whether or not the paper's experiments are "small scale" goes beyond my scope of knowledge, I am open to other reviewer's comments.

Thank you so much for your detailed and valuable review, and for generously noting the paper’s "convincing" experiments and "very solid" conclusion.

Regarding "when the model is trained much longer than what scaling laws recommend ... D2Z scheduling gives optimal performance"

Just to clarify: in all cases where we train exactly to the compute-optimal frontier (as recommended by scaling laws), D2Z does improve over 10x (e.g., Figure 11), and yet you are correct that when training beyond this level, the gains increase further.

Regarding "there is room for improvement in the writing: Section 3.1 and 3.2"

This is a very fair point, and we appreciate your help here. We will definitely improve these sections for the camera paper. We will explicitly develop the idea that decaying the LR to zero allows for more updates to be averaged, reducing the gradient variance; and we will note that variance reduction grows in importance in the later stages of optimization, particularly for higher-TPP training. We make points along these lines in Section 5.1, but adding discussion here will make the methods section much clearer. We will also motivate the differences between cosine D2Z and linear D2Z from a loss surface perspective, as suggested by Reviewer sx5Y, and then subsequently show the connection to the EMA coefficients; this should also improve the presentation here.

Regarding "the cyclic learning rate scheduling ... is used in experiments as a baseline but not explained or discussed"

Yes, we only briefly mention in Section 4.2 that "Appendix Figure 27 provides full LR curves and dual coefficients for all models in Figure 7", and this includes the Cyclic LR schedule (third row of Figure 27). We will add a separate figure that only shows the LR schedules (including Constant, Linear, Cosine, WSD, and Cyclic), and refer to this figure early in the main body of the paper, rather than only showing these schedules on the final page of the appendix. Good point.

Regarding "it would be nice if more details about the computational resources, model configuration ... can be revealed."

True, right now we only say, in Section 5.1, that "Appendix A has full experimental details." But we will revise this to note that "Appendix A has further experimental details, including details on the models (Appendix Table 2 providing the model width, number of layers, number of attention heads/head dimensionality, batch size, etc.), as well as details on the computational resources used during training." We will also expand Table 2 to explicitly provide the number of heads in the multi-head attention ($=d_{model}/d_{head}$). Please let us know if there's anything else that would be useful to include here.

Regarding "it would be nice if the experiment code can be released"

There is precedent in our organization for releasing code to accompany published papers. In particular, we should be able to share the code for calculating and visualizing the EMA coefficients for arbitrary LR schedules, and for reproducing our results with the NanoGPT codebase.

Regarding "modifying the weight initialization scale ... by 1 / lambda, such that the bias could potentially be reduced at larger lambda"

This is an interesting suggestion. Our overall philosophy has been to adopt the maximal update parameterization, where weights are initialized in order to ensure invariance to changes in model width. From this perspective, the ability of weight decay to enable variance reduction later in training, without impacting bias reduction, can perhaps be viewed as a feature rather than a bug – it gives us a tool to specifically impact the later stages of optimization. But this idea definitely merits further study, and the link to Wang & Aitchison’s Eq. 21 is worth noting in the paper, thanks!

Regarding "[why] weight decay is beneficial only used together with certain learning rate scheduling"

Thank you for the pointer to the prior work in Loshchilov & Hutter, 2017. As discussed in Section 5.1, we believe that without LR decay, a constant LR is averaging over too few updates later in training. Raising weight decay only further shortens the timescale over updates, so is not beneficial (although it might have been beneficial if much smaller constant LRs had been used). Meanwhile, with D2Z, we can use weight decay to fine-tune the average timescale during the decay period. This is more effective than adjusting the peak LR itself, as raising it too high causes instabilities, and too low affects movement from initial conditions. 10x decay (at an optimal LR) only mildly benefits from weight decay because it is somewhere in between these extremes. We would like to pursue further experiments along these lines in follow-up work, so any thoughts that you have on this topic would be very welcome! We will also revise this part of 5.1 to be more clear, and provide a forward-pointer to this discussion under Finding 4.

Thank you very much for your thoughtful and constructive feedback, and for noting the extensive and effective experiments. We are also grateful for all the reviewers scoring the paper highly in terms of presentation quality.

Regarding "a more rigorous mathematical framework that connects the empirical results to theoretical principles is needed"

We take your point regarding the lack of analytical results; there is always a tradeoff between the ability to derive formal theorems and the generality of the theory (i.e., requiring fewer assumptions). What attracted us to the EMA perspective is that it applies whenever AdamW is used as the optimizer. We view our work as a step forward in terms of the mathematics, as Wang & Aitchison’s prior work did not properly account for a dynamic LR schedule. Moreover, Equation (1), defined for SGD with convex loss as you note, is still useful as intuition about training having a phase where we move away from initial conditions, followed by a phase where we must primarily deal with gradient noise. Here, another contribution of our work is to note the EMA perspective does also recognize different phases. Indeed, we can quantify the dependence on the initial weights (bias), and we can show analytically how this is reduced via a larger LR, but not reduced via a larger weight decay. This is the point noted by reviewer f1Pq: "I really like the author's argument on why, given a fixed value of eta * lambda, modifying eta does not have the same effect as modifying lambda, since having a lambda too large would affect the bias reduction." This contribution reconciles the results in Figure 6, which otherwise are inconsistent with the original interpretation in Wang & Aitchison. Moreover, as we will note in the camera version of the paper, this contribution also suggests that for training to a higher tokens-per-parameter (training with a greater proportion of gradient noise), it may be more effective to reduce weight decay than to reduce the LR. This is a timely suggestion given that papers continue to appear that observe and recommend purely a decrease in the max LR for high-TPP training (e.g., Bjorck et al., 2024), which we view as counterproductive.

That being said, we agree with Reviewer f1Pq’s point that the presentation in 3.1 and 3.2 can be improved, and unified around the concept of variance reduction. We will revise these sections to explicitly describe why D2Z might improve over 10x decay from a variance reduction perspective (this is mainly discussed now in Section 5.1). We also believe motivating the differences between cosine D2Z and linear D2Z from a loss surface perspective, as suggested by Reviewer sx5Y, and then subsequently showing the connection to the EMA coefficients, will greatly improve readability here as well. Moreover, we will make clear where the theory is still lacking, and we will note our hope that, going forward, further theoretical principles can be developed to further explain our extensive findings. Thanks for your suggestions here.

Regarding "more comprehensive information on scaling law experiments"

This is a great point. First of all, as we noted in response to Reviewer sx5Y, we will revise the main paper to discuss these experiments (and our sweeping of max LR for 1.7B-parameter models), rather than only discussing these results in the appendices. In terms of further experimental information:

• While we provided information on model architecture for these experiments in Table 6, we neglected to include the batch size. We will do so for the camera-ready paper.
• Moreover, we did not include a table along the lines of Table 3, documenting the specific LR schedule (warmup, cooldown, and total tokens). We will do so for the camera-ready paper.
• We will also provide further details on our dataset mix, and describe in further detail the proprietary Arabic dataset that was partly used in the training of these models. For our submission, we suppressed a citation to this data in order to maintain anonymity.
• We will provide further details on how we obtained our power law fit. Essentially, we transformed our Loss-to-FLOPs equation $y=cx^m$ to a logarithmic form, $log(y)=log(c) + m log(x)$, and fit the slope and intercept parameters of this line using a least squares regression.
• Importantly, we will also provide the fitted power law exponents (-0.0608 for D2Z, -0.0593 for 10x decay), which strongly indicate that the gains from D2Z will continue to improve as we scale up models.
• Finally, we will add a subplot to Figure 24 to explicitly show the relative loss improvement of D2Z over 10x decay (from 0.51% at 256M, to 0.6% at 590/1.3B, to 1.21% at 2.3B).

If there are other specific details that you would like to see in this section, please let us know.

Based on the current response, I will maintain my rating of 5. From my perspective, Linear D2Z functions more as an engineering heuristic than a comprehensive algorithm. The authors mentioned in their response that they would provide additional theoretical analyses to establish its foundations, and I was waiting these in the revised manuscript. However, until now I have not observed any modifications addressing this aspect. Therefore, I will keep my rating unchanged.

Thank you for your very helpful review, and your kind words regarding the importance of the problem, the well-written and compelling story, and the convincing experiments on the SlimPajama dataset. We will address your main concerns by doing a better job of describing the larger-scale experiments and additional datasets in the main body of the paper (rather than only in the appendices). We will also provide additional results at the 1.7B scale.

Regarding "A base lr sweep … for the 1.7B model would benefit the paper greatly"

This is an important point. In fact, this sweep is shown in the appendix in Figure 11, right-hand-side. The figure shows that even when each LR schedule uses its optimal peak LR, the relative loss improvement of D2Z over 10x decay actually increases at 1.7B scale (from 0.8% at 617M to 1.2% at 1.7B). While the main paper focused on 617M models (where we could afford to perform more ablations), we unfortunately neglected to mention sweeping the base LR for some of our largest/most-expensive models (training 12 different 1.7B LLMs for this figure, each on 34.3B tokens!).

Regarding "the base learning rates for the small-scale models used in this paper are very high"

Note that MUP (Yang et al., 2022) is a technique that allows use of the same base LR (before adjustment) at different model scales. For all figures that sweep LR for MUP models, we provide this base LR, $\tilde{\eta}$, on the x-axis. The actual LR, $\eta$, is scaled down by a factor of 3 for 111M models ($d_{model} = 768$), and by 8 for the 617M and 1.7B models (both with $d_{model}=2048$), as prescribed by MUP, given our proxy-tuned model was 3x and 8x narrower than these models, respectively ($d_{model} = 256$). So, in Figure 11, the actual LRs are 8x lower for the bigger models. The MUP scaling of the LR is noted in the second paragraph of Section 2.2, and details for our proxy model are noted in Appendix A, but we will revise the main paper to note that D2Z continues to dominate 10x decay even as the actual adjusted LR, $\eta$, decreases substantially.

Regarding "this paper needs to show that the benefits of d2z do not completely vanish at large scales"

Continuing from the previous point, we experimented with a variety of other model widths, up to $d_{model} = 2560$ with a 2.7B-parameter model, in Appendix B.10. Here, the actual adjusted LR is an order of magnitude smaller than the MUP base LR, but the benefit of D2Z over 10x again grows as model scale increases (from 0.51% at 256M, to 0.6% at 590/1.3B, to 1.21% at 2.3B). We will revise to refer to these larger-scale results in the main paper. Currently we only provide a brief pointer to "results with different … model scales (Section B.3, B.10)" in the second paragraph of Section 5.3. Thanks a lot for raising this!

To further address your concerns, note we have also performed new experiments comparing D2Z and 10x decay by training 1.7B models to 80 tokens-per-parameter (i.e., 137B tokens, 4x beyond training-compute optimality). At the proxy-tuned base LR, we find D2Z achieves a quite significant 3.0% lower relative loss than 10x decay. We will add these results to Figure 4 and provide further experimental details in the paper.

Note that these 1.7B models trained to 80 TPP are not necessarily "toy" models by today’s standards; for example, we over-train models of this scale for use as proposal models in speculative decoding (to speed up inference). Moreover, smaller models like Gemma-2B and Phi-1.5 are used in many production settings. So, finding the right recipes for over-training models at this scale is an important research direction and a valuable supplementary contribution of this paper.

Regarding "the impact of the dataset is another thing that is not discussed"

We should have mentioned the conclusions of these tests in the main paper, as currently we only point the reader to "results with different … datasets (Section B.9, B.10)" in Section 5.3. We will revise the main paper to note that the benefits of D2Z over 10x do continue across different corpora and across different languages. Results in the main paper use the SlimPajama dataset, while results with NanoGPT use Open Web Text. Our scaling law experiments use a combination of the Pile and a proprietary Arabic dataset. In all cases, D2Z is superior to 10x decay under an optimal LR. That being said, we will also note in our Limitations section that our results may not hold across all datasets, and understanding the interactions between dataset quality and the optimal decay schedule is an avenue of future research. Again, good point.

3.2: Conceptual model: bias and variance in LLM training (continued)

Higher decay is preferable to simply reducing the peak LR because we also need to reduce bias, i.e., minimize the contribution of the initial random weights. In the EMA perspective, the contribution of the initial weights to parameters at step $t$ is $c_{t,1} = \prod_{j=2}^{t} (1 - \alpha_j)$. For a constant LR, $c_{t,1} = (1-\alpha)^{t-1}$. For a decaying LR, and where $\alpha_j \ll 1$, coefficient $c_{t,1}$ can be approximated as $c_{t,1} = (1-\bar{\alpha})^{t-1}$, where $\bar{\alpha}$ is the average $\alpha_j$ over the schedule (see Appendix A for a derivation). Exactly as in Eq. 1, bias therefore decreases exponentially in the absolute number of steps, $t$, with a rate of decrease that depends on the LR. Crucially, this means that as we train for more total steps (i.e., a higher tokens-per-parameter; TPP), there is a decrease in the fraction of steps required for $c_{t,1}$ to become negligible. At higher TPP, bias reduction becomes relatively less important than variance reduction.

3.3: Experimental Hypotheses

Hypothesis 1: As TPP increases, the relative benefit of D2Z over 10$\times$ decay will increase.

This hypothesis follows from the premise that gradient variance plays a larger role at higher TPP, and greater LR decay (as in D2Z) allows for more updates to be averaged, and thus greater variance reduction. A related hypothesis is that if we increase the batch size at each step, gradient variance will decrease, and so the benefits of D2Z over 10$\times$ decay should diminish. Note our conceptual framework does not say precisely at which TPP or batch size D2Z will first prevail; here we will rely on our empirical findings (Section 4) to fill in the theoretical gap.

Hypothesis 2: As TPP increases, the optimal peak LR will decrease, for all LR schedules.

Tuning the peak LR is about trading off movement from initial conditions (requiring a high LR) and mitigating variance (requiring a low LR). As TPP increases, and bias reduction plays a smaller role, we should observe the optimal peak LR to decrease. We hypothesize the decrease will be greater with a Constant schedule, as Constant does not use decay to balance the conflicting demands of bias and variance. Moreover, the optimal peak LR for other continuous LR schedules, like WSD and Cyclic, should also decrease with longer training durations. In this way, such "schedule-free" approaches are not truly schedule-free. This dependence is obvious when plotting the EMA coefficients for these schedules, as in appendix Figure 27.

This hypothesis also implies that when comparing LR schedules, the maximum LR could be a confounder. For this reason, in our experiments, we compare schedules when each is tuned to their optimal maximum LR.

Hypothesis 3: Linear D2Z will improve over Cosine D2Z.

While LR decay allows averaging over more weight updates, if the LR decays too quickly, the final weight updates may not contribute to the EMA. From a loss surface perspective, as the LR approaches zero, we take vanishingly small steps. Since Cosine reaches smaller step sizes faster than Linear (left side, Figure 2, Section 1), it will make less progress toward the optimum loss. From the EMA perspective, this is equivalent to Cosine having smaller $c_{t,i}$ coefficients as $i$ approaches $t$ (Figure 2, right). Note this problem is unique to Cosine D2Z and will not affect, e.g., Cosine 10$\times$ decay.

Hypothesis 4: A high LR (not weight decay), is needed to reduce bias and achieve optimal loss.

Note that weight updates have a coefficient of $\frac{1}{\lambda}$ in Eq. 4. So, while $\eta_t$ and $\lambda$ contribute equally to $\alpha_j$, increasing $\lambda$ to reduce bias is counterproductive as weight updates will be scaled down proportionally, reducing movement from initial conditions. However, if LR is in a high enough range, both LR and WD should equally affect variance reduction.

Since weight decay does not impact bias reduction, then at very high TPP, it should be more effective to reduce variance by lowering WD than by reducing the LR. This stands in contrast to very recent work that recommends purely decreases in LR for high-TPP training (Bjorck et al., 2024).