Your batch size should have a schedule, too

TL;DR. Batch size should be proportional to the current learning rate divided by the square root of the remaining area under the learning-rate curve, then normalized to keep the total token budget fixed. The resulting schedule depends on the LR shape, not its peak value.

We warm up the learning rate. We decay it. Yet we often keep the batch size fixed. Could the same training tokens do more useful work if we changed when we spent them?

In our paper, with Zhiqi Bu and Shiyun Xu, we derive a batch-size schedule from a planned learning-rate schedule, then test it on language and vision-language models. The idea is to redistribute a fixed budget across updates, rather than simply use more data.

The efficiency figures are reported in Figure 5. They are compute comparisons, not measured wall-clock speedups.

• 1. Look at both tokens and steps
• 2. Let the learning-rate shape choose the batches
• 3. Spend the budget where it reduces noise most
• 4. Scale the whole recipe
• 5. Check the downstream trade-offs
• 6. What to try in a training run
• Citation

1. Look at both tokens and steps

A larger batch processes more tokens per update. The two views of Figure 1 tell a complementary story: the ordinary dynamic schedule improves final loss, while a larger average batch reaches the reference loss in fewer optimizer updates at a similar token budget. Compare iterations on the left with tokens on the right.

2. Let the learning-rate shape choose the batches

The answer is not always an increasing ramp. Cosine calls for a rise and fall. WSD calls for a rise followed by a plateau during cooldown. Linear decay gives a constant batch size. Change the peak learning rate below: the batch-size shape stays the same.

3. Spend the budget where it reduces noise most

In the characterization, noise from different updates carries different weight in the final loss. A constant batch spends equally on unequal contributions. Move the slider from static to derived, then keep going: exaggerating the allocation can give back the benefit.

4. Scale the whole recipe

For a training horizon four times longer, the joint recipe halves both peak learning rate and weight decay. The normalized batch shape is unchanged; the budget sets its overall size. This combines batch scheduling with a transfer rule for the other hyperparameters.

5. Check the downstream trade-offs

Better validation loss should also be checked against the tasks we care about. In Table 1, the dense model improves on all eight reported benchmarks. The MoE model improves on seven, with a small ARC-Challenge regression. Inspect the exact scores below.

6. What to try in a training run

Keep the comparison simple: fix the model, sequence length, token budget, and learning-rate schedule, then compare static batches with the derived allocation. Implement feasible effective batches with rounding and appropriate gradient accumulation. Check tokens, FLOPs, wall time, and downstream quality separately.

The benefit is not universal. It shrinks as the baseline batch grows, and the paper finds that it can disappear in purely low-precision training. Those are reasons to test the recipe in your own stack, not to assume a headline gain will transfer unchanged.

The takeaway: batch size is a schedule worth designing alongside learning rate. The useful question is where each token does the most work.

Citation

If you find this work useful, please cite our paper:

@misc{li2026jointscaling,
  title         = {Towards joint scaling laws with optimal batch size schedules},
  author        = {Jiaxiang Li and Zhiqi Bu and Shiyun Xu},
  year          = {2026},
  eprint        = {2607.27731},
  archivePrefix = {arXiv},
  primaryClass  = {cs.LG},
  doi           = {10.48550/arXiv.2607.27731},
  url           = {https://arxiv.org/abs/2607.27731}
}