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

## Abstract

Abstract Finding the optimal learning rate for language model pretraining is a challenging task.
This is not only because there is a complicated correlation between learning rate, batch size, number of training tokens, model size, and other hyperparameters but also because it is prohibitively expensive to perform a hyperparameter search for large language models with Billions or Trillions of parameters.
Recent studies propose using small proxy models and small corpus to perform hyperparameter searches and transposing the optimal parameters to large models and large corpus.
While the zero-shot transferability is theoretically and empirically proven for model size related hyperparameters, like depth and width, the zero-shot transfer from small corpus to large corpus is underexplored.
In this paper, we study the correlation between optimal learning rate, batch size, and number of training tokens for the recently proposed WSD scheduler.
After thousands of small experiments, we found a power-law relationship between variables and demonstrated its transferability across model sizes.
Based on the observation, we propose a new learning rate scheduler, Power scheduler , that is agnostic about the number of training tokens and batch size.
The experiment shows that combining the Power scheduler with Maximum Update Parameterization ( μ ⁢ P 𝜇 P \mu\mathrm{P} italic_μ roman_P ) can consistently achieve impressive performance with one set of hyperparameters regardless of the number of training tokens, batch size, model size, and even model architecture.
Our 3B dense and MoE models trained with the Power scheduler achieve comparable performance as state-of-the-art small language models.
We open source these pretrained models at https://ibm.biz/BdKhLa .

## 1 Introduction

Learning rate is a critical hyperparameter for deep neural network training.
In the context of Large Language Models (LLMs), the cosine learning rate scheduler is the most commonly used strategy.
It has been shown to be effective across multiple state-of-the-art models, including Llama 3 (Dubey et al., ), Gopher (Rae et al., ), etc.
However, the cosine scheduler requires pre-defined training step counts to achieve the optimal loss.
This results in two main drawbacks:
1) the intermediate training checkpoints are suboptimal,
2) continual training of an existing language model becomes complicated.

MiniCPM (Hu et al., ) proposes the Warmup-Stable-Decay (WSD) learning rate scheduler (illustrated in Figure 1) to address these issues.
The WSD learning rate schedule is divided into three phases:
1) warmup phase, linearly increase the learning rate from 0 to peak;
2) stable phase, maintain the learning rate at peak value and training the model for most of the time;
3) decay phase, annealing the learning rate to 0 in a relatively short period.
The main advantage of this schedule is that specifying the number of training steps in advance is not required.
This is particularly convenient for large runs, as the decay can be applied at any time to observe model performance and decide whether to stop.
It also allows for continual learning by default, as training can be resumed from a stable phase checkpoint.
Moreover, the data mixture can be changed during the decay phase to increase the ratio of high-quality data.
This data curriculum is shown to be effective in several recent language models (Team et al., ; Hu et al., ; Dubey et al., ; Shen et al., ).

Figure: Figure 1: Illustration of learning rate curves for Cosine, WSD, and our Power schedulers.
Refer to caption: extracted/5849060/figures/lr.png

However, is WSD really agnostic to token count?
In our experiments, we found that although the WSD scheduler could, in theory, continue in the stable phase forever, the optimal learning rates are different for different amounts of training tokens.
In other words, the optimal learning rate scheduler is tied to the number of training tokens.
Thus, the WSD scheduler still faces the same issue as the cosine scheduler: the intermediate and continual training checkpoints are suboptimal if the number of tokens is too different from the original plan.

Furthermore, deciding the optimal learning rate is still challenging for large-scaling pretraining.
As our experiments will show, there is a complicated correlation between hyperparameters, including learning rate, model size, batch size, and number of training steps.
The size and training cost of modern LLMs make it impossible to do a hyperparameter search on the target model size and training flops.
Researchers have proposed using small proxy models to run hyperparameter searches and predict the optimal hyperparameters of large models from search results (Dey et al., ; Hu et al., ; Yang et al., ).
Among these methods, $\mu\text{Transfer}$ (Yang et al., ; ) is proposed to facilitate zero-shot hyperparameter transfer between different model sizes.
$\mu\text{Transfer}$ has been successfully applied to several language models, including Cerebras-GPT (Dey et al., ), miniCPM (Hu et al., ), and AFM (Gunter et al., ).

In this paper, we first combine the WSD scheduler and $\mu\text{Transfer}$ to study the learning rate transfer between proxy and large models.
Our extensive experiments show that $\mu\text{Transfer}$ does not provide direct zero-shot learning rate transferability across the numbers of tokens and batch sizes for the WSD optimizer.
Instead, the optimal learning rate $\eta_{\mathrm{opt}}$ satisfies a power-law relation with respect to batch size $\beta$ and number of tokens $T$:

$$ $\eta_{\mathrm{opt}}=\beta\cdot aT^{b}$ (1) $$

