Towards Joint Scaling Laws with Optimal Batch Size Schedules

Summary (Overview)

Core contribution: This paper derives a closed-form optimal batch size schedule for any prescribed learning rate schedule, based on a convex optimization framework that characterizes deep learning training dynamics as a sequence-to-sequence mapping from (learning rate, batch size) schedules to loss values.
Key theoretical result: The optimal batch size schedule $B_t^{\text{optim}} = \frac{K}{2} \frac{\eta_t}{\sqrt{\int_0^T \eta_k dk \int_t^T \eta_k dk}}$ depends only on the shape of the learning rate schedule and the total data/compute budget—not on peak learning rate, model size, architecture, or optimizer.
Empirical gains: Dynamic batch size schedules improve compute efficiency by 6–15% with zero computational overhead across Llama3 (dense), Qwen3 (MoE), and vision-language models up to 7B parameters, using both AdamW and Muon optimizers.
Joint scaling laws: The authors establish joint scaling laws where peak learning rate scales as $\eta^*(T) \propto 1/\sqrt{T}$ , weight decay scales as $\kappa(N,T) = \kappa^*(N_{\text{small}}, T_{\text{small}})/\sqrt{T/T_{\text{small}}}$ , and batch size follows the closed-form schedule—yielding $O(1/\sqrt{T})$ loss convergence.
Universal dynamics preserved: The optimal batch size schedule maintains universal training dynamics ("supercollapse"), where different model sizes share the same normalized loss curve under properly scaled hyperparameters, enabling reliable prediction of large-scale training from small-scale runs.

Introduction and Theoretical Foundation

Background and Motivation

Large-scale model training critically depends on hyperparameter selection, particularly learning rate (LR) and batch size (BS). While learning rate schedules (warmup + decay) are standard practice, batch size is typically held static throughout training. This is potentially suboptimal: a fixed batch size is only a constrained special case of the broader class of batch size schedules $\{B_t\}_{t=1}^T$ .

Key research question: What is the optimal batch size schedule and which factors determine it?

Theoretical Foundation

The work builds on stochastic convex optimization theory, which has been shown to capture deep learning training dynamics empirically even though neural network training is non-convex and non-smooth. The authors follow the line of work by Defazio et al. (2023), Bu et al. (2026), and Schaipp et al. (2025).

Theorem 1 (Learning rate to loss sequence). For convex loss and any learning rate sequence $\eta_t$ under SGD:

\mathbb{E} L(\mathbf{w}_\tau; \eta_t) \leq L_* + \frac{D^2}{2 \sum_{t=1}^\tau \eta_t} + \frac{1}{2} \sum_{t=1}^{\tau-1} \frac{\eta_t^2 \mathbb{E}\|\mathbf{g}_t\|^2}{\sum_{k=t+1}^\tau \eta_k} + \frac{1}{2} \eta_\tau \mathbb{E}\|\mathbf{g}_\tau\|^2

where $D = \|\mathbf{w}_t - \mathbf{w}_*\|$ and $\mathbf{g}_t$ is the mini-batch gradient at step $t$ .

Key Assumption on Batch Size

To incorporate batch size effects, the authors assume the gradient norm satisfies:

\mathbb{E}\|\mathbf{g}_t\|^2 \leq G^2 + X/B_t

where $G$ bounds the norm of the expected gradient $\|\mathbb{E}\mathbf{g}_t\|$ and $X$ bounds the trace of the per-example gradient covariance matrix.

Methodology

Sequence-to-Sequence Loss Prediction

The authors derive a general characterization (Theorem 2) that jointly predicts the loss from both learning rate and batch size sequences. This yields the following approximate equality (validated empirically):

\mathbb{E} L(\mathbf{w}_T) \approx L_* + \frac{D^2}{2 \sum_{t=1}^T \eta_t} + \frac{1}{2} \sum_{t=1}^{T-1} \frac{\eta_t^2 (G^2 + X/B_t)}{\sum_{k=t+1}^T \eta_k}

Optimal Batch Size Derivation

The optimization problem minimizes the variance term subject to a data/compute budget $K = B_{\text{static}}T$ :

\min_{B_t} \int_0^T \left(\frac{1}{B_t} \frac{\eta_t^2}{\int_t^T \eta_k dk}\right) dt, \quad \text{s.t.} \int_0^T B_t dt = K

Theorem 3 (Closed-form optimal batch size):

B_t^{\text{optim}} = \frac{K}{2} \frac{\eta_t}{\sqrt{\int_0^T \eta_k dk \int_t^T \eta_k dk}}

Optimal Batch Sizes for Common Schedules

Learning Rate	$\eta_t$ formula	$B_t^{\text{optim}}$ formula
constant	$\eta$	$\frac{B_{\text{static}}}{2} \sqrt{\frac{T}{T-t}}$
cosine	$\frac{\eta}{2}(\cos(\pi t/T) + 1)$	$\frac{B_{\text{static}} \cos^2(\pi t/2T)}{\sqrt{1 - t/T - \frac{1}{\pi}\sin(\pi t/T)}}$
linear	$\eta(1-t/T)$	$B_{\text{static}}$
WSD	$\begin{cases} \eta & t \leq cT \\ \eta\frac{T-t}{T-cT} & t > cT \end{cases}$	$\begin{cases} \frac{B_{\text{static}} T}{\sqrt{(T+cT-2t)(T+cT)}} & t \leq cT \\ \frac{B_{\text{static}}}{\sqrt{1-c^2}} & t > cT \end{cases}$

