提出非欧几里得梯度噪声尺度，实现自适应批量大小，提升训练效率。
Adaptive Batch Sizes Using Non-Euclidean Gradient Noise Scales for Stochastic Sign and Spectral Descent
- 基于双范数几何推导signSGD与specSGD的非欧梯度噪声尺度。
- 在1.6亿参数Llama模型上，减少66%训练步数且保持相同验证损失。
- 适合分布式训练中追求高效优化的开发者使用。
为最大化硬件利用率，现代机器学习系统通常采用大而固定的或手动调优的批量大小策略，依赖脆弱且成本高昂的启发式方法。现有基于梯度噪声尺度（GNS）的自适应策略提供了更合理的替代方案。然而，其假设SGD的欧几里得几何与基于广义范数的流行优化器（如signSGD/Signum（ℓ∞）和随机谱下降（specSGD）/Muon（S∞））存在根本性不匹配。本文推导了signSGD与specSGD在各自对偶范数几何下自然产生的梯度噪声尺度。为实际估算这些非欧度量，我们提出一种高效的方差估计方法，利用分布式数据并行系统中不同排名的局部小批量梯度。实验表明，使用非欧几里得GNS的自适应批量大小策略，在16000万参数Llama模型上，使Signum与Muon的训练步数减少高达66%，同时达到与固定批量基线相当的验证损失。
原文摘要 · Abstract (English)
To maximize hardware utilization, modern machine learning systems typically employ large constant or manually tuned batch size schedules, relying on heuristics that are brittle and costly to tune. Existing adaptive strategies based on gradient noise scale (GNS) offer a principled alternative. However, their assumption of SGD's Euclidean geometry creates a fundamental mismatch with popular optimizers based on generalized norms, such as signSGD / Signum ($\ell_\infty$) and stochastic spectral descent (specSGD) / Muon ($\mathcal{S}_\infty$). In this work, we derive gradient noise scales for signSGD and specSGD that naturally emerge from the geometry of their respective dual norms. To practically estimate these non-Euclidean metrics, we propose an efficient variance estimation procedure that leverages the local mini-batch gradients on different ranks in distributed data-parallel systems. Our experiments demonstrate that adaptive batch size strategies using non-Euclidean GNS enable us to match the validation loss of constant-batch baselines while reducing training steps by up to 66\% for Signum and Muon on a 160 million parameter Llama model.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。