where $a$ and $b$ are power-low coefficients.
Furthermore, our experiment confirms the zero-shot transferability across model sizes, different model sizes share very similar coefficients.
Inspired by this observation, we propose a new learning rate scheduler, PowerLR, that is agnostic to batch size and token number.
It allows direct transfer of the optimal learning rate scheduling across batch size, token numbers, and model size.
Thus, the expensive pretraining runs can be trained without specifying the number of training tokens, enabling early stop and continual pretraining without sacrificing convergence.

## 2 Background

### 2.1 Maximum Update Parametrization ( μ ⁢ P 𝜇 P \mu\mathrm{P} italic_μ roman_P )

Maximal Update Parameterization ($\mu\mathrm{P}$) (Yang & Hu, ; Yang et al., ; ) controls initialization, layer-wise learning rates, and activation magnitudes to ensure analytically stable training, independent of a model’s width and depth.
In addition to improving training stability, $\mu\mathrm{P}$ improves the transferability of training hyperparameters from small proxy models to large models, a technique called $\mu\text{Transfer}$.
The hyperparameter transferability of $\mu\mathrm{P}$ is theoretically justified for width (Yang et al., ) and depth (Yang et al., ).
Previous work also show empirical evidence for transferability across batch size, training steps, and training sequence length.

In this paper, we follow the $\mu\mathrm{P}$ config used in CerebrasGPT (Dey et al., ) to study the transferability of batch size and learning rate across different numbers of training tokens and model sizes.
Table 1 lists the $\mu\mathrm{P}$ changes we applied to model initialization, learning rate, and multipliers.

**Table 1: List of changes applied when using $\mu\mathrm{P}$. $m_{\text{width}}$ is width multiplier, defined as $d_{m}/d_{\text{base}}$, where $d_{\text{base}}$ is the embedding width, $d_{m}$ is the target model size.**
| Name | Function |
| --- | --- |
| Embedding multiplier | Multiply the embedding output with $m_{\text{emb}}$ |
| Residual Multiplier | Multiply the output of each attention and MLP layer with $m_{\text{res}}$ before adding to residual connection |
| Initialization std | Initialize internal weight matrices (excluding input and output embedding) with standard deviation $\text{init}_{\text{std}}/\sqrt{m_{\text{width}}}$ |
| Learning rate scaling | Set learning rate of internal weight matrices to $\eta/m_{\text{width}}$ |
| Attention logit scaling | Divide attention logits by $d_{\text{head}}$ |

### 2.2 Warmup-Stable-Decay (WSD) Scheduler

MiniCPM (Hu et al., ) proposes the WSD learning rate scheduler to divide pretraining into three stages: the warmup stage, the stable training stage, and the remaining decay stage.
The WSD scheduler is defined as:

$$ $\text{WSD}\left(n\right)=\begin{cases}\frac{n}{N_{\text{warmup}}}\cdot\eta& \text{if }n<N_{\text{warmup}}\\ \eta&\text{if }N_{\text{warmup}}<n\leq N-N_{\text{decay}}\\ f(n,N,N_{\text{decay}})\cdot\eta&\text{if }n>N-N_{\text{decay}}\end{cases},$ (2) $$

where $n$ is the current number of steps, $\eta$ is the stable learning rate, $N$ is the total number of steps, $N_{\text{warmup}}$ is the number of warmup steps, $N_{\text{decay}}$ is the number of decay steps, $f(n,N,N_{\text{decay}})$ is the learning rate decay function.
The main advantage of WSD scheduler over other schedulers (e.g. cosine and linear) is that specifying the number of training steps in advance is not required.
Since the learning rate in the stable training stage does not depend on the number of training steps, the WSD scheduler does not need to specify the number of training steps in advance.
This is particularly convenient for large runs, as the cooldown can be initiated at any time to observe model behavior and decide whether to stop (Hägele et al., ).
It also facilitates extended pre-training, which can be applied to the last checkpoint in the stable training phase.
Moreover, the data mixture can be changed during the cooldown phase to incorporate more high-quality data toward the end of the training.
This data mixture strategy has been proven effective in many recent language models (Hu et al., ; Dubey et al., ; Team et al., ).

## 3 Search Optimal Learning Rate for WSD scheduler with μ ⁢ P 𝜇 P \mu\mathrm{P} italic_μ roman_P

In this section, we focus on finding the optimal learning rate $\eta$ for a small proxy model.
Due to our use of $\mu\mathrm{P}$, we expect the optimal learning rate to be invariant across different model sizes, batch sizes, and training steps.
To verify this, we conducted extensive experiments to find the optimal learning rate for different model sizes and numbers of training tokens.
Table 2 shows our overall hyperparameter search configurations.
We use the WSD scheduler with 10% tokens in the decay phase following previous works (Hu et al., ; Hägele et al., ).
All models in this section and the next section are trained on the RedPajama (Computer, ) corpus and tested on a holdout test set.

