Abstract
Adaptive optimization methods (such as Adam) play a major role in LLM pretraining, significantly outperforming Gradient Descent (GD). Recent studies have proposed new smoothness assumptions on the loss function to explain the advantages of adaptive algorithms with structured preconditioners, e.g., coordinate-wise or layer-wise, and steepest descent methods w.r.t. non-euclidean norms, e.g., $\ell_\infty$ norm or spectral norm, over GD. However, it remains unclear how these smoothness assumptions manifest in language modelling tasks. In this work, we aim to analyze the benefit of $\ell_\infty$-norm descent (a.k.a. sign descent) directly from properties of the data distribution, namely, heavy-tailed class imbalance. We propose a minimal yet representative setting of next-token prediction, where we can provably show faster convergence of coordinate-wise algorithms such as Sign descent (steepest descent w.r.t. $\ell_\infty$ norm) over normalized GD (steepest descent w.r.t. to $\ell_2$ norm) in the presence of heavy tail class imbalance.
Abstract (translated)
自适应优化方法(如Adam)在大规模语言模型的预训练中扮演着重要角色,显著优于梯度下降法(GD)。近期研究提出了一些关于损失函数的新平滑假设,以解释具有结构化预条件器(例如坐标方向或层方向)的自适应算法以及相对于非欧几里得范数(如$\ell_\infty$范数或谱范数)最速下降方法的优势,相较于GD。然而,这些平滑性假设如何在语言建模任务中表现尚不清楚。 在这项工作中,我们的目标是从数据分布的特性,特别是长尾类别不平衡的角度分析$\ell_\infty$范数下降(即符号下降)带来的优势。我们提出了一种最小但具有代表性的下一个令牌预测设置,在这种设置下,当存在长尾类别不平衡时,我们可以严格证明坐标方向算法如符号下降(相对于$\ell_\infty$范数的最速下降方法)比归一化GD(相对于$\ell_2$范数的最速下降方法)收敛更快。
URL
https://arxiv.org/abs/2512.00763