Power Scheduler: A Batch Size and Token Number Agnostic Learning Rate Scheduler

The major limitation of the paper is only 12M 36M 121M models are trained which significantly weakens the liability of results. For example, the ones in [1] involves multiple scales that are larger than ~100M models.

The paper could benefit from a more in-depth discussion on the computational costs of this search process, particularly for very large models where even proxy model training can be expensive. Offering practical guidelines for efficiently setting or approximating these hyperparameters without extensive searches would improve the usability of the method. In particular roughly how many FLOPs required for hyper-parameter sweeps, especially for the kind of MoE architecture employed in the paper.

The writing is generally good, but the large table of hyper-parameters disrupts the flow. Tab 3 and 4 are better to be in Appendix.

While the experiments demonstrate strong performance on language modeling some downstream tasks, including more challenging or varied domains, such as reasoning-based tasks or multi-modal data, would better validate the generality of the Power Scheduler. This can be true when we use post-training data in the decay phase, as did in MiniCPM.

Some related works are missing: Scaling Optimal LR Across Token Horizons. This paper show that learning rate should be tuned down if we increase token size.

[1] Small-scale proxies for large-scale Transformer training instabilities

1. Limited Transferability Demonstrated: The paper's primary contribution relies on identifying specific coefficients for a power-law relationship among learning rate, batch size, and number of training tokens. However, the transferability of these coefficients to different models and datasets is not thoroughly explored. Except for the Mixture-of-Experts models based on the same language model used to uncover the power law, there is a lack of experimentation with other architectures or datasets. This gap leaves uncertainty about whether the proposed Power scheduler can be effectively applied to a broader range of models.

2. Insufficient Details on Hyperparameter Optimization: The methodology for determining the power-law coefficients lacks crucial details. It's unclear whether the authors employed grid search, random search, or another optimization strategy when tuning the coefficients $a$ and $b$ in Figure 5. Information about the number of configurations explored, the computational budget, and the overall search process is missing. This is particularly important if the coefficients are not universally transferable; practitioners aiming to apply the scheduler to different models would need guidance on how to find appropriate coefficients and understand the associated computational costs.

3. Missing Baseline Comparisons: The paper does not compare the proposed scheduler with simpler baselines, such as a warmup plus a stable learning rate scheduler. A basic scheduler that allocates a fixed percentage (e.g., $10$%) of training steps for warmup and then maintains a stable learning rate could also reduce hyperparameters and function well in continual learning setups. Including such baselines, possibly with a final cooldown period would strengthen the evaluation by providing a clearer context for the scheduler's performance improvements.

4. Absence of Learning Curves: Despite citing the Maximum Update Parameterization paper which focuses on stable training, the paper lacks learning curves (e.g., training loss versus number of tokens) that could illustrate training stability and dynamics. Including these curves would provide valuable insights into how the scheduler affects the training process over time and would support claims about improved stability or convergence rates.

5. Unclear Benefit Magnitude: While the experiments show that the optimal learning rate shifts across different token ranges and batch sizes, the practical significance of these shifts is not fully quantified. The performance improvements appear modest; for instance, in the experiments with the Dense model, the Power scheduler outperforms other methods in only 3 out of 7 tasks. A more detailed analysis of the benefits, including effect sizes and statistical significance, would help clarify the scheduler's impact.

6. Minor Typos and Missing Explanations: There are minor errors that need correction—Line 258 should use "T" instead of "x," and Figure 6 should label the coefficient as "b" instead of "a." Additionally, explanations for other models are missing and could be included in the appendix.

1. Lack of theoretical justifications and ablation studies for the design choice in Eq. (7). To solve the issue that Eq. (6) tends to provide a small learning rate for a large-scale training corpus, the authors propose a specific design choice in Eq. (7). However, it is unclear how this design choice aligns with the previous result built upon mu-P theory and whether there are better choices.

2. The improvement over WSD seems marginal. As shown in Table 5, measured by PPL, the gap between Cosine and WSD is obvious. However, the proposed method performs similarly to WSD, +0.1 and +0.2 for dense and MoE respectively.

I am open to raising my score if the authors address the above Weaknesses.

Major Concerns

1. The authors should cite theoretical and empirical work that has shown that the learning rate does depend on the batch size [1,2,3]. Indeed, these prior works make Fig 2b and Hypothesis 1 unnecessary. Most of the points on their grid agree with the square root scaling rule introduced in [1] and theoretically justified in [2] -- I suspect that if they also adjusted the $\beta_1,\beta_2$ hyperparameters in Adam, then it would exactly agree with the square root scaling rule. The bigger problem is that this prior work implies there is no way that the power scheduler can be agnostic to the batch size. Indeed, the main experiments simply pick a fixed batch size and thus don't substantiate the claim that this scheduler generalizes across batch sizes.
2. I am also concerned that there is not much mathematical thought put into any of the findings in the paper. Decisions are motivated heuristically, perhaps with incorrect plots of experiments, and this makes it hard to believe that the scheduler really will transfer robustly. It's not necessary to have formal proofs or anything like that in this paper, but I think it is necessary to provide intuition of some sort, either via math or careful analyses of the training curves.
3. The takeaway that I get from Fig 5 doesn't match the written takeaway. The plot shows that the choice of $a$ can drive changes of up to 10 points in the test perplexity, which is extremely significant. Yet the caption and the text in lines 376-377 say that the test perplexity is not sensitive to the choice of $a$.
4. Altogether, given the issues with Figs 5 and 6, I think the values chosen for $a$ and $b$ don't make any sense to me. The authors should have plotted $(a,b)$ together against the test perplexity. As it stands, there's no evidence to drive the choice of $a$ and $b$.

Minor Concerns

1. My above point here also raises the question: why didn't the authors study the effect of $\beta_1,\beta_2$?
2. The paper makes several slightly incorrect statements about $\mu$P. For example, at the start of Section 3, the authors say that "Due to our use of $\mu$P, we expect the optimal learning rate to be invariant across different model sizes, batch sizes, and training steps." However, Tensor Programs (the theoretical framework behind $\mu$P) can only comment on the optimal learning rate on a fixed dataset. I suggest the authors look for these slight mis-statements throughout their paper and correct them. These culminate in the unsupported claim that the hyperparameters for their power scheduler will transfer across model sizes, training tokens, and batch sizes -- there's no fundamental reason for this to be true.
3. Fig 3 is not made correctly. For example, why is the red curve between batch size 32 and 64 exhibiting a non-zero second-order derivative? It should just be a line, unless additional batch sizes between 32 and 64 were tested. This is misleading. Actually, I am not even sure why these points are connected instead of being presented as a scatter plot.
4. Fig 6 is also confusing. The caption says it is about the $a$ hyperparameter but the plot appears to be for the $b$ hyperparameter, since the values are all negative.
5. It would be nice to see training curves for the 1B runs with different schedulers. That can help people understand what is going on when you use the Power Scheduler.
6. The significance of the 1B results is not apparent to me. The improvement is marginal, and there are more hyperparameters to tune with the Power Scheduler. Also, given the issues with Section 4, I can't tell if the hyperparameters are even selected correctly.

[1] You et al., Large Batch Optimization for Deep Learning: Training BERT in 76 minutes (ICLR 2020).

[2] Malladi et al., On the SDEs and Scaling Rules for Adaptive Gradient Algorithms (NeurIPS 2022).

[3] Goyal et al., Accurate, Large Minibatch SGD: Training ImageNet in 1 Hour (2017)