Unified Neural Scaling Laws

1 INTRODUCTION

• Scaling laws aim to predict large-scale model behavior as compute, data, and parameters increase, since top-performing methods at small scales may fail at larger scales (Tolstikhin et al., 2021).The introduction argues that including more relevant predictors can only maintain or reduce conditional entropy, motivating a unified functional form using causally related predictors of performance.Source: §1 · p. 1
• UNSL extrapolations are considerably more accurate than those of other neural scaling forms on the evaluated set.The paper introduces UNSL as a functional form for modeling and extrapolating deep neural networks while multiple dimensions vary simultaneously.Source: §1 · p. 1
• UNSL models nonmonotonic transitions in overfitting and hyperparameters, including learning rate and initialization weight standard deviation, that other forms cannot express.These hyperparameters have a nonmonotonic relationship with the performance evaluation metric.Source: §1 · p. 1

2 THE FUNCTIONAL FORM OF UNIFIED NEURAL SCALING LAWS

• UNSL models an evaluation metric—such as prediction error or cross-entropy, upstream or downstream—as a function of multiple predictors, including model size, data size, training and inference steps, and hyperparameters.The unknown constants are estimated by fitting the functional form to predictor–metric data.Source: §2 · p. 2
• For multiple input dimensions, the MBNSL component defines n + 1 smoothly connected hyperplanes in multi-log space, with n hyperbreaks whose sharpness is controlled by f.The initial exponents correspond to the first hyperplane’s input gradients, while f controls hyperbreak sharpness: smaller |f| gives sharper breaks and less curvature; larger |f| gives wider breaks and more curvature.Source: §2 · p. 3
• Equation 3 combines a non-bottleneck component representing smoothly connected hyperplanes with summed components for performance limits caused by individual bottleneck dimensions.When |M| = 1, Equation 4 has the expressivity of Caballero et al.’s (2023) univariate BNSL, with the performance-limit term subtracted.Source: §2 · p. 3
• Overfitting and hyperparameters such as learning rate can exert oppositional forces on performance, while a0 and a2 represent limits on the metric’s best and worst values.The inner reciprocal models overfitting and can become non-negligible after more than one training epoch; S counts hyperparameter-force misperformance limits and is usually at most 1.Source: §2 · p. 3

2.1 THE ADDITIVE SYMMETRY

• For Equation 5, gradient changes across a hyperbreak are nonnegative and break sharpness depends on that change, whereas Equation 6 allows arbitrary gradient changes and decouples sharpness from them.The authors report that observed nonmonotonic transitions seem to follow Equation 5, which they use for phenomena such as overfitting and hyperparameter effects; they also use a specified version for transitions involving zero-gradient regions.Source: §2.1 · p. 4

2.2 DESIDERATA

• The desiderata include univariate scaling with Caballero et al.’s (2023) BNSL expressivity and break positions that shift through multiplication depending on other input dimensions.Source: §2.2 · p. 5
• Near the global optimum, a performance limit may depend on a constant or, when one dimension bottlenecks performance, on that dimension alone.For a constant optimum limit, the transition is characterized by summing a functional form with a constant.Source: §2.2 · p. 5
• Misperformance limits are constant when no input dimension bottlenecks misperformance, but depend only on the bottlenecking dimension when one does.The form also specifies constant misperformance limits through the inverse of a functional form plus a constant, and the stated examples include high loss or error from too-small training datasets that cause overfitting.Source: §2.2 · p. 5

• Rosenfeld et al. (2019) first described a multivariate form using model parameters and dataset size, while Kaplan et al. (2020) and Hoffmann et al. (2022) used training steps times batch size instead of dataset size in CF.Source: §3 · p. 5
• Muennighoff et al. (2023) introduced DC with model parameters, training steps times batch size, and dataset size; its displayed form includes exponential terms and a min-defined U_N.Source: §3 · pp. 5–6
• When the dataset is large enough that training lasts only one epoch, DC is mathematically identical to CF.Source: §3 · p. 6

4 EMPIRICAL RESULTS: FITS & EXTRAPOLATIONS OF FUNCTIONAL FORMS

• Plots distinguish fitting points (triangles), held-out extrapolation points (circles), and fitted curves (lines); every plot includes at least one held-out point.A point’s vertical distance from the same-color fitted line represents fitting or extrapolation error.Source: §4 · p. 6
• Extrapolation evaluations in tables marked with ↓ are reported as RMSLE ± root standard log error.The definitions are provided in Appendices 8 and 9.Source: §4 · p. 6

