Why Does Adaptive Batching Help LLM Pretraining? A Perspective from Unbounded Variance

📄 arXiv: 2610.02355v1 📥 PDF

作者: Arda Fazla, Antesh Upadhyay, Ege C. Kaya, M. Berk Sahin, Abolfazl Hashemi

分类: cs.LG, math.OC

发布日期: 2026-10-01

💡 一句话要点

提出自适应批处理调度器以优化LLM预训练中的方差控制

🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 自适应批处理 方差控制 大型语言模型 噪声模型 随机优化 训练效率 深度学习

📋 核心要点

  1. 现有的随机优化方法假设随机梯度方差是均匀有界的，但在实际应用中这一假设常常不成立，导致训练效果不佳。
  2. 论文提出了一种广义BG-$a$噪声模型，能够根据训练过程中的方差增长动态调整批处理大小，从而优化训练效果。
  3. 在预训练OLMo2模型的实验中，使用该自适应批处理调度器的验证损失低于小批量和大批量训练，同时迭代次数少于小批量训练的10%。

📝 摘要（中文）

在大型语言模型（LLM）预训练中，增加批处理大小是一种常见做法，但其成功的理论依据尚不明确。现有的随机优化分析通常假设随机梯度方差是均匀有界的，但在许多实际的非凸问题中，这一假设并不成立。本文提出了一种广义的BG-$a$噪声模型，能够更准确地描述实际噪声行为，并在此基础上提出了一种自适应批处理调度器，通过动态调整批处理大小来控制训练过程中的方差增长。实验结果表明，该调度器在预训练OLMo2模型时，能够在匹配的token预算下，取得比小批量和大批量训练更低的验证损失。

🔬 方法详解

问题定义：本文旨在解决在大型语言模型预训练中，随机梯度方差增长控制不足的问题。现有方法通常假设方差是有界的，但在实际应用中，这一假设往往不成立，导致训练效率低下。

核心思路：论文提出了一种广义BG-$a$噪声模型，允许方差根据距离初始化的程度进行调整，从而更好地反映实际噪声行为。通过动态调整批处理大小，控制方差的增长，进而提升训练效果。

技术框架：整体方法包括引入广义BG-$a$噪声模型、推导信息论下界、建立与批处理大小相关的上界，以及设计自适应批处理调度器。主要模块包括噪声模型的构建、复杂度分析和调度器的实现。

关键创新：最重要的创新在于提出了广义BG-$a$噪声模型，该模型在方差控制上比传统的有界方差假设更具灵活性，能够适应实际训练中的噪声变化。

关键设计：在设计中，调度器根据训练过程中的方差动态调整批处理大小，确保在迭代过程中有效控制方差增长。此外，模型的复杂度分析提供了理论支持，确保了方法的有效性。

🖼️ 关键图片

img_0
img_1
img_2

📊 实验亮点

实验结果显示，使用自适应批处理调度器的OLMo2模型在验证损失上优于小批量和大批量训练，且在匹配token预算下，迭代次数仅为小批量训练的10%以下，展现了显著的训练效率提升。

🎯 应用场景

该研究的潜在应用领域包括自然语言处理、机器翻译和对话系统等大型语言模型的训练。通过优化批处理调度，可以显著提高模型的训练效率和性能，降低计算资源的消耗，具有重要的实际价值和广泛的应用前景。

📄 摘要（原文）

Increasing the batch size during training is a common practice in large language model (LLM) pretraining, yet the theoretical justification behind its success is not well understood. Analyses of stochastic optimization often assume uniformly bounded stochastic gradient variance, yet recent evidence suggests that this assumption fails in many practical nonconvex problems. The Blum--Gladyshev (BG-$0$) noise model relaxes this assumption by allowing the variance to grow quadratically with the distance from initialization, suggesting that batch size schedulers can help by controlling the variance growth during training. However, this growth can be overly conservative in practice. We empirically investigate variance growth in LLM pretraining and observe that a generalized BG model with a tunable growth exponent provides a tighter description of practical noise behavior. Motivated by this observation, we introduce the generalized BG-$a$ noise model, which interpolates between bounded variance ($a=0$) and BG-$0$ noise ($a=2$). Under $L$-smoothness, we derive an information-theoretic lower bound with growth-dependent oracle complexity $Ω(ε^{-(4+a)})$ and establish a matching upper bound in $ε$-dependence by increasing the batch size as the iterates move away from initialization. Finally, we propose an adaptive batch scheduler that controls variance growth through dynamic batch size adjustments during training. In pretraining OLMo2 models of up to 1B parameters on C4, our scheduler achieves a lower validation loss than both small and large batch training under matched token budgets, while using less than 10\% of the iterations of small batch training.