Fast Catch-Up, Late Switching: Optimal Batch Size Scheduling via Functional Scaling Laws
Authors: Jinbo Wang
Binghui Li
Zhanpeng Zhou
Mingze Wang
Yuxuan Sun
Jiaqi Zhang
Xunliang Cai
Lei Wu
Venue: ICLR 2026 Poster (2026)
Keywords: Batch Size Scheduling; Training Dynamics
Primary area: learning theory
TL;DR: We study batch size schedule: uncover optimal batch size schedule, fast catch-up effect and later switch strategy from Functional Scaling Laws (FSL) theoretical framework, and bring our insights to LLM pre-training.

Abstract: Batch size scheduling (BSS) plays a critical role in large-scale deep learning training, influencing both optimization dynamics and computational efficiency. Yet, its theoretical foundations remain poorly understood.
In this work, we show that the **functional scaling law (FSL)** framework introduced in [Li et al. (2025a)](https://arxiv.org/abs/2509.19189) provides a principled lens for analyzing BSS. Specifically, we characterize the optimal BSS under a fixed data budget and show that its structure depends sharply on task difficulty. For easy tasks, optimal schedules keep increasing batch size throughout.  In contrast, for hard tasks, the optimal schedule  maintains small batch sizes for most of training and switches to large batches only in a late stage. 
To explain the emergence of late switching, we uncover a dynamical mechanism—the **fast catch-up** effect—which also manifests in large language model (LLM) pretraining. After switching from small to large batches, the loss rapidly aligns with the constant large-batch trajectory. Using FSL, we show that this effect stems from rapid forgetting of accumulated gradient noise, with the catch-up speed determined by task difficulty. Crucially, this effect implies that *large batches can be safely deferred to late training* without sacrificing performance, while substantially reducing data consumption. Finally,
extensive LLM pretraining experiments—covering both Dense and MoE architectures with up to **1.1B** parameters and **1T** tokens—validate our theoretical predictions. Across all settings, late-switch schedules consistently outperform constant-batch and early-switch baselines.

Published as a conference paper at ICLR 2026
FAST CATCH-UP, LATE SWITCHING: OPTIMAL BATCH SIZE SCHEDULING VIA FUNCTIONAL SCALING LAWS
Jinbo Wang' Binghui Li Zhanpeng Zhou², Mingze Wang', Yuxuan Sun',
Jiaqi Zhang*.,, Xunliang Cai & Lei Wu -2.6.
School of Mathematical Sciences, Peking University
"Center for Machine Learning Research, Peking University
3 School of Computer Science, Shanghai Jiao Tong University
"State Key Laboratory Cognitive Intelligence, University Science and Technology of China
$ Meituan, Beijing $ "A for Science Institute, Beijing
"wangjinbo@stu pku. edu. en, 1 lib wibinghuiiipii.. @pkuu. edu. cn
zhangjiacii 90me tt tnan. com, le i wwwmath.t..h.m pku edu. cn
ABSTRACT
Batch size scheduling (BSS) plays critical role in large-scale deep learning train- ing, influencing both optimization dynamics and computational efficiency. Yet, its theoretical foundations remain poorly understood. In this work, we show that the functional scaling law (FSL) framework introduced in Li et al. (2025a) provides principled lens for analyzing BSS. Specifically, we characterize the optimal BSS under a fixed data budget and show that its structure depends sharply on task diffi- culty. For easy tasks, optimal schedules keep increasing batch size throughout. In contrast, for hard tasks, the optimal schedule maintains small batch sizes for most of training and switches to large batches only in a late stage. To explain the emer- gence of late switching, we uncover dynamical mechanism the fast catch-up effect-which also manifests in large language model (LLM) pretraining. After switching from small to large batches, the loss rapidly aligns with the constant large-bate batch trajectory. Using FSL, we show that this effect stems from rapid for- getting of accumulated gradient noise, with the catch-up speed determined by task difficulty. Crucially, this effect implies that large batches can be safely deferred late training without sacrificing performance, while substantially reducing data consumption. Finally, extensive LLM pretraining experiments coviieintt-coverinng both Dense and MoE architectures with up to 1.1B parameters and 1T tokens-validatt our theoretical predictions. Across all settings, late-switch schedules consistently outperform constant-batte and early-switch baselines.
INTRODUCTION
Large language model (LLM) pretraining demands massive computational resources, making train- ing efficiency a central challenge. At scale, training efficiency depends critically on parallelism, and increasing the batch size directly improves hardware utilization and throughput (Goyal al., 2017: Brown et al., 2020; Hoffmann et al., 2022). Large-batch training has therefore become indispensable for scalable LLM pretraining.
However, using a constant large batch size throughout training is suboptimal in terms of sample efficiency (McCandlish al., 2018; Merrill al., 2025). From stochastic optimization perspective, the batch size determines the noise scale of stochastic gradients: each update can be viewed as the population gradient perturbed by noise whose variance decreases with the batch size. In the early stages of training, the optimization dynamics are signal-dominate, so aggressively reducing noise via large batches yields limited benefit while consuming more data. As training proceeds, signal weakens and the influence of gradient noise increases, making larger batches more effective
"quual contribution Corresponding author.

---

Published as a conference paper at ICLR 2026
1B MoE, val loss stop
linear regression, loss vs step
1.60
0.50
small batch
small botch
large batch
luge botth
4.50
small botch
barg batch 9.5k
umall butch- - large butch 3030
barge
0.40
small butch - large
validation 1ag
0.45
4.40
0.25
4.35
0.25
4.25
Fast Calcirug
Fast catch
4.15
0.15
10
20
30
50
0
2000
4:00
02000
99000
1000D 12000
num of steps housandss
steps
Figure 1: The fast catch-up effect when switching from a small to a large batch size. Left: Validation loss versus training steps for a 1B-parameter MoE model trained on approximately 0.4T tokens under four batch-size schedules: constant small batch, constant large batch, small-to-large with early switch, and small-to-large with late switch. Right: Validation loss versus training steps in the theoretical setting with s = 0.3 and B 1.5 (the hard-task regime), which demonstrates the same catch-up effect.
for improving iteration efficiency. This motivates batch size scheduling (BSS), i.e., dynamically increasing the batch size during training.
Indeed, BSS has become ubiquitous in industrial-scale LLM pretraining, adopted in models such as GPT-3 (Brown et al., 2020), PaLM (Chowdhery al., 2023), LLaMA-3 (Grattafiori et al., 2024), DeepSeek-V V3 (DeepSeek-Al et al., 2024b), MiniMax-01 (MiniMax et al., 2025), Nemotron-4 (Par- mar al., 2024; Nvidia et al., 2024), and GLM-4.5 (Zeng et al., 2025). This widespread adoption calls for principled understanding of how batch size scheduling shapes training dynamics and ef- ficiency. Yet existing analyses either focus on constant batch sizes (Ma et al., 2018; Zhang et al., 2025) or rely on empirical and heuristic insights (Smith et al., 2018; McCandlish et al., 2018; Mer- rill et al., 2025). As a result, current BSS design often depends on heuristic tuning or expensive large-scale experimentation.
The functional scaling law (FSL) framework introduced in Li et al. (2025a) provides a continuous- time modeling of how batch size and learning rate schedules affect loss dynamics. While originally derived for linear regression and kernel regression, FSL exhibits strong expressive power for mod- cling the loss dynamics of practical LLM pretraining. However, Li et al. (2025a) focuses solely on learning rate schedules. In this paper, we extend this framework to analyzing BSS. Our contributions are as follows.
Optimal batch size schedule. Under the FSL framework, we derive the optimal batch size schedule under a fixed data (or compute) budget. The optimal BSS depends critically on task difficulty: easy tasks favor a monotonically increasing schedule, while hard tasks require keep- ing batch sizes small for most of training, with growth deferred to late phase. This stable- growth strategy increases the number of optimization steps under a fixed data budget, which benefits hard tasks. Extending to practical few-stage schedules, we find that easy tasks again favor constant latge-batch training, whereas hard tasks demand a prolonged small-batch phase followed by a late switch to large batches.
The fast catch-up effect. To explain why large batches can be safely deferred for hard tasks, uncover striking and highly robust fast catch-up effect: when training switches from small to a large batch size, the loss rapidly collapses to that of the constant large-batch run. This phenomenon appears consistently across LLM pretraining experiments with diverse chitectures, model scales, and data regimes (see Figure 1). Using FSL, we further provide a theoretical explanation of this effect and quantitatively characterize how task difficulty governs the speed of catch-up.
Large-scale validation of late-switch superiority. The fast catch-up effect implies that large batches can be safely deferred to late training without sacrificing performance, while substan- tially reducing data consumption. We validate this principle through extensive LLM pretraining
2

---

Published as a conference paper at ICLR 2026
experiments spanning Dense and MoE architectures, model sizes from 50M to 1B parameters, and data scales from 10B to 1T tokens. Across all settings, stage-wise BSS with late switching consistently outperforms constant-baie batca and carly-switch baselines.
1.1 RELATED WORK
Neural scaling laws. Hestness al. (2017) first observed that the performance of deep learn- ing follows predictable power-law relationships with model and data size, a phenomenon later formalized as neural scaling laws (Kaplan et al., 2020). These laws have since become guiding principles for configuring large-scale training and been refined across architectures and training regimes (Henighan et al., 2020; Hoffmann et al., 2022; Kadra et al., 2023; Aghajanyan et al., 2023; Muennighoff et al., 2023; Tissue et al., 2024; Luo et al., 2025; Qiu et al., 2025), with parallel theo- retical efforts explaining their origins and mechanisms (Bordelon et al., 2024; Lin et al., 2024; Bahri et al., 2024; Paquette et al., 2024; Yan et al., 2025; Kunstner & Bach, 2025; Li et al., 2026a). In this work, we build on the framework of Li al. (2025a) to provide a scaling-law analysis of batch size scheduling.
Large-batch training and batch size scheduling. Large-batch training is essential for leverag- ing hardware parallelism at scale. Existing work largely focuses on static batch sizes, aiming to determine how large the batch size can be increased without sacrificing data efficiency, typically characterized by the critical batch size (McCandlish et al., 2018; Ma et al., 2018; Kaplan et al., 2020; Gray al., 2024; Zhang et al., 2025; Merrill et al., 2025). In practice, however, LLM pre- training routinely employs batch size schedules. Despite its prevalence, BSS has received far less theoretical attention than learning rate schedules (Defazio al., 2023; Hu et al., 2024; Hägele al., 2024). Existing analyses of BSS either rely on heuristic arguments (Smith et al., 2018; McCandlish al., 2018) framed as optimal control problems (Lee al., 2022; Zhao al., 2022; Perko, 2023), offering limited structural insight. In contrast, we develop a scaling- law-based theory of BSS that systematically explains empirical practice and yields new design principles.
One-pass SGD in kernel regression. The convergence of one-pass stochastic gradient descent (SGD) in kernel regression-often interpreted as high-dimensional linear regression-has been ex- tensively studied. In particular, Dieuleveut & Bach (2015); Mücke al. (2019) showed that aver- aged one-pass SGI) achieves the minimax-optimal rate D D=*0/(a3+1) in easy-task regimes and the rate D= in hard-task regimes. Subsequent work further established that the same rates can be at- tained by last iterate when combined with appropriate learning rate decay (Wu et al., 2022a; Lin et al., 2024; Li et al., 2026b). In contrast, we show that one-pass SGD with a constant learning rate, when coupled with properly designed BSS, achieves the same optimal rates.
2 PRELIMINARIES
Throughout the paper, the notation = indicates equivalence up to a constant factor, and S (resp. 2) indicates inequality up constant factor. For two nonnegative functions fig IR>0 -
Notation.
R>o. we write f(t x g() if there exist constants C1, C2 > 0, independent of t, such that Cif(t) S
g(t) S C2f(t), Vt 2
2.1 PEATURE-SPACE LINEAR REGRESSION
Let X' and denote the input domain and distribution, respectively. Labels are generated as y = fa( + with e ~ (0, 2²). We assume o Z and the target function fr is given by f*(x) := (0(x), 0*). Here, Φ X + N is a feature map and 0* E R (with N E N. u (2o)) is the unknown target parameter. We assume (fx) ~ (O H) with (1i)/1 denoting the eigenvalues of
H Eg (<))<(<)) in a decreasing order.
Assumption 2.1 (Power-law structures). The following two conditions hold:
(Capacity condition) X - eー for some 3 oo).
(Source condition) 10 マ -ー - (((+(()-1))) for some E 20).
3

---

