PaperlayerView exact revision on arXiv

Source-linked AI summary

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

arXiv:2602.14208v2cs.LGmath.OCstat.ML

TL;DR

Batch size scheduling lacks a principled theoretical foundation despite its importance for efficient large-scale training. The paper extends functional scaling laws to derive task-dependent optimal schedules and explains late switching through fast catch-up, with experiments showing late-switch schedules outperform constant-batch and early-switch baselines.

  • Problem

    Batch size scheduling is important for large-scale training, but its theoretical foundations and design principles remain limited.

    Source: §1 · pp. 1–2

  • Method

    The paper extends functional scaling laws to analyze optimal batch-size schedules under fixed data or compute budgets and explain fast catch-up.

    Source: §1 · p. 2, §4.1 · p. 9

  • Results

    Late-switch schedules consistently outperform constant-batch and early-switch baselines across the evaluated LLM pretraining settings.

    Source: §1 · pp. 2–3, §6 · p. 11

  • Takeaways & Limitations

    For hard tasks, large batches can be deferred to late training while preserving performance and reducing data consumption.

    Source: §1 · pp. 2–3, §6 · p. 11

    The analysis assumes standard SGD, whereas modern LLM training predominantly uses adaptive optimizers.

    Source: §6 · p. 11

Abstract

from arXiv · show

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) 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.

1 Introduction

The paper extends functional scaling laws to batch size scheduling, deriving task-dependent optimal schedules and identifying fast catch-up as the mechanism enabling late switching. Experiments across LLM settings support late-switch schedules over constant-batch and early-switch alternatives.

1 Introduction

  • Large batches improve hardware utilization and throughput, but constant large-batch training is suboptimal for sample efficiency, motivating dynamic batch-size schedules.The introduction links this trade-off to changing optimization signal and gradient-noise influence during training.Source: §1 · pp. 1–2
  • FSL supplies a principled framework for analyzing batch size scheduling, extending earlier work that focused primarily on learning-rate schedules.The framework is presented as a response to limited theoretical analysis and heuristic BSS design.Source: §1 · p. 2, §1.1 · p. 3
  • Hard tasks favor prolonged small-batch training followed by a late switch to large batches, whereas easy tasks favor increasing batch sizes throughout training.In stage-wise schedules, easy tasks favor constant large-batch training while hard tasks retain the late-switch structure.Source: §1 · p. 2
  • After switching from small to large batches, loss rapidly collapses to the constant large-batch trajectory across theoretical and LLM-pretraining settings.The effect is reported across diverse architectures, model scales, and data regimes and is illustrated in Figure 1.Source: Figure 1 · p. 2, §1 · p. 3
  • Across Dense and MoE models from 50M to 1B parameters and 10B to 1T tokens, late-switch schedules consistently outperform constant-batch and early-switch baselines.The reported validation spans diverse architectures, model scales, and data regimes.Source: §1 · p. 3

1.1 Related Work

  • The paper builds on scaling-law research and contrasts its BSS theory with prior work centered on static batch sizes, heuristics, or optimal-control formulations.Related work had given BSS less theoretical attention than learning-rate scheduling.Source: §1.1 · p. 3

2 Preliminaries

The preliminaries model one-pass SGD in feature space and then express its loss dynamics through a continuous-time stochastic differential equation and functional scaling law. The resulting decomposition separates signal learning from batch-size-dependent noise accumulation.

2.1 Feature-Space Linear Regression

  • The setup uses feature-space linear regression with Gaussian features, noisy labels, power-law assumptions, and one-pass SGD with a constant learning rate and batch-size schedule.Performance is measured by excess risk.Source: §2.1 · p. 4
  • The source exponent s measures target alignment and task difficulty, while capacity exponent β controls spectral decay and model capacity.Smaller s indicates harder learning problems, whereas smaller β corresponds to higher model capacity.Source: §2.1 · p. 4

2.2 Functional Scaling Laws

  • The discrete SGD dynamics are modeled by an Itô SDE with a continuous-time batch-size function b satisfying b(kη) = Bk.Continuous training time maps each discrete iteration k to time t = kη.Source: §2.2 · p. 4
  • The FSL maps the batch-size schedule to loss dynamics and separates signal learning from noise accumulation.The signal term is governed by source exponent s, while the noise term depends on how the schedule dissipates gradient noise.Source: §2.2 · p. 5
  • The forgetting kernel K(t − τ) describes how noise injected at time τ affects later loss, with higher-capacity models tending to forget noise more slowly.The kernel is specified as K(t) = (t + 1)^−(2−1/β).Source: §2.2 · p. 5

3 Theoretical Analyses via Functional Scaling Laws

