Training Neural Networks at Any Scale: Training Neural Networks at Any Scale

Training Neural Networks at Any Scale

Speakers: Leena Chennuru Vankadara, Volkan Cevher

Conference: ICML 2025

YouTube: https://slideslive.com/39043354

Overview

This detailed technical article delves into the intricacies of training neural networks, particularly focusing on the challenges and opportunities presented by scaling models to unprecedented sizes. Presented by Volkan Cevher and Leena Chennuru Vankadara at ICML 2025, the talk dissects the prevailing "scaling law" paradigm, where larger models, more compute, and extensive data are believed to inherently lead to superior performance. While acknowledging the success of this approach, the speakers critically examine its limitations, demonstrating how naive scaling can paradoxically degrade model performance.

The tutorial is structured into two complementary parts. Volkan Cevher leads the first segment, focusing on advanced optimization algorithms that enhance training efficiency and stability. He introduces a unifying framework based on Linear Minimization Oracles (LMOs) and the Frank-Wolfe (Stochastic Conditional Gradient) method, reinterpreting common optimizers like Adam and providing a principled approach to hyperparameter management through explicit constraints. Leena Chennuru Vankadara then takes the stage to address the theoretical underpinnings of scaling, introducing "principled scaling rules." Her segment explores concepts like stability, feature learning, and the Maximal Update Parameterization (μP), demonstrating how these theoretical insights can lead to predictable performance gains and crucial hyperparameter transferability from small to large models, ultimately reducing the prohibitive computational costs associated with large-scale AI development.

The talk is highly significant for anyone involved in developing or deploying large-scale machine learning models, especially large language models (LLMs). It moves beyond the simplistic "bigger is better" narrative, providing a sophisticated understanding of how to scale effectively. By presenting both algorithmic innovations and foundational scaling theories, the speakers offer a roadmap for building more efficient, stable, and predictable AI systems, addressing critical issues like hyperparameter tuning costs, generalization performance, and the interaction between model architecture, optimization, and hardware.

Background

▶ Watch: Introduction: Language models and scaling challenges (0:00)

The advent of models like ChatGPT marked a pivotal moment in AI, ushering in an era dominated by the "scaling law" hypothesis. This principle posits a direct correlation: more compute, larger models, and more training data invariably lead to improved performance. The exponential growth in NVIDIA's stock price, often cited by Volkan Cevher, underscores the industry's commitment to this paradigm. However, the talk immediately challenges the notion that "scale is everything," presenting a crucial caveat: incorrect scaling, such as a naive increase in the width of a Mamba architecture, can lead to performance degradation rather than improvement. This highlights the necessity for not only optimal algorithms but also sophisticated scaling rules that intelligently adjust model architectures, hyperparameters, and algorithms.

