🤖 machine learning

$μ$ pscaling small models: Principled warm starts and hyperparameter transfer

This paper introduces a principled width-based upscaling method that combines weight perturbation with theoretically grounded scaling laws to ensure functional equivalence and enable efficient hyperparameter transfer from small to large models, thereby reducing the cost of tuning for diverse inference budgets.

Original authors: Yuxin Ma, Nan Chen, Mateo Díaz, Soufiane Hayou, Dmitriy Kunisky, Soledad Villar

Published 2026-07-03

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Yuxin Ma, Nan Chen, Mateo Díaz, Soufiane Hayou, Dmitriy Kunisky, Soledad Villar

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you have a very talented, small apprentice chef who has spent months mastering a specific recipe. Now, you want to open a massive restaurant with a huge kitchen and a team of 100 chefs to cook that same dish for thousands of people.

The old way of doing this was to hire 100 brand-new chefs and teach them the recipe from scratch. This takes a long time and costs a fortune in training.

A newer, smarter way (called "upscaling") is to take your one talented apprentice, clone them 100 times, and put them all in the big kitchen. Since they all know the recipe perfectly, the big kitchen starts off with the same quality as the small one. However, there's a catch: if you just clone them exactly, they will all do the exact same thing at the exact same time. They won't learn to use the extra space or the new tools in the big kitchen; they'll just be 100 copies of the same person, which is inefficient.

This paper introduces a new, mathematically proven method to make this "cloning and expanding" process work perfectly. Here is how they do it, explained simply:

1. The Perfect Clone (Dynamic Equivalence)

First, the authors figured out the exact mathematical rules for cloning the small chef into a big team so that, initially, the big team acts exactly like the small one.

The Analogy: It's like taking a single musician playing a song and copying their sheet music 100 times. If everyone plays the exact same notes at the exact same speed, the sound is identical to the solo performance.
The Innovation: Previous methods could do this at the start, but they didn't know how to keep the big team in sync as they started learning and changing. This paper proves that if you adjust the "volume" (learning rate) and the "tempo" of the big team correctly, they will stay perfectly in sync with the original small chef throughout the entire training process.

2. The "Nudge" (Breaking the Symmetry)

Once the big team is perfectly synced, they are stuck in a rut. They are all doing the same thing, so they can't explore the new, larger kitchen to become even better.

The Analogy: Imagine 100 clones standing in a circle. If they all step forward at the same time, they just move as a block. To make them useful, you need to give them a tiny, random "nudge" so they step in slightly different directions.
The Innovation: The paper introduces a precise way to add this "noise" or nudge. It's not a random guess; it's a calculated amount of chaos. This nudge breaks the perfect symmetry, allowing the big team to start using their extra capacity to learn new things, while still keeping the core knowledge of the original chef.

3. The "Magic Map" (Hyperparameter Transfer)

The hardest part of training a big team is figuring out how much "nudge" to give and how fast they should learn (the learning rate). Usually, you have to test this on the massive, expensive big kitchen, which is a waste of money.

The Analogy: Imagine you want to know how much salt to put in a giant pot of soup. Instead of buying a giant pot and testing it, you use a tiny saucepan. You figure out the perfect amount of salt for the tiny pot, and then you use a "magic map" to know exactly how much salt to put in the giant pot without ever testing the giant one.
The Innovation: The authors proved that this "magic map" exists. You can train a tiny version of the upscaled model (a small team of clones), find the perfect settings there, and then transfer those settings directly to the massive model. This saves a tremendous amount of time and computing power.

Why This Matters

Speed: It allows us to take a small, already-trained model and turn it into a giant, powerful one much faster than starting from scratch.
Cost: It saves money because we don't need to run expensive experiments on the giant model to find the right settings; we do it on a cheap, small model first.
Reliability: The paper provides a mathematical guarantee that this process works, rather than just guessing based on trial and error.

What They Tested

The authors tested this method on three different types of "chefs" (neural networks):

MLPs: Simple networks used for tabular data (like predicting forest types).
ResNets: Networks used for image recognition (like identifying cats and dogs).
GPT-2: Large language models used for generating text.

In all cases, the "cloned and nudged" models learned faster and reached a better final result than models trained from scratch, while using the "magic map" to skip the expensive tuning phase.

In short: This paper gives us a rulebook for how to safely and efficiently grow a small, smart AI into a giant, smart AI without wasting time or money, ensuring the giant version doesn't just become a confused crowd of clones.

Technical Summary: µpscaling Small Models: Principled Warm Starts and Hyperparameter Transfer

Problem Statement

Modern deep learning workflows typically involve training families of models at multiple scales. Small models are used for rapid prototyping and scaling law analysis, while larger models are trained from scratch for final deployment to meet diverse hardware and latency constraints. This workflow is computationally wasteful, as it discards the knowledge acquired during the training of smaller models.

While "model upscaling" (initializing larger models from trained smaller ones) has been proposed to mitigate this waste, existing methods (e.g., Net2Net) face two critical limitations:

Hyperparameter Sensitivity: Upscaled models often require hyperparameters (particularly learning rates and noise levels) that differ significantly from those of the base model or models trained from scratch. Tuning these at the target large scale is prohibitively expensive.
Lack of Theoretical Grounding: Previous approaches rely on empirical heuristics. It remains unclear whether tuning hyperparameters on smaller models and extrapolating via scaling laws is valid in the context of upscaling, as the training dynamics of an upscaled model differ from both the base model and a scratch-trained large model.