The FSL framework turns optimal batch-size scheduling under a fixed data budget into an analytically tractable resource-constrained problem. Its solutions depend sharply on task difficulty, favoring increasing batches for easy tasks and delayed growth or switching for hard tasks.

3 Theoretical Analyses via Functional Scaling Laws

  • The analysis imposes a fixed data budget and hardware lower bound b(t) ⩾ Bmin, while treating the schedule as a resource-constrained variational problem.For a fixed model, data and compute budgets are equivalent because computational cost scales linearly with data size.Source: §3 · p. 5, §3.1 · p. 5

3.1 Optimal Batch Size Scheduling without Shape Constraints

  • For easy tasks, the optimal unconstrained schedule increases batch size throughout training, with peak size scaling as Bmax ≂ D^α for α > 0.Larger datasets therefore favor larger peak batches under the derived schedule.Source: §3.1 · p. 6
  • For hard tasks, the unconstrained optimum holds the batch size at Bmin for most training, then grows during a vanishingly small late phase.The growth phase occupies (T⋆−T⋆1)/T⋆ = oD(1), so extremely large batches are needed mainly near the end.Source: §3.1 · p. 6
  • The hard-task stable-growth structure trades fewer signal-learning errors for late noise reduction by preserving more optimization steps under a fixed data budget.Small batches increase the number of steps T, reducing the signal-learning term T^-s; late growth primarily reduces noise.Source: §3.1 · p. 6
  • The derived rates are D^-sβ/(1+sβ) for easy tasks and D^-s for hard tasks, matching the relevant minimax or one-pass-SGD optima.These schedules preserve data efficiency while reducing iteration complexity relative to learning-rate scheduling.Source: §3.1 · p. 6, Figure 2 · p. 7
  • Discrete-time SGD with the prescribed optimal schedule follows the predicted data scaling, and its constant-learning-rate performance is comparable to cosine and warmup–stable–decay schedules.Figure 2 reports the easy-task minimax-rate match and the hard-task D^-s match.Source: §3.1 · pp. 6–7, Figure 2 · p. 7

3.2 Stage-Wise Optimal Batch Size Scheduling

  • In the two-stage setting, easy tasks optimally use the large batch throughout, whereas hard tasks use the small batch for most training and switch only very late.The switching point follows a scaling law, enabling estimation from small-scale pilot experiments and extrapolation to large training runs.Source: §3.2 · pp. 7–8

4 The Fast Catch-Up Effect: A Bridge to LLM Pretraining

The fast catch-up effect explains why hard-task schedules can defer large batches until late training while preserving the large-batch trajectory. FSL attributes this behavior to rapid forgetting of noise accumulated during small-batch training, with harder tasks catching up faster.

4 The Fast Catch-Up Effect: A Bridge to LLM Pretraining

  • After a batch-size increase, the loss rapidly collapses onto the corresponding constant large-batch trajectory, even when small batches were used for most training.This fast catch-up repeatedly occurs across multiple switching times and stages.Source: §4 · p. 8
  • Fast catch-up consistently appears across Dense and MoE settings, including models up to 1.1B parameters and training runs up to 1T tokens.Figure 3 spans a 0.5B LLaMA model, a 1B MoE model, and a 1.1B MoE model; its middle panel includes a four-stage schedule.Source: §4 · p. 8, Figure 3 · p. 8
  • Late switching matches constant large-batch validation loss while reducing token consumption and computational cost without sacrificing performance.The strategy starts with small batches and defers the transition to a large batch until a late stage.Source: §4 · pp. 8–9

4.1 An Explanation via Functional Scaling Laws

  • At the switching point, the loss gap arises solely from different noise accumulation because both runs have undergone the same optimization time.FSL decomposes excess risk into signal learning and a batch-size-dependent noise term.Source: §4.1 · p. 9
  • After switching, catch-up progressively forgets the initial small-batch noise, while the forgetting exponent depends on capacity exponent β rather than task difficulty.The large-batch baseline evolves on the slower signal-learning time scale.Source: §4.1 · p. 9
  • For hard tasks, the switched trajectory relaxes on a much faster time scale than the large-batch baseline, and smaller task difficulty produces faster catch-up.The separation is expressed as δ_ϵ ≪ t⋆ for hard tasks.Source: §4.1 · p. 9

5 Validating Late Switching in LLM Pretraining

Experiments validate that late batch-size switching improves training outcomes across model scales and architectures. The optimal switch occurs late, and the predicted power-law relationship is supported empirically.

Fine-grained analysis of the switching time

  • At approximately 70% of training tokens, the optimal switching point yields improved performance.This result comes from the 200M LLaMA experiment summarized in Figure 4 (left).Source: §5 · p. 10, Figure 4 · p. 10

Power-law scaling

  • The fitted relationship between optimal switching point and data size closely follows the predicted power law.Experiments use a 50M-parameter model trained on C4 with token budgets from 1.3B to 5B.Source: §5 · p. 10

