$μ$LO: Compute-Efficient Meta-Generalization of Learned Optimizers

This paper proposes a meta-training recipe for learned optimizers using Maximal Update Parametrization ($\mu$P), demonstrating that $\mu$-parameterized learned optimizers ($\mu$LOs) significantly outperform standard parametrization counterparts in meta-generalizing to wider, deeper, and longer-horizon tasks under the same compute budget.

The Gist

Imagine you are teaching a robot to learn how to learn.

In the world of Artificial Intelligence, we usually hand-craft the "rules" the robot uses to improve itself. These rules are called optimizers (like the famous "Adam" or "SGD"). They are like a standard, one-size-fits-all recipe for baking a cake. It works okay for a small cake, but if you try to bake a giant wedding cake (a massive neural network), the standard recipe might make it collapse or taste terrible.

Learned Optimizers (LOs) are a newer, smarter idea. Instead of giving the robot a fixed recipe, we train a tiny AI (the "optimizer") to invent its own recipe for every new problem. It's like hiring a master chef who can taste the ingredients and decide exactly how much salt to add, rather than following a cookbook.

However, there's a big problem with these master chefs: They are bad at scaling.
If you train your chef on small cakes (small neural networks), and then ask them to bake a skyscraper-sized cake (a massive, wide network), they often panic. They might burn the cake, forget the ingredients, or just give up. They can't "generalize" to bigger tasks.

The Paper's Big Idea: "The Golden Rule of Scaling"

The authors of this paper, µLO, discovered a secret sauce to fix this. They applied a concept called Maximal Update Parametrization (µP).

Think of µP as a "Golden Rule" for how to build and train these neural networks.

• Without µP (Standard Parametrization): Imagine you are building a tower with blocks. If you make the tower twice as wide, you accidentally make the blocks twice as heavy. The tower becomes unstable and falls over. This is what happens to standard optimizers when networks get wider.
• With µP (The µLO approach): The Golden Rule says, "No matter how wide you make the tower, keep the weight of the blocks and the speed of the builders perfectly balanced." This ensures that whether you are building a small shed or a skyscraper, the physics of the construction remain stable.

What Did They Do?

1. Derived the Rules: They took two of the most advanced "master chef" optimizers (called VeLO and small_fc_lopt) and rewrote their internal code to follow this Golden Rule (µP).
2. The Training Recipe: They didn't just change the code; they also changed how they trained these chefs. Instead of training them on just one type of small cake, they trained them on a mix of small, medium, and large cakes.
3. The Result: They created µLOs (µ-Parameterized Learned Optimizers).

The Magic Results

The paper tested these new µLOs against the old ones and the standard hand-crafted optimizers. Here is what happened, using some analogies:

• The "Wider" Test: They asked the optimizers to train networks that were 8 times wider than anything they had ever seen before.
  • Old Optimizers: The tower collapsed immediately. The loss (error) went sky-high.
  • µLOs: They built the skyscraper smoothly. They didn't even break a sweat.
• The "Deeper" Test: They asked the optimizers to train networks that were 5 times deeper (more layers) than they were trained on.
  • Old Optimizers: They got confused and stopped learning.
  • µLOs: They handled the depth perfectly, even though the theory didn't strictly promise they would. It was a happy surprise!
• The "Longer" Test: They asked the optimizers to train for 25 times longer than their training sessions.
  • Old Optimizers: They got tired and started making mistakes (diverging).
  • µLOs: They kept going, stable and efficient, for the entire marathon.

Why Does This Matter?

Usually, to get a robot to handle giant tasks, you need massive amounts of computing power (like thousands of supercomputers running for months).

The µLO approach is a "hack" that costs zero extra money. By simply changing the mathematical "rules" of how the optimizer updates its weights (the Golden Rule), they got the same massive performance boost without needing more supercomputers.

The Takeaway

This paper is like discovering that the reason your car engine stalls when you try to drive a truck is that you were using a bicycle chain. By switching to a truck-sized chain (µP), the same engine can now pull a massive load effortlessly.

They proved that if you teach your "learning-to-learn" AI the right scaling rules, it can generalize to huge, unseen problems, saving us time, money, and energy in the future of AI development.

Technical Summary

1. Problem Statement

Learned Optimizers (LOs) are neural networks trained to optimize other neural networks. While they have shown potential to outperform hand-designed optimizers (like Adam or SGD) when meta-trained on massive scales (e.g., VeLO with 4000 TPU-months), they suffer from severe meta-generalization limitations. Specifically:

• Width Generalization: LOs trained on narrow networks often fail to optimize wider, unseen networks, frequently diverging or performing poorly.
• Depth and Horizon Generalization: They struggle to generalize to deeper architectures or training horizons (number of steps) significantly longer than those seen during meta-training.
• Cost: Achieving reasonable generalization often requires prohibitively expensive meta-training on thousands of diverse tasks, making it impractical for many research settings.

The core problem addressed is: How can we enable learned optimizers to generalize to wider, deeper, and longer training tasks without requiring massive-scale meta-training?