**Table 2: Hyperparameter search for finding the optimal learning rate and batch size combination across different numbers of training tokens.**
| Model size | 12M | 36M | 121M |
| --- | --- | --- | --- |
| Fixed hyperparameters |  |  |  |
| Attention head size | 64 | 64 | 64 |
| Number of layers | 32 | 32 | 32 |
| Number of attention heads | 2 | 4 | 8 |
| Embedding size | 128 | 256 | 512 |
| MLP hidden size | 320 | 640 | 1280 |
| Initialization std | 0.02 | 0.02 | 0.02 |
| Sequence Length | 4096 | 4096 | 4096 |
| Warmup tokens | 1B | 1B | 1B |
| $m_{\mathrm{width}}$ | 0.5 | 1 | 2 |
| $m_{\mathrm{emb}}$ | 1 | 1 | 1 |
| $m_{\mathrm{residual}}$ | 1 | 1 | 1 |
| Variable hyperparameters |  |  |  |
| Training tokens $T$ | 2B, 4B, 8B, 16B, 32B, 64B, 128B, 256B |  |  |
| Batch size $\beta$ | 16, 32, 64, 128, 256, 512 |  |  |
| Learning rate $\eta$ | 0.0002, 0.0004, 0.0008, 0.0016, 0.0032, 0.0064, 0.0128, 0.0256 |  |  |

### 3.1 Does Optimal Learning Rate Transfer?

Figure: (a) Batch size $\beta=128$
Refer to caption: x1.png

We conduct two controlled experiments to verify the optimal learning rate’s transferability.
First, we fix the batch size $\beta=128$ and swept across all combinations of learning rates $\eta$ and numbers of training tokens $T$.
Figure 2(a) shows that the optimal learning rate, $\eta_{\mathrm{opt}}$, does not transfer across a wide range of $T$.
$\eta_{\mathrm{opt}}$ tends to decrease with respect to the number of training tokens.
For example, $\eta_{\mathrm{opt}}$ is 0.0128 for 2B training tokens and 0.0008 for 256B.
Second, we fix the number of training tokens $T$ to 128 billion tokens and swept across all combinations of learning rates $\eta$ and batch sizes.

Figure 2(b) shows that the optimal learning rates do not transfer across a wide range of $\beta$, but tends to increase with respect to the batch size.
The optimal learning rate $\eta_{\mathrm{opt}}$ is 0.0002 for batch size 16 and 0.0064 for batch size 512.

Figure: Figure 3: Batch Size v.s. Learning Rate for different numbers of training tokens. The optimal learning rate increases with respect to the batch size.
Refer to caption: x3.png

### 3.2 What is the relationship between η opt subscript 𝜂 opt \eta_{\mathrm{opt}} italic_η start_POSTSUBSCRIPT roman_opt end_POSTSUBSCRIPT , β 𝛽 \beta italic_β and T 𝑇 T italic_T ?

To better understand the relationship between $\eta_{\mathrm{opt}}$, $\beta$, and $T$, we swept across all possible combinations of these three variables.
Figure 3 shows the optimal learning rate for each batch size and training token combination.
Like our previous observation, the optimal learning rate consistently decreases with respect to the number of training tokens across different batch sizes.
Furthermore, we also notice that the ratio between the optimal learning rate $\eta_{\mathrm{opt}}$ and batch size $\beta$ is relatively stable for each number of training tokens.
Based on this observation, we make our first hypothesis:

###### Hypothesis 1

The optimal learning rate $\eta_{\mathrm{opt}}$ for the WSD scheduler and a given pair of $(T,\beta)$ is proportional to the training batch size $\beta$.

Thus, we define $\gamma$ as the ratio between $\eta_{\mathrm{opt}}$ and $\beta$:

$$ $\gamma=\frac{\eta_{\text{opt}}}{\beta}$ (3) $$

Figure: (a) 36M
Refer to caption: x4.png

To verify Hypothesis 1, we conducted an extensive hyperparameter search for three model sizes: 12M, 36M, and 121M.
After finding the optimal learning rate $\eta_{\mathrm{opt}}$ for every combination of $(T,\beta,\text{model size})$, we only keep the three best batch sizes to focus on the optimal scenario.
The $\gamma$ of the three best batch sizes for each $T$ and model size are plotted in in Figure 4(a), Figure 4(b), and Figure 4(c).
These results show that, given a fixed number of training tokens $T$, $\gamma$ falls in a relatively small region.

Furthermore, we notice that $\gamma$ approximately follows a power-law relation with respect to the number of training tokens.
Thus, we make a second hypothesis:

