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

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

Speakers: Shikai Qiu, Google DeepMind; Lechao Xiao, Google DeepMind; Andrew Wilson, New York University; Jeffrey Pennington, Google DeepMind; Atish Agarwala, Google DeepMind

Conference: ICML 2025

YouTube: https://slideslive.com/39043901

Overview

This article delves into a groundbreaking discovery presented at ICML 2025 by Shikai Qiu and collaborators from Google DeepMind and New York University, titled "Scaling Collapse Reveals Universal Dynamics in Compute-Optimally Trained Neural Networks." The talk introduces the phenomenon of "scaling collapse," a novel approach to understanding and predicting the full training dynamics of large-scale machine learning models. The core idea is that when normalized appropriately, the loss curves of neural networks trained across vastly different scales collapse onto a single, universal trajectory, revealing a profound underlying predictability in complex training processes.

The motivation behind this work stems from the immense computational and financial costs associated with training increasingly large AI models. Just as physicists develop robust predictive models for complex systems like orbital dynamics to guide expensive space missions, the ML community urgently needs reliable models to predict the outcomes of large training runs. Traditional scaling laws, while valuable, typically only predict a limited set of observables, such as final loss or optimal token budget. This research pushes beyond these limitations, aiming to model the entire loss curve, thereby offering a more comprehensive and granular understanding of how models learn and scale. Such predictive capabilities are invaluable for optimizing training procedures, fine-tuning hyperparameters, and diagnosing deviations from expected behavior during critical, large-scale experiments.

By moving beyond hand-picked functional forms and simple power-law extrapolations, the researchers propose a data-first methodology rooted in the principle of scaling collapse. This approach not only offers higher fidelity and predictability for modeling full training curves but also provides a powerful diagnostic tool. The presence or absence of collapse can empirically indicate whether a model's scaling rules are "properly" configured, thus offering a new lens through which to evaluate and discover optimal scaling strategies for various training parameters. This work highlights the emergence of universal behaviors in neural network training, drawing parallels to universality observed in other fields of science and opening new avenues for theoretical exploration in AI.

Background

▶ Watch: Understanding training dynamics at scale: limitations of scaling laws (0:00)

The pursuit of understanding and predicting the behavior of large neural networks has largely been driven by the concept of scaling laws. Pioneering works, such as the Kaplan paper and later the Chinchilla paper, demonstrated that quantities like final loss often follow predictable power-law relationships with model size, dataset size, and compute budget. This led to a three-step procedure: collect small-scale data, identify trends (often power laws), and extrapolate to larger scales. While effective for a few specific observables—like final loss on the compute-optimal Pareto frontier, optimal token budget, or optimal batch size—this traditional approach faces significant limitations when attempting to model the full, intricate training dynamics of a neural network.

The primary challenge is that the assumption of simple, discoverable formulas breaks down for complex phenomena like entire loss curves. Early attempts, such as fitting sums of power laws to training loss curves, yielded only "okay" fits at best. The reason for this inadequacy lies in the rich, multi-dimensional dependencies of loss curves on various hyperparameters. Beyond model size and data tokens, factors like learning rate schedule, weight decay, batch size, and context length all profoundly influence the training trajectory. These "messy and important details" are difficult to capture with simple, hand-picked functional forms, making the task of modeling a multi-dimensional scaling law that accounts for all these variables incredibly complicated and intractable through mere guesswork and curve fitting.

Furthermore, the traditional approach often focuses on asymptotic behavior or final performance, making it difficult to detect suboptimal scaling choices during the early or intermediate stages of training. What appears to be a minor deviation at small scales or in final performance might mask significant inefficiencies in the underlying scaling rules. The need for a more robust, data-driven methodology that can reveal deeper, scale-invariant dynamics throughout the entire training process became evident, paving the way for the exploration of alternative approaches like scaling collapse. This method seeks to directly extrapolate observed trends from small-scale experiments without the need for pre-defined functional forms, offering a more general and higher-fidelity understanding of scaling phenomena.

