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

TL;DR

This paper analyzes how the heavy-tailed Zipf distribution of words affects the efficiency of gradient and sign descent in training linear bigram models, revealing that sign descent performs significantly better with large vocabularies.

Contribution

It derives new scaling laws for gradient and sign descent under Zipf's law, especially highlighting the challenges when the tail is heavy (α=1) and the advantages of sign descent.

Findings

Sign descent scales with the square root of dimension for Zipf data.

Gradient descent is most challenging at α=1, scaling almost linearly with dimension.

Sign descent outperforms gradient descent in heavy-tailed data scenarios.

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 $k$th most frequent word $\pi_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 $\pi_k \propto 1/k^\alpha$ parameterized by the exponent $\alpha > 0$. We derive optimization scaling laws for deterministic gradient descent and sign descent as a proxy for Adam as a function of the exponent $\alpha$. Existing theoretical investigations in scaling laws assume that the eigenvalues of the data decay as a power law with exponent $\alpha > 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 $\alpha = 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.