2. Methodology

The authors propose µLO, a framework that adapts the Maximal Update Parametrization (µP) to learned optimizer architectures.

A. Theoretical Derivation of µP for LOs

The paper derives the necessary scaling rules for two state-of-the-art LO architectures: small_fc_lopt and VeLO. To satisfy the µP desiderata (ensuring stable pre-activations and maximal updates in the infinite-width limit), the authors modify:

1. Optimizee Initialization: Hidden and input layer weights are initialized with variance $1/\text{FAN_IN}$, while output layers use variance 1.
2. Pre-activation Multipliers: Output layer pre-activations are scaled by $1/\text{FAN_IN}$.
3. Optimizer Update Scaling: The update rule $w_t = w_{t-1} - \alpha \lambda_1 d \exp(\lambda_2 m)$ is re-scaled. For hidden layers, the update is divided by $\text{FAN\_IN}$ to ensure the parameter update magnitude remains $\Theta(1)$ as width increases.

The authors prove theoretically (Propositions 4.1 and 4.2) that these modifications ensure that the input features to the LO remain $\Theta(1)$ (stable) regardless of the width of the optimizee network, satisfying the conditions for maximal update parametrization.

B. Meta-Training Recipe

Instead of training on thousands of diverse tasks, the authors propose a low-cost, multi-width meta-training recipe:

• Task Distribution: Meta-train on a small set of MLP tasks with varying widths (e.g., 128, 512, 1024) rather than a single width.
• Unroll Length: Meta-train for a fixed number of steps (e.g., 1000), but test on much longer horizons.
• Cost Efficiency: The proposed µLOs are meta-trained on a single GPU for roughly 100 hours, a fraction of the compute used by previous SOTA LOs.

3. Key Contributions

1. Theoretical Extension: First derivation of µP specifically for learned optimizer architectures (VeLO and small_fc_lopt), proving that standard LO input features can be made width-invariant through proper scaling.
2. Novel Training Recipe: A simple, compute-efficient meta-training strategy involving multi-width tasks that unlocks strong generalization capabilities.
3. Empirical Validation: Extensive experiments demonstrating that µLOs outperform both standard parametrization (SP) LOs and heavily tuned hand-designed optimizers (AdamW, µAdam) on out-of-distribution (OOD) tasks.

4. Experimental Results

The evaluation suite includes 35 tasks spanning MLPs, Vision Transformers (ViT), and Language Models (LM) with varying widths, depths, and training lengths.

A. Generalization to Wider Networks (Primary Focus)

• Performance: µLOs (specifically µLOM and µVeLOM) achieve significantly lower training loss on networks 8x to 64x wider than those seen during meta-training.
• Stability: While SP-LOs (standard parametrization) diverge immediately or fail to converge on wide networks (e.g., width > 2048), µLOs maintain smooth training dynamics.
• Ranking: In average rank comparisons across 5 diverse OOD tasks, µLOs consistently secured the 1st and 2nd best ranks, outperforming per-task-tuned AdamW and µAdam baselines.

B. Generalization to Deeper Networks (Unexpected Finding)

• Although µP is theoretically designed for width scaling, µLOs showed surprising generalization to depth.
• When tested on networks 5x deeper (16 layers vs. 3 layers) than meta-training, µLOs remained stable, whereas SP-LOs diverged immediately.
• Hypothesis: The authors attribute this to the stabilization of pre-activations provided by µP, which prevents the "blow-up" of logits in deep networks.

C. Generalization to Longer Training Horizons

• µLOs were tested on training runs 25x longer than the meta-training unroll length (25,000 steps vs. 1,000 steps).
• µLOs continued to decrease loss smoothly, while SP-LOs either diverged or became unstable after a few thousand steps.

D. Comparison to VeLO-4000

• The authors compared their µLOs (trained with ~0.004% of the compute budget of VeLO-4000) against the massive VeLO-4000 baseline.
• On in-distribution wide tasks, µLOs outperformed VeLO-4000. On far OOD tasks, they performed competitively, suggesting µP improves scalability more effectively than simply increasing compute.

5. Significance and Impact

• Democratization of LOs: This work demonstrates that high-performance learned optimizers do not require thousands of TPU-months of training. A simple change in parametrization and a multi-width training recipe allows for effective meta-generalization on a single GPU.
• Solving the Generalization Bottleneck: It addresses the primary failure mode of LOs (inability to scale to wider/deeper models), making them viable for real-world large-scale training.
• Stability Mechanism: The paper provides empirical evidence that pre-activation stability (a core tenet of µP) is the key driver for generalization not just in width, but also in depth and training duration.
• Future Direction: It opens the door for "low-cost" meta-learning of general-purpose optimizers that can be applied to arbitrary architectures without task-specific hyperparameter tuning.

In summary, µLO proves that by aligning learned optimizers with the theoretical principles of Maximal Update Parametrization, one can achieve robust, compute-efficient meta-generalization to unseen, large-scale neural network tasks.