Source-linked AI summary
Unified Neural Scaling Laws
Ethan Caballero, Priyank Jaini, David Krueger, Irina Rish
TL;DR
Forecasting neural-network performance across changing model, data, compute, and hyperparameter dimensions is important for selecting scalable methods and supporting AI safety. The paper proposes UNSL, a multivariate functional form combining broken scaling behavior with bottleneck and oppositional-force components. Across varied tasks, UNSL extrapolates more accurately than alternative functional forms and captures nonmonotonic behavior they cannot express.
Problem
Accurate forecasting is needed because methods that lead at small scale may fail to remain strong at larger scale, while scale forecasts also inform AI-safety considerations.
Method
UNSL models performance as multiple dimensions vary simultaneously, using multivariate broken scaling components alongside bottleneck and oppositional-force terms.
Results
UNSL achieves the lowest RMSLE on 88.89% of evaluated language tasks and 60.87% of evaluated vision tasks, outperforming alternative functional forms in extrapolation.
Takeaways & Limitations
UNSL provides a unified way to model and extrapolate scaling across varied architectures and upstream and downstream tasks, including nonmonotonic behaviors other forms cannot express.
Abstract
from arXiv · showhide
We present a functional form (that we refer to as a Unified Neural Scaling Law (UNSL)) that accurately models and extrapolates the scaling behaviors of deep neural networks as multiple dimensions all vary simultaneously (i.e. how the evaluation metric of interest varies as one simultaneously varies the number of model parameters, training dataset size, number of training steps, number of inference steps, amount of compute, and various hyperparameters) for various architectures and for each of various tasks within a varied set of upstream and downstream tasks. This set includes large-scale vision, language, math, and reinforcement learning. When compared to other functional forms for neural scaling, this functional form yields extrapolations of scaling behavior that are considerably more accurate on this set.
1 INTRODUCTION
Neural scaling laws must forecast performance while multiple factors change, because small-scale leaders may fail at larger scales and reliable forecasts support model selection and AI safety. UNSL addresses this need by modeling multivariate scaling, including nonmonotonic behaviors that other forms cannot express.
- Motivation: Accurate scaling forecasts help select methods likely to perform best at scale, where small-scale leaders may not remain competitive.The motivation also links forecasting scale-dependent behavior to predicting novel capabilities for responsible AI development.
- Motivation: Including more relevant predictors cannot increase conditional uncertainty about performance, motivating a unified model P(Y|X) over causally related variables.The stated principle is H(Y|X) ≤ H(Y).
- Contribution: UNSL models and extrapolates deep-network scaling while multiple dimensions vary simultaneously.The paper presents it as a unified functional form for multivariate scaling behavior.
- Contribution: UNSL captures nonmonotonic transitions involving overfitting and hyperparameters such as learning rate and initialization-weight standard deviation.These behaviors are described as inaccessible to other functional forms.
2 THE FUNCTIONAL FORM OF UNIFIED NEURAL SCALING LAWS
UNSL combines multivariate broken scaling components with bottleneck and oppositional-force terms to model performance across predictors in multi-log space. Its design encodes smooth hyperbreaks, nonmonotonic behavior, and desiderata inherited from univariate broken scaling laws.
- Functional form: The predictors may include model parameters, dataset size, training steps, inference steps, and hyperparameters, while y denotes an upstream or downstream performance metric.The formulation treats these quantities as the multivariate inputs whose joint variation determines performance.
- Functional form: UNSL is presented as a general functional form built from Q, R, K, oppositional-force terms, and bottleneck components.K is specifically defined as a Multivariate Broken Neural Scaling Law.
- Multi-log geometry: In multi-log space, MBNSL extends BNSL to multiple dimensions as smoothly connected hyperplanes separated by n hyperbreaks.The initial exponents define the first hyperplane’s gradient, offsets define output position, and f controls hyperbreak sharpness.
- Functional form: The K component represents smoothly connected hyperplanes, while summed components represent performance limits caused by bottlenecks in individual dimensions.These components provide distinct mechanisms for multivariate scaling and dimension-specific limits.
- Oppositional forces: Q and R incorporate performance limits and oppositional forces from hyperparameters and overfitting.Hyperparameter force reflects degradation when settings such as learning rate or initialization scale are too large; overfitting force becomes non-negligible beyond one epoch.
- Additive symmetry: Equation 5 and Equation 6 differ in how they control gradient changes and hyperbreak sharpness, with Equation 6 decoupling sharpness from gradient change.Equation 5 restricts gradient changes to be nonnegative under the stated positivity assumptions.
- Additive symmetry: Equation 5 is used for phenomena with empirically observed nonmonotonic relationships, including overfitting and learning-rate or initialization-scale effects.A related constrained version is used when transitions involve a zero gradient in an input dimension.
- Desiderata: Each univariate scaling behavior is a BNSL, preserving expressivity needed to decouple transition sharpness from gradient change for a significant subset of hyperplane transitions.The paper states that UNSL satisfies the listed desiderata, with supporting evidence discussed in appendices.
3 RELATED WORK
Earlier multivariate scaling laws modeled performance using limited combinations of model size, data, and compute-related variables. UNSL builds on this lineage while targeting broader multivariate behavior.
- Prior multivariate laws: Rosenfeld et al. introduced a bivariate form using model parameters and training dataset size, later used by Kaplan et al. and Hoffmann et al.The paper calls this functional form CF.
- Prior multivariate laws: Muennighoff et al. introduced DC with trivariate input involving model parameters, training compute, and dataset size.The cited DC expression combines these inputs through minimum and exponential terms.
- Relations among prior laws: When dataset size is sufficiently large that training lasts one epoch, DC is mathematically identical to CF.The paper refers readers to an appendix for additional related work.
4 EMPIRICAL RESULTS: FITS & EXTRAPOLATIONS OF FUNCTIONAL FORMS
The paper evaluates UNSL against ablation and alternative functional forms for multivariate scaling across vision and language tasks. UNSL generally achieves the strongest extrapolation accuracy while modeling varied scaling behaviors and dimensions.
- 4 EMPIRICAL RESULTS: FITS & EXTRAPOLATIONS OF FUNCTIONAL FORMS: Held-out circle points evaluate extrapolation, while triangle points fit each functional form; line colors encode the corresponding color-bar dimension.Vertical distance from a point to the same-colored line represents fitting or extrapolation error.
- 4 EMPIRICAL RESULTS: FITS & EXTRAPOLATIONS OF FUNCTIONAL FORMS: The evaluations report extrapolation quality using RMSLE and root standard log error.These metrics are used for table evaluations marked with the downward-arrow symbol.
- 4.0.1 ABLATION FUNCTIONAL FORMS: UNSL incorporates more additive symmetries than A1, A2, and A3, although all four forms have the same supremal expressivity.The reported extrapolation advantage is attributed to enforcing more desiderata rather than greater maximum expressivity.
- 4.1 VISION: 60.87% of vision tasks have the lowest RMSLE under UNSL, compared with 21.74% for the next-best functional form.Vision evaluations vary training dataset size, training steps, and model parameters across few-shot Birds 200, Cars 196, and ImageNet tasks.
- 4.1 VISION: UNSL extrapolates downstream vision performance across bivariate and trivariate scaling behaviors, including the dimensions of dataset size, training steps, and model parameters.The evaluated downstream metrics are test error rates on few-shot image-classification tasks.
- 4.2 LANGUAGE: UNSL is best by RMSLE on 88.89% of language tasks, while the next-best functional form is best on 11.11%.The language experiments include trivariate and bivariate upstream and downstream scaling behaviors.
- 4.2 LANGUAGE: The language evaluations vary model parameters, processed tokens, training-dataset tokens, and training steps or processed tokens, with constant and chinchilla regimes also represented.The constant regime holds learning rate fixed without a schedule; Table 5 includes a scenario where dataset size is effectively infinite and training lasts one epoch.
- 4.2 LANGUAGE: UNSL additionally extrapolates multivariate inference scaling, where inference denotes test-time scaling.This result is reported in Appendix 18.2.
5 THE LIMIT OF THE PREDICTABILITY OF SCALING BEHAVIOR
The paper examines how far scaling behavior can be predicted using sparse-parity experiments and noiseless UNSL simulations. Accurate extrapolation across hyperbreaks depends on having fitting data sufficiently near each transition.
- 5 THE LIMIT OF THE PREDICTABILITY OF SCALING BEHAVIOR: The noiseless simulation averages away run-to-run deviations to recover UNSL constants through curve fitting.The simulation shows the behavior expected with infinitely many training runs or seeds.
- 5 THE LIMIT OF THE PREDICTABILITY OF SCALING BEHAVIOR: Accurate extrapolation past each hyperbreak requires the nearest fitting-hull distance to that hyperbreak to be sufficiently small.This identifies proximity to transitions as the stated boundary on extrapolation predictability.
6 DISCUSSION
UNSL models and extrapolates neural-network performance while multiple dimensions vary simultaneously across diverse architectures and tasks. The paper positions this unified form as applicable to upstream and downstream evaluation settings.
- 6 DISCUSSION: UNSL models and extrapolates performance as model parameters, dataset size, training steps, inference steps, and hyperparameters vary simultaneously.The evaluation metric may be measured upstream or downstream across varied architectures and tasks.
APPENDIX
This appendix extends additive symmetry relations from a sum of two zero-hyperbreak MBNSLs to two MBNSLs with arbitrary numbers of hyperbreaks.
- Equation 5 sums two n = 0 versions of the MBNSL in Equation 4.
- Extending the relations requires obtaining the n = 0 version of each MBNSL, which is its tangent hyperplane in multi-log space.
10 EXPERIMENTAL DETAILS OF FITTING UNSL
UNSL fitting minimizes mean squared log error with regularized exponents, while hyperparameters are selected using extrapolation-oriented validation and seeds are selected by training error.
- UNSL is implemented in KFAC-JAX and fitted by minimizing mean squared log error.
- L2 regularization is applied to UNSL exponents with weighting λ relative to the MSLE loss term.
- The values of n, S, and λ are selected using the lowest validation error under a coordinate-wise extrapolation split, then validation points are restored for fitting.
- KFAC-JAX fitting uses 20000 training steps and 20 random-initialization seeds, with the selected seed minimizing training error rather than validation error.
11 EXPERIMENTAL DETAILS OF SECTIONS 5, 13, 17 (BESIDES FIGURE 10), AND 14
The experiments use a large fixed batch, Adam optimization, and minimal regularization, with specific exceptions for two figures and parameter-count variation implemented through width.
- The batch size is 80000, and no regularization is used because the training dataset is approximately infinite and models train for one epoch.
- Adam is used with β1 = 0 and β2 = 0, except in Figures 11 and 12 and Table 6, which use β1 = 0.9 and β2 = 0.999.
- In Figure 4, the number of model parameters is varied by changing width while holding depth constant.
12 OBTAINING THE COMPUTE-OPTIMAL VALUES OF THE INPUT DIMENSIONS
Compute-optimal input dimensions are obtained by solving a system based on which predictors contribute directly to training compute and which do not.
- The dimensions are divided into compute-contributing set 𝒟 and non-compute-contributing set ℋ, with C denoting training compute and C0 a constant.
- For a given training-compute value C, the input dimensions yielding the lowest performance metric y are obtained by solving the stated system of equations.
13 EFFECT OF VARYING THE NUMBER OF OBSERVED POINTS USED FOR FITTING UNSL FUNCTIONAL FORM
UNSL is evaluated as a multivariate scaling law across observed-point efficiency, simultaneous extrapolation, expressivity, desiderata, and several application settings. The supplied results report accurate extrapolation across these tests, including reinforcement learning, inference scaling, width-versus-depth variation, and batch-size variation.
- 13 EFFECT OF VARYING THE NUMBER OF OBSERVED POINTS USED FOR FITTING UNSL FUNCTIONAL FORM: 9e0 to 9e2 observed points still produce accurate UNSL extrapolations on the (n, k)-sparse parity task.The experiment varies the number of observed fitting points from 9e0 to 9e2.
- 14 UNSL ACCURATELY EXTRAPOLATING TO SCALES AN ORDER OF MAGNITUDE LARGER IN MULTIPLE DIMENSIONS SIMULTANEOUSLY: UNSL accurately extrapolates to scales an order of magnitude larger in multiple dimensions simultaneously.The result is reported for an MLP trained for one epoch on the (n, k)-sparse parity task, with means computed over more than 100 seeds.
- 15 SUPREMAL EXPRESSIVITY EQUIVALENCE OF A1, A2, A3, AND UNSL: UNSL and A1, A2, and A3 have identical supremal expressivity because each can generate positive continuous functions that the others can approximate or recover.A1 is a single-hidden-layer softplus network, and the other forms can be reduced to it through parameter choices.
- 16.1 EXPLANATION OF HOW UNSL FUNCTIONAL FORM SATISFIES DESIDERATUM 1: When only one input dimension varies, UNSL recovers the univariate broken neural scaling law by absorbing the other dimensions into fitted constants.This establishes the single-dimension behavior described by Desideratum 1.
- 16.2 EXPLANATION OF HOW UNSL FUNCTIONAL FORM SATISFIES DESIDERATUM 2: Multiplicative interactions among exponentiated input dimensions shift the locations of UNSL’s hyperbreaks.The hyperbreak condition defines where smooth transitions between hyperplanes occur.
- 16.3 EXPLANATION OF HOW UNSL FUNCTIONAL FORM SATISFIES DESIDERATUM 3: UNSL represents nonmonotonic scaling behavior and includes separate terms for performance limits and oppositional forces from overfitting or hyperparameters.The formulation uses reciprocal combinations of functional terms and misperformance limits, while hyperparameter terms can oppose performance.
- 17.1 EMPIRICAL EVIDENCE OF DESIDERATUM 1: UNSL empirically supports its desiderata: gradient-change sharpness must be decoupled, performance limits are modeled, and UNSL outperforms A3 in most instances.The evidence also reports lower RMSLE than an ablated baseline with the same number of learned parameters.
- 18.1 PLOTS OF REINFORCEMENT LEARNING EXTRAPOLATION RESULTS: UNSL accurately extrapolates multivariate scaling for reinforcement learning, inference scaling, simultaneous width and depth changes, and simultaneous batch-size and training-step changes.The corresponding figures cover StarPilot (hard), chain-of-thought length, MLP width and depth, and batch size with training steps.