arXiv:2602.14208cs.LGmath.OC2026-02中稿 · ICLR被引 6

根据任务难易度动态调整训练批次大小,可大幅节省数据且不损失性能。

Fast Catch-Up, Late Switching: Optimal Batch Size Scheduling via Functional Scaling Laws

  • 按任务难度设计批次大小调度:简单任务持续增大,复杂任务前期小批后期切换
  • 晚期切换可快速追上大批次损失轨迹,关键在于梯度噪声快速遗忘
  • 适用于大模型预训练,尤其适合参数量超10亿的密集和MoE架构

批量大小调度(BSS)在大规模深度学习训练中至关重要,影响优化动态与计算效率,但其理论基础仍不清晰。本文基于Li等(2025a)提出的函数缩放律(FSL)框架,揭示了固定数据预算下的最优BSS结构强烈依赖于任务难度。对于简单任务,最优策略持续增加批量大小;而对于困难任务,最优策略在训练早期保持小批量,仅在后期阶段切换至大批量。我们发现这一‘晚期切换’现象源于‘快速追上’动力学机制——从少量批转换为大批量后,损失迅速对齐恒定大批量轨迹。通过FSL分析表明,该效应源于累积梯度噪声的快速遗忘,追上速度由任务难度决定。关键启示是:大批量可安全延后,显著降低数据消耗而不牺牲性能。大量LLM预训练实验(涵盖密集与MoE架构,参数量达1.1B,数据量达1T tokens)验证了理论预测:所有设置下,晚切换策略均优于恒定批量和早切换基线。

原文摘要 · Abstract (English)

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.

批量调度大模型训练函数缩放律高效训练

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。