Predicting Large Model Test Losses with a Noisy Quadratic System

Abstract: We introduce a predictive model that estimates the pre-training loss of large models from model size (N), batch size (B) and number of weight updates (K). This is the first loss prediction model that can handle changing batch size. The model outperforms Chinchilla's loss model, a model of the test loss using the batch size and number of tokens, in terms of projecting the loss at extrapolated compute budgets (up to 1000 folds). A natural use of the model is to find optimal N, B, K configurations under explicit and compound resource constraints like time, memory and compute. In our experiments, the model-selected configurations are close to ground-truth optimal. Our work advocates for loss prediction as a better alternative to heuristic-based laws, which are growing in complexity. The implementation is available on https://github.com/chuningxdy/Noisy-Quadratic-System.

• The paper introduces a mechanistic Noisy Quadratic System (NQS) that models batch-induced noise and training dynamics to predict LLM test losses.
• It outperforms the Chinchilla model, maintaining low error even when extrapolating across compute gaps up to 1000×.
• The NQS framework enables resource-constrained optimization by accurately forecasting losses over varying model sizes, batch sizes, and weight updates.

Predicting Large Model Test Losses with a Noisy Quadratic System

Introduction and Motivation

The paper "Predicting Large Model Test Losses with a Noisy Quadratic System" (2605.09154) provides a mechanistic, parameterized predictive model—the Noisy Quadratic System (NQS)—to forecast the test loss of LLMs under varying pre-training configurations defined by model size ($N$), batch size ($B$), and number of weight updates ($K$). The prevailing approach epitomized by the Chinchilla method models test loss as a function of $N$ and dataset size $D$ (total tokens), neglecting batch size as an independent variable and thus obscuring crucial optimizer dynamics. As scaling laws become more complex and the number of relevant training parameters increases, there is urgent demand for predictive models accommodating a broader set of resource variables and supporting optimization under practical compute, memory, and time constraints.

The NQS is derived from recent theory on SGD dynamics (e.g., Noisy Quadratic Model), incorporates mini-batch-induced gradient noise, and supports closed-form, numerically efficient prediction. The model can directly inform optimal resource allocation, e.g., selection of $(N, B, K)$ under compound constraints. The authors empirically demonstrate that the NQS consistently outperforms the Chinchilla loss model, especially in extrapolation scenarios with compute budgets up to $1000\times$ beyond those seen in training, and is robust to changes in the training set.

Evaluation of Chinchilla Loss Model

The Chinchilla loss model (Method 3) posits pre-training loss as a function primarily of $N$ and $D$, following a power-law structured formula. While this model is descriptive in-distribution, its extrapolation ability across large compute gaps is substantially limited. Concretely, when fitted on low-compute training data and evaluated on high-compute holdouts, its predictive accuracy deteriorates substantially.

Figure 1: Chinchilla Method 3 matches in-distribution data but fails to extrapolate to higher compute budgets, with increasing error over large gaps between training and test compute.

This highlights a key deficiency: the Chinchilla method captures asymptotics for parameter and data scaling but ignores variance contributions that are critical when scaling batch size or training for significantly longer durations.

Noisy Quadratic System Model: Theoretical Framework

The NQS is founded on the following principles:

• Mechanistic modeling of test loss: Formulates test loss as risk from minimizing a quadratic loss surface using SGD, parametrized by effective spectrum exponents for data, noise, and model initialization.
• Explicit role of batch size: Models batch-induced noise via a distinct power-law spectrum, unifying influence of batch size with model size and number of weight updates.
• Closed-form solution: The expectation over the SGD trajectory under NQS admits a numerically tractable, low-dimensional sum with at most seven degrees of freedom after parameterization.

A further realistic adjustment incorporates the effect of LayerNorm, dynamically reducing the effective learning rate as a function of weight vector norm, particularly impactful for low-batch regimes.

Figure 2: Incorporating LayerNorm's effect by scaling the learning rate with inverse weight norm is crucial for accurate modeling at small batch sizes.

NQS parameters are estimated using gradient-based optimization over scaling datasets containing $(N, B, K, l)$, with validation on strict train/holdout splits. Calculation leverages both summations and continuous integral approximations for computational efficiency.