4.0.1 ABLATION FUNCTIONAL FORMS

• The ablations progressively incorporate more additive symmetries: A2 includes more of the discussed symmetries than A1, A3 more than A2, and UNSL incorporates all discussed symmetries, more than A3.The constant a0 corresponds to the limit on how far y can be reduced or maximized, even if all inputs reach the indicated values.Source: §4.0.1 · pp. 6–7

4.0.2 SUMMARY OF RESULTS

• A1, A2, A3, and UNSL share the same supremal expressivity, so UNSL’s better extrapolation than A1–A3 is attributed to enforcing more desiderata, including all additive symmetries.The summary comparison is reported in Table 1.Source: Table 1 · p. 7, §4.0.2 · p. 7

4.1 VISION

• UNSL accurately extrapolates reported multivariate scaling behavior for reinforcement learning, width and depth, batch size, and learning rate, initialization standard deviation, and training steps.The learning-rate, initialization, and training-step result is trivariate.Source: §4.1 · p. 8
• UNSL has the lowest RMSLE on 60.87% of vision tasks, while the next-best functional form performs best on 21.74%.The evaluated bivariate dimensions are training dataset size and training steps; the trivariate dimensions add model parameter count.Source: §4.1 · pp. 7–8, Table 2 · p. 8, Table 3 · p. 8

4.2 LANGUAGE

• UNSL has the lowest RMSLE on 88.89% of language tasks, versus 11.11% for the next-best functional form.The evaluation covers trivariate and bivariate scaling behavior; in the Table 5 setting, UNSL becomes A3 because the training dataset is effectively infinite and training lasts one epoch.Source: §4.2 · pp. 8–9, Table 4 · p. 9, Table 5 · p. 9

5 THE LIMIT OF THE PREDICTABILITY OF SCALING BEHAVIOR

• UNSL accurately extrapolates sparse-parity scaling despite no observable loss or error progress for the first few hundred training steps.Figure 4’s left plot uses an MLP trained for one epoch on the (n, k)-sparse parity task with n=40 and k=4; each point averages more than 100 seeds.Source: Figure 4 · p. 9, §5 · pp. 9–10
• A noiseless UNSL simulation models the idealized case of infinitely many training runs or seeds, averaging away run-to-run noise to recover the fitted constants as well as possible.Source: Figure 4 · p. 9, §5 · p. 10
• Accurate extrapolation past each hyperbreak requires the fitting points’ convex hull to be sufficiently close to that hyperbreak.Source: §5 · p. 10

• UNSL models and extrapolates scaling as model parameters, dataset size, training and inference steps, and hyperparameters vary simultaneously.Source: §6 · p. 10
• Its scope spans varied architectures and upstream and downstream tasks, including large-scale vision, language, math, and reinforcement learning.Source: §6 · p. 10
• UNSL yields considerably more accurate scaling-behavior extrapolations than other neural-scaling functional forms on this set.Source: §6 · p. 10

Extension of the additive symmetry relations to two MBNSLs with arbitrary hyperbreaks

• The extension covers sums of two MBNSLs, each with an arbitrary number of hyperbreaks n.The appendix frames this as extending the additive symmetry relations discussed in Section 2.1.Source: APPENDIX · p. 11
• Equation 5 sums two n = 0 versions of the MBNSL in Equation 4.For the extension, each arbitrary-hyperbreak MBNSL must be reduced to its n = 0 version.Source: APPENDIX · p. 11
• Each n = 0 version is the tangent hyperplane in multi-log space, characterized by values of wb and wci.The supplied passage introduces those values but does not provide them.Source: APPENDIX · p. 11

10 EXPERIMENTAL DETAILS OF FITTING UNSL