Key Findings

▶ Watch: Introducing scaling collapse: a data-first approach (3:30)

The central discovery of this research is the phenomenon of scaling collapse in compute-optimally trained neural networks. When the loss and compute of different models are normalized by their respective final values (i.e., loss as a fraction of final loss, and compute as a fraction of total compute), the training curves across various model sizes collapse into a single, universal curve. This remarkable convergence implies that the full training dynamics, beyond just the endpoints, are highly predictable and exhibit a universal shape.

A key aspect of this collapse is its exceptional precision. The researchers quantified this using a metric called collapse deviation, which measures the variance across the normalized curves. They found that this deviation is often lower than the noise floor, representing the inherent randomness from different random seeds, for a significant portion of the training process, especially when learning rate decay is employed. This precision underscores the robustness of the observed universality. The phenomenon is also shown to be highly general, observed across diverse scenarios including transformers on the C4-5M dataset, transformers on chess data, and MLPs trained on regression datasets. It holds whether model depth or width is scaled, suggesting a fundamental property of neural network training.

The researchers draw an interesting connection between scaling collapse and the concept of universality in other scientific fields, such as statistical physics and probability theory (e.g., the Central Limit Theorem). In these domains, systems of different sizes or even entirely different systems can exhibit identical behavior up to rescaling, often revealing deeper theoretical principles. While speculative, this connection suggests that similar underlying mechanisms might govern the observed universality in neural network training.

Crucially, scaling collapse serves as a powerful empirical indicator for the "correctness" of scaling rules. The talk demonstrates that if specific scaling rules, such as muP for learning rate scaling or Chinchilla-optimal training for token budgets, are not applied, the collapse breaks down. For instance, ablating muP scaling or training models beyond Chinchilla optimality leads to distinct, non-collapsing curves. This is particularly significant because these suboptimal scaling procedures are often difficult to discern by merely inspecting final performance or small-scale results. The presence of super collapse thus acts as an "integration test" for scaling, providing a novel methodology to empirically determine optimal scaling rules for other dimensions like batch size, momentum, or context length.

Finally, the work offers a partial theoretical explanation for scaling collapse. For loss curves that approximate a sum of power laws, collapse can be shown to occur when models are trained to be compute-optimal. More broadly, the researchers propose a simple yet effective model for the impact of learning rate schedules on real transformer loss curves. This model, derived from SGD on quadratic loss with noisy gradients, approximates the loss as a sum of a gradient flow term (independent of learning rate schedule when time is measured in gradient flow time) and an excess loss term (proportional to learning rate times the trace of the gradient covariance). This framework, with a single fitted parameter of 0.21, accurately predicts loss curves across different learning rate schedules, model sizes, and total token budgets, providing a mechanistic understanding of why collapse occurs even with varying schedules.

Technical Deep Dive

▶ Watch: Discovery of universal collapse by normalizing loss and compute (4:30)

The technical foundation of this work revolves around a data-first approach to understanding scaling trends, circumventing the limitations of traditional scaling laws that rely on hand-picked functional forms. The core methodology is scaling collapse, which transforms diverse training loss curves into a unified, universal trajectory.

The experimental setup primarily involved training transformers for next-token prediction on a subset of the C4 dataset, specifically C4-5M. To ensure that hyperparameters were well-specified across different scales and to mitigate confounding factors, the researchers employed muP (maximal update parameterization) for both initialization and learning rate scaling, which is known to preserve signal propagation and allow for scale-invariant training dynamics. Models were trained to be compute-optimal, adhering to the principles derived from the Chinchilla scaling law, which dictates the optimal balance between model size and training tokens for a given compute budget. Additionally, the learning rate was consistently decayed to zero by the end of training, a common practice for achieving good final performance.