Empirical Results: Loss Prediction and Extrapolation

Comprehensive experiments are conducted across Pythia and Llama model families. The NQS is evaluated not just on its fit to training data but, crucially, on its ability to extrapolate (predict loss at larger compute budgets and for unseen batch sizes).

Figure 3: NQS accurately predicts LLM test losses on holdout data across variations in model and batch size, surpassing Chinchilla especially as the compute scale moves far beyond the training regime.

Notably, the NQS:

• Generalizes across up to $B$0 compute gaps, maintaining low Huber and MSE error on holdouts where Chinchilla degrades rapidly.
• Handles batch size variation directly, letting practitioners reason about optimizer batch/loss tradeoffs using predicted curves.
• Maintains robustness under aggressive subsampling of the training set, producing tight confidence intervals (see Figures 8 and 9).

  Figure 4: NQS retains predictive accuracy under extreme extrapolation (compute gap up to $B$1 between train and test), where Chinchilla performance collapses.

  

  Figure 5: On Llama + LM1B, NQS predictions remain accurate across model scales and data budgets, replicating the generalization seen with Pythia + OWT2.

  

  Figure 6: Confidence intervals for NQS remain narrow even as the training set is randomly subsampled, exhibiting robustness to training data variation.

Resource-Constrained Optimization

The differentiability and closed-form nature of NQS predictions directly enable $B$2 optimization under multiple resource constraints (compute, memory, wall-clock time). Given a constraint set (e.g., total FLOPs, max training time, GPU memory), the optimal configuration is obtained by exhaustive prediction and selection.

Figure 7: NQS predictions support compound-constrained selection of model, batch, and update budget, matching or exceeding ground-truth-optimal configurations as resource constraints are varied.

Comparison of Model Complexity and Overfitting

Despite its higher parametric complexity (NQS: 7+1; Chinchilla: 5), the NQS does not overfit: training error is close to holdout error and, unlike Chinchilla, the extrapolation gap does not rapidly widen. This substantiates the necessity of mechanistic bias/variance modeling and explicit batch-size terms for reliable scaling.

Discussion: Implications and Future Directions

Theoretical implications:

The NQS framework is a step toward unifying empirical scaling law observations with mechanistic models of SGD dynamics, bridging insights from linear regression, noisy quadratic losses, and normalization-induced learning rate scheduling. Its performance supports theoretical predictions about the nontrivial role of variance (often neglected in classical scaling law expressions) and the necessity of including batch size as a first-class variable [see also nqm, paquette202543phasescomputeoptimalneural]. The ability to efficiently simulate “rolled-out” loss dynamics at scale could catalyze more accurate scaling theories, especially as models diverge from idealized asymptotic regimes.

Practical implications:

NQS provides practitioners with a precise, computationally efficient method for pre-training configuration under arbitrary resource constraints—significantly enhancing the fidelity and confidence of planning schedules for large-scale LLM training, especially in industrial and resource-constrained environments. The empirical superiority in extrapolation is particularly salient, as future training budgets continue to grow and current scaling laws become out-of-distribution.

Extensions and limitations:

Key avenues for future research include extending NQS-like models to capture learning rate/batch size schedule interaction, weight decay effects, and nontrivial data/architecture heterogeneity. Incorporating multi-epoch scenarios and data mixture scaling, as well as integrating richer optimizer models (Adam, Adagrad, etc.), are tractable due to the flexible structure of the NQS approach. Open questions remain on joint optimization of $B$3 (initial weight norm) and dynamic schedule parameters, as well as inference stability in more exotic data regimes.

Conclusion

The NQS loss prediction framework substantially advances the state of the art in mechanistic, scalable loss modeling for LLMs. By explicitly capturing the interplay of model size, batch size, and training length, and demonstrating robust extrapolation across vast compute gaps and training regimes, NQS enables principled design of training configurations under complex, real-world constraints. Its theoretical and empirical advancements suggest a paradigm shift in both the practice and study of large-scale neural scaling, motivating further extensions and the eventual unification of mechanistic and empirical scaling law research.