Published as a conference paper at ICLR 2026
The capacity exponent 3 controls the decay rate of the eigenvalues. Smaller 3 corresponds to larger effective rank of the spectrum and thus higher model capacity. The source exponent s measures the alignment of the target function with the kernel eigenstructure: smaller corresponds to harder learning problems, with more energy concentrated in high-frequency components. 'hese capacity and source conditions are standard in the analysis of kernel methods and have recently been adopted in scaling-law studies (Paquette et al., 2024; Lin et al., 2024; Bordelon et al., 2025; Li et al., 2025a). A more detailed interpretation of the above setup provided in Appendix A.I.
One-pass SGD. We learn the target function S using a student model (æ;;0 = (d(æ)), 0) by minimizing the population risk R(0) マ (B_ - y) via one-pass SGD. At each iteration 1 & l I K, SGD samples a mini-batch S = {( vk.s)) BR and performs the update
BB
Dx+1 = 02
7
Vo
(f( 02) - )
(1)
BR
where 17 is constant learning rate and (B1; B2, Bxk denotes the batch size schedule (BSS). Notably, the iteration (1) can be rewritten as
OK+1 02 7 (VR (2) + [x)
where En denotes the gradient noise that follows E [[ 0, E[ E [(0,)/Br, where (() represents the covariance of gradient noise at 0 with the batch size 1. The learning performance is
measured using the excess risk: E(0) == R()) 2 2110 0 l where || = HE.
2.2 FUNCTIONAL SCALING LAWS
We analyze the loss dynamics of SGD using continuous time stochastic differential equation (SDE) model. The discrete update (2) can be modeled by the following Itô SDE (Li al., 2019; Orvieto & Lucchi, 2019; Ankirchner & Perko, 2024):
do, - VR(O.) de +
? S(O)) aw,, b(t)
where W, E R. is N-dimensional Brownian motion, and C(R2z) is the continuous-time batch size schedule with b(kr) = By for all & E N. Here f = kh represents continuous training time, with each discrete iteration corresponding time kn..
For the SDE (3), Li et al. (2025a) derived a functional scaling law (FSL) that characterizes the loss dynamics in continuous training time: Theorem 2.2 (Functional Scaling Law). Under Assumptions f.1, sufficiently large t,
+ noo J: K((t dr,
[E(E(0.)
((T)
signal learning
moise accumulation
where K((t) = + 11-
The above theorem is simplification of Li et al. (2025a, Theorem 4.1) for constant learning rate. For completeness, we provide a self-contained derivation of the above FSL in Appendix A.2. This law establishes functional-level map from the BSS function to the loss at time i. and notably, the two terms exhibit a clean interpretation:
The signal-learning term corresponds to the learning under full-batch gradient descent, captur- ing the rate at which SGD extracts the signal f*. This rate is determined by the source exponent s.
The noise-accumulation term characterizes how the BSS shapes the dissipation of gradient noise. The forgetting kerne KC(1 r) characterizes how the noise injected time still affects the loss at time Due X:(t)= (t 1) (22-1/3) higher-capacity model (smaller 3) tends to forget noise more slowly.
While the FSL framework was introduced by Li et al. (2025a), their analysis was restricted to con- stant batch sizes and focused primarily the analysis of learning rate scheduling. We extend this framework by showing that FSL also provides a principled tool for analyzing how batch size scheduling influences optimization dynamics and training efficiency.
4

---

Published as a conference paper at ICLR 2026
3 THEORETICAL ANALYSES VIA FUNCTIONAL SCALING LAWS
We begin by asking the following question:
Given total data budget D, what is the optimal batch-size schedule (BSS) when the loss dynamies follows the FSL (4)?
For fixed model, the data budget is equivalent compute budget, since the computational cost scales linearly with data size. Determining the optimal BSS is challenging, as the final-step loss depends on the entire training trajectory. This is essentially an optimal control problem (Zhao al., 2022: Perko, 2023), which generally does not admit explicit solutions. However, the explicit char- acterization provided by FSL enables an analytical treatment this problem. We address the above question under two settings: (1) unconstrained schedules; and (2) stage-wise BSS motivated by practical constraints.
3.1 OPTIMAL BATCH SIZE SCHEDULING WITHOUT SHAPE CONSTRAINTS
Under the FSL framework, seeking the optimal BSS can formulated as solving the following resource-constrained variational problem:
-t) b(t)
min T>0, b(-)
ED[T,b] = T +
dt
S b(t) dt)d D,
(5)
s.t.
(data/compute constraint)
b(t) Buins Vt [0,
(hardware constraint).
Here, the integral constraint 1-" b(t) dt D comes from the available data budget. The pointwise constraint b(t) 2 Buia captures hardware limitations in data-parallel training: the global batch size must be no smaller than the number of parallel devices (Narayanan et al., 2021).
We denote by (-).) the optimal BSS for problem (5), and let T be the corresponding total training time. We further define the final-step loss as E == ET
Theorem 3.1 (Optimal batch size schedule). Assume and Buin are sufficiently large. Then:
Easy-task regime (s > 1 - 1/3). The optimal BSS satisfies
6a (t) Bmax (T" - -
0 OSEST
1/24+8
3 T R Driw Moreover,
with Brnax ~
Hard-task regime (s < 1 1/8). The optimal BSS exhibits two-phase stable-growth
structure:
Buinn
(1(1)
B Buax T* - -t 1) stst
where T D, T
TT"
D
1-1/8-1 2-113
Bmax D 2
Moreover.
D
The shape of the optimal BSS. the easy-task regime, the optimal BSS takes the form (*(t) 2 Bmax(I* -t 1)-? corresponding to progressively increasing batch size throughout training, as illustrated Figure 2 (left). The peak batch size scales with the data budget as Buax ~ D with > 0. indicating that larger datasets favor larger batch sizes. This provides a theoretical explanation for the empirical practice increasing batch size with dataset size (DeepSeck-At et al., 2024a; Zhang et al., 2025: Li et al., 2025b).
5

---

Published as a conference paper at ICLR 2026
In the hard-task regime, the optimal BSS exhibits a stabeegrower structure: it stays at the minimal batch size Buin for the first T steps, followed by growth phase with the same functional form as in the easy-task regime. Notably, the growth phase occupies only tiny fraction of the total training horizon, (T - / = oD (1), implying that extremely large batch sizes are required only near the end of training. See Figure 2 (left) for an illustration. Intuitively, for hard tasks (small s), maintaining small batch size allows more optimization steps (larger T') under a fixed data budget, thereby significantly reducing the signal-learning term T The late-stage batch growth primarily serves to noise reduction. This stable growth structure can be viewed as the batch-size analogue of the warmup- stable stable-decay learning rate schedule (Hu al., 2024; Hägele et al., 2024; Li et al., 2026b).
Same data efficiency, fewer iterations. In the easy-task regime, the excess risk rate D~ +0+((++))) matches the minimax optimal rate of this problem (Caponnetto & De Vito, 2007, Theorem 2). In the hard-task regime, the excess risk scales as D=* matching the best rate attainable by one-pass SGD (Diculeveut & Bach, 2016; Pillaud-Vivien et al., 2018). These suggest that, with a properly designed BSS, constant learning rate can achieve the same data efficiency as carefully tuned learn- ing rate schedules (Lin et al., 2024; Li et al., 2025a). The key distinction, however, lies in iteration complexity: batch-size scheduling significantly reduces the total number of iterations compared to learning rate scheduling. When coupled with modern GPU parallelism, this reduction directly trans- lates into shorter wall-clock training time. In short, batch-size scheduling preserves data efficiency while substantially reducing iteration complexity.
Numerical validation. Although the FSL (4) is derived from the continuous-time SDE (3), we empirically confirm that Theorem 3.1 accurately predicts the behavior of discrete- time SGD. Specif- ically, we run SGD (1) using the optimal BSS prescribed by Theorem 3.1 and report the final-step loss as a function of the data size in Figure 2 (middle, right); see Appendix B.2 for experimental details. The observed data scaling closely matches the theoretical predictions. Additionally, in Ap- pendix 1, we compare constant learning rates combined with the optimal BSS against popular learning rate schedules (cosine and warmup-stabll-decay). We find that the optimal BSS with a constant learning rate achieves comparable performance to these widely used learning rate sched- ules.
Optimal BSS comparison
Scaling behavior (a 100 = 2.0)
Scaling behavior (s 0.4,5 = 2.0)
Fily took ingien Bldd tak regine
rrdected
and std
20-
batch a0
750
4''30"
50-
9. 12
3900
11860
1500
20340
112
Figure 2: Optimal BSS experiments for the feature-space linear regression. Left: Illustration of the optimal
BSSs for the casy-task and hard-task regimes. Middle: In the easy-task regime (s 1.0, B 2.0), one- pass SGD with optimal BSS attains the predicted minimax rate D (1-0), Right: In the hard-task regime 0.4, 3 2.0). it matches the optimal rate D "ttainablle by one-pass SGD.
3.2 STAGE-WISH OPTIMAL BATCH SIZE SCHEDULING
Theorem 3.1 shows that the unconstrained optimal BSS follows a smoothly increasing schedule. In practice, however, batch sizes are discrete and constrained by hardware limitations. Moreover, changing the batch size during training incurs nontrivial system overhead, such data pipeline and communication reconfiguration. Consequently, practical schedules typically permit only a small number of stage-wise adjustments (DeepSeek-Al et al., 2024b; MiniMax et al., 2025).
In this section, we study the simplest nontrivial stage-wise setting: a two-stage schedule that begins with small batch size B, and later switches to larger batch size B2, In practice, B, and B2 are

---

Published as a conference paper at ICLR 2026
largely determined by hardware constraints, the key issue is to determine the optimal timing of this switch.
We denote by BB (t) the loss at time under this two-stage schedule. Let D be the total number of training samples and P E [0, DJ denote the number of samples processed before switching from B1 B2. The corresponding BSS B. (t), the switching time s, r, and the total training time Tp are defined as follows:
[BI. 0 S
P
D
b/ BB) -
T.P = B1' Tp
(6)
B2, T.p <t<
B1
B2
We denote by ED - (P) the expected final-step loss under this schedule.
Theorem 3.2 (Optimal two-stage batch size schedule). Let B1 < B2 be constants independent of D, and assume D is sufficiently large. Define PD arg mempeeinnnnei BB (P). Then:
. Ifs > 1 1/B, then Pn =
1-1/8- Ifs S - 1/B, then D-PD D 2-1/8
This theorem shows that, even within the restricted class two-stage BSS, the optimal strategy still depends sharply on task difficulty. For easy tasks, is optimal employ large-batch training throughout. In contrast, for hard tasks with s < 1/3, one should maintain a small batch size for most of training and switch consistent with the behavior in the unconstrained setting.
large batch only very late stage as (D PD) D oD(1),
Additionally, Theorem 3.2 shows that the optimal switching point obeys scaling law: D PT D for some exponent under the FSL framework. This suggests principled tuning strategy: one can estimate the scaling exponent via small-scale pilot experiments and extrapolate the resulting optimal switching point to large-scale training.
4 THE FAST CATCH-UP EFFECT: A BRIDGE TO LLM PRETRAINING
A central insight of the preceding analysis is that, for hard tasks, the optimal schedule switches to very large batch size only in a late stage of training. We now provide a dynamical perspective that explains this phenomenon and extends naturally to LLM pretraining. We refer to this mechanism as the fast catch-up effect, and regard as key insight for practical BSS designing.
The fast catch-up effeet. Figure reveals a consistent phenomenon, observed from linear regres- sion to LLM pretraining, when the batch size is increased from small to large:
Once the batch size increases. the loss rapidly collapses onto the trajectory of large-batch training.
In other words, although the model is trained with a small batch size for most of the trajectory, rapidly "catches up" to the performance of training with the larger batch throughout.
To evaluate the robustness of this phenomenon in realistic large-scale settings, we conduct exper- iments across multiple switching times, model architectures (Dense and MoE), and scales (1.1B parameters and training tokens). Figure 3 shows that fast catch-up consistently occurs all configurations. In particular, Figure 3 (middle) presents four-stage BSS (640 1280 1920 2560). After each stage transition, the loss trajectory rapidly collapses to that of training continu- ously the corresponding larger batch size. This repeated collapse across stages underscores the robustness of the fast catch-up effect.
The late-switch principle. The fast catch-up effect reveals a simple yet powerful principle for batch size scheduling:
The validation loss of constant large-batch training be matched by starting with small batch and deferring the transition to the large batch until a late stage.

---

Published as a conference paper at ICLR 2026
0.98 Dense
1B MoE
MoE
4.23
baich wholicie
25 bacch now from 3005 tcher
3.8-
8.0
batch IBc
basth whilmine
basch 2Bc
basth whileine
4.-
De bunth die from can
Zx banch dre from 20c0
De batch dis whok time
baich siz from S0CB
4.60
0.53
2b bacch vos from 5085 toher
burth size from 5ee
223
butch sia from 2018
mubiple
0.43
4.33
2.8
4.33
(Thousurds)
(Thousunds)
(Thoasands)
Figure 3: The fast catch-up effect across diverse model architectures, model and data scales. Left: A p.rB-parrmettr LLaMA model trained on the C4 datasct with a base batch size of 512. Middle: A 1B- parameter MoE model trained on approximately 0.4T tokens with base batch size of 640; the gray curve shows an additional 4-stage schedule beyond the two-stage runs, Right: 1.1B- 1.IB-parameter MoE model trained on ITT tokens with a base batch size of 1024.
Because the subsequent large-batch phase rapidly aligns with the corresponding large-batch trajec- tory, the final loss and total optimization steps remain unchanged. At the same time, the prolonged small-batch phase substantially reduces token consumption, lowering computational cost without sacrificing performance. We term this strategy late switching, and validate its effectiveness in real- istic LLM pretraining in Section 5.
4.1 AN EXPLANATION VIA FUNCTIONAL. SCALING LAWS
We now explain the fast catch-up effect using FSL. Specifically, consider the two-stage BSS (6) and denote by f. the switching time. Additionally, we denote by EB, (t) and Eg, (t) the losses under constant batch sizes B, and B2, respectively. By Theorem 2.2, for training with a constant batch size B and sufficiently large i, the excess risk admits the decomposition Et( < t= - no²/ /B, where the first term corresponds signal learning and the second term captures noise accumulation.
Loss gap at the switching point. At the switching point I,, the loss gap is given by
1102 B
1702 B₂
1 B2
G. = E (1+)) Eng = 1.
1.
=10
B
Thus, the loss gap at the switching point arises purely from the difference in noise accumulation. The signal-learning term is identical across the two runs, since both have undergone the same opti- mization time 4,.
Post-switch gap decay (catch-up dynamics). After switching to the larger batch size, FSL (4)
implies that after an additional interval 8, the gap decays as
(/- +8 K(L. + 1-1) 1)
4. K(l., + 8 -1)
ER,, (1+ + 6) Eng (1+ +)
dt
b (t) dt
bB, bB 1B (t)
1102
1122
/'8 K(t,, + 5 f) dt G. (-1-11)))))
B
B2
This indicates that the catch-up dynamics progressively forget the noise accumulated during the ini- tial amall-batch phase. Importantly, the forgetting exponent depends only on the capacity exponent 3 and is independent of the task difficulty s.
When is the catch-up fast? We quantify "fast" catch-up through a comparison of time scales. Define the catch-up time 8, as the smallest 6 such that
EB (t. + 5) S (1+c) + +0),
i.e., the switched trajectory lies within a (1 + E) factor of the large-batch baseline. Combining the
gap decay (7) with En, (4, + 6) = (l, + 8) - + n B2 yields
8, z 1 1-7
8

---

Published as a conference paper at ICLR 2026
In contrast, the large-batch loss evolves the time scale ò t., since the signal term (t, + 8) - changes appreciably only when 2 t.. For hard tasks with 8 < 1 1/B, we have dg < t, which establishes clear time-scale separation: the switched trajectory relaxes the fast scale 6,, whereas the large-hatch baseline evolves on the slow scale 4.. The fast catch-up phenomenon therefore direct consequence of this separation of time scales. Moreover, since &, decreases with 8, harder tasks (smaller s) exhibit faster catch-up.
5 VALIDATING LATE SWITCHING IN LLM PRETRAINING
We now examine how the preceding theoretical results manifest in practical LLM pretraining. The experimental setup is summarized below; further details are provided Appendix B.1.
Small-scale. For small-scale experiments, we adopt the popular NanoGPT codebase (Karpa- thy, 2022) and evaluate standard dense LLaMA architectures (Touvron et al., 2023) on the C4 dataset (Raffel et al., 2020). Following Chinchilla law (Hoffmann et al., 2022), the total number of training tokens set be approximately 20x the number of model parameters, convention commonly adopted in small-scale training studies. Concretely, we consider model sizes 50M, 200M, and 492M (+x 0.5B) parameters.
Large-scale. We conducted large-scale experiments using the widely adopted Megatron-LMM codebase (Shoeybi et al., 2019). Our models are based on sparse Mixture-of-Experts (MoE) architecture, specifically the shorteut-comneeted MoE proposed by Cai al. (2025). To bet- ter reflect real-world LLM pretraining, we train our models with latentto-parmmete ratios that substantially exceed the canonical 20:1 guideline, placing our experiments in beyond- Chinchillaooptimal regime (Sardana et al., 2024). We consider two model configurations: (i) 1001M (ss 1B) total parameters with 209M parameters activated per token, trained on approxi- mately 0.4T tokens; (ii) 1119M (<2 1.1B) total parameters with 291M parameters activated per token, trained on approximately 1T tokens.
Fine-grained analysis of the switching time. Figure (left) shows how the final-step loss varies with the switching time. The optimal switching point occurs at approximately 70% of the total train- ing tokens, corroborating the theoretical prediction in Section 3.2: late switching yields improved performance.
We next validate Theorem 3.2, which establishes a power-law relation between the optimal switching
point Pj and the total data size D: D P cDv, for some >> 0 and E (0.1). 1). Taking logarithms yields the linear relation log( Pp) > log log c. We conduct experiments with 50M-parameter model trained C4, with token budgets ranging from 1.3B to SB. For each D. we perform a grid search determine the optimal switching point P and fit log( - PD) against log D using least squares. As shown in Figure 4 (right), the fitted line indeed closely follows power-law relation.
0.2B Dense, different switch point
50M Dense, switch point scaling
3.20
8.24
D Pb (Billio)
3.20
10f
3.20
3.16
Rox 1/9
3xx
D.ID
D (Billion)
Figure 4: Left: Validation loss under different batch size switching points. The x-axis denotes the fraction of data processed before switching. Right: Power-law scaling between D Ph and D. A linear fit in log-log coordinates yields R?2 0.990, supporting the predicted power-law relation.

---

Published as a conference paper at ICLR 2026
McE, switch to 2x batch size
1B MoE, switch to 4x batch size
1.18 McE, switch to batch size
4:3
2300
Ze botes se fore 3038 toker
De bunch sie whale tine
tunth die from 100
Zu hoters cae foom 5339 riber
4.7
2: batch san from Sie
2n baich skn from 2006
4.7
baich from 200
4.60
barch des from 2000
validation love
valdatio
validationn
8:00
4.4
1.45
42-
4.2
1.00
4.1
4.20
num tokers (Bilioo)
rMm takenn (Bl0
mam tokens (Bilion)
Figure 5: Left: Validation loss versus training tokens under different switch points for IB MoË model trained on 0.4T tokens; batch size increases from 640 1280. Middle: Same IB MoE model and dataset; batch size increases from 512 to 2048. Right: I.1B MoE model trained on IT tokens; batch size increases from 1024 to 2048.
Larger-scale validation of late-switch superiority. We now turn to large-scale settings and demonstrate that late switching consistently outperforms early switching. The main results are pre- sented in Figure 5, with additional details provided Figure 6. Across different switching ratios, model architectures, and training scales, late switching yields consistently better performance than early switching. We further evaluate the late-switch principle in multi-stage batch-size schedules (Appendix B.5, Figures 7 and 8), where deferring batch-size increases continues to outperform early switching. Together, these results confirm that late switching is robust across model and data scales.
CONCLUSION
In this work, we demonstrate that the functional scaling law (FSL) provides a principled framework for analyzing batch size scheduling. We characterize the optimal batch size schedules in both the unconstrained and stage-wise settings, and show that the optimal structure depends sharply on task difficulty. In particular, hard tasks favor late switching: using small batch size for most of training and transitioning to a large batch only in a late stage. To explain this structure, we uncover the fast catch-up effect and show that it extends beyond the thcoretical setting to realistic LLM pretraining.
Several important directions remain for future work. First, the FSL framework is derived under stan- dard SGD, whereas modern LLM training predominantly relies on adaptive optimizers. Extending the analysis to adaptive methods is therefore an important open problem.
Second, for analytical clarity, our analysis focuses on constant learning rates. This assumption is meaningful in its own right, as widely used learning rate schedules such as warmup-stablc-dee maintains a constant learning rate throughout most of training (Zhai et al., 2022; Hu et al., 2024; Hägele et al., 2024). Nevertheless, understanding the joint effect of learning rate decay and batch- size mnneduling-particullae
how learning rate decay influences the fast catch-up effect and the resulting late-switch strategy-remains an important direction. As a preliminary exploration, we provide experiments with cosine learning-rate decay in Appendix B.6, which suggest that the fast catch-up effect continues to hold approximately. A systematic treatment of these interactions is left for future work.
ACKNOWLEDGMENT
Lei Wu is supported by the National Natural Science Foundation China INSF NSFC92470122, and NSFC12288101). Binghui Li is supported by the Elite Ph.D. Program in Applied Mathematics at Peking University. Mingze Wang is supported by Young Scientists (PhD) Fund of the National Natural Science Foundation of China (No. 124B2028). We also thank Weinan E. Zilin Wang, Shaowen Wang, Kairong Luo, Haodong Wen and Kaifeng Lyu for many helpful discussions, and the anonymous reviewers for their valuable feedback.
10

---

Published as a conference paper at ICLR 2026
REFERENCES
Armen Aghajanyan, Lili Yu, Alexis Conneau, Wei-Ning Hsu, Karen Hambardzumyan, Susan Zhang, Stephen Roller, Naman Goyal, Omer Levy, and Luke Zettlemoyer. Scaling laws for gen- erative mixed-modal language models. In International Conference on Machine Learning, pp- 265-279. PMLR, 2023. 3
Stefan Ankirchnerr and Stefan Perko. A comparison of continuous- approximations to stochastic gradient descent. Journal of Machine Learning Research, 25(13):1-555 2024.
Yasaman Bahri, Ethan Dyer, Jared Kaplan, Jachoon Lee, and Utkarsh Sharma. Explaining neural scaling laws. Proceedings of the National Academy of Sciences, 121(27))e2311878222 2024. 3
Blake Bordelon, Alexander Atanasov, and Cengiz Pehlevan. A dynamical model of neural scaling laws. In International Conference on Machine Learning, pp. 4345.4382. PMLR, 2024. 3
Blake Bordelon, Alexander Atanasov, and Cengiz Pehlevan. How feature learning can improve neural scaling laws. In International Conference Learning Representations, 2025. 4
Tom B. Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared Kaplan, Prafulla Dhari- wal, 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. Advances in neural information processing systems, 33:1 137-19911,,,, 2020. 1,2
Weilin Cai, Juyong Jiang, Le Qin, Junwei Cui, Sunghun Kim, and Jiayi Huang. Shorteut-connecter expert parallelism for accelerating mixture-of-experts. In International Conference on Machine Learning, pp. 621 1-6228. PMLR, 2025. 9, 26
Andrea Caponnetto and Ernesto De Vito. Optimal rates for the regularized least-squares algorithm. Foundations of Computational Mathematies, 7:331-368, 2007. 6
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 Lev- skaya, Sanjay Ghemawat, al. PaLM: Scaling language modeling with pathways. Journal of Machine Learning Research, 24(240):1-113 2023. 2
DeepSeck-AL, Xiao Bi, Deli Chen, Guanting Chen, Shanhuang Chen, Damai Dai, Chengqi Deng, Honghui Ding, Kai Dong, Qiushi Du, Zhe Fu, Huazuo Gao, Kaige Gao, Wenjun Gao, Ruiqi Ge, Kang Guan, Daya Guo, Jianzhong Guo, Guangbo Hao, Zhewen Hao, Ying He, Wenjie Hu, Panpan Huang, Erhang Li, Guowei Li, Jiashi Li, Yao Li, Y. K. Li, Wenfeng Liang, Fangyun Lin, A. X. Liu, Bo Liu, Wen Liu, Xiaodong Liu, Xin Liu, Yiyuan Liu, Haoyu Lu, Shanghao Lu, Fuli Luo, Shirong Ma, et al. DeepSeek LLM: Scaling open-source language models with longtermism. arXiv preprint ar Xixi::2001. 02954,, 2024a.
DeepSeek-AL, Aixin Liu, Bei Feng, Bing Xue, Bingxuan Wang, Bochao Wu, Chengda Lu, Cheng- gang Zhao, Chengqi Deng, Chenyu Zhang, Chong Ruan, Damai Dai, Daya Guo, Dejian Yang. Deli Chen, Dongjie Ji, Erhang Li, Fangyun Lin, Fucong Dai, Fuli Luo, Guangbo Hao, Guanting Chen, Guowei Li. H. Zhang, Han Bao. Hanwei Xu, Haocheng Wang. Haowei Zhang, Honghui Ding, Huajian Xin, Huazuo Gao. Hui Li, Hui Qu, J. L. Cai, Jian Liang, Jianzhong Guo, Jiaqi Ni, Jiashi Li, et al. DeepSeek- V3 technical report. arXiv preprint arXiv:24/2. 2t..24///.999.... 2024b. 2, 6
Aaron Defazio, Ashok Cutkosky, Harsh Mehta, and Konstantin Mishchenko. Optimal linear decay learning rate schedules and further refinements. arXiv preprint arXiv:2310.07831, 2023.
Aymeric Dieuleveut and Francis Bach. Non-purametric stochastic approximation with large step sizes. Annals of Statisties, 44(4), 2015.
11

---

Published as a conference paper at ICLR 2026
Aymeric Dieuleveut and Francis Bach. Nonparametric stochastic approximation with large step- sizes. The Annals of Statistics, 44(4):1363 1399, 2016.
Priya Goyal, Piotr Dollár, Ross Girshick, Pieter Noordhuis, Lukasz Wesolowski, Aapo Kyrola, An- drew Tulloch, Yangqing Jia, and Kaiming He. Accurate, large minibatch SGD: Training ImageNet in 1 hour. arXiv preprint arXiv: 02777700077777 2017.
Aaron Grattafiori, Abhimanyu Dubey, Abhinav Jauhri, Abhinav Pandey, Abhishek Kadian, Ahmad Al-Dahle, Aiesha Letman, Akhil Mathur, Alan Schelten, Alex Vaughan, Amy Yang, Angela Fan, Anirudh Goyal, Anthony Hartshorn, Aobo Yang, Archi Mitra, Archie Sravankumar, Artem Ko- renev, Arthur Hinsvark, Arun Rao, Aston Zhang, Aurelien Rodriguez, Austen Gregerson, Ava Spataru, Baptiste Roziere, Bethany Biron, Binh Tang, Bobbie Chern, Charlotte Caucheteux, Chaya Nayak, Chloe Bi, Chris Marra, et al. The Llama 3 herd of models. arXiv preprint 1xXiv:2407.217880 2024. 2
Gavia Gray, Aman Tiwari, Shane Bergsma, and Joel Hestness. Normalization layer per-example gradients are sufficient to predict gradient noise scale in Transformers. Advances in Neural Infor- mation Processing Systems, 37:93510 10-59399999,,. 2024. 3
Alexander Hägele, Elie Bakouch, Atli Kosson, Loubna Ben allal, Leandro Von Werra, and Martin Jaggi. Scaling laws and compute-optimal training beyond fixed training durations. Advances in Neural Information Processing Systems, 37:76232-76264 2024. 3, 6, 10, 29
Tom Henighan, Jared Kaplan, Mor Katz, Mark Chen, Christopher Hesse, Jacob Jackson, Heewoo Jun, Tom B. Brown, Prafulla Dhariwal, Scott Gray, Chris Hallacy, Benjamin Mann, Alec Radford, Aditya Ramesh, Nick Ryder, Daniel M. Ziegler, John Schulman, Dario Amodei, and Sam Mc- Candlish. Scaling laws for autoregressive generative modeling. arXiv preprint arXiv:2010.14701, 2020. 3
Joel Hestness, Sharan Narang, Newsha Ardalani, Gregory Diamos, Heewoo Jun, Hassan Kianinejad, Md Mostofa Ali Patwary, Yang Yang, and Yanqi Zhou. Deep learning scaling is predictable, empirically. arXiv preprint 171.vv1712.0000 2017. 3
Jordan Hoffmann, Sebastian Borgeaud, Arthur Mensch, Elena Buchatskaya, Trevor Cai, Eliza Rutherford, Diego de Las Casas, Lisa Anne Hendricks, Johannes Welbl, Aidan Clark, Tom Hen- nigan, 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. Advances in neural information processing systems, 35:30016-30038 2022. 1,3,9
Shengding Hu, Yuge Tu, Xu Han, Chaoqun He, Ganqu Cui, Xiang Long, Zhi Zheng. Yewei Fang, Yuxiang Huang, Weilin Zhao, Xinrong Zhang, Zheng Leng Thai, Kaihuo Zhang, Chongyi Wang, Yuan Yao, Chenyang Zhao, Jie Zhou, Jie Cai, Zhongwu Zhai, Ning Ding, Chao Jia, Guoyang Zeng, Dahai Li, Zhiyuan Liu, and Maosong Sun. MiniCPM: Unveiling the potential of small language models with scalable training strategies. In Conference on Language Modeling, 2024. 3. 6, 10, 29
Arlind Kadra, Maciej Janowski, Martin Wistuba, and Josif Grabocka. Power laws for hyperparam- eter optimization. arXiv preprint arXiv:2302.00441. 2023. 3
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.0836), 2020.
Andrej Karpathy. NanoGPT. https: wtt.th/ww../eeerppttt/tt/annnaaa 2022. 9
Frederik Kunstner and Francis Bach. Scaling laws for gradient descent and sign descent for linear bigram models under zipf's law. arXiv preprint arXiv:2505.19227. 2025.
Tim Tsz-Ki Lau, Han Liu, and Mladen Kolar. AdAdaGrad: Adaptive batch size schemes for adap- tive gradient methods. arXiv preprint arXiv:2402. 12/5, 2024.
12

---

Published as a conference paper at ICLR 2026
Tim Tsz- Kit Lau, Weijian Li, Chenwei Xu, Han Liu, and Mladen Kolar. Adaptive batch size sched- ules for distributed training of language models with data and model parallelism. In Proceedings of Conference on Parsimony and Learning, 2025.
Kiwon Lee, Andrew Cheng, Elliot Paquette, and Courtney Paquette. Trajectory of mini-batch mo- mentum: batch size saturation and convergence in high dimensions. Advances in Neural Infor- mation Processing Systems, 35:36944-36957, 2022.
Binghui Li, Fengling Chen, Zixun Huang, Lean Wang, and Lei Wu. Functional scaling laws in kernel regression: Loss dynamics and learning rate schedules. arXiv preprint 12xx:::509.191889 2025a. 1,2, 4, 18
Binghui Li, Kaifei Wang, Han Zhong, Pinyan Lu, and Liwei Wang. Muon in associative memory learning: Training dynamics and scaling laws, arXiv preprint arXiv:2602.05725 2026a. 3
Binghui Li, Zilin Wang. Fengling Chen, Shiyang Zhao, Ruiheng Zheng, and Lei Wu. Optimal learning-rai rate schedules under functional scaling laws: Power decay and warmup-stable-decay, arXiv preprint arKiv:2602.0077 2026b. 3, 6
Houyi Li, Wenzhen Zheng. Jingcheng Hu, Qiufeng Wang, Hanshan Zhang, Zili Wang, Shijie Xuyang, Yuantao Fan, Shuigeng Zhou, Xiangyu Zhang, and Daxin Jiang. Predictable scale: Part I I- optimal hyperparameter scaling law in large language model pretraining. arXiv preprint arXiv:2503.04715, 15, 2025b.
Qianxiao Li, Cheng Tai, and Weinan Stochastic modified equations and dynamics stochastic gradient algorithms I: Mathematical foundations. Journal of Machine Learning Research, 20(40):
1-47, 2019. 4
Licong Lin, Jingfeng Wu, Sham M Kakade, Peter L Bartlett, and Jason D Lee. Scaling laws in linear regression: Compute, parameters, and data, Advances in Neural Information Processing Systems, 37:60556-60666 2024. 3, 4,
Yinhan Liu, Myle Ott, Naman Goyal, Jingfei Du, Mandar Joshi, Danqi Chen, Omer Levy. Mike Lewis, Luke Zettlemoyer, and Veselin Stoyanov. RoBERTa: A robustly optimized BERT pre- training approach. arXiv preprint arXiv: 1692, 2019. 26
Meituan LongCat, Bayan, Bei Li, Bingye Lei, Bo Wang, Bolin Rong, Chao Wang, Chao Zhang, Chen Gao, Chen Zhang, Cheng Sun, Chengeheng Han, Chenguang Xi, Chi Zhang, Chong Peng, Chuan Qin, Chuyu Zhang, Cong Chen, Congkui Wang, Dan Ma, Daoru Pan, Defei Bu, Dengchang Zhao, Deyang Kong. Dishan Liu, Feiye Huo. Fengcun Li, Fubao Zhang. Gan Dong, Gang Liu, Gang Xu, Ge Li, Guoqiang Tan, Guoyuan Lin, Haihang Jing. Haomin Fu, Hao- Yan, Haoxing Wen, Haozhe Zhao, al. LongCat-Flash technical report. arXiv preprint arXiv:2509.013322. 2025. 26
Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization. International Confer- ence on Learning Representations, 2019. 26
Kairong Luo, Haodong Wen, Shengding Hu, Zhenbo Sun, Zhiyuan Liu, Maosong Sun, Kaifeng Lyu, and Wenguang Chen. A multi-power law for loss curve prediction across learning rate schedules. In International Conference on Learning Representations, 2025.
Siyuan Ma, Raef Bassily, and Mikhail Belkin. The power of interpolation: Understanding the effectiveness of SGD in modern over-parametrized learning. In International Conference on Machine Learning. pp. 3325-3333 3334. PMLR, 2018. 2,
Sam McCandlish, Jared Kaplan, Dario Amodei, and OpenAl Dota Team. An empirical model of latch-batch training. arXiv preprint urXiv:18/2.000
2018. 1, 2,
William Merrill, Shane Arora, Dirk Groeneveld, and Hannaneh Hajishirzi. Critical batch size visited: A simple empirical approach to large-batch language model training. arXiv preprint arXiv:2505.23971, 2025. 1, 2,
13

---

Published as a conference paper at ICLR 2026
MiniMax, Aonian Li, Bangwei Gong, Bo Yang, Boji Shan, Chang Liu, Cheng Zhu, Chunhao Zhang, Congchao Guo, Da Chen, Dong Li, Enwei Jiao, Gengxin Li, Guojun Zhang, Haohai Sun, Houze Dong, Jiadai Zhu, Jiaqi Zhuang. Jiayuan Song, Jin Zhu, Jingtao Han, Jingyang Li, Junbin Xie, Junhao Xu, Junjie Yan, Kaishun Zhang. Kecheng Xiao, Kexi Kang, Le Han, Leyang Wang. Lian- fei Yu, Liheng Feng, Lin Zheng, Linbo Chai, Long Xing, Meizhi Ju, Mingyuan Chi, Mozhi Zhang, et al. MiniMax-01: Scaling foundation models with lightning attention. arXiv preprint arXiv:2501.00 08311, 2025. 2. 6
Takashi Mori, Liu Ziyin, Kangqiao Liu, and Masahito Ueda. Power-law escape rate of SGD. In International Conference on Machine Learning, pp. 15959-15975. PMLR, 2022. 19
Nicole Mücke, Gergely Neu, and Lorenzo Rosasco. Beating SGD saturation with tail tail-averaging and minibatching, Advances in Neural Information Processing Systems, 32, 2019. 3
Niklas Muennighoff, Alexander Rush, Boaz Barak, Teven Le Scao, Nouamane Tazi, Aleksandra Piktus, Sampo Pyysalo, Thomas Wolf, and Colin A Raffel. Scaling cota-const language models. Advances in Neural Information Processing Systems, 36:50358-50376, 2023. 3
Deepak Narayanan, Mohammad Shoeybi, Jared Casper, Patrick LeGresley, Mostofa Patwary, Vi- jay Anand Korthikanti, Dmitri Vainbrand, Prethvi Kashinkunti, Julie Bernauer, Bryan Catanzaro, Amar Phanishayee, and Matei Zaharia. Efficient large-scale language model training on GPU clusters using Megatron-L.M. In Proceedings of the international conference for high perfor- mance computing, networking. storage and analysis, pp. 1-15, 2021. 5
Nvidia, Bo Adler, Niket Agarwal, Ashwath Aithal, Dong H. Anh, Pallab Bhattacharya, Annika Brundyn, Jared Casper, Bryan Catanzaro, Sharon Clay, Jonathan Cohen, Sirshak Das, Ayush Dattagupta, Olivier Delalleau, Leon Derezynski, Yi Dong, Daniel Egert, Ellie Evans, Alek- sander Ficek, Denys Fridman, Shaona Ghosh, Boris Ginsburg, Igor Gitman, Tomasz Grze- gorzek, Robert Hero, Jining Huang, Vibhu Jawa, Joseph Jennings, Aastha Jhunjhunwala, John Kamalu, Sadaf Khan, Oleksii Kuchaiev, et al. Nemotron-4 340B technical report. arXiv preprint arXiv:2406. 11704, 2024. 2
Antonio Orvieto and Aurelien Lucchi. Continuous-tine time models for stochastic optimization algo- rithms. Advances in Neural Information Processing Systems, 32, 2019. 4
Elliot Paquette, Courtney Paquette, Lechao Xiao, and Jeffrey Pennington. 4+3 phases of compute- optimal neural scaling laws. Advances in Neural Information Processing Systems, 37:16459- 16537, 2024. 3,
Jupinder Parmar, Shrimai Prabhumoye, Joseph Jennings, Mostofa Patwary, Sandeep Subramanian, Dan Su. Chen Zhu, Deepak Narayanan, Aastha Jhunjhunwala, Ayush Dattagupta, Vibhu Jawa, Jiwei Liu, Ameya Mahabaleshwarkar, Osvald Nitski, Annika Brundyn, James Maki, Miguel Martinez, Jiaxuan You, John Kamalu, Patrick LeGresley, Denys Fridman, Jared Casper, Ash- wath Aithal, Oleksii Kuchaiev, Mohammad Shoeybi, Jonathan Cohen, and Bryan Catanzaro. Nemotron-4 15B technical report. arXiv preprint arXiv:2402.1668) 2024. 2
Stefan Perko. Unlocking optimal batch size schedules using continuous-time control and perturba- tion theory. arXiv preprint arXiv:2312.01899 2023. 3,5
Loucas Pillaud-Viviin Alessandro Rudi, and Francis Bach. Statistical optimality of stochastic gra- dient descent on hard learning problems through multiple passes. Advances in Neural Information Processing Systems, 31, 2018.
Shikai Qiu, Lechao Xiao, Andrew Gordon Wilson, Jeffrey Pennington, and Atish Agarwala. Scaling collapse reveals universal dynamics in compute-optiman trained neural networks. arXiv preprint arXiv:2507.022119, 2025. 3
Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter J Liu. Exploring the limits of transfer learning with a unified text-to-text Transformer. The Journal of Machine Learning Research, 214555545--555, 2020. 9, 26
Nikhil Sardana, Jacob Portes, Sasha Doubov, and Jonathan Frankle. Beyond Chinchilla-optimal Accounting for inference in language model scaling laws. In International Conference on Ma- chine Learning, 2024. 9
14

---

Published as a conference paper at ICLR 2026
Mohammad Shoeybi, Mostofa Patwary, Raul Puri, Patrick LeGresley, Jared Casper, and Bryan Catanzaro. Megatron-LM: Training multi-billion parameter language models using model par- 2019. 9
allelism. arXiv preprint
Samuel Smith, Pieter-Jan Kindermans, Chris Ying, and Quoe V Le. Don't decay the learning rate, increase the batch size. In International Conference on Learning Representations, 2018. 2,
Jianlin Su, Murtadha Ahmed, Yu Lu, Shengfeng Pan, Wen Bo, and Yunfeng Liu. RoFormer: En-
hanced Transformer with rotary position embedding. Neurocomputing, 568:127063 2024. 26
Howe Tissue, Venus Wang, and Lu Wang. Scaling law with learning rate annealing. arXiv preprint arXiv:2408. 11029, 2024. 3
Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie-Anne Lachaux, Timothée Lacroix, Baptiste Rozière, Naman Goyal, Eric Hambro, Faisal Azhar, Aurelien Rodriguez, Ar- mand Joulin, Edouard Grave, and Guillaume Lample. LLaMA: Open and efficient foundation language models. arXiv preprint arXiv:2302. 13971, 2023. 9, 26
Jinbo Wang, Mingze Wang, Zhanpeng Zhou, Junchi Yan, Weinan E, and Lei Wu. The sharpness disparity principle in Transformers for accelerating language model pre-training. In International Conference on Machine Learning, pp. (8859-64879. PMLR, 2025a. 26
Mingze Wang and Lei Wu, A theoretical analysis of noise geometry in stochastic gradient descent.
arXiv preprint arXiv:2310.00692 2023. 19
Mingze Wang, Jinbo Wang, Jiaqi Zhang, Wei Wang, Peng Pei, Xunliang Cai, Weinan E, and Lei Wu. GradPower: Powering gradients for faster language model pre-training. arXiv preprint 2rXiv:2505.24275. 2025b. 26
Kaiyue Wen, Zhiyuan Li. Jason Wang. David Hall, Percy Liang, and Tengyu Ma. Understanding warmup-stable-deet learning rates: A river valley loss landscape perspective. International Conference on Learning Representations, 2025. 29
Jingfeng Wu, Difan Zou, Vladimir Braverman, Quanquan Gu, and Sham Kakade. Last iterate risk bounds of SGD with decaying stepsize for overparameterized linear regression. In International Conference on Machine Learning, pp. 14280-24311 PMLR, 2022a.
Lei Wu, Mingze Wang, and Weijie Su. The alignment property of SGD noise and how helps select flat minima: A stability analysis. Advances in Neural Information Processing Systems, 35: 4680-4693, 2022b. 19
Tingkai Yan, Haodong Wen, Binghui Li, Kairong Luo, Wenguang Chen, and Kaifeng Lyu. Larger datasets can be repeated more: A theoretical analysis of multi-epoch scaling in linear regression. arXiv preprint arXiv:2511.1342) 2025. 3
Aohan Zeng, Xin Lv, Qinkai Zheng, Zhenyu Hou, Bin Chen, Chengxing Xie, Cunxiang Wang, Da Yin, Hao Zeng, Jiajie Zhang, al. GLM-4.5: Agentic, reasoning, and coding (ARC) founda- tion models. arXiv preprint arXiv:2508.06477 2025. 2
Xiaohua Zhai, Alexander Kolesnikov, Neil Houlsby, and Lucas Beyer. Scaling Vision Transformers. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pp. 12104-10 12113, 2022. 10
Hanlin Zhang, Depen Morwani, Nikhil Vyas, Jingfeng Wu, Difan Zou, Udaya Ghai, Dean Foster, and Sham Kakade. How does critical batch size scale pre-training? International Conference on Learning Representations, 2025. 3,
Jiawei Zhao, Zhenyu Zhang, Beidi Chen, Zhangyang Wang, Anima Anandkumar, and Yuandong Tian. GaLore: Memory-efficient LLM training by gradient low-rank projection. International
Conference on Machine Learning, 2024. 26
Jim Zhao, Aurelien Lucchi, Frank Norbert Proske, Antonio Orvieto, and Hans Kersting. Batch size selection by stochastic optimal control. In Has Trained Yet? NeurIPS 2022 Workshop, 2022. 3, 5
15

---

Published as a conference paper at ICLR 2026
Hanqing Zhu, Zhenyu Zhang, Wenyan Cong, Xi Liu, Sem Park, Vikas Chandra, Bo Long, David Z Pan, Zhangyang Wang, and Jinwon Lee. APOLLO: SGD-like memory, AdamW-leve perfor- mance. In Conference Machine Learning and Systems, 2025. 26
16

---

Published as a conference paper at ICLR 2026
Appendix
Theoretical Setup and Proofs
18
A.1 Interpretation of the Source and Capacity Conditions
18
A.2 Proof of Theorem 2.2 (Self-Contained Derivation of FSL).
18
A.3 Proof of Theorem 3.1 "Shape-Unconstrained Optimal BSS)
21
A.4 Proof of Theorem 3.2 (Optimal Two-Stage BSS)
24
B Experimental Details and Additional Results
25
B.1 LLM Pretraining: Models, Data, and Training Setup
25
B.2 Linear Regression Experiments: Setup and Details
26
B.3 Additional Details of Fast Catch- Up Experiments
27
B.4 Additional Details of Switching Time Analysis Experiments
28
B.5
Additional Details and Results for Late-Switch Superiority Experiments
28
B.6 Extension: Interaction with Learning Rate Scheduling
29
B.6.1 Comparison with Cosine and WSD
29
B.6.2 Fast Catch-Up under Learning Rate Decay
30
Statement
30
C.I Ethics Statement
30
C.2 Reproducibility Statement
30
C.3 LLM Usage Statement
31
17

---

Published as a conference paper at ICLR 2026
A THEORETICAL SETUP AND PROOFS
A.1 INTERPRETATION OF THE SOURCE AND CAPACITY CONDITIONS
In this section, we provide detailed description of the parameters of feature-space linear regression (aka power-law kernel regression in Li et al. (2025a)) and their interpretation in the context of LLM pretraining. Let kg == 05/25² for j E [M], so that (Øg b forms an orthonormal basis of (D).
Model Capacity 3: A model of the form
N Σ 0 2 ~ 0g
N Σ 0
f(:0) = 0
8/2 2/23
1=1
shows that higher-index features are increasingly down-weighted by the factor j 8-8/2 As B in- creases, the spectrum decays more rapidly, causing the model effectively rely on fewer features.
Additionally, for fixed target function f", one can use different (potentially nonlinear) feature maps & (and consequently, different values of 3). The value of reflects the capacity of the chosen features. For instance, consider (bx) <aN(<; 0), where (-; denotes a neural network. In this case, (xx) corresponds to neural tangent features, and the associated kernel
Ks x') (xx) (xxz)
known as the neural tangent kernel (NTK). Here, the network depth and activation functions govern the spectral decay, determining the effective exponent B.
Task Difficulty s: The target function admits the expansion
0
1/2
1/2
(89+1)/2
Since } are orthonormal, this assumption implies that the spectral energy of f decays according to power law. The exponent a := sB thus quantifies the task's intrinsic difficulty, which depends only the target function itself and is independent of the model's spectrum. In contrast, measures the relative difficulty with respect to a model of capacity 8: for a fixed f* (and fixed a),, adopting higher-capacity model (smaller (3) increases words, the same task appears casier to a higher-capacity model.
a/3, making the task relatively easier. In other
Connection with LLM Pretraining. In the context of large language model (LLM) pretraining, the parameter 3 reflects the model architecture and determines its capacity. Specifically, , influenced by factors such as the depth of the model, the activation functions, and the choice of feature map. model with a larger capacity (smaller (3) has spectrum that decays more slowly, allowing it to utilize a broader range of features, whereas model with smaller capacity (larger 3) down-weights higher-index features more rapidly. On the other hand, the parameter reflects the difficulty of the task relative to the model architecture. It quantifies how challenging particular task for a given model capacity B. For a fixed target function f", increasing the model's capacity (reducing () leads to lower value of making the task easier. In other words, the same task will appear easier to a model with a higher capacity, because the model can better accommodate the complexity of the task due its architecture.
A.2 PROOF OF THROREM 2.2 (SELF-CONTAINEI DERIVATION OF FSI.)
key insight that makes the above SDE (3) analytically tractable is the anisotropic noise structure, which can be formalized follows:
Lemma A.1 (Anisotropic noise). For any 0 E N it holds that
(2E(0) + 22) H ( < (4E(0) + o²) H.
18

---

Published as a conference paper at ICLR 2026
Lemma A.1 demonstrates that the noise covariance Σ(0) approximately admits closed-form
pression: ( x H, as observed in (Mori et al., 2022: Wu et al., 2022b; Wang & Wu, 2023).
This flomed-form expression enables precise characterization of the noise dynamics, thus providing a framework for tracking the SGD training dynamics.
Proof. For a given data point z (xx. y), we define the point-wise risk as f(z: 0) 2 ((x)
v)?. By definition of {(2; 0) and R.))), we have
<((=; 0) (x))(())) (0 0-) Φ(xx),
VR(0) = E[V((2;;))) = H (0 - -0"). For the stochastic mini-batch gradient noise = V((2; 0) VR(O), we have
5 = ==<)))(<=) (0 0*) o(x)c H -0-00)
Hence, the covariance matrix Σ(0) = Elgg [0) satisfies
Σ(0) = (E (0())) (xx) uu o( (<)))())) Huu H) + o2H,
where u = 0 - 0" Let M == E (())) T uu Φ(z)δ()) and Mij be (i,j) entry of M. Calculating Mij using Wick's probability theorem
Mj Σ
Σ 2
un H + Ha).
Recognizing each term, we know
up H.H.HHHHHi (Huu H)
Σ
ug Hg (Hum H)
kd
Σ
H.ugHHHH -uu H.j k.l
Hence
M 2Huu H + (u Hu)H
E()) Huu H + (u Hu)H + H.
Noting that u Hu -0) for any vector æ with the same shape of we have
x (Huu") T H) = x) S (2, (u. Hu) Hæ.
Hence Huu H < (u. Hu)H and
(0) ž (u T Hu)H + o + +
(0) 3 2(u Hu)H+ (4E(0) + + H.
Now proceed to the main theorem.
Proof. For the SDE (3)
d0, VR(0.) + V b(t)
E(0.) dB..
For each coordinate j. we define p; = e, E(0))e, we have
d0,())= = -> 07) di +
V D(0) dB,
17
19