Key Properties

Decoupling from learning rate tuning: The optimal batch size is independent of peak learning rate, weight decay, model size, architecture, and optimizer.
Endpoint behavior: If $\eta_t \sim (T-t)^p$ as $t \to T$ , then $b_t \sim (T-t)^{(p-1)/2}$ , meaning linear/WSD schedules yield constant terminal batch sizes while cosine schedules drive $B_t^{\text{optim}} \to 0$ .

Empirical Validation / Results

Validation Approach

The authors validate their theory through linear regression fitting of the loss prediction (Eq. 2.4), demonstrating precise characterization of training dynamics across:

Optimizers: Muon-NSGD and AdamW
Models: Llama3 (dense), Qwen3 (MoE), SmoLM2-360M (VLM)
Scales: 0.1B to 7B parameters, up to 1e21 FLOPs

Key Empirical Findings

1. Compute Efficiency Improvement

The dynamic batch size schedule delivers 6–15% improvement in compute efficiency over static batch size baselines under fixed compute of 1e20 FLOPs, with the same learning rate and weight decay settings.

2. Loss Convergence

Dynamic batch size consistently outperforms static baselines across cosine, WSD, and constant learning rate schedules.
Relative improvement in best perplexity: 0.8% (cosine), 2.5% (WSD), and 14.6% (constant LR).

3. Evaluation Benchmarks (Llama3-7B, last-iterate)

Model	Batch Size	LAMBADA	HellaSwag	ARC-Easy	ARC-Challenge	WinoGrande	PIQA	OpenBookQA	BoolQ
Llama3	static	32.23	42.93	56.22	29.26	53.03	67.89	32.40	61.89
Llama3	dynamic	35.36	44.33	56.36	30.46	53.99	68.34	33.20	62.26
Qwen3	static	36.06	45.61	57.83	32.59	52.49	68.77	34.00	61.47
Qwen3	dynamic	37.71	46.80	59.09	32.25	55.33	69.97	36.20	61.59

4. Post-training: Vision-Language Models

Dynamic batch size outperforms static baseline on 7 of 9 VLM evaluation tasks, with no significant regression on any benchmark.

5. Post-training: Math Fine-tuning

On Qwen3-0.6B fine-tuned on orca_math with WSD schedule, dynamic batch size shows better validation loss and GSM8K accuracy trajectories throughout training.

Variance Reduction Result

The advantage of dynamic over static batch size is formalized via Jensen's inequality:

L_{\text{dynamic}}(T) - L_{\text{static}}(T) = \frac{X T}{2B_{\text{static}}} \left[ \left(\frac{\int_0^T Z_t dt}{T}\right)^2 - \frac{\int_0^T Z_t^2 dt}{T} \right] \leq 0

where $Z_t := \frac{\eta_t}{\sqrt{\int_t^T \eta_k dk}}$ , confirming dynamic schedules always reduce the loss upper bound.

Theoretical and Practical Implications

Theoretical Implications

Convex dominance in deep learning: The empirical accuracy of the convex-theory-based loss prediction across non-convex neural networks and adaptive optimizers suggests that deep learning training dynamics are dominated by convex behavior, extending findings from prior work.
$\eta_t/B_t$ is not a sufficient statistic: The authors show that matching $\eta_t/B_t$ trajectories does NOT produce equivalent training dynamics—realizing schedules through learning rate consistently outperforms batch size (by +0.017 to +0.743 validation loss). This challenges the common simplification in prior work (Smith et al., 2018; Goyal et al., 2017).
Joint scaling laws: The framework unifies scaling of learning rate ( $1/\sqrt{T}$ ), weight decay ( $1/\sqrt{T}$ ), and batch size (closed-form schedule) into a single principled framework with provable $O(1/\sqrt{T})$ convergence.

Practical Implications

Zero-cost improvement: The optimal batch size schedule requires no additional hyperparameter tuning—it is fully determined by the learning rate schedule shape and total compute budget.
Universal dynamics for prediction: The optimal batch size schedule preserves universal training dynamics ("supercollapse"), enabling reliable prediction of large-scale training behavior from small-scale experiments.
Generalization across domains: Benefits extend beyond pre-training to fine-tuning (math) and vision-language models, with no observed regressions.
Limitations: Advantages diminish as batch size grows (theoretically vanishing as $B \to \infty$ ) and may disappear under purely low-precision training.

Conclusion

This work establishes that the optimal batch size schedule is a principled, closed-form function of the learning rate schedule shape and total computational budget—decoupled from peak learning rate, model architecture, and optimizer choice. The derived schedule consistently outperforms static batch size baselines by 6–15% in compute efficiency with zero additional cost, while preserving the universal dynamics that enable reliable scaling predictions.

Future directions include: extending the framework to more general cost functions (e.g., wall-clock time), investigating interactions with other hyperparameters beyond learning rate and weight decay, and exploring whether the convexity-driven insights hold at even larger scales and more diverse architectures.