Scaling Collapse Reveals Universal Dynamics in Compute-Optimally Trained Neural Networks

Thank you for your feedback and supportive review. We have attached some additional figures here, and address your specific questions below.

On the limited range of widths in CIFAR-5M experiments. We acknowledge the limitations of this experiment, which was a combination of our computational constraints as well as the small size of the dataset itself. We set the maximum width to 2048 in the CIFAR-5M experiments to 1) limit the computational cost of our experiments, and 2) avoid excessive repetition of the training data. Based on the estimated compute-optimal training horizon, the width 2048 model already required training on the CIFAR-5M dataset (which only has 5M training images) for over 10 epochs. A larger dataset is necessary to train even wider models to compute-optimality without overfitting effects.

We note that our MLP experiment does not suffer from this data constraint, where we scaled up the width by 8x from 512 to 4096, corresponding to a 64x scale up in model size, and demonstrated supercollapse across all models.

In addition, we demonstrate in Figure 1 (middle, right) here that supercollapse occurs in two additional data modalities: chess game string representations and chemical structures, and along one more scaling dimension: depth, further establishing the generality of our observations.

On the utility of Equation 6. We emphasize that our claim regarding the dominance of two terms in Equation 54 is mathematically valid in the asymptotic limit of large $D$, not merely a convenient approximation. This generalized analysis of power law loss curves extends the connection between supercollapse and compute-optimal scaling beyond the simplified form in Equation 7. This extension is particularly important because recent theoretical work [1] demonstrates that this broader class of loss curves naturally emerges in high-dimensional regression models, whereas the two-term power law in Equation 7 fails to capture significant regions of the phase diagram (Figure 2, [1]).

To empirically validate our theoretical findings, Figure 6 attached via this link presents additional experiments using the Power Law Random Features (PLRF) model introduced in Section 3.1. Unlike our setup in Figure 4, we specifically selected values of $\alpha$ and $\beta$ to produce loss curves asymptotically described by three power laws (detailed in [1]), which cannot be accurately represented by Equation 6. These results confirm that supercollapse occurs exactly as predicted by our analysis in Appendix G. Without our generalized framework, we would lack the theoretical foundation to explain supercollapse in these more complex settings.

On determining the optimal decay schedule For a fixed $\hat{K}$, equation 10 suggests that decaying the learning rate as late as possible to maximize total gradient flow time is best. However, as we show in Figure 5 and Figure 9 in the paper, the best fit $\hat{K}$ has non-trivial dependence on the schedule and correlates with the preconditioned NTK trace, which tends to be higher for schedules that start decaying earlier. This gives a tradeoff between progress in gradient flow time (favoring later decay) and more favorable geometry (favoring early decay). Understanding this tradeoff requires a deeper understanding of $\hat{K}$, its dynamics, and its relationship to the schedule. We believe this is an exciting direction for future research on optimizing learning rate schedules.

Thank you again for your review. We hope we have addressed your questions. Please let us know if we can clarify anything further.

[1] Paquette et al. "4+3 phases of compute-optimal neural scaling laws."

Thank you for the constructive feedback! We provide several additional results and clarifications here and will include them in the updated paper. Corresponding figures are in the linked PDF.

Further evidence for the generality of supercollapse. We focused on width scaling as it's the most well-studied dimension with extensive literature [1,2,3,4] on optimal learning rate and initialization scaling, which made finding the proper hyperparameter scaling needed for supercollapse much easier. Inspired by your question, we conducted a depth-scaling experiment by adding residual connections to our MLP (creating a well-defined depth scaling limit) and using depth-µP [5] to scale learning rate and initialization. Figure 1 (left) in the linked PDF shows depth scaling also produces supercollapse, confirming it's not width-specific; this plot is consistent with our theory linking supercollapse to power law compute-loss pareto frontiers.

In addition, Figure 1 (middle, right) demonstrates supercollapse occurring in two additional data modalities: chess game string representations and chemical structures, further establishing its generality.

Realistic scaling recipes that don’t admit supercollapse. We expect supercollapse to break down when hyperparameters such as the learning rate are not properly scaled with the model. In Figure 3 of the linked PDF, we ablate µP while keeping learning rate constant for all widths, which makes the learning rate increasingly suboptimal as width scales [4]. Without µP the rescaled curves don’t collapse.

Practical relevance of supercollapse. Exploring supercollapse's practical applications is an exciting direction for future work! Our primary goal in this work is to establish supercollapse as a novel phenomenon across multiple data modalities and architectures and to show how studying loss curve universality enhances our understanding of scaling.

That said, we believe this work already provides compelling evidence of supercollapse's practical value. Figures 3(a) and 4 in the paper show how absent supercollapse reveals errors in estimated data exponent $\gamma$. Similarly, MLP without µP (Figure 3, linked PDF) shows comparable performance to MLP with μP on the loss-compute plot except for the largest model, but the absence of supercollapse reveals scaling errors. Supercollapse is a powerful diagnostic tool, making errors more detectable than when evaluating final losses alone.