---

Published as a conference paper at ICLR 2026
Applying Itô's formula to (83 03)² we obtain
B[O0 03) 22pt + 1-/ E -22,(t-=2)
((2) ps dz.
2 = Σ N1103p
2NSt
Σ X 1/ E
222((t==)) 1
b(z)
ps
dz.
j=1
j=1
By Lemma A.I, it is trivial that ps = ej E( ~ 2,(E(0,) + 0212). we have the following Volterra equation:
1'" E -22 "
2B(E(0,)| ~ Σ 2 2ASt
22yet + Σ X/
Py dz
b(2)
j=1
j=1
2 e((t) + S/¹
11
K(t 2) (B(E(0-)] + dz,
+zd)dz
b(z)
where
e(t) = Σ A3 1510719e-211 22ds = S/ u0-1e 2ud du.
3=1
20
130-20 21yt 2 /'' u a-- 2ut du.
k(t)= = Σ 5=1
Let f(t)= = E(E(0.)). = e(t) + 0² J0 K(tt - 2) n ((a
dz, and define the linear operator
Tf(t) = 1." K(t 2) x){-=)) b(z)
"
f()) dz.
With this notation, the Volterra equation admits the compact representation f= g + Tf. Formally, the solution can be expressed via the Neumann series expansion:
f=g+
f= -= T'g. g.
1=0
Note that K. * K(()) = 2000 K(I - 2)(( S 2K(112) /22 K(z) & K(1/2) $ K(1), by m//((t) Sn
Telt) [ 7 (1) 1. K(t x((-2)) 2) n
S
dz oz) g(t))
b(z)
Hence, we have
a(1)+T9(1 1()) (0050(0000000)00)000 +T()+ + D x 2 (() + 1- - Tg()).
As a result,
1 2 = -()= +T(()
= 1o +n f K(t r) dr.
1
E[R(0.)]
b(r)
20

---

Published as a conference paper at ICLR 2026
A.3 PROOF OF THEOREM 3.1 UNCONS-UNCONSRRRINNE OPTIMAL BSS)
Lemma A.2. Define the feasible region of BSS under data D:
Bn (T,b) T Z5o b & L. (O,,))) > a.e., [e b(1) (()dd=11 (0)dd=DD )
Consider the following optimal batch size scheduling problem:
K(T t)
min E[T,6] (T
8 J
dt.
b(t)
The optimal batch size schedule obeys
b(1) -T+ - 1) D2
(
with
T ~ Dirts E ~ D 1193.
Proof. We first minimize the second term of the loss under fixed T. By the Cauchy-Schwarg in- equality, we have
(1 12K(T--0a al) (1. (()))aa) (fo' VK(T- (I"VETT
b(()
Equality holds when
b(t) =
where C ensures So b(t) df = D. The minimizer must satisfy the above equality, combining with
fo b(t) dt = D. we have
b(t) ~ D (T-++))#-1
(T+ 1)
and consequently,
(Lo (T)
K(T - t)
(T + 1)
b(t) dt
So b(l) di
D
Consequently, define
/'
T K(T -
1
(T (1+1)
g(T)
= min (T,b)EB, T
Jo b(t)
dt
T
D
Minimizing the above risk with respect to T, we obtain the optimal T
Th D
Substituting T back, the minimum & satisfies
(T
E (T*)-*+
= 175 D 1427
T473
D
The corresponding optimal batch size schedule satisfies
(7(1)= ~ (T* - 1) D.
(T+ + +)
21

---

Published as a conference paper at ICLR 2026
Lemma A.3. Define the feasible region of BSS under data D:
BD (T,b) TERsob (0, b(t) s B2 a.e., [' b(t) dt =D]
Consider the following optimal batch size scheduling problem:
K(T - t)
min E[T,6] = T
dt,
(T,b)EBo
JO
b(t)
The optimal batch size schedule must take one of two possible forms: (i)
(²(1) = Ci CiV C2(T - + 1) for 0505,,
with b' (0) 2 B1. (ii)
B1, b²()= - { [ -
for 1 < Ti:
for 2 Tı:
where T₁ is determined by the boundary condition C2(T Ti 1) (21)) B1.
Proof. We now consider the constrained problem under fixed T with B1 S b(t) S B2. Since the integrand KC(T t)/b(t) is convex for b(t) > 0, and the constraints are linear, so Slater's condition holds. Consequently, any point satisfying the KKT conditions is a global minimizer. Consider the Lagrangian
b[ = !'"
[K(TT Ab(t) (B - b(t)) + E(+)((()) Ba) dt AD b(t)
with ((t)= 2 0 and E(t) 2 0. The stationarity condition is given by
K(T
+ + (ut) + E(t) = 0.
(8)
b(t)2
The complementary slackness conditions are p(()(B: b(()) 0, (())((()) B2) 0, B S b(1) < B2, ml() 2 0, E(L) 2 0.
(9)
Define
A: [t < b(t) < B2}. I = [+|b(t) = B1}. Iz := [t|b(t) = B2}.
In A, both constraints are inactive, thus u(t) {(t) = 0. The stationarity condition (8) yields
K(T t)
+
b(1)2
Solving for b(t), we obtain
b(t) K(T -t)
In Zi, we have b(() = B, and E(t) = O; In Iz, we have b(t) B2, n(t) = 0. By stationarity
(8) and complementary slackness(9), we have the following two relations on I₁ and Iz: For Iz
K((-1) + X - ptt) = 0 implies uut) A K(T-t)
B 2 0, which further yields b(t) = B1 2
K(T-t), For Iz. K(T- + {(t) implies ((t) K(T-1) x which in turn yields B
B2
B
b(1) = B2 <
K(T x
Therefore, the optimal batch size schedule is given by
(t) = clip (V V K(( 1) A (K(T--)
B, = -kip(CV -t, B1, Bz),
where clip(æ, a, b) max (a, min minf Exploiting the monotonicity of K(T - t), the schedule admits the following piecewise form:
B1:
for l < Tí, for Ti < 1 S T2
(t)
C( (0--))-
B2.
for t > Tz,
22

---

Published as a conference paper at ICLR 2026
where C' and C" are problem-dependene constants that depend on D,T, B, BL, B2-
(i) When Ti = 0, the schedule takes the first form
(²()) = Civ - C2 --) for SIST,
with (-(0)= 2 B1.
(ii) When Ti > the schedule takes the second form
(t) - {
B1,
for 1 < Ti
for T1.
where 2i is determined by the boundary-matching condition C2(T - Ti + 1)2 B1.
Now we proceed to the main theorem. For the main theorem, we only consider Lemma A.3 with B1 Buin and B2 oo.
Proof. (I) We now consider whether the optimal batch size schedule b* (t) in Lemma A.2 can satisfy the constraint b(t) 2 B1 under the easy-task regime. Since b" (t() non-decreasing. and
(0) D/(T" + = D 1734= 1.
It follows that under the easy task regime, the constraint bz b*(1) 2 B1 is automatically satisfied when D is sufficiently large. Consequently, we have
T Drfw, ED riso
We have (+) ~ C2(T" - + 1)x0 where the constant C2 is determined from the budget constraint
b²(t) dt C2 (" (++1))/(23)-1dr dt D,
c
Solving for C2, we obtain
C2 D2= = DD
1/21
(II) Under the hard task regime with 1 1/3 11-111 trivially similar), the unconstrained solution in Lemma A.2 is infeasible due to T 2 D. which trivially violates the constraint. We therefore analyze the constrained candidates in Lemma A.3 and determine which achieves lower objective value. We first consider the second form in Lemma A.3.
(*())= = B1, Bi
for t < Ti
(1
for T < T == Ti T2
(10)
The data- butaeb constraint implies
Ti = B
dt
(1I)
T2 +1
Let a = 1/(23) - 1. We consider the objective function
KC -t) t)
E T
dt
b(t)
Substituting (10) into the above objective yields
E- 1 (t+ 1)20 dt
1 B
(T2 (T2+ + 1)a (* 11+1)aat
B 1
Let these three parts be E,, &g, and Eg, respectively. Their derivatives with respect to are
--1 d7
sT -T
dTp
dT,
23

---

Published as a conference paper at ICLR 2026
dEq
dT
(T 1)2n
(T2 1)20
d72
d72
d&g
2a + 1
(72 + 1)2n
(72 + 1)a-1
d72
a + 1
a a+1 1
The optimal 72 must satisfy
3p
d&2
dEg
dT2 = 0
dT,
dT²
dT,
For the regime D 2 1, by Equation (11), have Ti D/B1- 23 (T2 + 1) (T2 1
this implies T 2 1. Moreover, we must have T2 Z 1; otherwise, the expression above would be
dominated by its first term and become negative when D 2 1. Keeping only the dominant terms gives
T (T2 + 1) ~ T 18-2
Hence,
T2 T 28-1,
Since T S D, follows that T2 D 23 and therefore
21 -1
(+++1) +1
dt = 23 (T2 + -Tb + 1))
D 2t12
By the hard-task regime condition, sB3 + B 23 1. Together with (11), this yields T1 ~ D, hence, T ~ D, which yields
s
201 and =
Tr R D 20
In particular, & is now dominated by the signal learning term with B,T, 2 (1 D. We next consider the first form Lemma A.3.
b(1) = 1) 1) ro S SIST. T.
Since b(0)= 2 B1, we have
( b(1) 2 B1 '" (++1
1/(28)-1
[++1)
dl
23B1 (T+ (T + 1)-- )
2 (23 B)T
which implies the intrinsic term satisfies
D
T (23 - c)
However, in the second form, is dominated by the signal learning term and satisfies T 2 (1 c)D/B: Br. Therefore, the signal-learning term under the first form is than that under the second form by at least a constant factor. Since in the second form signal-dominated, this constant- factor improvement carries over to the total error, implying that the second form strictly dominates the first and is therefore optimal. Finally, from
C2 -1 Bi,
we obtain C2 D #+1 which gives the desired scaling for C2.
A.4 PROOF OF THEOREM 3.2 (OPTIMAL TWO-STAGE BSS)
Proof. Recalling that
6 ())-B,() = B1: 0<t
P
D
P
T.P =
Tp
B2, T.p <t< Tp,
B
B
B2
24

---

Published as a conference paper at ICLR 2026
For clarity, omit the explicit dependence D in P(D). Following Theorem 2.2,
d EB. B Ba (Tp) (11
1
(((Tp)
KCI (((D-))BB))
sTp
ap
B₂
Bi
B2
Since B1 B2, we have 1/B2
1/B22 0. Note that under the two-stage batch schedule setting,
P
D - P
dTp
1
Tp B1 +
0.
B2
dP
B1
B2
Since Tp ~ D.
K(Tp)
D (2-
sT-51 -D
B
Bi
(I) Under the hard-task regime with < 1/3 (s = 1 1/8 is trivially similar), it is trivial to, since D (-(2-1/0) o(D===1) the minimizer P* must satisfy the stationary point condition:
d
EB, & (B,
p=p- 0.
dp
In particular,
(D-P-)
D PM
~ D-
B2
Trivially, K((D B))/Ba) + 0, By the monotonicity of Kc this implies D - p*
- co. We have
(D (P) (2211111/8)) D D-
D P =
We now verify that the stationary point Pa a minimizer and corresponding Tp.:
d2
1
(T(TP)
K'((D P)/ B2)
dp2 EB, (BT) =
s(s 1)T
"T, +
Bi
B2
B1
B3
Note that Tú 0 and K' < 0, it suffices to prove
(T((TP)
Tp K'((D - PP/B2)
>0.
B1
B3
P=P+
The above inequality is trivial since
= ol (D-P))/B:o((T) = (
(II) Under the easy-task regime, since D= =0 00D-(2-1//))) have
dp E8 - (Tp) 0.
for sufficiently large D and for any Tp E [D/ Bi), thus the optimum satisfies P 0 for sufficiently large D.
B EXPERIMENTAL DETAILS AND ADDITIONAL RESULTS
B. LLM PRETRAINING: MODELS, DATA, AND TRAINING SETUP
Unless otherwise specified, language model pretraining Sections 4 and uses the following set- tings.
To verify whether the observed phenomena are consistent across scales, we perform experiments under two distinct settings.
25

---

Published as a conference paper at ICLR 2026
Table Model configurations
Type Model Size
LLaMA 200M
MoE
50M
492M
1001M
1119M
Activated Size
209M
291M
model dEF
512
1024
1280
512
576
2048
4096
5120
1408
1152
dFF_Mon
1408
192
__head
8
16
20
8
6
head
8
16
20
depth
15
12
24
n_expert
64
224
activated _eppeet
16
Small-scale experiment settings.
Model. LLaMA (Touvron et al., 2023) is a dense, decoder-only Transformer architecture that integrates several modern design components, including Rotary Positional Encoding (RoPE) (Su et al., 2024), Swish- Gated Linear Units (SwiGLU), and Root Mean Square Layer Normalization (RMSNorm). We pretrain LLaMA models with parameter sizes ranging from 50M to 492M. A full list of model configurations is provided in Table 1.
Dataset. Colossal Clean Crawled Corpus (C4) (Raffel et al., 2020) is large-scale, pub- licly available language dataset widely adopted for LLM pretraining, including models such as RoBERTa (Liu et al., 2019) and T5 (Raffel et al., 2020). For tokenization, we employ the T5 tokenizer with vocabulary size of 32,100. Following the setup of Zhao et al. (2024); Zhu et al. (2025); Wang et al. (2025a;b), we train with a sequence length 256. We use 1,000 linear warmup steps.
Large-scale experiment settings.
Model. Shortcut-comes connected Mixture of Experts (ScMoE) (Cai et al., 2025) is novel MoE chitecture that addresses communication overheads in expert parallelism by introducing short- connections and overlapping parallelization strategy. ScMoE decouples the usual quential dependency between communication and computation, enabling up to 100% over- lap of those two processes, which has demonstrated notable gains in inference efficiency and throughput compared to models of a comparable scale (LongCat et al., 2025). A full list of model configurations is provided in Table
Dataset. We train private, real-world LLM dataset to ensure that our experiments closely reflect practical deployment scenarios. The tokenizer configured with a vocabulary size of 131,072, and training is performed with maximum sequence length of 8,192.
Optimizer. For both small-scale and large-scale experiments, we adopt the standard AdamW (Loshchilov & Hutter, 2019) optimizer the baseline. The baseline configuration fol- lows protocols from LLaMA pretraining (Touvron et al., 2023), using hyperparameters 31 = 0.9. 3g = 0.95, weight decay X 0.1, and gradient clipping threshold of 1.0.
B.2 LINEAR REGRESSION EXPERIMENTS: SETUP AND DETAILS
empirically validate that the optimal batch size schedule alone is sufficient to achieve the optimal rates attainable for one-pass SGD in both easy-task and hard-tase task regimes.
The easy-task regime. We consider task with parameters s = 1.0, B 2.0, and a 1.0. We set the learning rate 7 0.0005 and adopt the batch size schedule preseribed by Theorem 3.1 as follows. Recalling that the optimal schedule under the easy-task regime satisfies
6 (t) ~ (max((" - - + 0 s T with Bmax D 1/2
T D 14*8
Due to the discrete nature of batch sizes, we replace the continuous time variable t by the iteration index k. and the horizon T by total number of iterations K. We introduce introduce data-scale
26

---

Published as a conference paper at ICLR 2026
hyperparameter Do and scale constant >> 0 to control the target data scale. The discrete batch size schedule is then constructed as
BR
Do T
1/2+e0
(K (K-
k= == 1....... K
with K = (a T
with v >> 0 stabilizes the schedule near the terminal stage. Accordingly, the total data size is of the same order as the data-scale hyperparameter, i.c. D = D 2i=l Bx ~ Do. We fix a = 1000, v/ = 10 and Do to be 2, 4, 8, 16, and 32. The corresponding values of D are 6346, 13973, 30331, 64962, and 137693, As illustrated in Figure 2 (middle), the batch size schedule alone is sufficient to achieve the minimax optimal risk rate under the easy-task regime.
The hard-task regime. this case, we consider the task with s = 0.4, B = 2.0, and o = 1.0. We set learning rate 7 0.0005 and the batch size schedule is configured according to Theorem 3.1 as follows. Recalling that the optimal schedule under the hard-task regime satisfies
Bminn
<
(t)
Bmax (T - t 1) #- s
with
D,
T - T
Bmax 211
Due to the discrete nature of batch sizes, we replace the continuous time variable t by the iteration index k and the horizon T and ii by discrete training length K and K₁. We introduce introduce a data-scale hyperparameter Do and scale constant C1 > 0 to control the target data scale. The discrete batch size schedule is then constructed as
for h: 1,... K1; for k: K1 + 1.... K,
#
K-K+v K-K1-1
with
K = [aDo, =1= ai= D
where scale constant V/ > O stabilizes the schedule near the terminal stage. Accordingly, the total data size is of the same order as the data-scale hyperparameter, i.c. D : E B₁ x Do- We fix = 1, b 10 and Do to be 2000, 4000, 8000, 16000, and 32000. The corresponding values of D are 6346, 13973, 30331, 64962, and 137693. As illustrated by Figure 2 (right), batch size schedule matches the predicted best rate achievable by one-pass SGD with learning rate schedule under the hard-task regime.
B.3 ADDITIONAL. DETAILS OF FAST CATCH-UP EXPERIMENTS
We conduct fast catch-up experiments across multiple scales:
0.5B model. We train 492M (== 0.5B) LLaMA model with learning rate 5 10-4 using a two-stage batch size schedule, switching from 512 to 1024, 2048, 4096 at step and step 25,000 in training, with total 100,000 steps.
1B model. We train a 1001M (Nr IB) MoE model using a two-stage batch size schedule, switching from 640 1280, 2560 50B, 200B and 300B tokens in training. In addition, we evaluate a multi-stage schedule that progressively increases the batch size-from 640 to I 1280, then 1920, and finally 2560 at 100B, 150B and 200B tokens in training, with total 600,000 steps.
1.1B model. We train 1119M (s= 1.IB) MoE model using a two-stage batch size schedule, switching from 1024 to 2048 at 300B and 600B tokens, with total 50,000 steps.
27

---

Published as a conference paper at ICLR 2026
B.4 ADDITIONAL DETAILS OF SWITCHING-TIME ANALYSIS EXPERIMENTS
In Figure 4 (left), we train a 200M LLaMA model on 4B tokens with learning rate 1 X 10 using a two-stage batch size schedule, switching from 256 to 512 at different points in training. The total data size corresponding
the full large batch size training step is 30000. We switch batch size at different ratio {0. 1/16, 2/16, 3/16, 4/16, 5/16, 6/16, 7/16, 8/16, 9/16, 10/16, 11/16, 12/16, 13/16, 14/16, 15/16, 16/163. Each ratio is repeated multiple times to reduce variance in the results.
In Figure 4 (right), we train a 50M LLaMA model trained on the C4 dataset with learning rate 1 10-3 using small batch size of 128 and large batch size of 256. The total data sizes corresponding to the full large batch size training step are (20000, 25000, 30000, 35000, 40000, 45000, 50000, 55000, 60000, 65000, 70000, 75000). For each data size, we perform grid search to determine the optimal switching point D with a precision of D/32. Each configuration of D is repeated multiple times to reduce variance in the results.
B.5 ADDITIONAL DETAILS AND RESULTS FOR LATE-SWITCH SUPERIORITY EXPERIMENTS
We conduct late-switch experiments across multiple scales:
0.5B model. We train a 492M (s= 0.5B) LLaMA model on 4B tokens with learning rate 5x 10-4 using two-stage batch size schedule. Specifically, we switch the batch size from 512 to either 1024, 2048, or 4096 at step 25,000.
1B model. We train a 1001M (== IB) MoE model on 0.4T tokens using a two-stage batch size schedule, switching from 640 to either 1280 or 2560 at the 50B, 200B, or 300B token marks. In addition, we evaluate a multi-stage schedule that progressively increases the batch size from 640 to 1280, then 1920, and finally 2560 at 100B, 150B, and 200B tokens, respectively.
1.1B model. We train a 111 ( I.IB) MoE model on IT tokens using a two-stage batch size schedule, switching from 1024 to 2048 at either the 300B or 600B token mark.
0.58 Denso, switch to x2 batch size
0.58 Densc, switch to x4 batch size
5.5
wholetime
wholetime
3.2-
2x batch diz from 25k
3.2-
4x batch siee from 25k
validation loss
3.1
validation loss
3.3
3.0
0.0
2.9
2.9
2.8
277
2.7
num of tokons (Billion)
num of tokens (Billion)
Figure 6: Validation loss versus training tokens under different batch size switching times using 0.5B LLaMA model trained on around 10B tokens, switching batch size from 512 1024 (left) and 2048 (right), respectively.
Experimental results are shown in Figure 4 (left), Figure 5, and Figure Moreover, we compare multi-stage butch size scheduling strategies for 200M LLaMA model and 1.1B MoE model. For 1119M MoE model, we train on IT tokens using a four-stage batch size schedule, switching from 1024 to 2048, then 3072 and finally 4096 at different time steps. For 200M LLaMA model, we train on 4B tokens using a four-stage batch size schedule, switching from 128 to 256, then finally 512 at different time steps.
In Figure and Figure 8, the left panels show how batch size evolves with training tokens, while the right panels report the corresponding validation loss. Across both model scales, later switching consistently yields lower validation loss than earlier switching, validating the effectiveness late- switch superiority in multi-stage batch size scheduling regime.
28

---

Published as a conference paper at ICLR 2026
batch vs
val loss vs taken
4.7
exry
4.6
5072
batch site
2045
4.5
4.4
earty
100
300
1041
100
1000
2000
400
600
500
1000
tokens (Bilioo)
num of tokens (Billion)
Figure 7: Validation loss versus training tokens with four-stage batch size schedule using 1.1B MoE model trained on tokens. Left: batch size versus training tokens; Right: validation loss versus training tokens.
batch size vs token
val loss vs token
512
3.4
carls
ddde
batdl 5bc
2.4
109
carly
middi
32
3.1
num tokens (Billion)
num (Billion)
Figure 8: Validation loss versus training tokens with three-stage batch size schedule using 200M LLaMA model trained on 4B tokens. Left: batch size versus training tokens; Right: validation loss versus training tokens.
B.6 EXTENSION: INTERACTION WITH LEARNING RATE SCHEDULING
B.6.1 COMPARISON WITH COSINE AND WSD
learning rate vs token
batch size us token
val loss us token
3000
weli
banch Size
00022
0990
(188or)
Figure 9: Validation loss versus training tokens among batch size schedule, cosine decay learning rate schedule, warmup-stable-deca. learning rate schedule using 50M I.LaMA model trained on 1B tokens. Left: learning rate versus training tokens; Middle: batch size versus training tokens; Right: validation loss versus training tokens.
We conduct set of proof-of-concept experiments to evaluate whether constant learning rate with batch size schedule can perform on par with mainstream learning rate schedulers used in LLM pretraining, such as cosine decay and Warmup-Stable-Deaa (WSD) schedule (Hu et al., 2024; Wen et al., 2025). Following established conventions (Hägele et al., 2024), the cosine schedule decays the learning rate 10% of its maximum value, whereas the WSD schedule decays it to zero, with the ratio of decay phase as 20%. For all figures, the left panels show the evolution of the learning rate over training tokens, the middle panels show the batch size trajectory, and the right
29

---

Published as a conference paper at ICLR 2026
learning rote
batch
0100
0.44
0.70
300 F1j
Caniie utz
6.00
238
15:01
11000
12000
2002
0.040
(900
2:34)
10:00
1200
2000
-.03
08.0
1290
mam
sum
Figure 10: Validation loss versus training tokens among batch size schedule, cosine decay learning rate schedule, warmup-stable-decay learning rate schedule using linear regression model. Left: learning rate versus training tokens; Middle: batch size versus training tokens; Right: validation loss versus training tokens.
panels report the corresponding validation loss curves. We denote the constant learning rate with batch size schedule 'bs", the cosine schedule 'cos', and the WSD schedule as 'wsd".
Figure shows the comparison for LLM pretraining. For the batch size schedule, we begin with a base batch size and increase it in stage-wing wise manner: switching to 2x the base batch size 70% of training tokens, 4x at 80%, and 8x at 90%. The base batch size is set to 128 for the 50M model. We emphasize that this batch size schedule is determined heuristically and is not optimized.
Figure 10 shows the comparison for linear regression. We set 8 = 0.3, 3 1.5, = 2, 7) = 0.05, the exponent in 22/3 is the batch size schedule derived by 1/(2.) 1. With the explicit 8, design an optimal batch size schedule according to Theorem 3.1.
We observe that, across both LLM pretraining and linear regression, constant learning rate with an appropriately designed batch size schedule achieves performance comparable to widely adopted learning rate schedulers.
B.6.2 FAST CATCH- UP UNDER LEARNING RATE DECAY
In this section, to explore the influence of learning rate decay, we replicate the late-switch superior- ity experiments from Appendix B.5 on 50M and 0.5B LLaMA models using cosine learning rate schedule. As shown in Figure/ 11, the characteristic phenomena fast catch-up and later switch- ing-persist. Note that catch-up quantified in terms of the intrinsic time T. Under constant learning rate regime, T advances at uniform pace, whereas a cosine schedule causes it to advance more slowly toward the end of training. Consequently, the apparent merge speed decreases in the final stages. While current theoretical analysis focuses a constant LR, the FSL mechanism is general and naturally carries over to other LR schedules, making the theoretical extension to such settings straightforward.
C STATEMENT
C.I ETHICS STATEMENT
We have confirmed that this research was conducted in full compliance with the ICLR Code of Ethics. All experiments respect the principles of integrity, faimess, and transparency. No part of this work involves harm to humans, animals, or the environment, and we have taken care to ensure the responsible use of data, models, and computational resources.
C.2 REPRODUCIBILTY STATEMENT
We believe that all experimental results this work are reproducible. The paper specifies compre- hensive training and evaluation
hyperparameters, optimizer choices, and other relevant settings-in Section 5 and Appendix B. For small-scale experiments, we provide open- source code in the supplemental material, and all datasets used publicly available. For large-scale
30

---

Published as a conference paper at ICLR 2026
cos, 50M, val loss vs step
cos, 50M, val loss vs taken
3.9
3.90
2x batch von wholeime from
2x batch win wholecime
9.85
3.8
3.80
validation loss
3.755
validation
3.6
8.70
3.65
3.60
3.4
3.65
3.5
3.50
10
20
2.5
300
num of steps (Thousands)
num of tokens (Dillion)
(a) cosine learning schedule, 50M Dense
COS, 0.5B, val loss vs step
cos, 0.5B, loss vs token
luttc size wholetime 2x buitch size from 25
batch size wholetime batch vow from 25k
0.1
3.2
validetion loss
3.0
validation loss
2.8
2.0
27
24
2.7
20
40
60
90
100
num of (Thousands)
num of tokens (Billion)
(b) cosine learning schedule, 0.5B Dense
Figure 11: Two-stage batch size switching using 50M and 0.5B LLaMA model trained 3.2B and 10B tokens, respectively. Left: validation loss versus training steps; Right: validation loss versus training tokens.
experiments, we believe that employing comparable datasets and training pipelines will reproduce the same phenomena.
C.3 LLM USAGE STATEMENT
We used the LLM writing assistant during paper preparation. The model was used to identify and correct grammatical errors throughout the manuscript. It suggested ways to make our sentences clearer and smoother. The LLM helped polish the language while keeping our meaning intact. We limited LLM use to only language editing tasks. All research content and ideas came entirely from human work.
Beyond serving as tools, LLMs were themselves the subject of our study. We trained these models and analyzed their behavior uncover and explain novel phenomena. Importantly, this use of LLMs as research objects should not be misinterpreted as substantive contribution from the models to the work itself.
31


Reviews:
[{"review_id": "IkH43OriQu", "reviewer": "Reviewer_iVq9", "summary": "This paper studies the effect of batch size scheduling in LLM pretraining, focusing on a simplified case, where the batch size increases just once. They make two findings: (1) when the batch size increases, the test error drops quickly then stabilizes, and (2) after a while, the test error curve merges with the one induced by using a constant large batch size from the beginning. \n\nUsing a continuous approximation of online SGD applied to kernel linear regression, this paper proves the above two observed phenomena. Moreover, they calculated the optimal batch size schedule (in the simplified sense) using their theory. Finally, they conducted experiments in LLaMA (upto 492M parameters) and MoE models (upto 291M activated parameters) to verify their theory prediction.", "questions": "My major concern with the paper is regarding the theory component, which seems fabricated instead of fully rigorous. First, the problem setting is highly simplified: constant step size online SGD for kernel linear regression, with a potential two-stage batch size scheduler, and under the source and capacity conditions. \n\nThis problem, without the two-stage batch size complication, is very well studied in kernel linear regression literature. However, this paper chooses to take a weird approach by analyzing a continuous approximation of the discrete online SGD process. \n\nNote that in online SGD, the optimization time is tied with the data size. Hence, it can be dangerous extending insights from continuous process to the discrete process. I have noticed one such issue in Theorem 4.4. Detailed as the following question. \n\n1. Theorem 4.4 cannot be correct -- it is well known that one-pass SGD is suboptimal for certain power-law classes. Specifically, in Theorem 4.4, when $\\\\beta > 1+s\\\\beta$, the optimal time $T^*$ is greater than data size $D$, which violates the online nature of the algorithm. Therefore, the discussions in Section 4.3 are misleading. \n\nDespite this major issue, it is unclear why the paper focuses on constant step size setting. \n\nIn this setting, it is well known that the last iterate of SGD does not converge due to the additive variance error. Indeed, the drop of the loss by increasing batch size is exactly because increasing batch size decreases the variance error. However, in standard SGD literature, one should analyze the averaged iterates or the last SGD iterate but with a decaying step size scheduler. The latter is also what practitioners do. This leads to my second question:\n\n2. For online SGD with a reasonable step size scheduler, e.g., exponentially decaying one or cosine, would increasing batch size also cause a sudden drop of the test error? \n\nThose issues could have been avoided by using the well known rigorous tools developed by prior kernel linear regression literature, instead of using the heuristic continuous approximation. \n\n\nBesides those theoretical questions, I also feel the experiments are not on a sufficiently large scale. However, I am not an expert here, so I will also wait to see other reviewers’ comments on this. \n\nOverall, I am not fully convinced by the sudden drop of the loss story. But even if I choose to believe it, I do not quite see the implications of this observation:\n\n3. Is the two-stage batch size schedule of any practical importance? When considering more comprehensive batch size schedulers, would the sudden drop still hold to some extent? What's the motivation for studying two-stage batch size schedule?\n\n\nA minor issue.\n\n4. Lemma 3.3 relies on a certain fourth moment hypercontractivity condition on the feature, however, this is not explicitly mentioned. The constants 2 and 4 seem to suggest they need the feature to be exactly Gaussian.", "limitations": "", "rating": 2, "confidence": 3, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "See below.", "weaknesses": "See below.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "Cr4kiFC9sv", "reviewer": "Reviewer_wCEH", "summary": "The paper studies two-stage batch-size schedules for LLM pre-training under a constant learning rate: start with a small batch, then switch once to a larger batch. Empirically, the authors report two robust phenomena across dense and MoE models up to 1.1B parameters and up to 1T tokens: (i) a “sudden drop” in loss at the switch, and (ii) “final merge,” where the loss trajectory measured in steps converges toward the always-large-batch curve. They analyze these behaviors using the Functional Scaling Law (FSL) within a power-law kernel (PLK) teacher–student framework, proving the sudden-drop and final-merge effects, and deriving a “later-switch” rule in data-limited regimes along with a power-law scaling of the optimal switch point with total data size. Experiments at several scales corroborate the theory and show later switches tend to yield better final loss under a fixed token budget.", "questions": "Do sudden-drop/final-merge and the later-switch rule hold under the cosine decay method?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. Clear empirical phenomena distilled. The paper cleanly isolates and names two behaviors (“sudden drop,” “final merge”) and shows them across architectures/scales, which aids practitioner understanding.\n2. Theory that matches practice. The FSL-based analysis explains both phenomena and yields a concrete later-switch rule and a power-law prediction for the optimal switch point, predictions borne out in experiments (including a strong log–log fit).\n3. Breadth of evidence. Results span dense (LLaMA-style) and MoE models, billions of tokens, and multiple switch ratios/timings; figures are easy to digest.\n4. Actionable takeaway. Under a fixed token budget, “switch later rather than earlier” is a simple guideline practitioners can trial. The paper also discusses hardware constraints motivating staged schedules.", "weaknesses": "1. Theory and most experiments assume constant LR, while real LLM training typically uses warmup + cosine/linear decay.\n2. Large-scale runs use a private dataset, which limits external reproducibility.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "pFR0Sb8C0A", "reviewer": "Reviewer_A634", "summary": "The paper studies batch size scheduling via a power-law regression theoretical model. It provides an explanation behind two observed effects when switching batch size from low to high: sudden loss drop and eventual matching of the loss curve with the high batch size loss curve. It also proposes an optimal batch size schedule from the theoretical model.", "questions": "1. Does Lemma 3 hold generally? Can the authors provide a proof for the same?\n\n2. Is there a practical way of implementing (or obtaining) the optimal batch size scheduling scaling law? Can it be verified that its performance matches the performance of cosine decay in practice?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The paper studies batch size scheduling from a theoretical perspective of a power-law regression model and provides insights related to ‘sudden drop’ and ‘final merge’ of the loss values.\n\n2. It proposes scaling law for optimal switching time from small to large batch in a training run and also empirically verifies that practical settings obey a scaling law.\n\n3. It also proposes an optimal batch size scheduling algorithm for the power-law model.", "weaknesses": "1. The paper only studies a constant learning rate schedule, which deviates from practice.\n\n2. Although the paper proposes an optimal batch size scheduling algorithm as a power law, it provides no way of actually developing a practical optimal scheduling algorithm.\n\n3. I don't think Lemma 3 holds for any arbitrary $\\theta$, but only for local minimizers as the expected gradient has to be zero for this to hold.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "AzGxOPPHqt", "reviewer": "Reviewer_nk2G", "summary": "This paper provides a theoretical foundation for the common practice of batch size scheduling in LLM pre-training. Motivated by two robust empirical observations of a \"sudden drop\" in loss and a \"final merge\" of loss curves upon switching batch sizes, the authors apply the Functional Scaling Law (FSL) framework from Li et al. (2025) to explain these dynamics. The theoretical analysis successfully explains these phenomena and yields several key predictions: a \"later is better\" rule for switching in data-limited regimes, a quantitative power law relating the optimal switch point to the total data budget, and the functional form of a minimax-optimal continuous schedule. These predictions are supported by experiments on models up to 1.1B parameters, bridging the gap between theory and practice.", "questions": "1. How should FSL parameters be interpreted for LLMs? The FSL framework relies on a power-law kernel characterized by the capacity parameter $\\beta$, which governs eigenvalue decay and drives all theoretical predictions. However, the paper doesn't clarify what $\\beta$ represents in language model pre-training. What aspects of the data distribution, architecture, or task does it capture?\n2. Can the FSL framework predict optimal joint schedules? Li et al. (2025) used FSL to analyze learning rate schedules, while this paper analyzes batch size schedules. The natural next step is unification: can FSL simultaneously optimize both schedules?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "1. **Concrete and Predictive Theory:** This work transforms batch size scheduling from heuristic practice into principled science by applying the FSL framework to derive quantitative predictions. The theorems rigorously characterize loss dynamics (sudden drop and final merge) and provide actionable guidance for schedule design.\n2. **Strong Empirical Corroboration:** The theoretical predictions are robustly validated across diverse settings. The core phenomena are shown to be robust across different model architectures (LLaMA and MoE) and scales. It is particularly compelling to see the near-perfect empirical fit to the predicted power-law scaling of optimal switch points, which shows the theory's predictive power and practical utility.", "weaknesses": "1. **Constant Learning Rate Assumption:** The analysis assumes a constant learning rate throughout training to isolate batch size effects. This is a major limitation, as modern large-scale pre-training universally employs learning rate schedules (e.g., cosine decay with warmup). \n2. **Missing Empirical Validation of the Optimal Schedule:** Theorem 4.4 derives a continuous, increasing batch size schedule that is provably optimal within the FSL framework. This is arguably one of the paper's most significant theoretical contribution, yet it is not empirically tested. The authors mention practical hardware and software constraints that make continuous schedules difficult to implement, but a proof-of-concept using a discretized approximation of the optimal schedule would have strengthened the paper's conclusions.\n3. **Asymptotic Nature of the Theory:** The core theoretical results rely on asymptotic analysis where training time or data budget approaches infinity. While experiments demonstrate reasonable approximation quality for large finite runs, the theory cannot formally guarantee performance in the finite regimes where practitioners actually operate.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)