Reproduce “Predicting Large Model Test Losses with a Noisy Quadratic System” on OpenResearch

About this paper

Chuning Li, Chris J. Maddison

NQS predicts language-model training loss across model size, batch size, and update count. It treats training as noisy optimization through a curved loss landscape, with batch size controlling noise, and fits parameters to runs from a model family. On held-out Pythia and OpenWebText2 IsoFLOPs tests, average log-Huber loss was 9.0 × 10⁻⁵ for Chinchilla and 2.5 × 10⁻⁵ for NQS; evidence covers families and datasets.

Key insights

• A mechanistic loss model, grounded in theoretical understanding of neural network training dynamics, can offer greater extensibility and interpretability than purely empirical scaling laws.
• Explicitly modeling the variance error component's dependence on batch size is critical for comprehensive and accurate loss prediction across diverse training configurations.
• The influence of normalization techniques like LayerNorm on the effective learning rate, particularly for small batch sizes, is a necessary consideration for robust predictive modeling.

Summary by alphaXiv.