###### Hypothesis 2

The Learning rate to batch size ratio $\gamma$ has a power-law correlation with $T$:

$$ $\gamma=aT^{b}$ (4) $$

when using $\mu\mathrm{P}$, the correlation can be transferred across model sizes.

To verify the Hypothesis 2, we compute the average $\gamma$ of the best batch sizes to estimate the real $\gamma$ for each number of tokens $\beta$.
Figure 4(d) shows that all three model sizes share a similar power-law correlation.
After aggregating the results from three sizes, we get the following correlation:

$$ $\gamma=4.6T^{-0.51}$ (5) $$

where $x$ is the number of training tokens.

In other words, the relation between optimal learning rate $\eta_{\mathrm{opt}}$ for the WSD scheduler, batch size $\beta$, and the number of training tokens $T$ can be approximately modeled with the following equations:

$$ $\eta_{\mathrm{opt}}=\beta\cdot aT^{b}$ (6) $$

Equation 6 provides an easy way to predict the optimal learning rate given the number of training tokens and training batch size for the WSD scheduler.
As in prior work, $a$ and $b$ can be easily estimated through hyperparameter search on a small proxy model.

However, this relationship between the number of training tokens and the optimal learning rate constrains the number of training steps to be specified before training.
We also noticed that the equation tends to provide a small learning rate for a large-scale training corpus.
For example, given a 10 trillion token corpus and a batch size 1024, the predicted optimal learning rate is $0.0011$.
If we consider that $\mu\mathrm{P}$ will divide the learning rate by a factor of $m_{\text{width}}=d_{\text{model}}/d_{\text{base}}$ for matrices inside the model, the actual learning rate will be very small for large models.
While a small learning rate is good for final convergence, it is also likely to cause insufficient exploration at the beginning of training.
To solve these issues, we propose a novel power learning rate scheduler in the next section.

###### Hypothesis 1

###### Hypothesis 2

## 4 Power Scheduler

Inspired by previous observations, we propose a new power learning rate based on the observation from the previous section:

$$ $\eta_{\text{power}}\left(n\right)=\min\left(\eta_{\text{max}},\beta\cdot an^{b }\right)$ (7) $$

where $\beta$ is the batch size, $n$ is the number of tokens already trained, $a$ is the amplitude of the learning rate, $b$ is a power-law exponent for decaying the learning with respect to the number of trained tokens, and $\eta_{\text{max}}$ is the learning rate upper bound that rejects very large learning.
Like the constant learning rate, the power learning rate also does not require specifying the number of training steps or total training tokens beforehand, since the learning rate only depends on the current number of training tokens.

Like the WSD scheduler, we can combine the power learning rate with warmup and decay to get the final Power scheduler:

$$ $\text{Power}\left(n\right)=\begin{cases}\frac{n}{N_{\text{warmup}}}\cdot\eta_{ \text{power}}\left(N_{\text{warmup}}\right)&\text{if }n<N_{\text{warmup}}\\ \eta_{\text{power}}\left(n\right)&\text{if }N_{\text{warmup}}<n\leq N-N_{\text {decay}}\\ f(n,N,N_{\text{decay}})\cdot\eta_{\text{power}}\left(N-N_{\text{decay}}\right) &\text{if }n>N-N_{\text{decay}}\end{cases}$ (8) $$

**Table 3: Hyperparameter search config for Power scheduler. The search range for $a$ and $b$ is decided based on our observation in Section 3.**
| Model size | 36M |
| --- | --- |
| Fixed hyperparameters |  |
| Attention head size | 64 |
| Number of layers | 32 |
| Number of attention heads | 4 |
| Embedding size | 256 |
| MLP hidden size | 640 |
| Initialization std | 0.1 |
| Sequence Length | 4096 |
| Warmup tokens | 1B |
| $m_{\text{width}}$ | 1 |
| $m_{\text{emb}}$ | 12 |
| $m_{\text{residual}}$ | 0.26 |
| Max learning rate $\eta_{\text{max}}$ | 0.02 |
| Variable hyperparameters |  |
| Batch size $\beta$ | 32, 64, 128, 256, 512 |
| Training tokens $T$ | 2B, 4B, 8B, 16B, 32B, 64B, 128B |
| $a$ | $[3,5]$ |
| $b$ | $[-0.6,-0.4]$ |

The Power scheduler requires three separate hyperparameters $(\eta_{\mathrm{max}},a,b)$, instead of only one hyperparameter required by the cosine and WSD schedulers.
However, we expect these hyperparameters to transfer across different model sizes, number of training tokens, and batch sizes.
Additionally, since Figure 1 shows that the $\eta_{\mathrm{max}}$ has a very limited impact on the overall learning rate curve, we can simply set it to a large enough value.
In this paper, we set $\eta_{\mathrm{max}}$ to $0.02$.
Then, we conduct another hyperparameter search for $a$ and $b$ to find the optimal learning rate hyperparameter.
Table 3 shows the configuration used in the Power scheduler hyperparameter search.

