Towards joint scaling laws with optimal batch size schedules

2026-07-30

AI Summary

This work addresses the limitations of static batch size training, which overlooks the coupled influence of learning rate and batch size on training dynamics, thereby constraining model performance. From a convex optimization perspective, the paper establishes a theoretical characterization of the loss function under joint scheduling of learning rate and batch size, and for the first time derives a closed-form solution for the optimal dynamic batch size applicable to arbitrary learning rate schedules. Furthermore, it introduces a joint scaling law that generalizes across optimizers and model architectures. Experimental results demonstrate that the proposed dynamic batch size scheduling significantly outperforms fixed batch size baselines in large language model training, confirming its effectiveness and broad applicability.

Abstract

Modern deep learning typically keeps the batch size static throughout training, thus overlooking the joint effect of learning rate and batch size on the training dynamics. In this paper, we study the deep learning dynamics through the lens of convex optimization and derive a joint characterization of loss in terms of both schedules, applicable to general optimizers and model architectures. This characterization yields a closed-form optimal batch size schedule for any prescribed learning rate schedule, and further leads to joint scaling laws that consistently outperform static batch size baselines, highlighting the significance of dynamic batch size schedule in large language model training.

Problem

• batch size schedule
• learning rate
• training dynamics
• scaling laws
• large language models

Innovation

• joint scaling laws
• optimal batch size schedule
• learning rate schedule
• training dynamics
• convex optimization