The foundational problem addressed is the classical machine learning task of finding a functional mapping (often a neural network) that minimizes an empirical risk estimate based on observed data and labels. This "deep learning problem" is fundamentally an optimization challenge. The historical context of optimization is crucial, beginning with the Stochastic Gradient Descent (SGD) algorithm, formalized in 1951. SGD, in its basic form, uses Euclidean geometry to iteratively update parameters based on a stochastic gradient. Enhancements like mini-batching (reducing variance), momentum estimators (e.g., Polyak's heavy ball method, Nesterov acceleration), and more advanced variance-reduced estimators (e.g., STORM, SPIDER) have evolved to improve SGD's efficiency and convergence properties. SGD's enduring popularity stems from its ability to avoid traps in non-convex landscapes, its generalization capabilities, and its proven convergence properties.

The talk then transitions to two conceptual improvements for performance gains: on-the-fly adaptation and a priori adaptation. Adam is presented as the quintessential example of on-the-fly adaptation, an optimizer that adapts its step size and pre-conditioner based on observed gradients during training. While immensely popular (over 220k citations), Adam's hyperparameter tuning complexity for large-scale models is a significant challenge. A priori adaptation, on the other hand, involves methods like SMUG (Smooth Until Guilty), CO-SMUG (Coordinate-wise SMUG), and HELLA-SMUG (Hessian-based SMUG), which attempt to adapt step sizes or capture loss landscape geometry more proactively. This sets the stage for a deeper exploration into how these adaptive mechanisms can be unified and improved, especially in the context of large-scale neural network training.

Key Findings

▶ Watch: Stochastic Gradient Descent (SGD) fundamentals (4:00)

The talk presents several pivotal findings across both algorithmic innovation and theoretical scaling rules:

From Volkan Cevher's segment on Algorithms:

• Unifying Optimizer Framework: Many adaptive optimizers (e.g., Normalized SGD, LARS, LAMB, RProp, SignSGD, Signum, Lion, Muon) can be understood as instances of LMO-based steepest descent or the Stochastic Conditional Gradient Method (Frank-Wolfe). This provides a powerful, unifying template for analysis and extension.
• Weight Decay Reinterpretation: Traditional weight decay can be rigorously reinterpreted not just as a decoupled regularization term, but as an explicit constraint within the Frank-Wolfe framework. This perspective allows the "weight decay parameter" to be seen as a scaling parameter for explicit constraints, facilitating its transferability.
• Operator Norm Constraints for Transferability: By modeling layer-wise constraints using operator norms (OPN) (e.g., spectral norm for auxiliary layers, sign spectral-sign for input/output layers in image domains, spectral-spectral for image domains), the proposed Stochastic Conditional Gradient Method with Operator Norms enables hyperparameter transferability. This means that optimal step sizes and constraint radii tuned on smaller proxy models (e.g., 100 million parameters) can reliably transfer to much larger models (e.g., 3 billion parameters), drastically reducing the computational cost of tuning.
• Enhanced Stability and Efficiency: Explicitly constraining matrix norms provides perfect control over the spectral norms of weight matrices, which is critical for stability during longer training runs. These constrained algorithms also exhibit better behavior with increasing batch sizes, leading to higher hardware utilization and efficiency.
• Reduced Normalization Layers: The use of explicit operator norm constraints can, in some cases, make traditional normalization layers like LayerNorm or RMSNorm redundant, simplifying model architectures while maintaining stability.

From Leena Chennuru Vankadara's segment on Scaling Rules:

• Predictability Crisis in Scaling Laws: While empirical scaling laws exist, naive scaling often leads to unpredictable performance, with models sometimes performing worse at larger scales (e.g., Mamba example). Principled scaling rules are essential for achieving both optimality and predictability.
• Limitations of Standard Parameterization (SP): Under standard practices like He or LeCun initialization (SP), networks fail to admit feature learning in the infinite-width limit and exhibit instability for learning rates more aggressive than O(1/M) (where M is width). This contradicts observed empirical success.
• Controlled Divergence for Cross-Entropy Loss: For cross-entropy (CE) loss, SP can operate in a "benignly unstable" or "controlled divergence" regime. Here, the output layer might diverge, but hidden layers continue to learn features with a more aggressive learning rate scaling (O(1/√M)), outperforming Mean Squared Error (MSE) loss in similar settings.
• Maximal Update Parameterization (μP) for Optimal Scaling: μP is identified as a unique scaling rule that achieves effective feature learning in every layer. This means both effective and propagating updates remain O(1), making feature dynamics width-independent. μP leads to monotonic performance improvement with width, better generalization, and critically, hyperparameter transferability (e.g., optimal learning rates remain constant across different model widths).
• μP-squared (μP²) for Sharpness-Aware Minimization (SAM): For algorithms like SAM which perturb weights, μP² is a unique parameterization that ensures both effective feature learning and effective perturbations in every layer. This prevents the "perturbation collapse" seen in standard SAM under μP, where perturbations become effective only in the last layer, leading to suboptimal generalization.
• Complexities of Depth Scaling: When scaling both width and depth (e.g., in ResNets), the order and ratio of scaling significantly impact outcomes. Different "alpha" values for residual block scaling can lead to degeneracies like linearization within blocks or lack of feature diversity, affecting hyperparameter transferability.

Technical Deep Dive

▶ Watch: Common SGD estimators: Mini-batch and momentum (4:50)

The technical depth of the talk spans both advanced optimization theory and the theoretical foundations of neural network scaling.

Part 1: Algorithms for Efficient Training (Volkan Cevher)

The discussion begins with Stochastic Gradient Descent (SGD), its basic update rule ($X_{k+1} = X_k - \eta_k D_k$, where $D_k$ is the stochastic gradient and $\eta_k$ is the step size), and its dual feedback interpretation. Key variants include:

• Mini-batching: Averaging gradients over a batch to reduce variance.
• Momentum: Exponentially moving averages of gradients (Nesterov acceleration, Polyak's heavy ball method).
• Variance-reduced estimators: Such as STORM and SPIDER, offering theoretically superior complexity but often higher per-iteration cost.

Adam is presented as the dominant "on-the-fly" adaptive optimizer. It adapts to a Mahalanobis geometry by building a pre-conditioner from exponentially moving averages of first-order gradient information (squared gradients). This has roots in online learning algorithms like AdaGrad. Adam typically employs weight decay (keeping weights close to the origin) and learning rate schedulers (warm-up, constant phase, linear ramp-down, cosine decay) to manage training dynamics.

A crucial point is the challenge of applying dense pre-conditioners in large models. While AdaGrad and its "friends" explored denser Hessian approximations, their storage and computational costs become prohibitive (e.g., $2NM \times 2NM$ matrices for a two-layer network). This led to the development of structured Hessian approximations:

• Adafactor (Noam Shazeer): Uses row and column preconditioners, reducing degrees of freedom to $O(MN)$ (sublinear).
• Shampoo: Decouples layers and applies diagonal adaptation, $O(M^2 + N^2)$ degrees of freedom.
• SOAP: Incorporates Kronecker approximations and subspace rotations.

Volkan then introduces Stochastic Spectral Descent, an approach that treats weight matrices directly as matrices and uses spectral norms (Schatten infinity norms). This is framed using the concept of a Linear Minimization Oracle (LMO), where updates are derived by minimizing an inner product with the dual feedback over a norm-constrained space. The LMO for spectral norm involves computing the matrix sign function of the dual feedback, which can be done efficiently using randomized linear algebra (e.g., $O(MNR)$ for rank R) or Newton-Schulz / QR dynamically weighted Halley method (matrix multiplications, $O(MN^2)$). The overhead of these operations can vanish relative to backpropagation costs in sufficiently large GPU setups.

The LMO-based descent framework is presented as a unifying template. It shows that algorithms like Normalized SGD (which normalizes gradients to unit length to combat exploding/vanishing gradients) are instances of this, minimizing a linearized objective within a fixed distance (non-Euclidean geometry).

The "unsung hero" is the Stochastic Conditional Gradient Method (Frank-Wolfe). This algorithm explicitly treats training as a constrained optimization problem. The update $X_{k+1} = (1-\eta_k)X_k + \eta_k \text{LMO}(D_k)$ ensures iterates remain within a convex constraint set if $\eta_k \in [0,1]$. A key insight is that weight decay in this context can be understood as the Frank-Wolfe step size, with the constraint radius being the weight decay parameter. This allows operator norms (OPN) to define layer-wise constraints (e.g., spectral norm for $L_2 \to L_2$ stability, max norm for $L_1 \to L_\infty$ to limit extreme feature values). Specific OPN choices are recommended: spectral norms for auxiliary layers, and sign spectral-sign or spectral-spectral for input/output layers in image domains. This framework, particularly with the right scaling, yields Muon as a special case.

Part 2: Principled Scaling Rules and Theory (Leena Chennuru Vankadara)

Leena's part focuses on achieving optimal and predictable scaling through theory-driven rules. A scaling limit is defined as the limiting model obtained by extrapolating a scaling curve to infinite compute.

Stability at Initialization: For a multi-layer perceptron (MLP) with width $M$, activations $H$ are sums of $M$ i.i.d. random variables. To keep pre-activations stable (order one), the variance of weights must scale as 1/fan-in. This is the basis for He and LeCun initializations. Two key parameterizations are introduced:

• Standard Parameterization (SP): Weights sampled i.i.d. from Gaussian with variance $\sigma^2 / \text{fan-in}$.
• Neural Tangent Parameterization (NTP): Trainable weights scaled by $1/\sqrt{\text{fan-in}}$, then sampled from Gaussian. Equivalent at initialization but different during training.

Stability During Training:

• NTP: If learning rate $\eta$ grows faster than $O(1)$, unstable. If $\eta < O(1)$, stable but in the kernel regime (lazy regime), where training is equivalent to kernel regression with a fixed Neural Tangent Kernel (NTK). This regime lacks feature learning (updates to last-layer representations $\Delta H$ remain $O(1)$).
• SP: Maximally stable learning rate is $O(1/M)$. Faster rates lead to instability, slower rates to a kernel regime (a renormalized NTK). SP's last-layer initialization ($O(1/\sqrt{M})$ RMS norm) is fundamentally incompatible with feature learning, causing propagating updates to diverge.

Controlled Divergence Regime (for Cross-Entropy Loss): A critical nuance for CE loss (not MSE) is the existence of a "benignly unstable" regime. Here, the output logits diverge, but other dynamical quantities (activations, hidden layer gradients) remain stable, allowing hidden layers to learn features. The edge of this regime corresponds to a learning rate scaling of $O(1/\sqrt{M})$, which aligns with practical observations.

Maximal Update Parameterization (μP): To achieve both stability and effective feature learning in every layer (where both effective updates $\Delta W \cdot X$ and propagating updates $W_0 \cdot \Delta X$ remain $O(1)$), μP is introduced. This involves a specific, layer-wise scaling of initialization variances and learning rates. For signSGD, the learning rate should scale as $1/\text{fan-in}$. μP ensures all layers learn at a width-independent rate, making the maximally stable learning rate uniform across layers and improving predictability.

μP-squared (μP²) for SAM: For Sharpness-Aware Minimization (SAM), which perturbs weights, standard μP leads to a collapse of perturbations to only the last layer in the infinite-width limit. μP² is a unique scaling rule that ensures effective perturbations (avoiding collapse) in every layer alongside effective feature learning, by correctly scaling the perturbation radius (e.g., $O(1/\sqrt{M})$) and learning rates.

Depth Scaling: For ResNets (where $F_L$ is scaled by $L^{-\alpha}$), different depth limits (width then depth, or simultaneous) exist. For width $\to \infty$ then depth $\to \infty$ under μP:

• $\alpha=1$ (ODE-type limit): All features (including within residual blocks) are updated, but lack feature diversity (nearby layers learn similar features).
• $\alpha=0.5$ (SDE-type limit): Layers within residual blocks are linearized, but maximal feature diversity is maintained.

Both limits present challenges for hyperparameter transferability due to remaining discrepancies between finite and infinite models.

Experimental Setup & Results

▶ Watch: Advanced SGD methods: Nesterov, STORM, SPIDER (6:00)

The talk provides compelling empirical evidence to support its theoretical claims and algorithmic advancements.

Part 1: Algorithms (Volkan Cevher)

• LLM Speedrunning with Hyperparameter Transfer: A core demonstration involves training large language models. Using the proposed Stochastic Conditional Gradient Method with Operator Norms (effectively Muon with specific scaling rules), optimal step sizes and constraint radii were tuned on a small 100 million parameter proxy model. These hyperparameters were then directly transferred to a 3 billion parameter model, achieving state-of-the-art speedrunning performance with reliable transfer. This significantly reduces the computational expense of hyperparameter search for large models, which conventionally requires tuning at scale. The method also allows for the use of low-precision gradients and shows improved behavior with increasing batch sizes, critical for system efficiency and maximum FLOPs utilization.
• Image Classification with Data-efficient Transformers: The method was applied to Data-efficient Transformers for ImageNet classification. The proposed approach achieved substantial speed improvements over a heavily tuned Adam baseline. For a base model, it required 30% fewer epochs, resulted in 40% better wall-clock time, and 80% fewer iterations. Similar to LLMs, tuning on a tiny model could successfully transfer to a base model, demonstrating broad applicability of the transferability concept.
• Spectral Norm Control: The explicit constraint mechanism in the proposed algorithm provides perfect control over the spectral norms of weight matrices, which was shown to be critical for better performance in longer training runs and aligns with observations from production settings (e.g., Kimi AI).

Part 2: Scaling Rules (Leena Chennuru Vankadara)

• Mamba Architecture Scaling: An illustrative example showed a Mamba model where naive scaling of model width (using standard initialization and tuning learning rate at each width) led to a deterioration of performance. In contrast, applying "principled scaling" (to be discussed in the theory section) consistently improved performance with increasing width.
• SP vs. μP with MSE/CE Loss (MNIST, CIFAR-10): Experiments on MLPs trained on MNIST and CIFAR-10 compared the performance of Standard Parameterization (SP) and Maximal Update Parameterization (μP) under both Mean Squared Error (MSE) and Cross-Entropy (CE) losses. Under SP, CE loss consistently and considerably outperformed MSE loss. However, under μP, the performance gap between MSE and CE vanished, with MSE sometimes even slightly outperforming CE at scale, validating the theoretical prediction about μP's robustness to loss function choice.
• SP Learning Rate Scaling and Predictability: For MLPs and Transformers, empirical analysis showed that under SP, the "maximally stable learning rate" scaled as O(1/√M), while the "optimal learning rate" was closer to O(1). This discrepancy meant that extrapolating optimal learning rates from smaller models would lead to divergence at larger scales, demonstrating strong unpredictability and non-monotonic performance with width, even with optimal tuning.
• μP Performance and Predictability (Transformer on WikiText): A Transformer model trained on WikiText under μP exhibited monotonic improvement in performance with increasing width. Crucially, the loss curves for different widths strongly overlapped, and the optimal learning rate remained the same across all widths. This provides strong empirical validation for μP's claims of predictability and hyperparameter transferability.
• μP vs. μP² for SAM (MLPs on CIFAR-10): Experiments with Sharpness-Aware Minimization (SAM) on MLPs for CIFAR-10 showed that μP² (the parameterization designed for SAM) achieved better test accuracy compared to μP when the perturbation radius was properly tuned. Furthermore, under μP², the joint optima of learning rate and perturbation radius were transferable across scales, unlike under μP where both hyperparameters exhibited significant shifts.

Practical Implications

▶ Watch: Why SGD is popular: Non-convex optimization benefits (7:00)

The insights from this talk carry profound practical implications for practitioners, infrastructure teams, model builders, and deployers in the AI/ML landscape.

1. Drastic Reduction in Hyperparameter Tuning Costs: The most immediate and impactful implication is the ability to reliably transfer optimal hyperparameters (like learning rates and constraint radii) from small, computationally inexpensive proxy models to large, production-scale models. This eliminates the need for costly and time-consuming hyperparameter sweeps at the largest scales, which can otherwise run into millions of dollars and weeks of GPU time. This directly addresses one of the biggest bottlenecks in training foundation models.

1. Enhanced Training Efficiency and Hardware Utilization: Algorithms like the proposed Frank-Wolfe variant, which behave better with increasing batch sizes, allow for more efficient utilization of modern GPU and TPU clusters. Infrastructure teams can push hardware to its limits, maximizing FLOPs utilization without sacrificing model stability or performance.

1. Improved Model Stability and Generalization: Explicit control over matrix norms, particularly spectral norms, through the Frank-Wolfe framework ensures greater stability during long training runs. This can lead to more robust models that generalize better, especially in scenarios where traditional optimizers might struggle with exploding or vanishing gradients.

1. Principled Model Design and Architecture Choices: The theoretical framework of principled scaling rules (like μP and μP²) guides model builders away from arbitrary scaling. Instead of simply making models wider or deeper, they can make informed decisions about parameterization, initialization, and learning rate schedules that guarantee predictable performance and feature learning across all layers. This means models are not just "bigger" but "better designed."

1. Simplified Architectures and Reduced Normalization Layers: The ability of explicit operator norm constraints to sometimes replace or reduce the need for normalization layers (e.g., LayerNorm, RMSNorm) offers avenues for simplifying model architectures. This could lead to models that are not only easier to reason about but potentially faster during inference due to fewer operations.

1. Informed Loss Function Selection: The finding that cross-entropy loss can enable a "benignly unstable" regime for standard parameterization, leading to better feature learning than MSE, provides practical guidance on selecting appropriate loss functions, especially when working with prevalent but theoretically suboptimal parameterizations.

1. Algorithm-Specific Parameterization: The necessity of μP² for effective Sharpness-Aware Minimization across all layers highlights that the optimal parameterization can be algorithm-dependent. Practitioners must consider the interplay between their chosen optimizer and the model's scaling rules to unlock full potential.

Tradeoffs and Limitations:

While revolutionary, the framework acknowledges several open challenges:

• Data Modeling: The current theoretical framework largely abstracts away the data distribution. How optimal scaling rules interact with diverse or imbalanced datasets remains an open research question.
• Transferability of Regularization: Regularization hyperparameters like weight decay and dropout are not yet reliably transferable across different scaling dimensions (width, depth, data size, training time) as they are often highly data-dependent.
• Multi-Dimensional Scaling: The talk primarily focuses on width and depth scaling. Optimally allocating resources across multiple scaling dimensions simultaneously (e.g., width, depth, data size, training time, compute budget) is a complex, unsolved problem.
• Architectural Nuances: While the framework applies broadly, highly specialized architectures (e.g., Mamba) or specific hardware constraints might require tailored analysis and potentially different optimal scaling rules.

Despite these limitations, the presented work provides a powerful, principled foundation for navigating the complexities of large-scale neural network training, offering concrete strategies to overcome significant practical hurdles.

Key Takeaways

• Scaling neural networks, particularly large language models, demands principled approaches beyond simply increasing model size and compute; naive scaling can degrade performance and predictability.
• Many popular adaptive optimizers can be unified under the Linear Minimization Oracle (LMO) and Stochastic Conditional Gradient (Frank-Wolfe) framework, allowing for a more rigorous understanding and extension of their properties.
• Weight decay can be reinterpreted as an explicit constraint within the Frank-Wolfe method, enabling its "radius" to be transferred across different model scales alongside the learning rate.
• By using operator norms for layer-wise constraints within the Frank-Wolfe framework, it is possible to achieve hyperparameter transferability, allowing optimal settings from small proxy models (e.g., 100M parameters) to be reliably applied to much larger models (e.g., 3B parameters), drastically reducing tuning costs.
• The Maximal Update Parameterization (μP) is a critical theoretical discovery that ensures effective feature learning in every layer as models scale. This leads to monotonic performance improvement with width, better generalization, and consistent optimal learning rates across different model sizes.
• For advanced optimizers like Sharpness-Aware Minimization (SAM), a specialized parameterization called μP-squared (μP²) is necessary to ensure effective perturbations across all layers, preventing suboptimal "perturbation collapse" seen with other parameterizations.
• The choice of loss function (e.g., cross-entropy vs. MSE) and the precise rules for depth scaling (e.g., specific alpha values for ResNets) significantly impact training stability, feature learning, and hyperparameter transferability, requiring careful consideration.
• While theoretical frameworks offer invaluable guidance for building efficient, stable, and predictable large-scale ML systems, open challenges remain in areas such as comprehensive data modeling, the transferability of all types of hyperparameters (e.g., regularization), and optimal resource allocation across multi-dimensional scaling.

About the Speaker(s)

Volkan Cevher is a distinguished researcher and professor, known for his work in machine learning and optimization. During the talk, he mentioned teaching a class on training LLMs with Çağlar Gülçehre, highlighting his expertise in large-scale model development. His research encompasses adaptive methods for optimization, including algorithms like AccelGrad, UniXGrad, and UnderGrad, and he has made significant contributions to the understanding of spectral descent methods.

Leena Chennuru Vankadara is a prominent figure in the theoretical aspects of machine learning, with a specific focus on understanding and developing principled scaling rules for neural networks. Her contributions, as demonstrated in the talk, delve into the fundamental properties of initialization, parameterization, and their impact on stability, feature learning, and the predictability of model performance at scale.