Figure: Figure 5: Hyperparameter search results for 32B training tokens. Left: Test perplexity vs $a$. The test perplexity is not sensitive to different choices of $a$ within the range of $[3,5]$. Right: Test perplexity vs $b$. The test perplexity is more sensitive to the choices of $b$. The optimal $b$ is between $-0.52$ and $-0.51$.
Refer to caption: x8.png

Figure: Figure 6: Optimal $a$ vs Number of training tokens. The optimal $a$ is consistent across different numbers of training tokens.
Refer to caption: x10.png

Figure 5 shows the hyperparameter search results for 32 billion training tokens.
We observe that the test perplexity is not sensitive to different choices of $a$ within our search range of $[3,5]$.
In contrast, the test perplexity is more sensitive to the choices of $b$.
The optimal $b$ is between $-0.52$ and $-0.51$.
A similar observation has been made across different numbers of training tokens.
Figure 6 shows that the optimal $b$ for different $\beta$ and $T$ consistently falls within the range of $[-0.49,-0.53]$ with few exceptions.
Combining this and previous observations, we select $a=4$ and $b=-0.51$ for the rest of this paper.

## 5 Pre-training Experiments

This section compares the Power scheduler with the WSD and Cosine scheduler across different scenarios.
In the first part, we conduct a controlled experiment to train 1B transformer models and 1B mixture-of-experts (MoE) models with different learning rate schedulers to show that the Power scheduler is comparable to or better than other schedulers.
In the second part, we take a more realistic setting, training a 3B transformer language model and a 3B MoE language model with high-quality data in the decay phase to compare with strong open-source language models.
For all the Power scheduler experiments in this section, we will use the hyperparameters from Section 4, $a=4$, $b=-0.51$, and $\eta_{\mathrm{max}}=0.02$.

### 5.1 1B Controlled Experiment

**Table 4: Hyperparameter for 1B and 3B models.**
| Model | 1B Dense | 1B MoE | PowerLM-3B | PowerMoE-3B |
| --- | --- | --- | --- | --- |
| Embedding size | 1536 | 1024 | 2304 | 1536 |
| Number of layers | 40 | 24 | 40 | 32 |
| Attention head size | 64 | 64 | 64 | 64 |
| Number of attention heads | 24 | 16 | 36 | 24 |
| Number of KV heads | 24 | 8 | 36 | 8 |
| MLP hidden size | 4096 | 512 | 9216 | 512 |
| MLP activation | swiglu | swiglu | swiglu | swiglu |
| Number of Experts | – | 32 | – | 40 |
| MoE TopK | – | 8 | – | 8 |
| Initialization std | 0.1 | 0.1 | 0.1 | 0.1 |
| Sequence Length | 4096 | 4096 | 4096 | 4096 |
| $m_{\text{width}}$ | 6 | 4 | 8 | 6 |
| $m_{\text{emb}}$ | 12 | 12 | 12 | 12 |
| $m_{\text{residual}}$ | 0.22 | 0.28 | 0.22 | 0.22 |
| #Parameters | 1.2B | 1.3B | 3.5B | 3.3B |
| #Active Parameters | 1.2B | 377M | 3.5B | 800M |
| #Training tokens | 1T | 1.1T | 1.25T | 3T |

We train a series of 1B parameter dense and Mixture-of-Experts (MoE) transformer models using WSD, cosine, and Power schedulers.
All models are trained with 1T tokens.
We use the optimal hyperparameters proposed in MiniCPM (Hu et al., ) and adapt them to our 1B setting with $\mu\text{Transfer}$.
All models are trained with batch size 1024.
For WSD and cosine scheduler, we set $\eta=0.01$, following the optimal learning rate from MiniCPM.
For the Power and WSD scheduler, we exponentially decay the learning rate to 0 for the last 100B tokens.
The MoE models are implemented with ScatterMoE (Tan et al., ).
More model details can be found in Table 4.

