Universal One-third Time Scaling in Learning Peaked Distributions

TL;DR

The paper addresses why neural scaling laws exhibit power-law loss decay in training time and proposes a mechanism intrinsic to the softmax and cross-entropy nonlinearities when learning peaked distributions. Through a minimal one-layer toy model and gradient-flow analysis, it derives a universal exponent $1/3$ for the loss decay in the low-temperature regime, supported by numerical and LLM-based experiments. The key finding is that $L(\tau) \sim \tau^{-1/3}$ in the intermediate regime, largely independent of data structure, offering a mechanistic explanation for observed training dynamics and suggesting avenues to improve LLM training efficiency. Empirical validation on Pythia models shows the predicted scaling in practice, reinforcing the practical relevance of the mechanism for real-world large-scale language modeling.

Abstract

Training large language models (LLMs) is computationally expensive, partly because the loss exhibits slow power-law convergence whose origin remains debatable. Through systematic analysis of toy models and empirical evaluation of LLMs, we show that this behavior can arise intrinsically from the use of softmax and cross-entropy. When learning peaked probability distributions, e.g., next-token distributions, these components yield power-law vanishing losses and gradients, creating a fundamental optimization bottleneck. This ultimately leads to power-law time scaling of the loss with a universal exponent of $1/3$. Our results provide a mechanistic explanation for observed neural scaling and suggest new directions for improving LLM training efficiency.

Figures (27)

• Figure 1: The toy model exhibits power-law training behaviors at low temperatures. At high temperatures (small $\beta^*$), the loss vs. step $t$ curves are concave on the log-log plot, similar to exponential decay. At low temperatures (large $\beta^*$), the loss vs. step $t$ curves converge to a line on the log-log plot, like power laws with the same exponent. Details in \ref{['app:adamTLR']}.
• Figure 2: Actual training inspires the aligned student ansatz (student weight is proportional to the teacher's). We project rows of teacher weight $W^*_i \in \mathbb{R}^m$ and the corresponding student weight rows $W_i$ to a 2-dimensional space. The student weight is initialized at zero. The arrows are fixed teacher weight rows. The colored dots are student rows at different steps. Details in \ref{['app:adamTLR']}.
• Figure 3: Numerical loss and gradient agree with the theory. The vertical dashed line is $\sqrt{2\ln n} \approx c_0$. Both loss (panel a) and gradient (panel b) tend to be power-law with $\beta$ in the intermediate regime. With increasing $\beta^*$, the power-law regions are clearer, and the exponents seem to converge. Fitting of the curves under $\beta^* = 100$ yields that the loss exponent is $1.15\pm0.04$ and the gradient exponent is $2.26\pm0.01$. Details in \ref{['app:neumerics']}.
• Figure 4: Aligned student ansatz can describe actual training results. Initializing student weights at zero and scanning learning rates $\eta$, we found that the loss under the optimal learning rate (lowest loss with the same step) has clear $1/3$ scaling (panel a), and the corresponding inverse temperature grows with an algebraic exponent $0.38$ close to the theoretical value $1/3$ (panel b). Here, $\tau \propto t$, and power-law fitting with $t$ or $\tau$ does not change the exponent. The red dashed line means $\beta=c_0$. Larger learning rates cannot follow gradients to align the student. Smaller learning rates follow the same dynamics (gradient flow) whose curves collapse with the dynamic time $\tau$ (panel c), some of which do not show $1/3$ scaling clearly as their $\tau$ is small and are not deep in the low-temperature regime. More details in \ref{['app:adamTLR']}.
• Figure 5: The $1/3$ time scaling of loss can be true beyond the aligned student ansatz. (a and b) Initializing the student at different scales, the aligned student ansatz is wrong at first, but can be correct later, yielding the $1/3$ time scaling. Details in \ref{['app:Adaminit']}. (c and d) By adding weight decay, we observe that under the optimal learning rate, $1/3$ time scaling of loss is correct, while the aligned student ansatz is wrong since $\beta$ does not follow the theory. Beyond norm growth, rotation of parameters in the low-temperature regime may also lead to the $1/3$ time scaling of loss. Details in \ref{['app:wd']}.