Methodology for fitting optimal training compute. We estimate optimal compute by training models for a fixed number of steps (7e5 for transformers, 1e6 for MLPs) and identifying the compute-optimal point. We find models with lowest loss at each compute value on a logarithmic grid to create the compute-loss Pareto frontier (Figure 2). Our discrete model sizes mean each model is optimal for a range of compute values, producing multiple y-values per x-value in Figure 7 (Appendix C.1). We fit a power law to these points using log-log linear regression to estimate optimal compute per model size, similar to Chinchilla’s Figure 1 [6]. We use constant schedule to efficiently sweep compute/token budgets, rather than running separate experiments for each budget as in Chinchilla. This reduces experimental costs, a strategy also used in [7] for LLM scaling law studies.

Additional clarifications

• $\theta$ refers to the fraction of optimal compute.
• Figure 4 in linked PDF shows collapse deviation std dev across 5 trials. Colors indicate width (now in the legend).
• $\Delta < \sigma$ for an $O(1)$ interval isn't trivial as $\Delta$ can grow quickly away from $\theta=1$, as we showed for constant schedule in Figure 6 of the paper.
• SGD noise scales as $\Theta(\eta/B)$ (see Eq 33, Appendix E). Annealing LR thus reduces noise similar to increasing batch size.
• Intuition for L415-420: variance of noise difference between timepoints is smaller than individual variances when sufficiently correlated.
• We expanded $P_0$ range for sensitivity analysis in Figure 5 (linked PDF).
• We suppressed $s$ in Eq 2 for brevity.
• Outliers in Figure 9(a) occur when $\Delta\eta$ approaches 0 in multi-cycle cosine schedule, making $\hat{K}$ divergent and not well-defined.

[1] Yang et al. "Feature learning in infinite-width neural networks."

[2] Yang et al. "Tensor programs ivb: Adaptive optimization in the infinite-width limit."

[3] Bordelon et al. "Self-consistent dynamical field theory of kernel evolution in wide neural networks."

[4] Everett et al. "Scaling exponents across parameterizations and optimizers."

[5] Yang et al. "Feature learning in infinite-depth neural networks."

[6] Hoffmann et al. "Training compute-optimal large language models."

[7] McLeish, Sean, et al. "Gemstones: A Model Suite for Multi-Faceted Scaling Laws."

Thank you for your careful reading of our draft and supportive review. We will update the paper to ensure that all notations are clearly defined in the main text and fix the typos you identified. Specifically regarding the definition of $k_i$, we first sample a scalar $s_i$ from the power-law distribution defined in L140 and then sample a random unit vector $v_i$ and define $k_i$ as $s \cdot v_i.$ This will be clarified in the text.

We have attached a PDF here with additional experiments requested by other reviewers to further demonstrate the generality of supercollapse. If you are interested, the details can be found in our other rebuttals. Let us know if we can clarify anything further.

Thank you for the thoughtful review.

Your point about the limited diversity of the datasets is well taken. To address this point, we conducted additional experiments on two non-image domains: the Lichess chess games dataset dataset of 22M chess games recorded in algebraic chess notation, and SMILES-250M, a large collection of chemical structure representations. In Figure 1 (middle, right) of the linked PDF, we show the next-token prediction loss curves of compute-optimally trained transformers and their rescaled versions. Our results demonstrate that supercollapse occurs consistently across both these datasets. In addition, we showed that scaling depth also gives supercollapse in the MLP setting similar to scaling width (Figure 1 left). Combined with our original findings on CIFAR-5M, MLP regression, and high-dimensional linear models, these results provide compelling evidence that supercollapse is a general phenomenon spanning diverse data modalities and model architectures.

Regarding LLM pre-training, computational constraints prevented us from conducting meaningful experiments. The smallest model in the Chinchilla scaling law [1] study has approximately 80M parameters, similar to our largest tested model. Our limited computation meant that we could not perform high-quality scaling studies in this setting. We believe that our consistent observation about supercollapse across multiple architectures (transformer, MLP, linear model) and data modalities (images, chess, chemical structures, regression), as well as our theoretical analysis establishing the connection between supercollapse and sum-of-power-law loss curves motivates future studies on LLM pre-training, which we hope to pursue in followup work.

While we recognize the potential practical impact of our findings for LLM pre-training, we would like to emphasize that the primary contribution of this work is establishing supercollapse as a novel and intriguing phenomenon across diverse settings, and demonstrating that studying universality in loss curves offers a promising avenue for advancing our scientific understanding of scaling. We will revise the paper to articulate this scientific focus more clearly.

Thank you again for your constructive feedback, which has helped improve the clarity and positioning of our paper. We will also correct the typos you mentioned in the revised version of the paper. We hope you will consider raising your score in light of our response and additional experiments.

[1] Hoffmann et al. "Training compute-optimal large language models."