• The UNSL is implemented in KFAC-JAX and fitted by minimizing mean squared log error (MSLE).Source: §10 · p. 11
• The fitting objective uses 𝜖= 10−16 and L2 regularization on UNSL exponents, weighted by 𝜆 relative to MSLE.Source: §10 · p. 11
• Values of 𝑛, 𝑆, and 𝜆 are selected by validation error, with training points restricted so none is larger than a validation point in every 𝑥 dimension.Source: §10 · p. 12
• After selecting 𝑛, 𝑆, and 𝜆, validation points return to training, while the circle-shaped hold-out points remain reserved for extrapolation RMSLE evaluation.Source: §10 · p. 12
• In practice, 𝑆≤1 unless the target scaling behavior is extravagant, similar to Figure 9 of Appendix 17.5.Source: §10 · p. 12
• Fitting uses 20000 training steps and 20 random-initialization seeds; the selected seed minimizes training error, unlike hyperparameters selected by validation error.Parameters use JAX default “LeCun Normal” initialization.Source: §10 · p. 12

• Training uses batch size 80000, no regularization because the dataset is ∼infinite and training spans a single epoch, and Adam.Adam uses 𝛽1 = 0 and 𝛽2 = 0, except in Figures 11 and 12 and Table 6 of Section 17.7.Source: §11 · p. 12
• Adam uses 𝛽1 = 0 and 𝛽2 = 0 except for Figures 11 and 12 and Table 6 in Section 17.7, which use 𝛽1 = 0.9 and 𝛽2 = 0.999.Source: §11 · p. 12
• Learning rate and weight-initialization standard deviation are held constant unless a figure explicitly varies them.Source: §11 · p. 12
• In Figure 4, model parameter count is varied by changing width while holding depth constant.Source: §11 · p. 12

12 OBTAINING THE COMPUTE-OPTIMAL VALUES OF THE INPUT DIMENSIONS

• Compute-driving dimensions are indexed by 𝒟, while non-compute dimensions such as learning rate and initialization-weight standard deviation are indexed by ℋ.C denotes training compute; C0 is a constant, with Hoffmann et al. (2022) given as an example where C0 = 6.Source: §12 · p. 12
• For a given compute amount C, the method seeks input-dimension values that yield the lowest y by solving a system of equations.Source: §12 · p. 12

13 EFFECT OF VARYING THE NUMBER OF OBSERVED POINTS USED FOR FITTING UNSL FUNCTIONAL FORM

• UNSL is reported to extrapolate accurately both with few observed fitting points and to scales an order of magnitude larger in multiple dimensions.Figure 5 varies the number of observed points from 9e0 to 9e2; Figure 6 shows the larger-scale extrapolation on sparse parity.Source: §13 · p. 13, Figure 5 · p. 13, §14 · p. 13

15 SUPREMAL EXPRESSIVITY EQUIVALENCE OF A1, A2, A3, AND UNSL

• A1, A2, A3, and UNSL have identical supremal expressivity, with the A1 family dense in positive continuous functions on compact domains.The stated argument uses the universal approximation theorem for softplus networks and shows that the other forms can approximate A1, and vice versa.Source: §15 · p. 14

16.1 EXPLANATION OF HOW UNSL FUNCTIONAL FORM SATISFIES DESIDERATUM 1

• UNSL recovers a univariate broken neural scaling law when only one input dimension varies, while multiplicative interactions shift hyperbreak locations across dimensions.The appendix explains the single-dimension reduction and identifies multiplicative interactions between exponentiated inputs as shifting hyperbreaks.Source: §16.1 · p. 14, §16.2 · p. 15

16.5 EXPLANATION OF HOW UNSL FUNCTIONAL FORM SATISFIES DESIDERATUM 5

• The functional form permits multiple distinct misperformance limits, with S greater than 1 used when different conditions produce multiple limits.The examples distinguish overfitting from nonoptimal hyperparameters and describe different limits at large versus small training-step counts.Source: §16.5 · p. 16

17.1 EMPIRICAL EVIDENCE OF DESIDERATUM 1

• Empirical tests report that sharpness must be decoupled from gradient-change magnitude for accurate fitting and extrapolation, and that Desideratum 2 also holds empirically.The first test attributes the needed extra expressivity to f in Equations 4 and 6.Source: §17.1 · p. 18, §17.2 · p. 18

17.8 EMPIRICAL EVIDENCE OF DESIDERATUM 8

• In the Greydanus–Kobak setting, UNSL outperforms A3 in the majority of instances; separate tests report misperformance limits under specified learning-rate and training-step conditions.The limits include cross-entropy 2.3 at learning rates greater than 3 and more than 100 steps, and a larger-than-plotted limit at learning rates significantly greater than 20 and fewer than 2 steps.Source: §17.5 · p. 19, Figure 9 · p. 19, §17.8 · p. 23

