Paper deep dive
Towards joint scaling laws with optimal batch size schedules
Jiaxiang Li, Zhiqi Bu, Shiyun Xu
Intelligence
Status: succeeded | Model: Gemma-4-26B-A4B | Prompt: intel-v1 | Confidence: 95%
Last extracted: 8/1/2026, 2:06:23 AM
Summary
This paper introduces a method for optimizing deep learning training dynamics by deriving a closed-form optimal batch size schedule that varies over time, coupled with learning rate schedules. Using convex optimization theory, the authors characterize the loss function in terms of both learning rate and batch size sequences. They demonstrate that this dynamic approach yields joint scaling laws that consistently outperform static batch size baselines in large language model training, improving compute efficiency by 6-15% while maintaining universal training dynamics.
Entities (9)
Relation Signals (6)
Dynamic Batch Size Schedule → appliedto → Llama3
confidence 95% · We validate the advantage of our batch size schedules... on Llama3 (dense)
Dynamic Batch Size Schedule → appliedto → Qwen3
confidence 95% · We validate the advantage of our batch size schedules... on Qwen3 (MoE)
Joint Scaling Laws → incorporates → Dynamic Batch Size Schedule
confidence 95% · We build joint scaling laws that employ... our scale-free batch size schedule
Dynamic Batch Size Schedule → outperforms → Static Batch Size
confidence 95% · joint scaling laws that consistently outperform static batch size baselines
Dynamic Batch Size Schedule → derivedfrom → Convex Optimization
confidence 90% · we study the deep learning dynamics through the lens of convex optimization and derive a joint characterization
Dynamic Batch Size Schedule → improves → compute efficiency
confidence 90% · with an improvement on compute efficiency by 6∼15%
Cypher Suggestions (0)
No Cypher suggestions yet.
Abstract
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.
Tags
Links
- Source: https://arxiv.org/abs/2607.27731v1
- Canonical: https://arxiv.org/abs/2607.27731v1
Trouble viewing inline? Open PDF directly →
Full Text
75,749 characters extracted from source content.
Expand or collapse full text
Towards joint scaling laws with optimal batch size schedules Jiaxiang Li 1 , Zhiqi Bu 1 , Shiyun Xu 2 1 Meta, 2 Independent researcher 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. 1 Introduction Large-scale model training has been central to recent progress in artificial intelligence (Kaplan et al., 2020; Hoffmann et al., 2022a), but its effectiveness depends critically on robust choices of hyperparameters, particularly the learning rate and batch size (BS). The learning rate controls the magnitude of parameter updates: overly small values slow optimization, whereas overly large values can destabilize training and cause loss divergence (Yang et al., 2021; Li et al., 2025). The batch size controls a trade-off between model quality and systems efficiency: larger batches can improve throughput by exploiting data parallelism, amortizing communication, and increasing hardware utilization (Shallue et al., 2019; Goyal et al., 2017; Rajbhandari et al., 2020); however, beyond a critical regime, increasing the batch size can degrade model performance under fixed data or compute budgets, because it reduces the number of optimization steps and changes the gradient-noise profile (McCandlish et al., 2018; Golmant et al., 2018; Keskar et al., 2017). Consequently, hyperparameter selection is essential for efficient model training, yet exhaustive tuning is often prohibitively computationally expensive at frontier scale. This motivates principled methods for transferring or scaling hyperparameters across model sizes and training horizons, so that small-scale experiments can reliably inform large-scale ones (Kaplan et al., 2020; Li et al., 2025; Hu et al., 2024; Zhang et al., 2025; Bergsma et al., 2025a; Shuai et al., 2024; Bi et al., 2024). For the learning rate, the mainstream practice is to use a scheduleη t T t=1 rather than a fixed constant. Common schedules typically combine warmup, which mitigates early-training instability, with decay, which improves late-stage optimization (Vaswani et al., 2017; Godbole et al., 2023; Defazio et al., 2023; Schaipp et al., 2025). Under maximal update parametrization (μP), optimal learning rates can transfer across model sizes when the training horizon is fixed (Yang et al., 2021, 2024; Dey et al., 2026). Furthermore, when the training horizonTchanges, the optimal peak learning rate must be rescaled, often following a power law of the form η ⋆ (T )∝ T −α (Bi et al., 2024; Bjorck et al., 2025; Bu et al., 2026). In contrast, batch size has been less explored systematically as a scheduled hyperparameter. Existing approaches often use fixed batch sizes or hand-designed batch-size schedules coupled to the learning-rate schedule (Smith et al., 2018; Balles et al., 2017; Meterez et al., 2025). See Section E for a summarization of batch size schedules for recent LLMs. This is potentially suboptimal: a fixed batch size is only a constrained special case of a broader class of batch-size schedulesB t T t=1 . This gap motivates a principled study of batch-size scheduling and its transfer across model sizes and training horizons. In this work, we study the fundamental question: 1 arXiv:2607.27731v1 [cs.LG] 30 Jul 2026 What is the optimal batch size schedule and which factors determine it? Our approach builds on the stochastic convex optimization theory, capturing the training dynamics of deep learning (Defazio et al., 2023) which is non-convex and non-smooth. To be specific, we develop a principled and generalizable mapping in(2.6)that characterizes the loss sequence from both the learning rate sequence and the batch size sequence. This sequence-to-sequence characterization is empirically accurate and analytically tractable to formalize a constrained optimization problem: the objective is the last iteration of the loss sequence, the domain is batch size sequences, and the constraint is data or compute budgets. We solve this problem for the optimal batch size schedule in Theorem 3. In summary, our answer is that the optimal batch size schedule is a function of total computational budget and learning rate schedule shape (cosine, linear, WSD), but independent of peak learning rate, weight decay, model size, model architecture, or optimizer. As a result, the batch size becomes a principled scaling dimension rather than a separately tuned hyperparameter. Contributions: 1.Loss prediction by learning rate and batch size schedules. We derive an any-iteration loss characterization in(2.5)from any learning rate and batch size schedules. This characterization is empirically precise for general optimizers and model architectures, providing evidence for the convex dominance in deep learning. 2.Closed-form optimal batch size schedule. We derive the closed-form optimal batch size for any learning rate schedule in Theorem 3, given the data, compute or general budgets. Notably, the optimal batch size schedule is independent of the peak learning rate, thus decoupled from the learning rate tuning. 3.Joint scaling laws and universal dynamics. We build joint scaling laws that employ 1/ √ T-scaling of learning rate and weight decay, as well as our scale-free batch size schedule, with an improvement on compute efficiency by 6∼15% and zero computational overhead. Crucially, we show that the optimal batch size schedule preserves universal training dynamics, also known as supercollapse, rendering a robust and high-performing recipe for scaling. 4.Empirical ablations in various regimes. We validate the advantage of our batch size schedules, in pre-training and fine-tuning regimes, on Llama3 (dense), Qwen3 (MoE) and vision-language models up to 7B parameters, with AdamW and Muon optimizers, across different scales and configurations in Section C. Table 1 Dynamic batch size generally outperforms static batch size (marked in green) by lm-evaluation-harness v0.4 on the last-iterate. See Figure 9 for the evaluation on intermediate iterates. ModelBatch SizeLAMBADAHellaSwagARC-EasyARC-ChallengeWinoGrandePIQAOpenBookQABoolQ Llama3 (dense) static32.2342.9356.2229.2653.0367.8932.4061.89 dynamic35.3644.3356.3630.4653.9968.3433.2062.26 Qwen3 (MoE) static36.0645.6157.8332.5952.4968.7734.0061.47 dynamic37.7146.8059.0932.2555.3369.9736.2061.59 2Discrete loss prediction by sequences of batch size and learning rate We start with the stochastic gradient descent (SGD) as w t+1 = w t − η t g t , where w t is the parameters,η t is the learning rate, g t = g(w t ) is the mini-batch gradient with batch sizeB t , and 1≤ t≤ TwithTtotal iterations. In what follows, we denote w ∗ as the minimizer of loss L, and D =∥w t − w ∗ ∥. 2.1 From learning rate sequence to loss sequence The prediction from learning rate to loss value has been investigated under convexity or strong convexity loss, Lipschitz continuity or smoothness, and averaged or current iterate. We follow the work of Defazio et al. 2 0200040006000800010000 Iteration 2.8 3.0 3.2 3.4 3.6 3.8 4.0 Test loss -62.5% iter Static BS (baseline) Dynamic BS (baseline) Dynamic BS (large batch) 0.500.751.001.251.501.752.002.252.50 Tokens 1e9 2.8 2.9 3.0 3.1 3.2 3.3 3.4 3.5 Test loss Static BS (baseline) Dynamic BS (baseline) Dynamic BS (large batch) 01000200030004000500060007000 Iteration 2.6 2.7 2.8 2.9 3.0 3.1 3.2 Test loss -37.5% iter Static BS (baseline) Dynamic BS (baseline) Dynamic BS (large batch) 0.40.60.81.01.21.4 Tokens 1e10 2.6 2.7 2.8 2.9 3.0 Test loss Static BS (baseline) Dynamic BS (baseline) Dynamic BS (large batch) Figure 1 Dynamic batch size outperforms static batch size under fixed compute of 1e20 FLOPs on Fineweb-edu dataset: (I) for fixed training horizon, dynamic batch size schedule (same BS) achieves better loss than static baselines; (I) for the same loss, dynamic batch size schedule (larger BS) uses significantly fewer iterations than static baselines, with the same learning rate. Upper row: Llama3-7B (dense). Lower row: Qwen3-A1B (MoE). See the evaluation in Table 1 and Figure 9. (2023); Bu et al. (2026); Schaipp et al. (2025) and study the sequence-to-sequence prediction under convex loss and SGD in Theorem 1 (see proof in Section A.1). Theorem 1. For convex loss and any learning rate sequence η t , under SGD optimizer, we have EL(w τ ;η t )≤ L ∗ + D 2 2 P τ t=1 η t + 1 2 τ−1 X t=1 η 2 t E∥g t ∥ 2 P τ k=t+1 η k + 1 2 η τ E∥g τ ∥ 2 (2.1) We emphasize that although the theoretical basis for Theorem 1 is limited to convex loss and SGD optimizer, the empirical applicability generalizes to non-convex neural networks (e.g. LLM and VLM) and adaptive optimizers (e.g. AdamW and Muon) as extensively investigated in Bu et al. (2026). 2.2 From batch size sequence to loss sequence To understand the effect of batch size on the loss, we simply assume E∥g t ∥ 2 ≤ G 2 + X/B t (2.2) whereGupper bounds the norm of expected gradient∥Eg t ∥, andXupper bounds the trace of per-example gradient covariance matrix. We derive (2.2) and extend it to adaptive optimizers in Section A.2. 3 Theorem 2. Assuming (2.2), for convex loss and any learning rate sequence η t , we have EL(w τ ;η t ,B t )≤ L ∗ + D 2 2 P τ t=1 η t + 1 2 τ−1 X t=1 η 2 t (G 2 + X/B t ) P τ k=t+1 η k + η τ 2 (G 2 + X/B τ ) (2.3) Proof of Theorem 2. Substituting (2.2) to Theorem 1. 2.3 Precise characterization under linear regression We validate Theorem 2 and thus justify the assumptions(2.1)and(2.2)on the various sequences of learning rate and batch size in Figure 2, where we construct three variables x 1,τ = 1 2 P τ t=1 η t ,x 2,τ = 1 2 τ−1 X t=1 η 2 t P τ k=t+1 η k + η τ 2 ,x 3,τ = 1 2 τ−1 X t=1 η 2 t /B t P τ k=t+1 η k + η τ /B τ 2 then solve a linear regression 1 with L = [L 1 ,...,L T ]: min L ∗ ,D,G,X L ∗ + D 2 x 1 + G 2 x 2 + Xx 3 − L 2 2 (2.4) constant batch size, dynamic learning rate 0500010000150002000025000300003500040000 Iteration 0.000 0.002 0.004 0.006 0.008 0.010 0500010000150002000025000300003500040000 Iteration 0.000 0.002 0.004 0.006 0.008 0.010 0500010000150002000025000300003500040000 Iteration 0.000 0.002 0.004 0.006 0.008 0.010 0500010000150002000025000300003500040000 Iteration 0.000 0.002 0.004 0.006 0.008 0.010 0500010000150002000025000300003500040000 Iteration 2.50 2.75 3.00 3.25 3.50 3.75 4.00 Loss L * + D 2 x 1 L * + G 2 x 2 + Xx 3 Fitted 0500010000150002000025000300003500040000 Iteration 2.50 2.75 3.00 3.25 3.50 3.75 4.00 0500010000150002000025000300003500040000 Iteration 2.25 2.50 2.75 3.00 3.25 3.50 3.75 4.00 0500010000150002000025000300003500040000 Iteration 2.6 2.8 3.0 3.2 3.4 3.6 3.8 4.0 constant learning rate, dynamic batch size 0500010000150002000025000300003500040000 Iteration 250000 255000 260000 265000 270000 275000 0500010000150002000025000300003500040000 Iteration 100000 200000 300000 400000 500000 0500010000150002000025000300003500040000 Iteration 100000 200000 300000 400000 500000 0500010000150002000025000300003500040000 Iteration 100000 150000 200000 250000 300000 350000 400000 0500010000150002000025000300003500040000 Iteration 2.8 3.0 3.2 3.4 3.6 3.8 4.0 Loss L * + D 2 x 1 L * + G 2 x 2 + Xx 3 Fitted 0500010000150002000025000300003500040000 Iteration 2.6 2.8 3.0 3.2 3.4 3.6 3.8 4.0 0500010000150002000025000300003500040000 Iteration 2.25 2.50 2.75 3.00 3.25 3.50 3.75 4.00 0500010000150002000025000300003500040000 Iteration 2.8 3.0 3.2 3.4 3.6 3.8 4.0 Figure 2 Sequence-to-sequence prediction by Theorem 2 for Llama3-1B model on Fineweb-edu data, trained by Muon-NSGD optimizer. Top two rows:B t is fixed (B= 256∗1024) andη t is dynamic. Bottom two rows:η t is fixed (η = 0.01) and B t is dynamic. Similar patterns for AdamW optimizer can be observed in Figure 10. 1 IfB t is constant, then the features(x 2 ,x 3 ) are collinear and perfectly correlated, rendering the coefficients(G, X) indistin- guishable, although the linear regression is solvable. 4 As shown in Figure 2, for all cases, a precise characterization emerges after an initial transient phase. This indicates that the training dynamics become dominated by convexity, and that the inequality in Theorem 2 may be replaced by equality in practice 2 : EL(w τ )≈ L ∗ + D 2 2 P τ t=1 η t + 1 2 τ−1 X t=1 η 2 t (G 2 + X/B t ) P τ k=t+1 η k + η τ 2 (G 2 + X/B τ ) (2.5) Specifically, if the learning rate decays to zero (η T = 0), then for the last iterate, EL(w T )≈ L ∗ + D 2 2 P T t=1 η t + 1 2 T−1 X t=1 η 2 t (G 2 + X/B t ) P T k=t+1 η k (2.6) 3 Continuous loss prediction and optimal batch size schedule Now we optimize the batch size with respect to the last-iterate loss in(2.6), subject to some constraints on the data size, compute budget, wall-clock training time, etc. To derive closed-form solutions, we will work in the continuous regime and replace the summations in (2.6) by integrals: L ∗ + D 2 2 R T 0 η t dt + G 2 2 Z T 0 η 2 t R T t η k dk ! dt + X 2 Z T 0 η 2 t /B t R T t η k dk ! dt(3.1) For static batch size B t = B static , we get L static (T ) := L ∗ + D 2 2 R T 0 η t dt + G 2 2 Z T 0 η 2 t R T t η k dk ! dt + X 2B static Z T 0 η 2 t R T t η k dk ! dt(3.2) 3.1 Data-optimal and compute-optimal batch size schedule We optimize the batch size given a fixed amount of data budgetK=B static T, while the learning rate schedule and the number of total iterationsTare fixed. This is equivalent to optimizing the batch size given a fixed compute budget C = 6NK (FLOPs), where N is model size. It suffices to only minimize the last term of (3.1) with respect to B t : min B t Z T 0 1 B t η 2 t R T t η k dk ! dt, s.t. Z T 0 B t dt = K(3.3) We solve this problem in Theorem 3 with a proof in Section A.3. Theorem 3. The closed-form solution to (3.3) is B optim t = K· b t R T 0 b t dt = K 2 η t q R T 0 η k dk R T t η k dk where b t := η t q R T t η k dk (3.4) Importantly, the optimal batch size schedule only depends on the shape of the learning rate schedule, not the peak learning rate. Thus,B optim t is fully decoupled from learning rate tuning and scale-free in terms of model size (see Table 2 for the independence on η t , and Figure 11 for ablations on peak learning rate). To give some concrete examples, we summarize the optimal batch size schedules corresponding to different learning rate schedules, which are visualized in Figure 3 3 . The endpoint behavior ofB optim t is governed by 2 X/G 2 is essentially the critical batch sizeB simple defined in McCandlish et al. (2018), which we treat as time-independent in line with Merrill et al. (2025); Golmant et al. (2018). 3 We observe no empirical difference whether the learning rate warmup is calculated in B optim t or not. 5 the order at which the learning rate approaches zero: ifη t ∼(T − t) p ast→ T, thenb t =η t / q R T t η k dk ∼ (T − t) (p−1)/2 , so schedules with a nonzero terminal slope (p= 1, e.g. linear and the WSD cooldown) yield a constant terminal batch size, whereas schedules that flatten to zero (p= 2, e.g. cosine withη t ∝(T − t) 2 ) drive B optim t → 0. Remark 3.1. Our characterization departs from the common simplification that the training dynamics are governed by the ratioη t /B t alone (Smith et al., 2018; Goyal et al., 2017; McCandlish et al., 2018), under which learning rate decay could be freely traded for batch size increase. We test this directly in Section C.5: for configurations with matchedη t /B t trajectories and identical token budgets, realizing the schedule through the learning rate consistently outperforms realizing it through the batch size (by +0.017 to +0.743 validation loss; see Table 4). Henceη t /B t is not a sufficient statistic for the training dynamics, and we justify that optimizing the batch size schedule is an independent dimension rather than as a re-parameterization of the learning rate. Table 2 Optimal batch size schedules for different learning rate schedules. learning rateη t formulab t formulaB optim t formula constantη q η T−t B static 2 q T T−t cosine η 2 (cos( πt T ) + 1) p η 2 1+cos(πt/T) √ T−t−T sin(πt/T)/π B static cos 2 ( πt 2T ) √ 1− t T − 1 π sin πt T linearη(1− t/T ) q 2η T B static WSD ( ηif t≤ cT η T−t T−cT if t > cT q 2η T+cT−2t if t≤ cT q 2η T−cT if t > cT B static T √ (T+cT−2t)(T+cT) if t≤ cT B static √ 1−c 2 if t > cT 0200040006000800010000 Iteration 0.000 0.002 0.004 0.006 0.008 0.010 Learning rate 0200040006000800010000 Iteration 0.000 0.002 0.004 0.006 0.008 0.010 Learning rate 0200040006000800010000 Iteration 0.0096 0.0098 0.0100 0.0102 0.0104 Learning rate 0500010000150002000025000300003500040000 Iteration 0 200 400 600 800 1000 1200 Batch size 0500010000150002000025000300003500040000 Iteration 600 800 1000 1200 1400 1600 Batch size 0500010000150002000025000300003500040000 Iteration 0 2500 5000 7500 10000 12500 15000 17500 20000 Batch size Figure 3 Optimal batch size schedules (blue) by Theorem 3 and static schedules (orange) for cosine, WSD, and constant learning rate schedules. 3.2 Optimal batch size leverages variance reduction To demonstrate the theoretical advantage of the optimal batch size schedule, we firstly derive the last-iterate loss. Corollary 3.2. Substituting (3.4) back to (3.1), we have L dynamic (T ) := L ∗ + D 2 2 R T 0 η t dt + G 2 2 Z T 0 η 2 t R T t η k dk ! dt + X 2B static T Z T 0 η t q R T t η k dk dt 2 Secondly, we compareL dynamic in Corollary 3.2 toL static in(3.2), where the only difference is in the last 6 terms. This difference explains the advantage of optimal batch size through the lens of variance reduction, as we show in Corollary 3.3. Corollary 3.3. Denoting Z t := η t √ R T t η k dk , we show that the difference is always non-positive, L dynamic (T )− L static (T ) = XT 2B static R T 0 Z t dt T ! 2 − R T 0 Z 2 t dt T ≤ 0 Proof of Corollary 3.3. Jensen’s inequality, or Cauchy-Schwarz inequality applied to Z t and 1. Additionally, we derive three facts from Table 3. •Variance reduction: Our batch size schedule reduces the gradient covariance term with X B static . For example, dynamic batch size improves 1.061 X B static → X B static for cosine learning rate, and 2.47 X B static → 1.9 X B static for WSD learning rate with 10% cooldown. •Scaled learning rate: The optimal peak learning rate for our batch size schedule is 1/ √ T-scaled for qualified learning rate, by minimizing L(T )− L ∗ over η in the third column. •Optimal loss convergence rate: Our batch size schedule consequently achievesO(1/ √ T) loss conver- gence in the last column. Table 3 Loss prediction by dynamic and static batch size schedules for different learning rate schedules. learning ratebatch sizeL(T )− L ∗ optimal L(T )− L ∗ constant dynamic D 2 2Tη + η 2 G 2 lnT + 2ηX B static D q G 2 lnT T + 4X/B static 1 T static D 2 2Tη + η 2 G 2 lnT + ηX 2B static lnTD q G 2 lnT T + X/B static lnT T cosine dynamic D 2 Tη + 1.061ηG 2 + Xη B static 2D p 1.061G 2 + X/B static q 1 T static D 2 Tη + 1.061ηG 2 + 1.061 Xη B static 2D p 1.061G 2 + 1.061X/B static q 1 T WSD dynamic D 2 (1+c)Tη + ηG 2 h 1 + 1 2 ln 1+c 1−c i + ηX/B static (1 + c) 2D √ (1+c)T r h 1 + 1 2 ln 1+c 1−c i G 2 + (1 + c) X B static static D 2 (1+c)Tη + η(G 2 + X/B static ) h 1 + 1 2 ln 1+c 1−c i 2D √ (1+c)T q G 2 + X B static r 1 + 1 2 ln 1+c 1−c 3.3 General computational cost-optimal batch size schedule Theoretically, we can derive the optimal batch size for any general computational cost functionf, as an extension of (3.3), e.g. the wall-clock training time. min B Z T 0 η 2 t B t 1 R T t η s ds ! dt, s.t. Z T 0 f (B t )dt = K(3.5) Notice that whenf(B) =B, the following reduces to the data-optimal case in Theorem 3. We obtain the closed form solution to (3.5) as B ∗ t (λ) in Lemma 4 (see proof in Section A.5). Lemma 4. Denote g −1 (x) = x 2 f ′ (x), then we can determine λ and B ∗ t through Z T 0 f g η 2 t R T t η s ds /λ !! dt = K, B ∗ t = g η 2 t R T t η s ds /λ ! 7 0500010000150002000025000300003500040000 Iteration 2.65 2.70 2.75 2.80 2.85 2.90 2.95 3.00 Theoretical loss bound Dynamic BS Static BS 0500010000150002000025000300003500040000 Iteration 2.65 2.70 2.75 2.80 2.85 2.90 2.95 3.00 Theoretical loss bound Dynamic BS Static BS 0500010000150002000025000300003500040000 Iteration 2.70 2.75 2.80 2.85 2.90 2.95 3.00 Theoretical loss bound Dynamic BS Static BS 0500010000150002000025000300003500040000 Iteration 2.65 2.70 2.75 2.80 2.85 2.90 2.95 3.00 Test loss Dynamic BS Static BS 0500010000150002000025000300003500040000 Iteration 2.65 2.70 2.75 2.80 2.85 2.90 2.95 3.00 Test loss Dynamic BS Static BS 0500010000150002000025000300003500040000 Iteration 2.70 2.75 2.80 2.85 2.90 2.95 3.00 Test loss Dynamic BS Static BS 0.00.20.40.60.81.0 Tokens 1e10 2.65 2.70 2.75 2.80 2.85 2.90 2.95 3.00 Test loss Dynamic BS Static BS 0.00.20.40.60.81.0 Tokens 1e10 2.65 2.70 2.75 2.80 2.85 2.90 2.95 3.00 Test loss Dynamic BS Static BS 0.00.20.40.60.81.0 Tokens 1e10 2.70 2.75 2.80 2.85 2.90 2.95 3.00 Test loss Dynamic BS Static BS Figure 4 Dynamic batch size consistently outperforms static batch size for Llama3-1B model on Fineweb-edu data, trained by Muon-NSGD optimizer. Left to right: cosine, WSD, and constant learning rate. Relative improvement of best perplexity is 0.8%, 2.5%, and 14.6%, respectively. 4 Joint scaling laws and universal dynamics 4.1 Hyperparameter scaling With our batch size schedule in Theorem 3, we present the joint scaling laws of multiple hyperparameters, which extends the single scaling law of peak learning rate in Bu et al. (2026): withη ∗ (N,T) andκ ∗ (N,T) being the optimal peak learning rate and weight decay for model size N and horizon T, peak learning rate: η ∗ (N,T ) = η ∗ (N small ,T small )/ p T/T small optimal batch size: B optim t (T ) = total data 2 η ∗ t q R T 0 η ∗ k dk R T t η ∗ k dk weight decay: κ(N,T ) = κ ∗ (N small ,T small )/ p T/T small (4.1) We theoretically validate(4.1)in Theorem 5, showing that 1/ √ Tscaling of peak learning rate and weight decay leads to asymptotically optimal loss convergence. Theorem 5. Assuming(2.2), convex loss and bounded parameters, for SGD with weight decay w t+1 = (1− η t κ)w t − η t g t and for qualified schedule with η ref / √ T peak learning rate, we have EL(w T )− L ∗ = O(1/ √ T ) under static batch size and optimal batch size schedule if κ = O(1/ √ T ). We demonstrate on the Llama3 and Qwen3 MoE families that the optimal batch size schedule is highly beneficial, on a wide range of model sizes (from 0.1B to 7B) and compute budgets (up to 1e21 FLOPs). We visualize two experiments in Figure 5, one with fixed compute and the other with fixed TPP. For the static batch size, we determine the optimal learning rate by following Bu et al. (2026): the peak learning rate is scaled asη ref / √ T, independent of model size, whereη ref is tuned on small-scale runs. Weight decay is determined in the same manner. We emphasize that our batch size schedule directly uses the same learning rate and weight 8 decay tuned for the static batch size, rendering the comparison fair but conservative. Nevertheless, in both experiments, the optimal batch size significantly outperforms the static batch size, with 15% improvement of compute efficiency on Llama3 as well as 6% improvement on Qwen3 MoE (see the middle column). Especially, the performance is consistent for most iterations during training. 10 2 10 3 Model size (M) 14 15 16 17 18 19 20 21 Validation perplexity Dynamic BS (R 2 =0.998) Static BS (R 2 =0.999) 10 20 Compute FLOPs (C) 2.50 2.55 2.60 2.65 2.70 2.75 2.80 2.85 Validation loss 1.15x Dynamic BS (R 2 =0.999): 36.933 * C 0.058 Static BS (R 2 =1.000): 34.184 * C 0.056 0.00.51.01.52.02.5 Tokens 1e10 2.4 2.5 2.6 2.7 2.8 2.9 3.0 3.1 3.2 Validation Loss Dynamic BS Static BS 10 3 4 × 10 2 6 × 10 2 Model size (M) 13 14 15 16 17 Validation perplexity Dynamic BS (R 2 =1.000) Static BS (R 2 =0.999) 10 20 2 × 10 20 3 × 10 20 4 × 10 20 Compute FLOPs (C) 2.4 2.5 2.6 2.7 2.8 Validation loss 1.06x Dynamic BS (R 2 =0.999): 244.963 * C 0.097 Static BS (R 2 =0.998): 246.244 * C 0.097 02468 Tokens 1e10 2.4 2.5 2.6 2.7 2.8 2.9 3.0 3.1 3.2 Validation Loss Dynamic BS Static BS Figure 5 Dynamic batch size significantly outperforms static batch size for Llama3 (top row) and Qwen3 MoE (bottom row). Left: perplexity under fixed compute of 1e20 FLOPs. Quadratic functions are fitted between validation loss and log(model size) following Hoffmann et al. (2022b). Middle & Right: TPP=50 for Llama3 and TPP=100 for Qwen3 MoE. Power laws are fitted between validation loss and compute FLOPs. 4.2 Universal dynamics Crucially, our joint scaling laws lead to the universal dynamics (Qiu et al., 2025; Bergsma et al., 2025b), where different models share the same normalized loss curve when (I) TPP is fixed (governing depth, width, and training horizon). (I) Hyperparameters are scaled (governing learning rate and weight decay). In Figure 6, each loss curve is normalized on the training horizonx=t/Tand on the validation loss F(x) = L(w xT )−L ∗ L(w T )−L ∗ . In addition, we observe another universal dynamics when the loss is normalized with respect to tokensx= R t 0 B k dk R T 0 B k dk , besides to iterations. Notice that these two types of normalization are equivalent for static batch size but different for dynamic batch size. As a result, our joint scaling laws enable the robust prediction of future scaling and demonstrate the advantage of optimal batch size. 5 Post-training experiments We further test our batch size schedule on two post-training experiments—vision-language model (VLM) and math-domain fine-tuning. See additional experiment details in Section D. 5.1 Vision-language model We train vision-language models (VLM) with SmoLM2-360M-Instruct as language backbone and Siglip2-large as vision backbone, following nanoVLM codebase (Wiedmann et al., 2025) on the Cauldron dataset (Laurençon 9 0.00.20.40.60.81.0 Normalized iterations 1.0 1.1 1.2 1.3 1.4 1.5 Normalized loss 0.00.20.40.60.81.0 Normalized tokens 1.0 1.1 1.2 1.3 1.4 1.5 Normalized loss 0.00.20.40.60.81.0 Normalized iterations 1.0 1.1 1.2 1.3 1.4 1.5 Normalized loss Figure 6 Dynamic batch size has universal dynamics with respect to normalized iterations and normalize tokens, similar to the static batch size for which two types of normalization are equivalent. Each plot shows 4 runs of Qwen3 MoE with TPP=100 from 0.4 to 2B parameters, with lighter color being smaller model size . et al., 2024). We train with Muon-NSGD optimizer. The static batch size is 0.26M tokens for text and 512*512 pixels for image. MMMU ScienceQA TextVQA DocVQA InfoVQA OCRBench MME perception MME cognition 0.75 0.85 0.95 1.00 Dynamic BSStatic BS Coarse perception Fine-grained perception Instance reasoning Logical reasoning Math Science 0.75 0.85 0.95 1.00 Dynamic BSStatic BS Figure 7 Normalized radar comparison of dynamic vs. static batch size across VLM benchmarks. Left: MMStar subtasks. Right: other benchmarks (MMMU, ScienceQA, TextVQA, DocVQA, InfoVQA, OCRBench, and MME). Each axis is normalized bymax(Dynamic BS, Static BS) so that mixed-scale metrics are visually comparable. See Table 6 for absolute scores. Figure 7 and Table 6 compare the dynamic and static batch sizes across various VLM evaluation tasks, where the dynamic batch size outperforms the static baseline on 7 of 9 tasks, while the static baseline can be slightly better on some MMStar subtasks. Overall, we see that dynamic batch size is competitive or superior for VLM tasks, with no significant regression on any benchmark. 5.2 Math-domain fine-tuning We fine-tune Qwen3-0.6B (Team, 2025) onorca_mathdataset (Mitra et al., 2024) with AdamW optimizer and WSD learning rate schedule. The peak learning rate is selected via a sweep with a constant batch size. In Figure 8, dynamic batch size shows better overall trajectories on both validation loss and GSM8K accuracy than static batch size. 10 0100200300400500 Iteration 0.50 0.55 0.60 0.65 0.70 0.75 0.80 Validation loss Dynamic BS Static BS 0100200300400500 Iteration 39 40 41 42 43 44 45 46 GSM8K Accuracy (%) Dynamic BS Static BS 02468 Tokens 1e6 0.50 0.55 0.60 0.65 0.70 0.75 0.80 Validation loss Dynamic BS Static BS 012345678 Tokens 1e6 39 40 41 42 43 44 45 46 GSM8K Accuracy (%) Dynamic BS Static BS Figure 8 Dynamic batch size consistently outperforms static batch size on math-domain fine-tuning. Left: validation loss and 5-shot GSM8K accuracy over iterations. Right: the same metrics plotted against total tokens consumed. 6 Conclusion In this work, we validated that the convex optimization theory can characterize the training dynamics of deep learning in(2.5), from the perspectives of learning rate and batch size schedules. This allows us to reverse-engineer and derive an optimal batch size schedule, which is decoupled from the choice of learning rate (no hyperparameter tuning) and consistently outperforms static batch size. As a result, we build joint scaling laws that inform the scaling of hyperparameters, and lead to better scaling of model performance. Throughout the paper, we ablate the most representative configurations to test the optimal batch size schedule: optimizers (Muon-NSGD and AdamW), model families (Llama3 and Qwen3 MoE), model sizes and training horizons in Figure 5, and learning rate schedules in Figure 2 and Figure 4. We provide additional ablations in the Section C: (I) We observe that our batch size schedule always outperforms the static batch size for any learning rate, even the non-optimal ones. (I) While our batch size always outperforms the static baseline, the advantage vanishes in theory asB → ∞yet is practically significant even when the base batch size is increased 8×. (I) However, the advantage may vanish under purely low-precision training. 11 References Lucas Balles, Javier Romero, and Philipp Hennig. Coupling adaptive batch sizes with learning rates. In Proceedings of the Thirty-Third Conference on Uncertainty in Artificial Intelligence, UAI 2017, Sydney, Australia, August 11-15, 2017. AUAI Press, 2017. http://auai.org/uai2017/proceedings/papers/141.pdf. 1 Shane Bergsma, Nolan Simran Dey, Gurpreet Gosal, Gavia Gray, Daria Soboleva, and Joel Hestness. Power lines: Scaling laws for weight decay and batch size in LLM pre-training. In The Thirty-ninth Annual Conference on Neural Information Processing Systems, 2025a. https://openreview.net/forum?id=bFXbLQzRoZ. 1 Shane Bergsma, Bin Claire Zhang, Nolan Dey, Shaheer Muhammad, Gurpreet Gosal, and Joel Hestness. Scaling with collapse: Efficient and predictable training of llm families. arXiv preprint arXiv:2509.25087, 2025b. 9 Xiao Bi, Deli Chen, Guanting Chen, Shanhuang Chen, Damai Dai, Chengqi Deng, Honghui Ding, Kai Dong, Qiushi Du, Zhe Fu, et al. Deepseek llm: Scaling open-source language models with longtermism. arXiv preprint arXiv:2401.02954, 2024. 1 Johan Bjorck, Alon Benhaim, Vishrav Chaudhary, Furu Wei, and Xia Song. Scaling optimal lr across token horizons. In International Conference on Learning Representations, volume 2025, pages 83640–83657, 2025. 1 Valentyn Boreiko, Zhiqi Bu, and Sheng Zha. Towards understanding of orthogonalization in muon. In Tiny Titans: The next wave of On-Device Learning for Foundational Models (TTODLer-FM), 2025. 21 Tom B. Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, Sandhini Agarwal, Ariel Herbert-Voss, Gretchen Krueger, Tom Henighan, Rewon Child, Aditya Ramesh, Daniel M. Ziegler, Jeffrey Wu, Clemens Winter, Christopher Hesse, Mark Chen, Eric Sigler, Mateusz Litwin, Scott Gray, Benjamin Chess, Jack Clark, Christopher Berner, Sam McCandlish, Alec Radford, Ilya Sutskever, and Dario Amodei. Language Models are Few-Shot Learners. In Advances in Neural Information Processing Systems, volume 33, pages 1877–1901, 2020.https://proceedings.neurips.c/paper/2020/ hash/1457c0d6bfcb4967418bfb8ac142f64a-Abstract.html. 27 Zhiqi Bu, Shiyun Xu, and Jialin Mao. Convex dominance in deep learning i: A scaling law of loss and learning rate. In The Fourteenth International Conference on Learning Representations, 2026.https://openreview.net/forum?id= dSdLqg02tx. 1, 3, 8 Aakanksha Chowdhery, Sharan Narang, Jacob Devlin, Maarten Bosma, Gaurav Mishra, Adam Roberts, Paul Barham, Hyung Won Chung, Charles Sutton, Sebastian Gehrmann, Parker Schuh, Kensen Shi, Sasha Tsvyashchenko, Joshua Maynez, Abhishek Rao, Parker Barnes, Yi Tay, Noam Shazeer, Vinodkumar Prabhakaran, Emily Reif, Nan Du, Ben Hutchinson, Reiner Pope, James Bradbury, Jacob Austin, Michael Isard, Guy Gur-Ari, Pengcheng Yin, Toju Duke, Anselm Levskaya, Sanjay Ghemawat, Sunipa Dev, Henryk Michalewski, Xavier Garcia, Vedant Misra, Kevin Robinson, Liam Fedus, Denny Zhou, Daphne Ippolito, David Luan, Hyeontaek Lim, Barret Zoph, Alexander Spiridonov, Ryan Sepassi, David Dohan, Shivani Agrawal, Mark Omernick, Andrew M. Dai, Thanumalayan Sankaranarayana Pillai, Marie Pellat, Aitor Lewkowycz, Erica Moreira, Rewon Child, Oleksandr Polozov, Katherine Lee, Zongwei Zhou, Xuezhi Wang, Brennan Saeta, Mark Diaz, Orhan Firat, Michele Catasta, Jason Wei, Kathy Meier-Hellstern, Douglas Eck, Jeff Dean, Slav Petrov, and Noah Fiedel. PaLM: Scaling Language Modeling with Pathways. Journal of Machine Learning Research, 24(240):1–113, 2023. https://jmlr.org/papers/v24/22-1144.html. 27 DeepSeek-AI. DeepSeek-V4: Towards Highly Efficient Million-Token Context Intelligence, 2026.https://huggingface. co/deepseek-ai/DeepSeek-V4-Pro/blob/main/DeepSeek_V4.pdf. Technical report. 27 Aaron Defazio, Ashok Cutkosky, Harsh Mehta, and Konstantin Mishchenko. Optimal linear decay learning rate schedules and further refinements. arXiv preprint arXiv:2310.07831, 2023. 1, 2, 15 Nolan Dey, Bin Zhang, Lorenzo Noci, Mufan Li, Blake Bordelon, Shane Bergsma, Cengiz Pehlevan, Boris Hanin, and Joel Hestness. Don’t be lazy: Completep enables compute-efficient deep transformers. Advances in Neural Information Processing Systems, 38:137707–137739, 2026. 1 Varun Godbole, George E. Dahl, Justin Gilmer, Christopher J. Shallue, and Zachary Nado. Deep learning tuning playbook, 2023. http://github.com/google-research/tuning_playbook. Version 1.0. 1 Noah Golmant, Nikita Vemuri, Zhewei Yao, Vladimir Feinberg, Amir Gholami, Kai Rothauge, Michael W Mahoney, and Joseph Gonzalez. On the computational inefficiency of large batch sizes for stochastic gradient descent. arXiv preprint arXiv:1811.12941, 2018. 1, 5 12 Priya Goyal, Piotr Dollár, Ross Girshick, Pieter Noordhuis, Lukasz Wesolowski, Aapo Kyrola, Andrew Tulloch, Yangqing Jia, and Kaiming He. Accurate, large minibatch sgd: Training imagenet in 1 hour. arXiv preprint arXiv:1706.02677, 2017. 1, 6, 22 Aaron Grattafiori, Abhimanyu Dubey, Abhinav Jauhri, Abhinav Pandey, Abhishek Kadian, Ahmad Al-Dahle, Aiesha Letman, Akhil Mathur, Alan Schelten, Alex Vaughan, et al. The llama 3 herd of models. arXiv preprint arXiv:2407.21783, 2024. 27 Jordan Hoffmann, Sebastian Borgeaud, Arthur Mensch, Elena Buchatskaya, Trevor Cai, Eliza Rutherford, Diego de Las Casas, Lisa Anne Hendricks, Johannes Welbl, Aidan Clark, et al. Training compute-optimal large language models. arXiv preprint arXiv:2203.15556, 2022a. 1 Jordan Hoffmann, Sebastian Borgeaud, Arthur Mensch, Elena Buchatskaya, Trevor Cai, Eliza Rutherford, Diego de Las Casas, Lisa Anne Hendricks, Johannes Welbl, Aidan Clark, Tom Hennigan, Eric Noland, Katie Millican, George van den Driessche, Bogdan Damoc, Aurelia Guy, Simon Osindero, Karen Simonyan, Erich Elsen, Jack W. Rae, Oriol Vinyals, and Laurent Sifre. Training Compute-Optimal Large Language Models. In Advances in Neural Information Processing Systems, volume 35, 2022b.https://proceedings.neurips.c/paper_files/paper/2022/file/ c1e2faff6f588870935f114ebe04a3e5-Paper-Conference.pdf. 9, 27 Shengding Hu, Yuge Tu, Xu Han, Chaoqun He, Ganqu Cui, Xiang Long, Zhi Zheng, Yewei Fang, Yuxiang Huang, Weilin Zhao, et al. Minicpm: Unveiling the potential of small language models with scalable training strategies. arXiv preprint arXiv:2404.06395, 2024. 1 Jared Kaplan, Sam McCandlish, Tom Henighan, Tom B Brown, Benjamin Chess, Rewon Child, Scott Gray, Alec Radford, Jeffrey Wu, and Dario Amodei. Scaling laws for neural language models. arXiv preprint arXiv:2001.08361, 2020. 1 Nitish Shirish Keskar, Dheevatsa Mudigere, Jorge Nocedal, Mikhail Smelyanskiy, and Ping Tak Peter Tang. On large-batch training for deep learning: Generalization gap and sharp minima. In International Conference on Learning Representations, 2017. 1 Kimi Team. Kimi K2: Open Agentic Intelligence, 2025. https://arxiv.org/abs/2507.20534. 27 Hugo Laurençon, Léo Tronchon, Matthieu Cord, and Victor Sanh. What matters when building vision-language models?, 2024. 9 Houyi Li, Wenzhen Zheng, Jingcheng Hu, Qiufeng Wang, Hanshan Zhang, Zili Wang, Shijie Xuyang, Yuantao Fan, Shuigeng Zhou, Xiangyu Zhang, et al. Predictable scale: Part i–optimal hyperparameter scaling law in large language model pretraining. arXiv e-prints, pages arXiv–2503, 2025. 1 Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization. In International Conference on Learning Representations, 2019. 21 Sam McCandlish, Jared Kaplan, Dario Amodei, and OpenAI Dota Team. An empirical model of large-batch training. arXiv preprint arXiv:1812.06162, 2018. 1, 5, 6, 22 William Merrill, Shane Arora, Dirk Groeneveld, and Hannaneh Hajishirzi. Critical batch size revisited: A simple empirical approach to large-batch language model training. In The Thirty-ninth Annual Conference on Neural Information Processing Systems, 2025. 5 Alexandru Meterez, Depen Morwani, Jingfeng Wu, Costin-Andrei Oncescu, Cengiz Pehlevan, and Sham M. Kakade. Seesaw: Accelerating training by balancing learning rate and batch size scheduling, 2025. 1 Arindam Mitra, Hamed Khanpour, Corby Rosset, and Ahmed Awadallah. Orca-math: Unlocking the potential of slms in grade school math, 2024. 10 Shikai Qiu, Lechao Xiao, Andrew Gordon Wilson, Jeffrey Pennington, and Atish Agarwala. Scaling collapse reveals universal dynamics in compute-optimally trained neural networks. In International Conference on Machine Learning, pages 50697–50720. PMLR, 2025. 9 Jack W. Rae, Sebastian Borgeaud, Trevor Cai, Katie Millican, Jordan Hoffmann, Francis Song, John Aslanides, Sarah Henderson, Roman Ring, Susannah Young, et al. Scaling Language Models: Methods, Analysis & Insights from Training Gopher, 2021. https://arxiv.org/abs/2112.11446. 27 Samyam Rajbhandari, Jeff Rasley, Olatunji Ruwase, and Yuxiong He. Zero: memory optimizations toward training trillion parameter models. In Proceedings of the International Conference for High Performance Computing, Networking, Storage and Analysis, pages 1–16, 2020. 1 13 Fabian Schaipp, Alexander Hägele, Adrien Taylor, Umut Simsekli, and Francis Bach. The surprising agreement between convex optimization theory and learning-rate scheduling for large model training. In International Conference on Machine Learning, 2025. 1, 3 Christopher J Shallue, Jaehoon Lee, Joseph Antognini, Jascha Sohl-Dickstein, Roy Frostig, and George E Dahl. Measuring the effects of data parallelism on neural network training. Journal of Machine Learning Research, 20 (112):1–49, 2019. 1 Xian Shuai, Yiding Wang, Yimeng Wu, Xin Jiang, and Xiaozhe Ren. Scaling law for language models training considering batch size. arXiv preprint arXiv:2412.01505, 2024. 1 Samuel L. Smith, Pieter-Jan Kindermans, and Quoc V. Le. Don’t decay the learning rate, increase the batch size. In International Conference on Learning Representations, 2018.https://openreview.net/forum?id=B1Yy1BxCZ. 1, 6, 22 Shaden Smith, Mostofa Patwary, Brandon Norick, Patrick LeGresley, Samyam Rajbhandari, Jared Casper, Zhun Liu, Shrimai Prabhumoye, George Zerveas, Vijay Korthikanti, et al. Using deepspeed and megatron to train megatron-turing nlg 530b, a large-scale generative language model. arXiv preprint arXiv:2201.11990, 2022. 27 Qwen Team. Qwen3 technical report, 2025. https://arxiv.org/abs/2505.09388. 10 Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Łukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural information processing systems, 30, 2017. 1 Luis Wiedmann, Aritra Roy Gosthipaty, and Andrés Marafioti. nanovlm.https://github.com/huggingface/nanoVLM, 2025. 9 An Yang, Anfeng Li, Baosong Yang, Beichen Zhang, Binyuan Hui, Bo Zheng, Bowen Yu, Chang Gao, Chengen Huang, Chenxu Lv, et al. Qwen3 technical report. arXiv preprint arXiv:2505.09388, 2025. 27 Ge Yang, Edward Hu, Igor Babuschkin, Szymon Sidor, Xiaodong Liu, David Farhi, Nick Ryder, Jakub Pachocki, Weizhu Chen, and Jianfeng Gao. Tensor programs v: Tuning large neural networks via zero-shot hyperparameter transfer. Advances in Neural Information Processing Systems, 34:17084–17097, 2021. 1 Greg Yang, Dingli Yu, Chen Zhu, and Soufiane Hayou. Tensor programs vi: Feature learning in infinite depth neural networks. In International Conference on Learning Representations, volume 2024, pages 55099–55150, 2024. 1 Aohan Zeng, Xiao Liu, Zhengxiao Du, Zihan Wang, Hanyu Lai, Ming Ding, Zhuoyi Yang, Yifan Xu, Wendi Zheng, Xiao Xia, Weng Lam Tam, Zixuan Ma, Yufei Xue, Jidong Zhai, Wenguang Chen, Peng Zhang, Yuxiao Dong, and Jie Tang. GLM-130B: An Open Bilingual Pre-trained Model. In The Eleventh International Conference on Learning Representations, 2023. https://openreview.net/forum?id=-Aw0rrrPUF. 27 Hanlin Zhang, Depen Morwani, Nikhil Vyas, Jingfeng Wu, Difan Zou, Udaya Ghai, Dean Foster, and Sham M Kakade. How does critical batch size scale in pre-training? In The Thirteenth International Conference on Learning Representations, 2025. 1 14 Appendix A Theoretical results A.1 Theorem 1 Proof of Theorem 1. Theorem 10 in Defazio et al. (2023) gives an upper bound of the loss EL(w τ )≤ L ∗ + D 2 2 P T t=1 η t + 1 2 P τ t=1 η 2 t E∥g t ∥ 2 P τ t=1 η t + τ−1 X k=1 η k P τ t=k+1 η t P τ t=k η 2 t E∥g t ∥ 2 P τ t=k η t ! (A.1) To simplify the notation, we write x t = η 2 t E∥g t ∥ 2 ,A k = η k P τ t=k+1 η t P τ t=k η t = 1 P τ t=k+1 η t − 1 P τ t=k η t so that (A.1) becomes L ∗ + D 2 2 P τ t=1 η t + 1 2 1 P τ t=1 η t τ X t=1 x t + 1 2 τ−1 X k=1 A k τ X t=k x t ! =L ∗ + D 2 2 P τ t=1 η t + 1 2 1 P τ t=1 η t τ−1 X t=1 x t + 1 2 x τ P τ t=1 η t + 1 2 τ−1 X k=1 A k τ X t=k x t ! =L ∗ + D 2 2 P τ t=1 η t + 1 2 1 P τ t=1 η t τ−1 X t=1 x t + 1 2 x τ 1 P τ t=1 η t + 1 2 x τ τ−1 X k=1 A k + 1 2 τ−1 X k=1 A k τ−1 X t=k x t ! We exchange the double sum to yield P τ−1 k=1 A k P τ−1 t=k x t = P τ−1 t=1 x t P t k=1 A k , and obtain L ∗ + D 2 2 P τ t=1 η t + 1 2 x τ 1 P τ s=1 η s + τ−1 X k=1 A k ! + 1 2 τ−1 X t=1 x t 1 P τ s=1 η s + t X k=1 A k ! which can be further simplified to L ∗ + D 2 2 P τ t=1 η t + 1 2 x τ η τ + 1 2 τ−1 X t=1 x t 1 P τ s=t+1 η s because the telescoping sum gives P t k=1 A k = 1 P τ s=t+1 η s − 1 P τ s=1 η s . A.2 Equation (2.2) According to the mean-covariance decomposition, E∥g t ∥ 2 =∥Eg t ∥ 2 + tr(Cov(g t )).(A.2) LetGbe the upper bound of∥Eg t ∥andXbe the upper bound of trace of Σ t :=Cov ξ (g(w t ;ξ)) which is the per-example stochastic gradient at iterationt, whereξis a data sample drawn from the data distribution. Since Cov(g t ) = Cov 1 B t B t X i=1 g(w t ;ξ t,i ) ! = 1 B t Σ t , it gives tr(Cov(g t )) = tr( 1 B t Σ t )≤ X B t . 15 Therefore, E∥g t ∥ 2 =∥Eg t ∥ 2 + tr(Cov(g t ))≤ G 2 + X B t .(A.3) Remark A.1. For adaptive gradient descent w t+1 = w t − η t φ(g t ), suppose some upper boundsG,X, then we obtain E∥φ(g t )∥ 2 ≤ G 2 + X/B t + O p (1/B 2 t ) A.3 Theorem 3 Proof of Theorem 3. Using the notation of b t , we rewrite the minimization problem (3.3) as min B t Z T 0 b 2 t B t dt s.t. Z T 0 B t dt = K. The Lagrangian functional associated with the equality constraint is, with λ∈ R, Z T 0 b 2 t B t dt + λ Z T 0 B t dt− K ! = Z T 0 b 2 t B t + λB t dt− λK. Denote the integrand as ℓ t (B t ) := b 2 t B t + λB t . The first-order condition gives ∂ℓ t ∂B t =− b 2 t B 2 t + λ = 0 =⇒ B optim t = b t √ λ , Next, we determine λ from the constraint: Z T 0 B optim t dt = 1 √ λ Z T 0 b t dt = K =⇒ 1 √ λ = K R T 0 b t dt =⇒ B optim t = K· b t R T 0 b t dt . To make B optim t explicit on η t , notice that Z T 0 b t dt = 2 s Z T 0 η k dk. A.4 Table 2 and Table 3 Here we deriveb t ,B optim t andL(T). We omit the optimalL(T) because it is trivial to minimize the formula L ∗ + a/η + b· η. constant learning rate We start from η t = η. Firstly, we write b t := η t q R T t η s ds = r η T − t Secondly, we get Z T 0 b t dt = 2 p Tη Thirdly, we get B t = BT · b t R T 0 b t dt = BT 2 p T (T − t) Lastly, substituting B t and η t into the loss bound: L(T ) = L ∗ + D 2 2Tη + η 2 G 2 lnT + 2ηX B 16 linear learning rate We start from η t = η(1− t/T ). Firstly, we write b t := η t q R T t η s ds = η(1− t T ) q (T−t) 2 2T = r 2 T Secondly, we get B t = B Lastly, substituting B t and η t into the loss bound: L ∗ + D 2 Tη + η(G 2 + X/B) WSD We start from η t = ( ηif t≤ cT η T−t T−cT if t > cT . Firstly, we write b t := η t q R T t η s ds where Z T t η s ds = ( η 2 (T + cT − 2t) if t≤ cT η 2 (T−t) 2 (T−cT) if t > cT so that b t = q 2η T+cT−2t if t≤ cT q 2η T−cT if t > cT Secondly, we get Z T 0 b t dt = p 2η(T + cT ) Thirdly, we get B t = BT · b ∗ t R T 0 b ∗ t dt = BT √ (T+cT−2t)(T+cT) if t≤ cT BT √ (T−cT)(T+cT) = B √ 1−c 2 if t > cT Lastly, substituting B t and η t into the loss bound: L ∗ + D 2 η(1 + c)T + η 2 2 + ln 1 + c 1− c G 2 + η(1 + c) X B . cosine learning rate We start from η t = η 2 (1 + cos ( πt T )). Firstly, we write b t := η t q R T t η s ds = cos 2 πt 2T s η T 2 1− t T − 1 π sin πt T . Secondly, we get B t = BT · b t R T 0 b t dt = B cos 2 πt 2T q 1− t T − 1 π sin πt T , Thirdly, substituting B t and η t into the loss bound: L ∗ + D 2 Tη + 1.061ηG 2 + ηX B 17 A.5 Lemma 4 Proof of Lemma 4.Similar to Theorem 3, using the notation ofb t , we rewrite the minimization problem(3.5) as min B t Z T 0 b 2 t B t dt s.t. Z T 0 f (B t )dt = K. The Lagrangian functional associated with the equality constraint is, with λ∈ R, Z T 0 b 2 t B t dt + λ Z T 0 f (B t )dt− K ! = Z T 0 b 2 t B t + λf (B t ) dt− λK. Denote the integrand as ℓ t (B t ) := b 2 t B t + λf (B t ). The first-order condition gives ∂ℓ t ∂B t =− b 2 t B 2 t + λf ′ (B t ) = 0 =⇒ g −1 (B t ) = B 2 t f ′ (B t ) = b 2 t λ =⇒ B ∗ t = g b 2 t λ , Next, we determine λ from the constraint: Z T 0 f (B ∗ t )dt = K =⇒ Z T 0 f g η 2 t R T t η s ds /λ !! dt = K =⇒ B ∗ t = g η 2 t R T t η s ds /λ ! A.6 Theorem 5 Proof of Theorem 5.DenoteL κ (w) =L(w) + κ 2 ∥w∥ 2 , then SGD with weight decay onLis equivalent to SGD without weight decay onL κ , where the bounded gradient condition becomes∇L κ =∇L+κw =⇒ ∥∇L κ ∥ =∥∇L∥ + κ∥w∥≤ G + κW. It is not hard to obtain (c.f. (2.6)) EL κ (w T )− L κ,∗ ≤ D 2 2 P T t=1 η t + 1 2 T−1 X t=1 η 2 t ((G + κW ) 2 + X/B t ) P T k=t+1 η k as long as the learning rate decays to zero for any fixed T. Notice that weight decay does not impact the last termX/B t , rendering the same minimization problem as (3.3) and hence the same solution B optim t , i.e. our optimal batch size schedule. For qualified learning rate schedule and scaled peak learning rate, we get Θ(1/ √ T) loss by Table 3, where the only difference is on G. Therefore, under either static batch size or optimal batch size, we get EL κ (w T )− L κ,∗ = O(1/ √ T ) To translate the gap between regularized loss L κ to the gap between vanilla loss L, we need EL(w T )− L ∗ ≤ EL κ (w T )− L ∗ ≤ EL κ (w T )− L κ,∗ + L κ,∗ − L ∗ . The first difference has already been derived. The second difference is obvious by contradiction: L κ,∗ − L ∗ ≤ κ 2 ∥w ∗ ∥ 2 ≤ κ 2 W 2 All in all, we need EL(w T )− L ∗ ≤ O(1/ √ T ) + κ 2 W 2 which requires κ = O(1/ √ T ). 18 B Experiment details B.1 Figure 1, Table 1 We train with WSD learning rate schedule and cooldown at 0.9; peak learning rate is 0.02. The optimizer is Muon-NSGD. Default initialization in Huggingface. Llama3 model has 48 layer, hidden size 3072, intermediate size 3072*4, batch size 256, and sequence length 1024. Qwen3 MoE model has 3B total parameters, 32 layers, hidden size 2048, intermediate size 2048*3, moe intermediate size 2048/4, 16 experts, batch size 512, and sequence length 4096. 0200040006000800010000 Iteration 0 5 10 15 20 25 30 35 LAMBADA Dynamic BS Static BS 0.00.51.01.52.02.5 Tokens 1e9 0 5 10 15 20 25 30 35 LAMBADA Dynamic BS Static BS 0200040006000800010000 Iteration 25.0 27.5 30.0 32.5 35.0 37.5 40.0 42.5 45.0 HellaSwag Dynamic BS Static BS 0.00.51.01.52.02.5 Tokens 1e9 25.0 27.5 30.0 32.5 35.0 37.5 40.0 42.5 45.0 HellaSwag Dynamic BS Static BS 0200040006000800010000 Iteration 35 40 45 50 55 ARC-Easy Dynamic BS Static BS 0.00.51.01.52.02.5 Tokens 1e9 35 40 45 50 55 ARC-Easy Dynamic BS Static BS 0200040006000800010000 Iteration 22 24 26 28 30 ARC-Challenge Dynamic BS Static BS 0.00.51.01.52.02.5 Tokens 1e9 22 24 26 28 30 ARC-Challenge Dynamic BS Static BS 0200040006000800010000 Iteration 50 51 52 53 54 WinoGrande Dynamic BS Static BS 0.00.51.01.52.02.5 Tokens 1e9 50 51 52 53 54 WinoGrande Dynamic BS Static BS 0200040006000800010000 Iteration 54 56 58 60 62 64 66 68 PIQA Dynamic BS Static BS 0.00.51.01.52.02.5 Tokens 1e9 54 56 58 60 62 64 66 68 PIQA Dynamic BS Static BS 0200040006000800010000 Iteration 24 26 28 30 32 34 OpenBookQA Dynamic BS Static BS 0.00.51.01.52.02.5 Tokens 1e9 24 26 28 30 32 34 OpenBookQA Dynamic BS Static BS 0200040006000800010000 Iteration 48 50 52 54 56 58 60 62 BoolQ Dynamic BS Static BS 0.00.51.01.52.02.5 Tokens 1e9 48 50 52 54 56 58 60 62 BoolQ Dynamic BS Static BS Figure 9 Evaluation of intermediate iterates for Llama3-8B with dynamic or static batch size. 19 B.2 Figure 2, Figure 4 Llama3 model has 28 layer, hidden size 1024, intermediate size 3072, batch size 256, and sequence length 1024. Figure 10 is analogous to Figure 2 but uses AdamW optimizer. The losses are smoothed by time weighted EMA. constant batch size, dynamic learning rate 0500010000150002000025000300003500040000 Iteration 0.0000 0.0005 0.0010 0.0015 0.0020 0500010000150002000025000300003500040000 Iteration 0.0000 0.0005 0.0010 0.0015 0.0020 0500010000150002000025000300003500040000 Iteration 0.0000 0.0005 0.0010 0.0015 0.0020 0500010000150002000025000300003500040000 Iteration 0.0000 0.0005 0.0010 0.0015 0.0020 0500010000150002000025000300003500040000 Iteration 2.50 2.75 3.00 3.25 3.50 3.75 4.00 0500010000150002000025000300003500040000 Iteration 2.25 2.50 2.75 3.00 3.25 3.50 3.75 4.00 Loss L * + D 2 x 1 L * + G 2 x 2 + Xx 3 Fitted 0500010000150002000025000300003500040000 Iteration 2.25 2.50 2.75 3.00 3.25 3.50 3.75 4.00 0500010000150002000025000300003500040000 Iteration 2.50 2.75 3.00 3.25 3.50 3.75 4.00 constant learning rate, dynamic batch size 0500010000150002000025000300003500040000 Iteration 250000 255000 260000 265000 270000 275000 0500010000150002000025000300003500040000 Iteration 100000 200000 300000 400000 500000 0500010000150002000025000300003500040000 Iteration 100000 200000 300000 400000 500000 0500010000150002000025000300003500040000 Iteration 100000 150000 200000 250000 300000 350000 400000 0500010000150002000025000300003500040000 Iteration 2.8 3.0 3.2 3.4 3.6 3.8 4.0 Loss L * + D 2 x 1 L * + G 2 x 2 + Xx 3 Fitted 0500010000150002000025000300003500040000 Iteration 2.6 2.8 3.0 3.2 3.4 3.6 3.8 4.0 0500010000150002000025000300003500040000 Iteration 2.8 3.0 3.2 3.4 3.6 3.8 4.0 0500010000150002000025000300003500040000 Iteration 2.6 2.8 3.0 3.2 3.4 3.6 3.8 4.0 Figure 10 Sequence-to-sequence prediction by Theorem 2 for Llama3-1B model on Fineweb-edu data, trained by AdamW optimizer. Top two rows:B t is fixed (B= 256∗1024) andη t is dynamic. Bottom two rows:η t is fixed (η = 0.002) and B t is dynamic. B.3 Figure 5, Figure 6 We use WSD learning rate schedule with 10% cooldown and Muon-NSGD optimizer. For Figure 5 (left column), the peak learning rate is 0.01; for other figures, the peak learning rate is scaled as 2/ √ T. TPP is based on active parameters for MoE and total parameters for dense models. For Llama3 models, we experiment with • total parameters 111M, 8 layers, 512 hidden size • total parameters 274M, 8 layers, 1024 hidden size • total parameters 426M, 10 layers, 1280 hidden size • total parameters 507M, 10 layers, 1440 hidden size • total parameters 631M, 12 layers, 1536 hidden size • total parameters 898M, 14 layers, 1792 hidden size • total parameters 1.236B, 16 layers, 2048 hidden size 20 • total parameters 2.163B, 20 layers, 2560 hidden size • total parameters 3.569B, 28 layers, 2880 hidden size • total parameters 6.585B, 48 layers, 3072 hidden size The intermediate size is set to 4 times hidden size, batch size 256, and sequence length 1024. For Qwen3 MoE models, we experiment with • active parameters 438M, 8 layers, 1280 hidden size • active parameters 606M, 16 layers, 1536 hidden size • active parameters 826M, 24 layers, 1792 hidden size • active parameters 1.127B, 32 layers, 2048 hidden size The intermediate size is set to 3 times hidden size, MoE intermediate size is 1/4 hidden size, 16 experts, batch size 512, and sequence length 4096. C Ablation results C.1 Optimizers Our main optimizer is Muon-NSGD (Boreiko et al., 2025). We additionally experiment on AdamW optimizer (Loshchilov and Hutter, 2019) to confirm the applicability of Theorem 2 in Figure 10. We also trained VLM in Section 5.1 with AdamW. C.2 Peak learning rate We have ablatedO(1/ √ T) learning rate in Section 4, which is empirically the optimal learning rate. Further- more, we show in Figure 11 that optimal batch size always outperforms static batch size for any learning rate, larger or smaller than the optimal learning rate (which is 0.01 in this case). This holds because our batch size derived in Theorem 3 is decoupled from the peak learning rate. 0500010000150002000025000300003500040000 Iteration 2.7 2.8 2.9 3.0 3.1 3.2 3.3 3.4 Test loss Dynamic BS (LR=0.003) Static BS (LR=0.003) 0500010000150002000025000300003500040000 Iteration 2.7 2.8 2.9 3.0 3.1 3.2 3.3 3.4 Test loss Dynamic BS (LR=0.01) Static BS (LR=0.01) 0500010000150002000025000300003500040000 Iteration 2.7 2.8 2.9 3.0 3.1 3.2 3.3 3.4 Test loss Dynamic BS (LR=0.03) Static BS (LR=0.03) 0.00.20.40.60.81.0 Tokens 1e10 2.7 2.8 2.9 3.0 3.1 3.2 3.3 3.4 Test loss Dynamic BS (LR=0.003) Static BS (LR=0.003) 0.00.20.40.60.81.0 Tokens 1e10 2.7 2.8 2.9 3.0 3.1 3.2 3.3 3.4 Test loss Dynamic BS (LR=0.01) Static BS (LR=0.01) 0.00.20.40.60.81.0 Tokens 1e10 2.7 2.8 2.9 3.0 3.1 3.2 3.3 3.4 Test loss Dynamic BS (LR=0.03) Static BS (LR=0.03) 13.013.514.014.515.015.5 Test Perplexity Dynamic BS (LR=0.003) Static BS (LR=0.003) Dynamic BS (LR=0.01) Static BS (LR=0.01) Dynamic BS (LR=0.03) Static BS (LR=0.03) -4.1% -2.5% -1.4% Figure 11 Impact of batch size schedules for Llama3 1B model on Fineweb-edu data. Left to right: peak learning rate=0.003,0.01 (optimal), 0.03. C.3 Training precision By default, we use mixed-precision training. We also observe strong performance of optimal batch size schedule when using full-precision training. However, with half-precision training, the advantage of our batch size vanishes in comparison to the static batch size. 21 0500010000150002000025000300003500040000 Iteration 2.7 2.8 2.9 3.0 3.1 Test loss fp32 (Dynamic BS (fp32)) fp32 (Static BS (fp32)) 0500010000150002000025000300003500040000 Iteration 2.7 2.8 2.9 3.0 3.1 Test loss fp32/bf16 (Dynamic BS (fp32/bf16)) fp32/bf16 (Static BS (fp32/bf16)) 0500010000150002000025000300003500040000 Iteration 2.7 2.8 2.9 3.0 3.1 Test loss bf16 (Dynamic BS (bf16)) bf16 (Static BS (bf16)) 0.00.20.40.60.81.0 Tokens 1e10 2.7 2.8 2.9 3.0 3.1 Test loss fp32 (Dynamic BS (fp32)) fp32 (Static BS (fp32)) 0.00.20.40.60.81.0 Tokens 1e10 2.7 2.8 2.9 3.0 3.1 Test loss fp32/bf16 (Dynamic BS (fp32/bf16)) fp32/bf16 (Static BS (fp32/bf16)) 0.00.20.40.60.81.0 Tokens 1e10 2.7 2.8 2.9 3.0 3.1 Test loss bf16 (Dynamic BS (bf16)) bf16 (Static BS (bf16)) 13.013.514.014.515.015.5 Test Perplexity Dynamic BS (fp32) Static BS (fp32) Dynamic BS (fp32/bf16) Static BS (fp32/bf16) Dynamic BS (bf16) Static BS (bf16) -2.6% -2.5% 0.1% Figure 12 Impact of batch size schedules for Llama3 1B model on Fineweb-edu data. Left to right: full-precision (fp32), mixed precision (fp32/bf16), and half-precision (bf16). C.4 Base batch size We show in Figure 13 and Figure 14 that optimal batch size always outperforms static batch size for any base batch size, under fixed compute and fixed iterations. Note that the advantage vanishes asBapproaches infinity, because the advantage is inversely proportional to base batch size by Corollary 3.3. HereB=1X is 128k tokens. 01000020000300004000050000600007000080000 Iteration 2.7 2.8 2.9 3.0 3.1 3.2 Test loss Dynamic BS (B=1X) Static BS (B=1X) 0500010000150002000025000300003500040000 Iteration 2.7 2.8 2.9 3.0 3.1 Test loss Dynamic BS (B=2X) Static BS (B=2X) 02500500075001000012500150001750020000 Iteration 2.6 2.7 2.8 2.9 3.0 Test loss Dynamic BS (B=4X) Static BS (B=4X) 0200040006000800010000 Iteration 2.60 2.65 2.70 2.75 2.80 2.85 2.90 Test loss Dynamic BS (B=8X) Static BS (B=8X) 0.00.20.40.60.81.0 Tokens 1e10 2.7 2.8 2.9 3.0 3.1 3.2 Test loss Dynamic BS (B=1X) Static BS (B=1X) 0.00.20.40.60.81.0 Tokens 1e10 2.7 2.8 2.9 3.0 3.1 Test loss Dynamic BS (B=2X) Static BS (B=2X) 0.00.20.40.60.81.0 Tokens 1e10 2.6 2.7 2.8 2.9 3.0 Test loss Dynamic BS (B=4X) Static BS (B=4X) 0.00.20.40.60.81.0 Tokens 1e10 2.60 2.65 2.70 2.75 2.80 2.85 2.90 Test loss Dynamic BS (B=8X) Static BS (B=8X) 1011121314151617 Test Perplexity Dynamic BS (B=1X) Static BS (B=1X) Dynamic BS (B=2X) Static BS (B=2X) Dynamic BS (B=4X) Static BS (B=4X) Dynamic BS (B=8X) Static BS (B=8X) -2.5% -2.5% -0.9% -1.1% Figure 13 Impact of batch size schedules for Llama3 1B model on Fineweb-edu data, under fixed compute (larger batch uses fewer iterations). 0500010000150002000025000300003500040000 Iteration 2.8 2.9 3.0 3.1 3.2 3.3 Test loss Dynamic BS (B=1X) Static BS (B=1X) 0500010000150002000025000300003500040000 Iteration 2.7 2.8 2.9 3.0 3.1 Test loss Dynamic BS (B=2X) Static BS (B=2X) 0500010000150002000025000300003500040000 Iteration 2.6 2.7 2.8 2.9 3.0 Test loss Dynamic BS (B=4X) Static BS (B=4X) 0500010000150002000025000300003500040000 Iteration 2.5 2.6 2.7 2.8 2.9 Test loss Dynamic BS (B=8X) Static BS (B=8X) 012345 Tokens 1e9 2.8 2.9 3.0 3.1 3.2 3.3 Test loss Dynamic BS (B=1X) Static BS (B=1X) 0.00.20.40.60.81.0 Tokens 1e10 2.7 2.8 2.9 3.0 3.1 Test loss Dynamic BS (B=2X) Static BS (B=2X) 0.00.51.01.52.0 Tokens 1e10 2.6 2.7 2.8 2.9 3.0 Test loss Dynamic BS (B=4X) Static BS (B=4X) 01234 Tokens 1e10 2.5 2.6 2.7 2.8 2.9 Test loss Dynamic BS (B=8X) Static BS (B=8X) 1011121314151617 Test Perplexity Dynamic BS (B=1X) Static BS (B=1X) Dynamic BS (B=2X) Static BS (B=2X) Dynamic BS (B=4X) Static BS (B=4X) Dynamic BS (B=8X) Static BS (B=8X) -1.7% -2.5% -1.2% -1.4% Figure 14 Impact of batch size schedules for Llama3 1B model on Fineweb-edu data, under fixed iterations of 40000. C.5 The η/B ratio does not determine training dynamics A natural hypothesis, suggested by the linear scaling rule (Goyal et al., 2017; Smith et al., 2018) and the analysis in McCandlish et al. (2018), is that training dynamics are governed by the ratioη/Brather than by ηandBindividually. If true, one could freely trade learning rate schedules for batch size schedules as long as η(t)/B(t) remains unchanged. We design a controlled experiment to test this hypothesis directly. Setup. We train Qwen3-0.6B from scratch on FineWeb-EDU using AdamW with batch size 256 and sequence length 1024. Each run trains for 10,000 steps, consuming 2.6B tokens. We sweep over three learning rate schedules (WSD, cosine, linear) and three base learning rates (3×10 −4 , 10 −3 , 3×10 −3 ), withη min = 0.1η max . 22 Design. For each (schedule,η max ) pair, we run two configurations that produce the sameη/Btrajectory at every step: • Type 1 (LR schedule): η(t) follows the schedule, B = G is constant. The ratio is η(t)/G. •Type 2 (BS schedule):η=η const is fixed,B(t)∝1/η ref (t) so thatη const /B(t) =η ref (t)/G. Here η const = T/ P T−1 t=0 1/η ref (t) is chosen to match the total token budget exactly. Both configurations consume identical total tokens by construction. Ifη/Bis the governing quantity, Type 1 and Type 2 should yield the same loss. Results. Figure 15 and Figure 16 show the validation loss curves versus iteration and total tokens, respectively. Across all 9 configurations, the LR-schedule run (Type 1, blue) consistently outperforms the BS-schedule run (Type 2, orange), often by a significant margin. Table 4 summarizes the final validation loss: the gap ranges from +0.017 to +0.74 in favor of Type 1. The effect is largest at small learning rates and for cosine/linear schedules, where the BS-schedule run must use a very small constantη const to satisfy the token budget constraint. These results demonstrate thatη/Balone does not determine training dynamics. The learning rate schedule provides benefits beyond its effect on theη/Bratio—likely through its interaction with the optimizer state and loss landscape curvature. 300040005000600070008000900010000 Iteration 3.0 3.2 3.4 3.6 3.8 4.0 4.2 4.4 Test loss WSD, LR=0.003 LR schedule + constant BS Constant LR + inverse BS (same /B) 300040005000600070008000900010000 Iteration 3.0 3.2 3.4 3.6 3.8 4.0 4.2 4.4 Test loss cosine, LR=0.003 LR schedule + constant BS Constant LR + inverse BS (same /B) 300040005000600070008000900010000 Iteration 3.0 3.2 3.4 3.6 3.8 4.0 4.2 4.4 Test loss linear, LR=0.003 LR schedule + constant BS Constant LR + inverse BS (same /B) 300040005000600070008000900010000 Iteration 3.0 3.2 3.4 3.6 3.8 4.0 4.2 4.4 Test loss WSD, LR=0.001 LR schedule + constant BS Constant LR + inverse BS (same /B) 300040005000600070008000900010000 Iteration 3.0 3.2 3.4 3.6 3.8 4.0 4.2 4.4 Test loss cosine, LR=0.001 LR schedule + constant BS Constant LR + inverse BS (same /B) 300040005000600070008000900010000 Iteration 3.0 3.2 3.4 3.6 3.8 4.0 4.2 4.4 Test loss linear, LR=0.001 LR schedule + constant BS Constant LR + inverse BS (same /B) 300040005000600070008000900010000 Iteration 3.0 3.2 3.4 3.6 3.8 4.0 4.2 4.4 Test loss WSD, LR=0.0003 LR schedule + constant BS Constant LR + inverse BS (same /B) 300040005000600070008000900010000 Iteration 3.0 3.2 3.4 3.6 3.8 4.0 4.2 4.4 Test loss cosine, LR=0.0003 LR schedule + constant BS Constant LR + inverse BS (same /B) 300040005000600070008000900010000 Iteration 3.0 3.2 3.4 3.6 3.8 4.0 4.2 4.4 Test loss linear, LR=0.0003 LR schedule + constant BS Constant LR + inverse BS (same /B) Figure 15 Validation loss vs. iteration for Type 1 (LR schedule + constant BS, blue) and Type 2 (constant LR + inverse BS, orange), which share the sameη/Btrajectory. Columns: WSD, cosine, and linear schedules. Rows top to bottom:η max = 3×10 −3 , 10 −3 , 3×10 −4 . The LR-schedule run consistently achieves lower loss despite identicalη/B at every step. Curves smoothed with a 50-iteration running average. 23 0.751.001.251.501.752.002.252.50 Tokens 1e9 3.0 3.2 3.4 3.6 3.8 4.0 4.2 4.4 Test loss WSD, LR=0.003 LR schedule + constant BS Constant LR + inverse BS (same /B) 0.51.01.52.02.5 Tokens 1e9 3.0 3.2 3.4 3.6 3.8 4.0 4.2 4.4 Test loss cosine, LR=0.003 LR schedule + constant BS Constant LR + inverse BS (same /B) 0.51.01.52.02.5 Tokens 1e9 3.0 3.2 3.4 3.6 3.8 4.0 4.2 4.4 Test loss linear, LR=0.003 LR schedule + constant BS Constant LR + inverse BS (same /B) 0.751.001.251.501.752.002.252.50 Tokens 1e9 3.0 3.2 3.4 3.6 3.8 4.0 4.2 4.4 Test loss WSD, LR=0.001 LR schedule + constant BS Constant LR + inverse BS (same /B) 0.51.01.52.02.5 Tokens 1e9 3.0 3.2 3.4 3.6 3.8 4.0 4.2 4.4 Test loss cosine, LR=0.001 LR schedule + constant BS Constant LR + inverse BS (same /B) 0.51.01.52.02.5 Tokens 1e9 3.0 3.2 3.4 3.6 3.8 4.0 4.2 4.4 Test loss linear, LR=0.001 LR schedule + constant BS Constant LR + inverse BS (same /B) 0.751.001.251.501.752.002.252.50 Tokens 1e9 3.0 3.2 3.4 3.6 3.8 4.0 4.2 4.4 Test loss WSD, LR=0.0003 LR schedule + constant BS Constant LR + inverse BS (same /B) 0.51.01.52.02.5 Tokens 1e9 3.0 3.2 3.4 3.6 3.8 4.0 4.2 4.4 Test loss cosine, LR=0.0003 LR schedule + constant BS Constant LR + inverse BS (same /B) 0.51.01.52.02.5 Tokens 1e9 3.0 3.2 3.4 3.6 3.8 4.0 4.2 4.4 Test loss linear, LR=0.0003 LR schedule + constant BS Constant LR + inverse BS (same /B) Figure 16 Same comparison as Figure 15, plotted against total tokens consumed. Since both types consume exactly the same number of tokens, the x-axes are aligned. The LR-schedule advantage persists when controlling for compute budget. Table 4 Final validation loss for Type 1 (LR schedule) vs. Type 2 (BS schedule) across all configurations. ∆ is Type 2 minus Type 1; positive values indicate the LR schedule wins. η max Schedule Type 1 (LR sched.) Type 2 (BS sched.)∆ 3× 10 −3 WSD3.0153.033+0.019 3× 10 −3 Cosine3.0503.082+0.031 3× 10 −3 Linear3.0413.059+0.017 10 −3 WSD2.9953.035+0.039 10 −3 Cosine3.0423.093+0.051 10 −3 Linear3.0353.063+0.028 3× 10 −4 WSD3.0773.212+0.136 3× 10 −4 Cosine3.2213.964+0.743 3× 10 −4 Linear3.2393.655+0.416 24 D Further experiment details Table 5 summarizes the hyperparameters used in the post-training experiments of Section 5. Table 5 Hyperparameters for the VLM and math fine-tuning experiments in Section 5. For the LR sweep, we train with constant batch size and select the best LR by validation loss and downstream accuracy, then fix it for the dynamic-vs.-static comparison. HyperparameterVLMMath ModelSmoLM2+Siglip2 Qwen3-0.6B DatasetCauldronOrcaMath OptimizerMuon-NSGDAdamW LR scheduleWSDWSD Peak LR2e-31e-5 Min LR00 Weight decay0.010.01 Context length4096512 Base batch size6432 Total steps40000500 Table 6 reports the absolute scores for VLM benchmark tasks under dynamic and static batch size schedules. Table 6 Benchmark scores for dynamic and static batch sizes on VLM evaluation (§5.1). Bold text indicates the better value per benchmark. BenchmarkDynamic Static MMMU (acc)0.2411 0.2311 ScienceQA0.3954 0.3924 TextVQA0.3556 0.3397 DocVQA0.1847 0.1704 InfoVQA0.1099 0.1067 OCRBench0.2810 0.2720 MME-Perception704.03 653.82 MME-Cognition169.29 171.79 MMStar (average)0.2361 0.2377 Coarse Perception0.2782 0.2819 Fine-grained Perception 0.2287 0.2317 Instance Reasoning0.2457 0.2392 Logical Reasoning0.2335 0.2341 Math0.2082 0.2369 Science0.2222 0.2024 25 E How existing paper use batch size? In this section we summarize the learning rate and batch size schedules used in training recent open source models. We summarize the schedules each model uses in Table 7, also visualize them in Figure 17 and 18. It could be seen that, most of the recent models are using a non-constant batch-size schedule, however they mostly employ a step function batch size schedule heuristically. 26 ModelLearning-rate scheduleBatch-size scheduleWhere specified Llama 3 405BPeak LR 8×10 −5 . Linear warmup for 8,000 steps; cosine decay to 8× 10 −7 over 1.2M steps; final long-context annealing linearly decays LR to 0 over the last 40M tokens. 4M tokens/update initially; 8M after 252M tokens; 16M after 2.87T tokens. (Grattafiori et al., 2024, Sec. 3.4, p. 14–15). Qwen3 / Qwen3.5 / Qwen3.6 Numeric LR schedule not disclosed in the public report/pages I found. Qwen3 states that predicted optimal LR/batch strategies are obtained by scaling laws and that LR decay is accelerated in Stage 2. Numeric batch-size schedule not disclosed. Publicly stated stages: > 30T tokens at 4K context, then ∼ 5T higher-quality tokens at 4K, then hundreds of billions of tokens at 32K. (Yang et al., 2025, Sec. 3.2, p. 4). DeepSeek-V4 Flash / Pro Flash: 2,000-step linear warmup; LR 2.7× 10 −4 for most training; near the end, cosine decay to 2.7× 10 −5 . Pro: same strategy, peak 2.0× 10 −4 , end 2.0× 10 −5 . Exact decay-start token count not disclosed. Flash: batch size increases from a small size to 75.5M tokens/update, then stays there. Pro: maximum batch 94.4M tokens/update. Exact ramp endpoint not disclosed. (DeepSeek-AI, 2026, Sec. 4.2.2, p. 24–25). Kimi-K2MuonClip with WSD. 500-step warmup to 2× 10 −4 , i.e. about 500× 67M = 33.5B tokens; constant 2× 10 −4 until 10T tokens; cosine decay to 2× 10 −5 over the next 5.5T tokens; terminal/long-context phase decays from 2×10 −5 to 7×10 −6 , with decay shape not explicitly specified. Constant 67M tokens/update.(Kimi Team, 2025, Sec. 2.5, p. 9). GPT-3 175BPeak LR 0.6× 10 −4 = 6× 10 −5 . Linear warmup for 375M tokens; cosine decay to 10% of peak over 260B tokens; constant at 10% of peak through 300B tokens. Final batch 3.2M tokens/update. Batch linearly increases from 32K tokens/update to final size over the first 4–12B tokens, depending on model size; the 175B-specific ramp length is not separately stated. (Brown et al., 2020, Table 2.1, p. 8; App. B, p. 42). MT-NLG 530BPeak LR 5×10 −5 . Linear warmup for 1B tokens; cosine decay targeting 10% of peak over 340B tokens. Training stops at 270B tokens, so the LR does not reach the 10% target. Batch starts at 32 sequences and increases by 32 sequences until 1,920 sequences over first 12B tokens. With sequence length 2,048, this is 65,536→ 3.93M tokens/update. (Smith et al., 2022, Sec. 3.2, p. 10). PaLM 540BAdafactor. LR 10 −2 for first 10,000 steps, then inverse-square-root decay proportional to 1/ √ k, where k is the step number. Sequence length 2,048. Batch size 512 until step 50K, 1,024 until step 115K, and 2,048 until step 255K; equivalently about 1.05M, 2.10M, and 4.19M tokens/update. (Chowdhery et al., 2023, Sec. 5, p. 10). Chinchilla 70B Peak LR 1× 10 −4 . Uses Gopher training setup except listed changes: warmup from 10 −7 to peak over 1,500 steps, followed by 10× cosine decay. Batch 1.5M → 3M tokens/update, doubled midway. Since Chinchilla trains on 1.4T tokens, this is about 700B tokens at each batch size. (Hoffmann et al., 2022b, Sec. 4.1/Table 4, p. 9); (Rae et al., 2021, Sec. 3.2, p. 6). Gopher 280B Peak LR 4× 10 −5 . Warmup from 10 −7 to peak over first 1,500 steps, followed by 10× cosine decay. Batch 3M → 6M tokens/update. The Chinchilla paper states the batch is doubled midway; for 300B training tokens this gives roughly 150B tokens at each batch size. (Rae et al., 2021, Sec. 3.1– 3.2/Table 1, p. 6); (Hoffmann et al., 2022b, Table 4, p. 9). GLM-130BPeak LR 8× 10 −5 . Warmup from 10 −7 to peak over first 0.5% of samples, then 10× cosine decay to 8× 10 −6 . For 400B tokens, 0.5%≈ 2B tokens. Sequence length 2,048. Batch warmup from 192 to 4,224 sequences over first 2.5% of samples, then constant; equivalently 0.393M → 8.65M tokens/update, with the ramp lasting about 10B tokens. (Zeng et al., 2023, p. 5; Table 11, p. 48). Table 7 Pre-training learning-rate and batch-size schedules. Batch size is expressed as tokens per optimizer update unless otherwise noted. 27 LRLR normalizedBatch sizeBatch size normalized Llama 3 405B 0.02.55.07.510.012.515.0 Tokens (T) 0 2 4 6 8 LR 1e5 Llama 3 405B: LR 0.00.20.40.60.81.0 Normalized tokens 0.0 0.2 0.4 0.6 0.8 1.0 Normalized LR Llama 3 405B: LR normalized 051015 Tokens (T) 5.0 7.5 10.0 12.5 15.0 Batch (M tokens) Llama 3 405B: batch 0.00.20.40.60.81.0 Normalized tokens 0.0 0.2 0.4 0.6 0.8 1.0 Normalized batch Llama 3 405B: batch normalized Kimi-K2 051015 Tokens (T) 0.00000 0.00005 0.00010 0.00015 0.00020 LR Kimi-K2: LR 0.00.20.40.60.81.0 Normalized tokens 0.0 0.2 0.4 0.6 0.8 1.0 Normalized LR Kimi-K2: LR normalized 051015 Tokens (T) 64 66 68 70 Batch (M tokens) Kimi-K2: batch 0.00.20.40.60.81.0 Normalized tokens 0.0 0.2 0.4 0.6 0.8 1.0 Normalized batch Kimi-K2: batch normalized DeepSeek-V4 Pro (approximate) 0102030 Tokens (T) 0.00000 0.00005 0.00010 0.00015 0.00020 LR DeepSeek-V4 Pro approx: LR 0.00.20.40.60.81.0 Normalized tokens 0.0 0.2 0.4 0.6 0.8 1.0 Normalized LR DeepSeek-V4 Pro approx: LR normalized 0102030 Tokens (T) 20 40 60 80 Batch (M tokens) DeepSeek-V4 Pro approx: batch 0.00.20.40.60.81.0 Normalized tokens 0.0 0.2 0.4 0.6 0.8 1.0 Normalized batch DeepSeek-V4 Pro approx: batch normalized GPT-3 175B 0.000.050.100.150.200.250.30 Tokens (T) 0 2 4 6 LR 1e5 GPT-3 175B: LR 0.00.20.40.60.81.0 Normalized tokens 0.0 0.2 0.4 0.6 0.8 1.0 Normalized LR GPT-3 175B: LR normalized 0.000.050.100.150.200.250.30 Tokens (T) 0 1 2 3 Batch (M tokens) GPT-3 175B: batch 0.00.20.40.60.81.0 Normalized tokens 0.0 0.2 0.4 0.6 0.8 1.0 Normalized batch GPT-3 175B: batch normalized MT-NLG 530B 0.000.050.100.150.200.25 Tokens (T) 0 1 2 3 4 5 LR 1e5 MT-NLG 530B: LR 0.00.20.40.60.81.0 Normalized tokens 0.0 0.2 0.4 0.6 0.8 1.0 Normalized LR MT-NLG 530B: LR normalized 0.000.050.100.150.200.25 Tokens (T) 0 1 2 3 4 Batch (M tokens) MT-NLG 530B: batch 0.00.20.40.60.81.0 Normalized tokens 0.0 0.2 0.4 0.6 0.8 1.0 Normalized batch MT-NLG 530B: batch normalized Figure 17 Learning-rate and batch-size schedules used in large-scale language model pre-training, part 1. The data are collected from the corresponding existing publications specified in Table 7. Each row corresponds to one model. The four columns show, from left to right: learning rate in original coordinates, learning rate in normalized coordinates, batch size in original coordinates, and batch size in normalized coordinates. In normalized plots, both axes are scaled to lie in [0, 1]. 28 LRLR normalizedBatch sizeBatch size normalized PaLM 540B 0.00.20.40.60.8 Tokens (T) 0.002 0.004 0.006 0.008 0.010 LR PaLM 540B: LR 0.00.20.40.60.81.0 Normalized tokens 0.0 0.2 0.4 0.6 0.8 1.0 Normalized LR PaLM 540B: LR normalized 0.00.20.40.60.8 Tokens (T) 1 2 3 4 Batch (M tokens) PaLM 540B: batch 0.00.20.40.60.81.0 Normalized tokens 0.0 0.2 0.4 0.6 0.8 1.0 Normalized batch PaLM 540B: batch normalized Chinchilla 70B 0.000.250.500.751.001.25 Tokens (T) 0.00000 0.00002 0.00004 0.00006 0.00008 0.00010 LR Chinchilla 70B: LR 0.00.20.40.60.81.0 Normalized tokens 0.0 0.2 0.4 0.6 0.8 1.0 Normalized LR Chinchilla 70B: LR normalized 0.000.250.500.751.001.25 Tokens (T) 1.5 2.0 2.5 3.0 Batch (M tokens) Chinchilla 70B: batch 0.00.20.40.60.81.0 Normalized tokens 0.0 0.2 0.4 0.6 0.8 1.0 Normalized batch Chinchilla 70B: batch normalized Gopher 280B 0.000.050.100.150.200.250.30 Tokens (T) 0 1 2 3 4 LR 1e5 Gopher 280B: LR 0.00.20.40.60.81.0 Normalized tokens 0.0 0.2 0.4 0.6 0.8 1.0 Normalized LR Gopher 280B: LR normalized 0.000.050.100.150.200.250.30 Tokens (T) 3 4 5 6 Batch (M tokens) Gopher 280B: batch 0.00.20.40.60.81.0 Normalized tokens 0.0 0.2 0.4 0.6 0.8 1.0 Normalized batch Gopher 280B: batch normalized GLM-130B 0.00.10.20.30.4 Tokens (T) 0 2 4 6 8 LR 1e5 GLM-130B: LR 0.00.20.40.60.81.0 Normalized tokens 0.0 0.2 0.4 0.6 0.8 1.0 Normalized LR GLM-130B: LR normalized 0.00.10.20.30.4 Tokens (T) 0 2 4 6 8 Batch (M tokens) GLM-130B: batch 0.00.20.40.60.81.0 Normalized tokens 0.0 0.2 0.4 0.6 0.8 1.0 Normalized batch GLM-130B: batch normalized Figure 18 Learning-rate and batch-size schedules used in large-scale language model pre-training, part 2. The data are collected from the corresponding existing publications specified in Table 7. Each row corresponds to one model. The four columns show, from left to right: learning rate in original coordinates, learning rate in normalized coordinates, batch size in original coordinates, and batch size in normalized coordinates. In normalized plots, both axes are scaled to lie in [0, 1]. 29