The key insight for achieving scaling collapse is a specific normalization procedure. Instead of directly comparing raw loss curves, which span different ranges of loss and compute, each curve is normalized by its own final values. Specifically, the loss at any point in training is expressed as a fraction of the model's final loss, and the compute utilized up to that point is expressed as a fraction of the total compute allocated for that model. Mathematically, if $L(t)$ is the loss at time $t$ and $C(t)$ is the compute at time $t$, then the normalized loss is $L(t)/L_{final}$ and normalized compute is $C(t)/C_{total}$. When this normalization is applied, the curves across different model sizes become "almost indistinguishable," collapsing onto a single, universal trajectory. This means that while the absolute values of loss and compute vary greatly with scale, their relative progress through training follows a universal pattern.

The precision of this collapse is rigorously quantified by comparing the collapse deviation (the variance across the normalized curves) against the noise floor (the relative fluctuations within a single model's loss curve due to random seeds). The findings show that the collapse deviation is consistently lower than the noise floor for much of the training process, particularly when learning rate decay is employed. This suggests that the observed universality is not merely an approximation but a highly precise phenomenon, often exceeding the inherent stochasticity of the training process itself. The robustness of this collapse extends to various learning rate schedules, including "exotic" ones, and is observed across different datasets (C4-5M, chess data) and architectures (transformers, MLPs), whether scaling model depth or width.

While the empirical observation of collapse is powerful, the researchers also provide a theoretical framework to explain its occurrence. They demonstrate that if loss curves can be approximated as a sum of power laws, then collapse naturally emerges when models are trained to be compute-optimal. However, real-world loss curves are significantly influenced by learning rate schedules. To account for this, a more sophisticated model is introduced, inspired by Stochastic Gradient Descent (SGD) on quadratic loss with noisy gradients. This model posits that the loss curve can be approximated as the sum of two terms: a gradient flow term and an excess loss term.

The gradient flow term describes the deterministic behavior of training in the absence of noise. Crucially, this term becomes independent of the specific learning rate schedule when time is re-parameterized into gradient flow time, defined as the integral of the learning rate over wall-clock time. This re-parameterization allows for a cleaner separation of concerns. The excess loss term, on the other hand, accounts for the additional loss incurred due to the noise in the gradients and is proportional to the learning rate ($\eta$) multiplied by the trace of the gradient covariance. By allowing a single, slightly adjusted scaling factor (found to be 0.21) for this excess loss term, the model achieves "quite excellent agreement" with real experimental results across different learning rate schedules, model sizes, and total training horizons. This unified model, combined with observations on how the trace of the gradient covariance scales with the loss itself, successfully explains the observed collapse across diverse learning rate schedules, extending beyond the simple power-law approximation.

The discussion also briefly touches upon a potential connection to renormalization group theory, pioneered by Ken Wilson in the 1970s. This framework provides a unifying perspective for universality phenomena by relating effective descriptions of systems at different scales. While speculative, the researchers hint that similar deeper theoretical principles might underpin the observed scaling collapse in neural networks, suggesting a rich area for future theoretical exploration.

Experimental Setup & Results

▶ Watch: Generality of scaling collapse across datasets and architectures (7:00)

The experimental core of this work involved a systematic study of neural network training dynamics across varying scales, primarily focusing on transformers. The models were trained on the C4-5M dataset for a next-token prediction task, a standard setup for language modeling. To ensure comparability and proper scaling across different model sizes, the researchers meticulously configured hyperparameters. A crucial element was the use of muP (maximal update parameterization) for both network initialization and the scaling of the learning rate. This technique is known to facilitate more stable and predictable scaling behavior by maintaining signal propagation properties across different model widths and depths. All models were trained to be compute-optimal, following the empirically derived principles of Chinchilla-like scaling laws, which balance model size and the number of tokens seen during training to maximize performance for a given compute budget. The Adam optimizer was used, and the learning rate was consistently decayed to zero by the end of training for all runs.

The primary "results" observed were the loss curves during training. When these raw curves, collected from models of different sizes, were normalized—specifically, by dividing the instantaneous loss by the final loss and the instantaneous compute by the total compute—they exhibited a striking collapse onto a single, universal curve. This visual evidence, presented graphically in the talk, showed that curves across different model sizes became "almost indistinguishable," demonstrating the core phenomenon of scaling collapse.

To quantify the precision of this collapse, the researchers introduced collapse deviation, which measures the variance of the normalized loss curves around their mean. This was compared against the noise floor, defined as the relative fluctuations in a single model's loss curve across multiple random seeds, representing the inherent stochasticity of the training process. A key finding was that the collapse deviation was often lower than the noise floor for a significant portion of the training, especially when learning rate decay was active. This indicates that the universal dynamics revealed by scaling collapse are not merely approximate but remarkably precise.

The generality of scaling collapse was demonstrated across several dimensions:

• Architectures: Observed in transformers (on C4-5M and chess data) and MLPs (on regression datasets).
• Scaling Dimensions: Held true whether scaling the depth or the width of the models.
• Learning Rate Schedules: The collapse was robust across various learning rate schedules, including "exotic" ones, particularly towards the end of training when the learning rate decayed to zero.

Crucial ablations were performed to highlight the conditions under which scaling collapse occurs, revealing its utility as a diagnostic tool:

• Ablating muP Learning Rate Scaling: When muP was not used for learning rate scaling, and a constant learning rate was applied, the normalized loss curves did not collapse. Instead, they showed distinct, non-overlapping trajectories, indicating a lack of scale-invariant dynamics.
• Non-Chinchilla-Optimal Training: When models were trained beyond what is considered Chinchilla-optimal (e.g., by altering the exponent for the data versus compute power law), the normalized curves again failed to collapse. A "small but consistent shift over model size" was observed, signaling that the scaling rule for the token budget was suboptimal.

These ablation results underscore a critical insight: traditional scaling laws, which often focus on final performance, can struggle to differentiate between optimal and suboptimal scaling procedures, especially at small scales. The "asymptotically suboptimal" may not be obvious from small-scale final performance alone. Scaling collapse, however, provides an immediate and clear empirical signal, acting as an "integration test" for the appropriateness of scaling rules. If everything is scaled properly, the training dynamics should be scale-invariant, leading to collapse.

Finally, the talk presented results from a simple yet effective theoretical model explaining the impact of learning rate schedules on loss curves. This model, which approximates loss as a sum of a gradient flow term and an excess loss term, demonstrated "quite excellent agreement" with real experiments. A single fitted parameter, 0.21, successfully captured the behavior across different learning rate schedules, model sizes, and total token budgets, validating the model's predictive power and providing a mechanistic understanding of collapse under diverse training conditions.

Practical Implications

▶ Watch: Connecting scaling collapse to universality in other sciences (7:50)

The discovery of scaling collapse carries significant practical implications for practitioners, infrastructure teams, model builders, and deployers in the AI/ML landscape. At its core, scaling collapse offers a powerful predictive model for the full training dynamics of large neural networks. By running smaller, more affordable experiments and observing the universal curve, practitioners can potentially predict the entire loss curve of an expensive, large-scale training run before launching it, provided the final loss can still be extrapolated (e.g., using traditional scaling laws). This capability can dramatically reduce wasted compute and time, treating large model training more like a scientific experiment with robust prior predictions.

One immediate application is in optimizing training procedures and hyperparameters. If the universal curve is known, deviations from it in early stages of a large run could signal issues with hyperparameter choices or the training setup. This allows for proactive adjustments rather than discovering problems late in a costly experiment. Furthermore, scaling collapse provides a novel diagnostic tool: if a model's normalized loss curves fail to collapse, it serves as a strong empirical indicator that the underlying scaling rules (e.g., for learning rate, data budget, batch size, momentum, or context length) are suboptimal or incorrectly specified. This "integration test" for scaling can help diagnose problems that are otherwise hard to detect from simple final performance metrics or small-scale experiments, which often mask suboptimal asymptotic behavior.

The research opens the door to empirically discovering new optimal scaling rules. If the hypothesis holds that proper scaling across any axis leads to scale-invariant training dynamics (i.e., collapse), then practitioners can systematically test different scaling strategies for previously unstudied hyperparameters like batch size or context length. By observing which strategies lead to collapse, they can identify empirically sound scaling rules, reducing reliance on intuition or trial-and-error.

Moreover, scaling collapse can potentially improve the accuracy of traditional scaling law fits. Instead of relying solely on endpoint measurements (like final loss) to fit scaling law parameters, the entire training trajectory can be leveraged. By incorporating more data points from the full loss curve, the fits for scaling law parameters could become more robust and precise, leading to better extrapolations for future large-scale models.

However, there are also practical tradeoffs and limitations. While scaling collapse provides high fidelity for the shape of the loss curve, it still requires knowledge of the final loss for normalization. This final loss might still need to be predicted using traditional scaling law extrapolations, meaning that scaling collapse doesn't entirely eliminate the need for those methods but rather complements them. The theoretical models explaining collapse, while effective, are described as "simple" rather than a deep, unifying principle, suggesting that a complete mechanistic understanding is still evolving. For instance, the discussion around joint scaling limits with dataset size, particularly concerning data repetition and overfitting, highlights areas where the universality might break down or require further nuance. Practitioners would need to be mindful of these conditions, ensuring that training regimes (e.g., avoiding excessive data repetition) align with the assumptions under which collapse is observed. Despite these nuances, the ability to gain such precise and general insights into training dynamics, often achievable without "huge compute budgets" for the initial collapse experiments, makes this a highly valuable and feasible scientific pursuit for improving the efficiency and predictability of large-scale AI development.

Key Takeaways

• Universal Training Dynamics: When appropriately normalized, the full training loss curves of neural networks across different scales collapse into a single, universal trajectory, revealing predictable dynamics beyond just final performance.
• Precision and Generality: This scaling collapse is remarkably precise, often exhibiting less variance than the noise introduced by random seeds, and is general across diverse architectures (transformers, MLPs), datasets (C4-5M, chess), and scaling dimensions (model depth, width).
• Empirical Indicator for Scaling Rules: The presence of scaling collapse acts as a powerful empirical "integration test." Its breakdown signals suboptimal scaling rules for hyperparameters like learning rate (e.g., without muP) or token budget (non-Chinchilla-optimal training), offering a new diagnostic tool.
• Data-First Predictive Power: Scaling collapse provides a data-first methodology for predicting full loss curves of expensive large-scale runs from smaller experiments, offering higher fidelity and predictability than traditional scaling laws that rely on hand-picked functional forms.
• Theoretical Explanation (Partial): Simple theoretical models, combining power-law behavior with a gradient noise model (approximating loss as gradient flow plus excess loss due to noise), can effectively explain the observed collapse, even across varied learning rate schedules.
• Future Directions and Practical Value: The phenomenon opens exciting avenues for discovering new optimal scaling rules for various hyperparameters (batch size, momentum, context length), improving existing scaling law fits, and exploring deeper theoretical principles of universality in machine learning, all with significant practical benefits for efficient AI development.

About the Speaker(s)

The work presented was a collaborative effort by a team of researchers from Google DeepMind and New York University. The talk was delivered by Shikai Qiu, who highlighted the significant contributions of his collaborators: Lechao Xiao, Andrew Wilson, Jeffrey Pennington, and Atish Agarwala. Specifically, Shikai Qiu gave special credit to Lechao Xiao for initially discovering the phenomenon of "super collapse," which served as the genesis of this impactful project. Their collective expertise spans advanced machine learning research, with a focus on understanding the fundamental dynamics and scaling properties of neural networks.