**Table 5: Language Modeling and Zero-shot performance of 1B models. $acc_{n}$ means accuracy with average log probability.**
| Task | Wiki | ARC | BoolQ | Hellaswag | OBQA | PIQA | WinoGrande | Average |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| Metric | $ppl$ | $acc_{n}$ | $acc$ | $acc_{n}$ | $acc_{n}$ | $acc_{n}$ | $acc$ |  |  |
| Dense | Cosine | 14.5 | 44.6 | 63.5 | 63.6 | 37.6 | 75.3 | 61.3 | 57.7 |
| WSD | 13.9 | 43.0 | 62.8 | 64.8 | 38.0 | 74.5 | 62.1 | 57.6 |  |
| Power | 13.8 | 44.3 | 65.6 | 64.6 | 37.0 | 76.0 | 61.8 | 58.2 |  |
| MoE | Cosine | 15.7 | 40.6 | 63.6 | 59.7 | 35.2 | 74.0 | 57.2 | 55.1 |
| WSD | 14.8 | 43.9 | 59.7 | 61.0 | 34.4 | 76.0 | 59.3 | 55.7 |  |
| Power | 14.6 | 44.0 | 60.4 | 61.6 | 36.0 | 74.8 | 58.3 | 55.9 |  |

We evaluate all 1B models on language model tasks and multiple-choice tasks from LM evaluation Harness (Gao et al., ).
The multiple-choice tasks include grade-school science questions (ARC, Clark et al. ()), yes/no questions (BoolQ, Clark et al. ()), common sense reasoning (Hellaswag, Zellers et al. ()), open book question answering (OpenBookQA, Mihaylov et al. ()), physical questions (PIQA, Bisk et al. ()), and Winograd schema task (Winogrande, Sakaguchi et al. ()).
Table 5 shows the performance.
The Power scheduler provides consistently better or comparable performance for both language modeling and downstream tasks.
Surprisingly, even though we only performed a hyperparameter search on a small dense model, the hyperparameters still performed well on the MoE model.

### 5.2 3B Realistic Experiment

