Softmax and Cross-entropy Lead to a Universal One-third Time Scaling of Loss When Learning Peaked Distributions

This repository contains code to reproduce the experiments in the paper Universal One-third Time Scaling in Learning Peaked Distributions, Yizhou Liu, Ziming Liu, Cengiz Pehlevan, and Jeff Gore, ICML 2026.

Overview

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$.

An interesting analogy can be drawn to statistical physics, where peaked Boltzmann distributions require low temperatures, leading to power-law free energy and internal energy (with inverse temperature).

Our results provide a mechanistic explanation for observed neural scaling and suggest new directions for improving LLM training efficiency.

Experiment-code Correspondence

All toy model experiments can be found in folder ./exp, where the names of scripts correspond to the type of experiment being run and are explained in Appendix of the paper.

LLM experiments are in ./LLMs folder, we use the script for Pythia-12B as an example. The scripts for other LLMs are similar. The code for fitting and analysis is in the notebook.

Experiment	Where in Paper	Code
Adam Scanning Temperatures and Learning Rates	Figure 1 and 4	exp-1
Adam Scanning Temperatures	Figure 2	exp-0
Numerics of Theory	Figure 3	test-2
Adam Scanning Initialization Scales	Figure 5, a and b	exp-2
Adam with Weight Decay	Figure 5, c and d	exp-3-2
LLM Temperature	Figure 6a	test-6
LLM Evaluation	Figure 6, b, c, and d	LLMs folder
SGD Scanning Temperatures and Learning Rates	Appendix D.2	exp-1-1
Adam with Learning Rate Schedule	Appendix D.4	exp-2-1
Toy Model with Power Laws in Data	Appendix D.6	exp-5
Generalized Toy Model with Residual Layers	Appendix D.7	exp-9
MNIST Classification with a MLP	Appendix D.8	exp-4-2

Citation

@article{liu2026universal,
  title={Universal One-third Time Scaling in Learning Peaked Distributions},
  author={Liu, Yizhou and Liu, Ziming and Pehlevan, Cengiz and Gore, Jeff},
  journal={arXiv preprint arXiv:2602.03685},
  year={2026}
}

Interested in Other Neural Scaling Laws?

• Width Scaling Due to Limited Representation: Superposition Yields Robust Neural Scaling (paper link, code link)
• Depth Scaling Due to Limited Transformation: Inverse Depth Scaling From Most Layers Being Similar (paper link, code link)