Larger-scale validation

  • Across switching ratios, architectures, and training scales, late switching consistently outperforms early switching.The comparison includes the larger-scale settings shown in Figure 5 and additional multi-stage schedules.Source: §5 · pp. 10–11, Figure 5 · p. 11

Experimental scope

  • The experiments cover 1B and 1.1B MoE models trained on 0.4T and 1T tokens with batch-size increases from 640 to 1280, 512 to 2048, and 1024 to 2048.Source: §5 · p. 10, Figure 5 · p. 11

6 Conclusion

The paper frames FSL as a principled framework for batch-size scheduling, explains late switching through task difficulty and noise dynamics, and identifies boundaries for the analysis. Its theoretical results characterize optimal schedules under data-budget constraints and motivate the empirical findings.

6 Conclusion

  • FSL provides a principled framework in which hard tasks favor small batches for most training before a late transition to large batches.The conclusion connects this schedule structure to the fast catch-up effect observed in realistic LLM pretraining.Source: §6 · p. 11
  • The analysis assumes standard SGD and constant learning rates, leaving adaptive optimizers and joint learning-rate decay with batch scheduling for future work.Cosine-decay experiments suggest fast catch-up continues approximately, but systematic treatment remains open.Source: §6 · p. 11

A.1 Interpretation of the Source and Capacity Conditions

  • The task difficulty parameter depends on target-function regularity relative to model capacity, so higher-capacity models make the same task appear easier.The source and capacity exponents determine this relative difficulty relationship.Source: §A.1 · p. 19

A.2 Proof of Theorem 2.2 (Self-Contained Derivation of FSL)

  • The FSL proof uses anisotropic noise structure and a closed-form noise covariance to characterize stochastic-gradient dynamics.The derivation represents the dynamics through an SDE and tracks noise using the resulting covariance structure.Source: §A.2 · p. 20

The optimal batch size schedule obeys

  • The optimal two-stage schedule is characterized by a switching point, with hard-task analysis showing that the optimal switch is delayed into a late phase.The proof compares candidate forms and establishes the relevant minimizer through the stationary-point conditions.Source: The optimal batch size schedule obeys · p. 23

A.3 Proof of Theorem 3.1 (Shape-Unconstrained Optimal BSS)

  • Under a fixed data budget, optimal unconstrained schedules differ by task difficulty: easy tasks support increasing batches, whereas hard tasks require constrained growth and late batching.The hard-task proof shows that the late-growth form dominates the alternative by a constant-factor improvement in signal learning.Source: The optimal batch size schedule obeys · pp. 22, 23, 24, 26

A.4 Proof of Theorem 3.2 (Optimal Two-Stage BSS)

The two-stage BSS proof analyzes the switching point under fixed data and shows that hard tasks favor a delayed transition from the smaller to the larger batch.

A.4 Proof of Theorem 3.2 (Optimal Two-Stage BSS)

  • For hard tasks, the optimal two-stage schedule has a switching point determined by stationary conditions, with the transition delayed rather than made early.The proof verifies that the stationary point is a minimizer and compares the candidate schedule against alternatives.Source: §A.4 · pp. 26–27

B.1 LLM Pretraining: Models, Data, and Training Setup

The experiments use separate dense and MoE LLM configurations at small and large scales. They specify the model architectures, datasets, token budgets, and baseline optimizer used for evaluating batch-size scheduling.

Training setup

  • Both small- and large-scale experiments use AdamW with β1 = 0.9, β2 = 0.95, weight decay λ = 0.1, and gradient clipping at 1.0.Source: §B.1 · pp. 27–28

Small-scale experiment settings

  • Small-scale experiments pretrain dense, decoder-only LLaMA models with 50M–492M parameters on the C4 dataset.Source: §B.1 · p. 28

Large-scale experiment settings

  • Large-scale experiments use Shortcut-connected Mixture of Experts (ScMoE), a novel MoE architecture with shortcut connections and overlapping parallelization, and train on a private, real-world LLM dataset.The tokenizer has a vocabulary size of 131,072, and training uses a maximum sequence length of 8,192.Source: §B.1 · p. 28

B.2 Linear Regression Experiments: Setup and Details

The linear-regression experiments test theoretically prescribed batch-size schedules in easy- and hard-task regimes under controlled data scales. The schedules achieve the predicted optimal rates, with batch-size scheduling alone attaining minimax-optimal risk for easy tasks and the best one-pass SGD rate for hard tasks.

Experimental setup

  • The experiments evaluate schedules derived from Theorem 3.1 for easy and hard tasks using fixed learning rate η = 0.0005.The easy task uses s = 1.0, β = 2.0, and σ = 1.0, while the hard task uses s = 0.4, β = 2.0, and σ = 1.0.Source: §B.2 · pp. 28–29

