Scaling Laws for Gradient Descent and Sign Descent for Linear Bigram Models under Zipf's Law

2025/05/25 by Kunstner, Frederik, Bach, Francis · 6 citations
#FOS: Computer and information sciences #FOS: Mathematics #Machine Learning (cs.LG) #Optimization and Control (math.OC)

paper · doi:10.48550/arxiv.2505.19227

Abstract

Recent works have highlighted optimization difficulties faced by gradient descent in training the first and last layers of transformer-based language models, which are overcome by optimizers such as Adam. These works suggest that the difficulty is linked to the heavy-tailed distribution of words in text data, where the frequency of the kth most frequent word π_k is proportional to 1/k, following Zipf's law. To better understand the impact of the data distribution on training performance, we study a linear bigram model for next-token prediction when the tokens follow a power law π_k ∝ 1/k^α parameterized by the exponent α> 0. We derive optimization scaling laws for deterministic gradient descent and sign descent as a proxy for Adam as a function of the exponent α. Existing theoretical investigations in scaling laws assume that the eigenvalues of the data decay as a power law with exponent α> 1. This assumption effectively makes the problem ``finite dimensional'' as most of the loss comes from a few of the largest eigencomponents. In comparison, we show that the problem is more difficult when the data have heavier tails. The case α= 1 as found in text data is ``worst-case'' for gradient descent, in that the number of iterations required to reach a small relative error scales almost linearly with dimension. While the performance of sign descent also depends on the dimension, for Zipf-distributed data the number of iterations scales only with the square-root of the dimension, leading to a large improvement for large vocabularies.

Cited by

• Provable Benefit of Sign Descent: A Minimal Model Under Heavy-Tailed Class Imbalance
• Fast Escape, Slow Convergence: Learning Dynamics of Phase Retrieval under Power-Law Data
• Kernel ridge regression under power-law data: spectrum and generalization
• Theory of Scaling Laws for In-Context Regression: Depth, Width, Context and Time
• Inductive Bias and Spectral Properties of Single-Head Attention in High Dimensions
• Scaling Laws and Spectra of Shallow Neural Networks in the Feature Learning Regime

Related

• Provable Benefit of Sign Descent: A Minimal Model Under Heavy-Tailed Class Imbalance
• Language Modeling at Scale
• Unifying Learning Dynamics and Generalization in Transformers Scaling Law
• Fira: Can We Achieve Full-rank Training of LLMs Under Low-rank Constraint?
• Scaling Law Phenomena Across Regression Paradigms: Multiple and Kernel Approaches