Methodology

The authors propose a theoretically grounded framework for µpscaling, which combines function-preserving weight transformations with the Maximal Update Parametrization (µP) to enable efficient hyperparameter transfer.

1. Dynamic Equivalence

The core theoretical contribution is the extension of "static equivalence" (where a widened model produces the same output as a base model at initialization) to dynamic equivalence.

Mechanism: By duplicating weights and applying specific width-dependent rescalings (based on the ⊗ Kronecker product with matrices of ones), a wider model can be constructed that parametrizes the exact same function as the base model.
Optimizer Coordination: Crucially, the authors prove that if optimizer hyperparameters (learning rate, weight decay, and auxiliary parameters like $\epsilon$ ) are rescaled according to specific rules derived from the homogeneity of the update function, the wider model follows an identical training trajectory in function space as the base model.
Generality: This dynamic equivalence holds for general architectures (expressible in the Ne⊗or⊤ program framework, including MLPs, ResNets, and Transformers) and general entrywise optimizers (SGD, Adam, AdamW).

2. The Upscaling Algorithm

To utilize the extra capacity of the wider model, the authors introduce a "copy-and-perturb" procedure:

Widening: Construct a dynamically equivalent wider model by duplicating and rescaling weights and optimizer states.
Symmetry Breaking: Inject a small amount of noise ( $\Delta$ ) into the parameters. Without noise, the model remains confined to the low-dimensional subspace of the base model. The noise allows the model to escape this subspace and exploit the increased capacity.
µP Scaling: The noise magnitude and optimizer hyperparameters for the upscaled model are scaled according to µP rules. This ensures that the training dynamics in the infinite-width limit are "optimal," preserving the signal from the base model while ensuring the injected noise contributes meaningfully without vanishing or exploding.

3. Hyperparameter Transfer

The framework enables zero-shot hyperparameter transfer. Because the infinite-width limit of the upscaled training dynamics is rigorously characterized using modified Tensor Programs, the optimal hyperparameters (learning rate and noise level) tuned on a small "proxy" upscaled system ( $n_0 \to kn_0$ ) transfer directly to the target large system ( $n \to kn$ ).

Key Contributions

Dynamic Equivalence Theory: The paper establishes that, under specific weight and hyperparameter rescalings, an upscaled model is dynamically equivalent to its base model throughout the entire training process, not just at initialization. This extends existing theoretical underpinnings to encompass the full training trajectory.
Principled Upscaling Algorithm: A unified procedure for upscaling that integrates weight duplication, symmetry-breaking noise injection, and µP-based hyperparameter scaling. This provides a rigorous alternative to heuristic methods like Net2Net.
Hyperparameter Transfer Guarantees: The authors provide the first rigorous framework for transferring hyperparameters in the upscaling setting. By extending the Tensor Program machinery to accommodate mid-training upscaling, they prove that optimal hyperparameters can be tuned on small, inexpensive proxy systems and transferred to large-scale targets, drastically reducing tuning costs.
Infinite-Width Analysis: A modified Tensor Program framework is introduced to characterize the infinite-width limit of upscaled training, showing that pre-activations converge to i.i.d. blocks with non-trivial covariance structures, distinct from standard non-upscaled training.

Experimental Results

The method was evaluated on three architectures (MLPs, ResNets, GPT-2) across various datasets (Forest Cover Type, CIFAR-100, FineWeb) and optimizers (SGD, AdamW).

Performance: Across all settings, upscaled models converged faster and reached lower final training loss compared to models trained from scratch at the same width.
Hyperparameter Transfer: Experiments verified that optimal learning rates and noise levels tuned on small systems (e.g., width $n_0=100$ ) transferred robustly to much larger systems (e.g., width $n=500$ or $n=1024$ ).
Efficiency: By performing hyperparameter tuning at a small scale rather than the target scale, the method achieved significant computational speedups in tuning costs:
- 23.6× speedup for MLPs ( $n/n_0 = 5$ ).
- 16.0× speedup for ResNets ( $n/n_0 = 4$ ).
- 48.2× speedup for GPT-2 ( $n/n_0 = 10$ ).
Generalization Caveat: While training loss improved, the upscaled ResNet on CIFAR-100 showed worse generalization (validation accuracy) than the scratch-trained model. The authors note this is consistent with their theory, which governs training dynamics rather than generalization, and is a known property of overparameterized CNNs on limited data.

Significance and Claims

The paper claims to provide the first rigorous theoretical framework for hyperparameter transfer in model upscaling. Its significance lies in:

Bridging Theory and Practice: Moving upscaling from empirical heuristics to a principled approach grounded in infinite-width limit theory and Tensor Programs.
Cost Reduction: Enabling the tuning of hyperparameters for massive models on small proxy systems, making the exploration of large-scale architectures more computationally feasible.
Universality: The framework applies to standard architectures and optimizers, unlike previous methods that were often architecture-specific.

The authors remain modest regarding generalization, acknowledging that their theory focuses on training dynamics and that generalization under upscaling remains an open question. They also note that the method currently applies to width-wise upscaling, with depth-wise upscaling left for future work.