To compare performance against strong open-source language models, we pretrain two 3B language models: 1) PowerLM-3B, a 3B dense language model, and 2) PowerMoE-3B, a 3B MoE language model(^1^11PowerLM-3B and PowerMoE-3B are open-sourced at [https://ibm.biz/BdKhLa](https://ibm.biz/BdKhLa).).
We pretrain these two models using the two-stage training schema in Hu et al. () and Power scheduler.
Stage 1 linearly warms up the learning rate and then applies the power decay.
The training corpus is a mix of large-scale, medium-quality open-source datasets with permissive licenses.
PowerLM-3B is trained on 1T tokens, and PowerMoE-3B on 2.5T tokens.
Stage 2 exponentially decays the learning rate to zero.
The training corpus is a mix of stage 1 data and a small amount of high-quality open-source/synthetic corpora with permissive licenses.
PowerLM-3B is trained on 250B tokens, and PowerMoE-3B on 500B tokens.
The training batch size is 1024.

**Table 6: Zero-shot performance on multiple-choice tasks.**
| Task | ARC | BoolQ | Hellaswag | OBQA | PIQA | WinoGrande | Average |
| --- | --- | --- | --- | --- | --- | --- | --- |
| Metric | $acc_{n}$ | $acc$ | $acc_{n}$ | $acc_{n}$ | $acc_{n}$ | $acc$ |  |
| Qwen1.5-4B | 50.6 | 77.7 | 71.5 | 39.8 | 77.0 | 64.4 | 63.5 |
| Gemma2-2B | 65.0 | 72.8 | 73.0 | 41.4 | 79.2 | 68.9 | 66.7 |
| PowerLM-3B | 60.5 | 72.0 | 74.6 | 43.6 | 79.9 | 70.0 | 66.8 |
| Qwen1.5-1.8B | 47.1 | 66.3 | 60.9 | 34.2 | 74.0 | 60.8 | 57.2 |
| Qwen2-1.5B | 48.1 | 72.5 | 65.4 | 36.4 | 75.3 | 66.2 | 60.7 |
| SmolLM-1.7B | 59.9 | 66.1 | 65.6 | 42.0 | 75.7 | 60.5 | 61.7 |
| PowerMoE-3B | 58.1 | 65.0 | 71.5 | 41.0 | 79.1 | 65.0 | 63.3 |

**Table 7: Performance on Math, Coding, and MMLU. The MATH 5-shot, exact matching setting is from Hendrycks et al. (), and the 4-shot, sympy matching setting is from Lewkowycz et al. ().**
| Task | GSM8k | MATH | HumanEval | MBPP | MMLU |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- |
| Prompt | 5-shot | 8-shot,cot | 5-shot | 4-shot | 0-shot | 0-shot | 5-shot |
| Metric | exact | exact | exact | sympy | pass@1 | pass@1 | acc |
| Qwen1.5-4B | 53.1 | 54.8 | 9.76 | 12.8 | 7.3 | 30.6 | 55.2 |
| Gemma2-2B | 23.9 | 25.5 | 9.80 | 15.4 | 19.5 | 27.8 | 53.0 |
| PowerLM-3B | 34.9 | 49.7 | 9.90 | 15.2 | 26.8 | 33.6 | 49.2 |
| SmolLM-1.7B | 6.6 | 6.8 | 4.86 | 3.0 | 20.7 | 30.6 | 28.8 |
| Qwen1.5-1.8B | 34.7 | 37.3 | 5.66 | 9.6 | 4.30 | 21.6 | 45.6 |
| Qwen2-1.5B | 58.3 | 58.6 | 4.38 | 23.3 | 16.5 | 31.0 | 55.9 |
| PowerMoE-3B | 25.9 | 38.4 | 9.26 | 14.8 | 20.1 | 32.4 | 42.8 |

Table 6 shows the multi-choices performance of our model and state-of-the-art models.
Table 7 shows MMLU and generative performance in the math and code domain.
The results show that, despite being trained with relatively fewer tokens, our PowerLM-3B still achieves comparable performance with state-of-the-art 2B to 4B language models.
Furthermore, our PowerMoE-3B uses only 800M active parameters but performs similarly to state-of-the-art 1B to 2B dense models.

## 6 Conclusion

In this paper, we systematically study the relationship between optimal learning rate, batch size, and number of training tokens.
We observed a power-law relation between these variables while using the WSD learning rate scheduler.
Inspired by this observation, we propose a new learning rate scheduler, the Power scheduler, that is invariant with respect to the number of tokens and batch size.
The experiment shows that combining the Power scheduler with Maximum Update Parameterization ($\mu\mathrm{P}$) can consistently achieve impressive performance with one set of hyperparameters regardless of the number of training tokens, batch size, model size, and even model architecture.

## Peer Reviews

confidence: 5
contribution: 3
initial_rating: 3
presentation: 2
questions: See above.
soundness: 2
strengths: 1. It is interesting that the authors find that the learning rate of the WSD scheduler is dependent on the number of tokens trained. The motivation of the scheduler is that one should be able to train in the stable (S) phase for an arbitrary period of time and decay (D) whenever training is completed. However, I do see substantial issues with their experiments that make me question the validity of this finding.
2. Designing methods that transfer across training budgets is of great interest, since it's hard to ablate all crucial design choices at scale.
summary: This paper proposes a new learning rate scheduler for language model pre-training that can accommodate varying numbers of pre-training examples. The authors identify a power-law relationship between the ratio LR / batch size and the number of training tokens. They then use this to design a scheduler that will decay the LR exponentially as training proceeds -- importantly, the scheduler improves over other options (cosine schedule) because it doesn't require knowing the number of tokens in advance.

Review summary: The problem is interesting, but there is a huge chunk of significant prior work (on scaling learning rates and batch sizes) that is missing. The plots of experiments appear to be incorrect, and most of the time, the takeaways do not match what is shown in the actual plot. Therefore, although the idea is interesting and may hold some merit, I am recommending rejection because the paper does not constitute a sound and valid scientific study.
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.
5. 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. 
7. 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.
8. 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)

---

confidence: 4
contribution: 2
initial_rating: 3
presentation: 3
questions: 1. Can you provide evidence of the power-law coefficients' transferability to other models and datasets?
2. What was your methodology and budget for determining the power-law coefficients $a$ and $b$?
3. What guidelines can you provide for practitioners on determining coefficients and required computational resources?
4. Have you compared the Power scheduler with simpler baselines like a (warmup plus) stable learning rate (+ final cooldown)?
5. Could you include learning curves to show training stability when using the Power scheduler?
6. How do you quantify the practical benefits of the Power scheduler given the modest improvements observed?
7. Have you tested the Power scheduler in continual learning scenarios with evolving dataset sizes?
8. Does the Power scheduler impact training efficiency regarding convergence speed or resource utilization?
soundness: 2
strengths: 1. Strong Motivation and Practical Relevance: The paper addresses the complex challenge of optimizing learning rates in language model pretraining, where hyperparameters like learning rate, batch size, number of training tokens, and model size are deeply intertwined. Given that hyperparameter optimization at the target scale is computationally prohibitive, their effort to reduce the number of hyperparameters is both timely and valuable.

2. Overcoming Limitations of Existing Schedulers: The authors tackle the drawbacks of current learning rate schedulers, such as the cosine scheduler's requirement to specify the total number of training tokens—a constraint that limits its use in scenarios like continual learning. The proposed Power scheduler is agnostic to the number of training tokens and batch size, offering greater flexibility and applicability.

3. Empirical Discovery of a Scaling Law: Through extensive experiments, they demonstrate that the optimal learning rate shifts across different ranges of tokens and batch sizes, indicating it doesn't transfer directly across scales. By uncovering a power-law relationship among these variables, they provide an empirical foundation for their scheduler.

4. Extensive Experimentation and Practical Impact: The significant amount of experimentation not only validates their findings but also showcases the scheduler's effectiveness on 3B-parameter dense and MoE models. This contributes to simplifying hyperparameter tuning for large-scale language models, which has substantial practical implications for the field.
summary: The paper addresses the challenge of optimizing the learning rate for language model pretraining, complicated by intricate correlations among learning rate, batch size, number of training tokens, and model size, as well as the prohibitive cost of hyperparameter optimization for large-scale models. Leveraging zero-shot transfer techniques, the authors explore dataset size as a fidelity parameter. They extend the WSD scheduler to focus on learning rate, batch size, and number of tokens, conducting extensive small-scale experiments that reveal a power-law relationship among these variables. This leads to the proposal of the Power scheduler, a learning rate scheduler agnostic to the number of training tokens and batch size. Key contributions include:

* Demonstrating that dataset size can effectively serve as a fidelity parameter for zero-shot transfer in hyperparameter tuning.
* Identifying a power-law relationship between learning rate, batch size, and number of training tokens.
* Proposing the Power scheduler to simplify hyperparameter tuning by reducing dependency on training tokens and batch size.
weaknesses: 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.

---

confidence: 3
contribution: 3
initial_rating: 5
presentation: 3
questions: 1. Besides model size and architectures, can the numbers a=4 and b=0.51 transfer to other optimizers or tokenizers? Do we need to search them again on the corresponding proxy models?

2. How to interpret the results in Table 6 and Table 7? If I understand correctly, there exist many differences between the PowerLM and the compared models, not solely the Power Scheduler. How could we attribute the performance gain to the proposed method?
soundness: 3
strengths: 1. The paper is well-written and easy to understand.

2. The research topic on learning rate scheduler is timely and important for the LLM community.

3. Theoretically, the proposed method is built upon the mu-P framework and thus has a good foundation.

4. Empirically, the effectiveness of the proposed method is verified on up to 3B dense and MoE models, which is comprehensive and convincing.
summary: The paper introduces the Power Scheduler, a novel learning rate scheduler for large language model (LLM) training that is invariant to batch size and the number of training tokens. The authors first explore the limitations of traditional schedulers, particularly the cosine and WSD schedulers, which are dependent on fixed training steps, leading to suboptimal checkpoints for models trained with varying token counts.

The key contributions are as follows:
1. Through theoretical and empirical studies, the paper demonstrates that the optimal learning rate follows a power-law relationship with batch size and token count, enabling a consistent learning rate setting across different model sizes and token volumes.

2. The scheduler is evaluated on various language models, including dense and MoE models, and shows that it performs comparably to or better than existing schedulers, even with different model sizes and architectures. The experiments further validate that Power Scheduler can transfer hyperparameters across proxy models to large-scale models, offering a cost-effective solution for extensive LLM training.
weaknesses: 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.

---

confidence: 3
contribution: 3
initial_rating: 6
presentation: 3
questions: Why are the curves in fig. 3 very non-monotonic?

In fig. 2(a) training longer requires a larger LR in general, but the other paper found the opposite is true in a larger scale, what cause the difference? [1]

[1] Scaling Optimal LR Across Token Horizons
soundness: 2
strengths: The Power Scheduler addresses a critical gap in learning rate scheduling for LLMs by eliminating the need to specify batch size or token counts in advance. This is a significant improvement over traditional approaches like the cosine and WSD schedulers, which require predefined training steps and can be suboptimal in continual learning scenarios. The paper introduces a well-motivated, novel concept of using a power-law relation to generalize the learning rate across different training setups, showing strong theoretical justification and empirical validation. The experimental results are demonstrating that the Power Scheduler achieves comparable or superior performance on both dense and MoE models. Its flexibility across model sizes and architectures makes it highly relevant for scaling deep learning workloads efficiently.

Another strength of the paper is its comprehensive experimental setup, which tests the Power Scheduler against state-of-the-art models in controlled environments and realistic training scenarios. The authors thoroughly explore hyperparameter search and transferability, particularly in the context of large-scale pretraining with billions of parameters. the integration with Maximum Update Parameterization (*µP*) shows that the scheduler can maintain stable training dynamics across various model sizes, a critical feature for improving the transferability of hyperparameters from small proxy models to large-scale models.
summary: Power Scheduler is designed to be independent of the batch size and number of training tokens. The motivation for this work is the inefficiencies in current schedulers like the Warmup-Stable-Decay (WSD) scheduler, which still ties the optimal learning rate to the number of training tokens, making continual training and early stopping suboptimal for large language models (LLMs). By identifying a power-law relationship between optimal learning rates, batch size, and the number of training tokens, the authors introduce a new scheduling method that generalizes across model sizes, batch sizes, and token counts. The results show that this approach improves model performance while maintaining flexibility in training large models, including both dense and Mixture of Experts (MoE) models.
weaknesses: 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