The easy-task regime

  • Batch size scheduling alone achieves the minimax optimal risk rate under the easy-task regime.The experiment varies D0 over 2, 4, 8, 16, and 32, producing total data sizes from 6346 to 137693.Source: §B.2 · p. 29

Discrete schedule construction

  • Discrete schedules replace continuous training time with iteration index k and control data scale through D0 and a positive scale constant.The resulting total data size is of the same order as the data-scale hyperparameter.Source: §B.2 · p. 29

The hard-task regime

  • The hard-task schedule matches the predicted best rate achievable by one-pass SGD with a learning-rate schedule.For D0 values 2000, 4000, 8000, 16000, and 32000, the corresponding total data sizes are 6346, 13973, 30331, 64962, and 137693.Source: §B.2 · p. 29

B.3 Additional Details of Fast Catch-Up Experiments

The appendix details fast-catch-up and switching-time experiments across multiple model scales, architectures, datasets, and switching configurations. These experiments vary switching points and batch-size trajectories to test how late switching affects training dynamics.

B.3 Additional Details of Fast Catch-Up Experiments

  • Fast-catch-up experiments span 0.5B LLaMA, 1B MoE, and 1.1B MoE models with two-stage and multi-stage batch-size schedules.The settings include switches from 512 to 1024–4096, from 640 to 1280–2560, and from 1024 to 2048 at different training times.Source: §B.3 · p. 29

B.4 Additional Details of Switching-Time Analysis Experiments

  • Switching-time analyses vary the switch ratio from 0 to 16/16 for a 200M LLaMA model and grid-search the optimal switch point for a 50M LLaMA model.The experiments repeat configurations multiple times to reduce variance.Source: §B.4 · p. 30

B.5 Additional Details and Results for Late-Switch Superiority Experiments

Additional experiments test late-switch superiority across multi-stage schedules, learning-rate baselines, and decayed learning rates. Late switching remains favorable in multi-stage settings, while batch-size scheduling performs comparably to mainstream learning-rate schedules and fast catch-up persists under cosine decay.

B.5 Additional Details and Results for Late-Switch Superiority Experiments

  • Late-switch experiments span a 492M (≈0.5B) LLaMA trained on 4B tokens, a 1001M (≈1B) MoE on 0.4T tokens, and a 1119M (≈1.1B) MoE on 1T tokens.The schedules switch batch sizes at specified training steps or token marks, including a multi-stage schedule increasing from 640 to 1280, 1920, and 2560.Source: §B.5 · p. 30
  • The late-switch experiments use two-stage schedules across 0.5B LLaMA, 1B MoE, and 1.1B MoE models at scales up to 1T training tokens.The 1.1B MoE setting switches from 1024 to 2048 at either 300B or 600B tokens.Source: §B.5 · p. 30
  • Later switching consistently yields lower validation loss than earlier switching in multi-stage schedules for both 200M LLaMA and 1.1B MoE models.The corresponding figures show batch-size trajectories on the left and validation-loss curves on the right.Source: §B.5 · p. 31, Figure 7 · p. 31, Figure 8 · p. 31

B.6.1 Comparison with Cosine and WSD

  • A constant learning rate with an appropriately designed batch-size schedule achieves performance comparable to cosine decay and Warmup–Stable–Decay schedules in LLM pretraining and linear regression.The LLM batch-size schedule uses heuristic increases at 70%, 80%, and 90% of training tokens, while the linear-regression schedule is designed from Theorem 3.1.Source: §B.6.1 · p. 32

B.6.2 Fast Catch-Up under Learning Rate Decay

  • Fast catch-up and the benefit of later switching persist under cosine learning-rate decay, although apparent late-stage merging slows because intrinsic time advances more slowly.The theoretical analysis focuses on constant learning rate, while the authors state that the FSL mechanism carries over to other learning-rate schedules.Source: §B.6.2 · p. 32, Figure 11 · p. 33

C.2 Reproducibility Statement

The reproducibility statement reports detailed training and evaluation specifications, open-source code for small-scale experiments, and publicly available datasets. The paper also states that LLMs were used only for language editing and were not contributors to the research content.

C.2 Reproducibility Statement

  • The paper provides training and evaluation details, open-source code for small-scale experiments, and publicly available datasets.The authors state that comparable datasets and training pipelines should reproduce the same phenomena at large scale.Source: §C.2 · p. 33

C.3 LLM Usage Statement

  • LLMs were used only to identify grammatical errors and polish language, while the research content and ideas came from humans.The trained LLMs were research objects, not substantive contributors to the work.Source: §C.3 · pp. 33–34
Loading 2602.14208v2…