18.1 PLOTS OF REINFORCEMENT LEARNING EXTRAPOLATION RESULTS

• UNSL is reported to extrapolate multivariate scaling for reinforcement learning, inference-time scaling, width and depth, and batch size; downstream-vision plots show UNSL alongside DC, A1, and A2 forms.The reinforcement-learning example is StarPilot (hard), the inference example is test error on MAA, and the width/depth and batch-size examples use an MLP.Source: §18.1 · p. 24, §18.2 · p. 24, §18.3 · p. 24

18.5.1 BIVARIATE

• The bivariate downstream vision visualization shows extrapolation results for the A3 functional form.The caption points to Section 4.1 for further details but supplies no numerical comparison.Source: Figure 21 · p. 30

18.5.2 TRIVARIATE

• The trivariate downstream vision visualizations compare UNSL, Muennighoff et al.’s “DC,” and A1–A3 functional forms.The supplied captions establish the extrapolation setting, not which functional form performs best.Source: Figure 22 · p. 31, Figure 23 · p. 31, Figure 24 · p. 32

18.6.1 TRIVARIATE

• The trivariate language visualizations compare UNSL, Muennighoff et al.’s “DC,” and A1–A3 functional forms.Each set contains 20 slices of one fitted trivariate scaling behavior, with model parameters in plot titles, training steps times batch size on the x-axis, and training dataset size encoded by color.Source: Figure 27 · p. 34, Figure 28 · p. 35, Figure 29 · p. 36

18.6.2 BIVARIATE

• The bivariate downstream and upstream language visualizations cover UNSL, Hoffmann et al.’s “CF,” and A1/A2 functional forms.The captions identify the scaling behavior as language performance but do not provide numerical comparisons.Source: Figure 32 · p. 39, Figure 33 · p. 40, Figure 34 · p. 41

19 ADDITIONAL RELATED WORK

• Additional related work spans alternative scaling settings, forecasting and hyperparameter methods, non-monotonic behavior, and emergent-ability interpretations.Examples include compute/model/dataset scaling (Henighan et al., 2020), transfer and fine-tuning (Hernandez et al., 2021), data curation (Sorscher et al., 2022; Ayed & Hayou, 2023), compute-optimal scaling (Alabdulmohsin et al., 2023), architecture heuristics (Tan & Le, 2019), RL (Neumann & Gros, 2023), and diffusion models (Liang et al., 2024). Forecasting connects to learning-curve extrapolation and HPO (Swersky et al., 2014; Falkner et al., 2018; Snoek et al., 2012), while cited work also treats large-batch scaling (Goyal et al., 2017; Smith et al., 2017; You et al., 2017; McCandlish et al., 2018), μP (Yang et al., 2022), double descent (Nakkiran et al., 2019), phase transitions/jamming, grokking (Power et al., 2022; Liu et al., 2022; 2023), and emergent abilities (Wei et al., 2022).Source: §19 · p. 43

20 IMPLEMENTATION OF UNSL

• UNSL implementation notes say to pass the training dataset for log-space z-normalization, which speeds convergence and reduces seed needs without changing expressivity.The notes recommend at least a dozen CPUs for quickly finding global optima, and say a second-order optimizer makes training faster.Source: §20 · pp. 43, 53
• Bivariate plots compare model predictions across log-spaced input grids while holding another input at a selected value.The plotted grid contains 4096 points spanning 10^-13 to 10^13.Source: §20 · pp. 54–55
• Bivariate visualizations distinguish training observations with triangle markers and test observations with circle markers.Marker color encodes another input dimension.Source: §20 · p. 55
• The trivariate plotting code reports RMSLE for all, training, and test data and prepares unique-coordinate counts and train/test splits.The metrics are computed against the training data used to fit the model.Source: §20 · p. 56
• For trivariate plots, the code groups training and test observations by the second input and iterates over its values as plot slices.Within each slice, it retrieves corresponding input and outcome values for both splits.Source: §20 · pp. 56–57
• Trivariate slice plots draw model-prediction curves while varying the first or third input and holding the second input fixed.Curve colors encode the third or first input, respectively.Source: §20 · p. 57
• Trivariate plots overlay training and test observations, use logarithmic input and outcome axes, and label each panel by its fixed second-input value.The plotting code saves the figure under a distinct trivariate filename.Source: §20 · pp. 57–58