$\mu$LO: Compute-Efficient Meta-Generalization of Learned Optimizers

Authors: Benjamin Thérien, Charles-Étienne Joseph, Boris Knyazev, Edouard Oyallon, Irina Rish, Eugene Belilovsky
Primary area: transfer learning, meta learning, and lifelong learning
Keywords: Meta Learning, Learned Optimizers, Pre-training

Abstract:
Learned optimizers (LOs) can significantly reduce the wall-clock training time of neural networks, substantially reducing training costs. However, they can struggle to optimize unseen tasks (meta-generalize), especially when training networks much larger than those seen during meta-training. To address this, we derive the Maximal Update Parametrization ($\mu$P) for two popular learned optimizer architectures and propose a simple meta-training recipe for $\mu$-parameterized LOs ($\mu$LOs). Our empirical evaluation demonstrates that LOs meta-trained with our recipe substantially improve meta-generalization to wider unseen tasks when compared to LOs trained under standard parametrization (e.g., as they are trained in existing work). When applying our $\mu$LOs, each trained for less than 250 GPU-hours, to large-width models we are often able to match or exceed the performance of pre-trained VeLO, the most performant publicly available learned optimizer, meta-trained with 4000 TPU-months of compute. We also empirically observe that learned optimizers trained with our $\mu$LO recipe also exhibit substantially improved meta-generalization to deeper networks ($5\times$ meta-training) and remarkable generalization to much longer training horizons ($25\times$ meta-training).

Decision: reject
Reviewer scores: 5, 6, 6, 6

Reviews and discussion:
<thread>
creation date:
1729335108586

thread_type:
official_review

<thread_object>
creator: reviewer_ohhy

<content>
summary:
In this paper, the authors aim to improve on the challenging task of improving the meta-generalization of Learned Optimizers (LOs) to longer optimization trajectories and larger base models while maintaining compute efficiency. The authors derive Maximal Update Parametrization for two popular learned optimizer architectures, enabling zero-shot hyperparameter transfer from smaller models to larger ones. A simple training recipe for µ-parameterized LO is proposed which leads to substantial improvement in meta-generalization to wider networks, deeper networks, and longer training horizons with orders of magnitude less computational cost.

strengths:
• Overall, I think the paper is well-written and addresses a very meaningful problem in the meta-learned optimizer. 

• The empirical performance of the proposed algorithm is very good (Figures 4, 5, and 6), and it is particularly impressive that they were able to get such good results using so little compute relative to VeLO (which is extremely expensive …).

weaknesses:
• The proposed method is quite simple, essentially, describing a way to construct a model such that the meta-learned optimizer can generalize to larger-width networks, deeper networks, and longer training trajectories. As far as I am aware the novelty of this paper is relatively limited since most of the findings are originally derived from [1] where Maximal Update Parametrizations was originally proposed. Therefore, from my perception, this appears to be a straightforward application of that work. Could the authors please elaborate on the technical novelty, of this paper? If the technical novelty is more than just an application of Maximal Update Parametrizations, I am happy to increase my score.

• The paper could be made more self-contained. I had to read [1] before really understanding much of the terminology and language being used in this paper.

[1] Yang, Greg, et al. "Tensor programs v: Tuning large neural networks via zero-shot hyperparameter transfer." NeurIPS2021

questions:
• From my understanding [1] shows how to initialize the network such that you can scale the width of the network while maintaining optimal hyper-parameters. This paper seems to suggest that using this strategy also allows you to scale the depth of the network and the optimization trajectory. Why is that the case? Can the authors please add some discussion about that and clarify more theoretically why the results are the way they are.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730485250446

thread_type:
official_review

<thread_object>
creator: reviewer_umkw

<content>
summary:
This paper proposed new optimizers to overcome the generalizability issue of the existing hand-crafted learned optimizers. It showed extensive experimental results so that the proposed optimizers improve meta-generalization to wider unseen tasks in most cases.

strengths:
*	Despite some cases where the proposed optimizers do not show the best performance, the “overall” results show that the proposed optimizers show better performance and better stability. 


*	It is interesting and informative that it shows a stable performance in the longer horizon.

weaknesses:
•	Why don’t you show results on ResNet and Plain ResNet (i.e., without residual connections)? Earlier, it has been shown that the reason why ResNet (with residual connections) works well by loss landscape – ResNet shows a smoother landscape while Plain ResNet shows a bumpy loss landscape in [1]. If the proposed optimizers can show stable performance even in Plain ResNets, I believe this work will be shown to be very effective. 		

•	In a nutshell, the technical addition of this work on top of the standard LP is adding 1/FAN_IN coefficient on the right term (the updated amount for the weight) in (2) if the weight is part of a hidden layer. In that sense, it is unclear how much technical contribution this work would contain. I wonder if the authors’ claim is that it shows a great advancement by adding such a small change. 

•	The result in Fig 4 is confusing with regard to the goal of the work. The goal of this paper is to develop a better-generalizable optimizer. However, it did not show the best performance on the most “out-of-distribution” tasks (transformer language modeling and ViT image classification, Fig. 4.) (Although the authors argued that their LOs are meta-trained with only 0.004% of VeLO-4000.) Hence, the contribution of this work is diluted. The proposed optimizers do not show the best performance for Transformer LM, in Fig 5 either. 

•	The legend in Fig 4 needs to be fixed. Either of VeLOM or LOM needs to be dotted line. Also, either of μLOM or μVeLOM needs to be a dotted line (by looking at (c), μVeLOM is a dotted line). 

[1]: Hao Li, Zheng Xu, Gavin Taylor, Christoph Studer, Tom Goldstein, “Visualizing the Loss Landscape of Neural Nets,” NeurIPS 2018.

questions:
•	Why did you select only MLP and Transformers (for language and vision) to show results on?

•	I am wondering how the loss landscapes of the approaches on ResNet and Plain ResNet would look.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730654662491

thread_type:
official_review

<thread_object>
creator: reviewer_s7qy

<content>
summary:
The paper investigates the limitations of current learned optimizers (LOs)—small, data-driven neural networks designed to optimize other neural models—in their ability to generalize to longer optimization horizons and larger models than those encountered during meta-training. The authors derive the Maximal Update Parametrization (muP) for two popular LO architectures and propose a simple meta-training recipe for mu-parameterized learned optimizers (muLOs). Empirical evaluations demonstrate that muLOs can match or exceed the performance of VeLO—a state-of-the-art LO trained with 4000 TPU-months of compute—when applied to large-width models.

strengths:
1. Effective Identification of Limitations in Learned Optimizers: The authors clearly pinpoint the generalization issues of current LOs, especially their generalization issues with larger models and longer training horizons.
2. Innovative use of Maximal Update Parametrization: By applying muP to LOs, they leverage theoretical guarantees for zero-shot hyperparameter tuning, going further to transfer optimizers. 
3. Practical Solution to Costly Hyperparameter Optimization: The proposed muLOs offer a way to avoid expensive HPO on each new architecture and task, addressing a significant practical challenge.
4. Significance in the Field: By enhancing the generalization capabilities of LOs and making them more practical for large-scale applications.

weaknesses:
1. Lack of Information on Computational Cost: The paper does not provide detailed information about the overall computational resources required for the method, spent on hyperparameter tuning for the mu-based methods. This is crucial for practitioners who may need to pre-train the optimizer themselves and assess the feasibility of implementing the method. 
2. High Pre-Training Computational Cost: According to Line 356, the pre-training of an LO with already tuned hyperparameters amounts to 16 TPU-months. This is a substantial computational expense, prohibiting practitioners from pre-training an LO themselves. This means that a SoTA method would be preferred, and hence the comparison to VeLO seems more critical.
3. Unclear Method Overhead: There is insufficient information regarding the computational overhead introduced by using the LOs compared to standard optimizers. Specifically, the paper does not quantify how much slower the training process becomes when employing the neural network-based optimizer. For practical applications, understanding this overhead is essential for adjusting compute budgets and training timelines accordingly.
4. Limited Clarity in Results Presentation: Throughout the paper, the primary metric presented is the training loss. However, details about how this metric is calculated and presented are lacking. For instance, in Figure 2, it is unclear whether the loss is measured exactly at the presented iteration if the error bars represent standard deviation across different seeds, or how the results would differ if showing the best loss achieved so far or using a moving average. 
5. Absence of downstream Evaluation Metrics.
6. Incomplete Presentation of Experimental Results: In Section 4.1.1, the authors mention collecting data for 1,000 training steps, but the plots presented display only 50% of this range (Figure 3 + appendix). Additionally, information about the number of seeds used and the variability across runs is missing. 
7. Ambiguities in Generalization Evaluation (Figure 4):
* Convergence Similarities: In Figures 4a and 4b, the loss curves for muAdam and the proposed muLOs converge to similar values. This suggests that the primary benefit is improved early performance (anytime performance), but without considering the method's overhead, it's unclear how significant this benefit is in practical applications.
* Behavior of muAdam in Figure 4c: The loss trajectory of muAdam has an unexpected shape in Figure 4c without a clear explanation, making it difficult to understand the underlying reasons.
* Confusing Legend and Color Coding: The legend in Figure 4 uses the same markers for VeLO and LO, and for muVeLO and muLO, making it challenging to distinguish between the different methods.
8. Terminology Clarification Needed: The paper refers to muAdam, but it is unclear whether this refers to muAdam or muAdamW. 
9. Inconsistent Performance Across Tasks (Figure 5): In Figure 5a, muLO appears to diverge, and across the tasks in Figures 5a and 5b, the ranking between muLO and muVeLO changes significantly. This inconsistency makes it difficult to determine which optimizer performs better overall and under what conditions, complicating the decision-making process for practitioners.
10. Questionable Claims About Stability: The claim in Line 516 that "muVeLO_M stably decrease training loss over time for each task" seems inaccurate. In Figure 6a, the loss begins to diverge, contradicting this statement. This raises concerns about the reliability of the method and the validity of the conclusions drawn.
11. High Variance in muVeLO Performance: The figures indicate that muVeLO exhibits a higher standard variance in its performance, while other methods show minimal variance. Understanding why muVeLO has higher variability is important for assessing its reliability and for guiding its practical adoption.

questions:
1. Can you provide detailed information on the computational resources required for your method, specifically the hyperparameter tuning of the mu-based methods?
2. Given that retraining of both the LO and VeLO is infeasible, how do you justify LO use compared to the SoTA method?
3. How much additional computational overhead does using the LO introduce compared to standard optimizers/ SoTA, and how does this affect training speed?
4. Could you clarify how the training loss is calculated and presented throughout the paper, and whether the error bars in Figure 2 represent standard deviation across different seeds?
5. Could you include any validation metric, to assess the practical impact of your method?
6. In Figure 3, why do the plots display only 50% of the collected data, and could you provide information about the number of seeds used and variability across runs?
7. Is muAdam in your experiments referring to muAdam or muAdamW, and could you clarify this terminology?
8. Given the inconsistent performance between muLO and muVeLO across tasks (as seen in Figures 5a and 5b), can you provide guidance on which optimizer performs better under specific conditions?
9. How do you reconcile the claim that "muVeLO_M stably decreases training loss over time for each task" (Line 516) with the divergence observed in Figure 6a?
10. Why does muVeLO exhibit higher variance in its performance compared to other methods, and what are the implications for its practical adoption?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730964461641

thread_type:
official_review

<thread_object>
creator: reviewer_i2y6

<content>
summary:
This paper proposes $\mu$LO, a new approach, inspired by hyperparameter transfer with the maximal update parameterization ($\mu$P), for training learned optimizers. Using $\mu$P for a model's initialization and weight updates has been shown to enable transferring high quality hyperparameter settings for the optimizer to networks of greater width, depth and longer training horizons.  Similarly, $\mu$LO addresses the limited ability of previous learned-optimizers to generalize to wider networks and longer training horizons by deriving the learned optimizer update required by $\mu$P.  Experiments show that $\mu$LO generalizes better than learned optimizers trained with standard parameterizations when applied to networks of larger width, depth, and longer training horizon as was observed for $\mu$P in the standard hyperparameter transfer settings.

strengths:
- The experimental results are quite compelling and show $\mu$ LO addresses the generalization issues of previous incantations of learned optimizers.
- This is an interesting application of $\mu$ P and could potentially reduce the cost of model tuning dramatically.  However, the bar is quite high and will require additional work to compare $\mu$ LO with newer optimizers like Shampoo and Modula.

weaknesses:
- While the experimental results are strong, they do not back up the claimed benefit that $\mu$ LO addresses the three issues of $\mu$ -transfer by avoiding the need for hyperparameter tuning for reach new architecture and task.
- The experimental results are missing comparisons to VeLO-4000 in Tables 5 and 6.  I think another reasonable baseline is per-task tuned Adam and $\mu$ Adam in the resource constrained setting to see if $\mu$ LO obvious the need for task specific optimizer tuning.
- The proposed $\mu$ LO does not take into account followup work on $\mu$ P that extends the original theoretical justification for transferring across width to depth (see [Tensor Program VI](https://arxiv.org/abs/2310.02244)).  Hence, the generalization across depth is more consequential then intended and may be improved with an explicit maximal update parameterization for depth.
- The legend for Figure 4 does not distinguish between dashed and solid lines.

questions:
- According to Table 1, $\mu$ LO_M and $\mu$ VeLO_M are trained on ImageNet classification using 3-layer MLP with 3 different width.  Is that correct?  Does that mean transformer and ViT architectures studied in Figure 4 d and e were not seen during meta-training and are there to show generalization to unseen architectures?
- Please include VeLO-4000 in Figures 5 and 6.
- Why is $\mu$ LO better than $\mu$ VeLO on average despite VeLO being a more powerful model?
- How does $\mu$ LO compare to a per task tuned Adam/ $\mu$ Adam?
- What is the overhead of using the learned optimizer per gradient step relative to Adam?
</content>
</thread_object>

</thread>
<thread>
creation date:
1732312496394

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
General reply to all reviewers

comment:
**General comment:**  We would like to thank all reviewers for their thoughtful insights and constructive criticism about our paper. We are pleased that reviewers believe we effectively highlight generalization issues with current learned optimizers [s7Qy] and that $\mu$ learned optimization ($\mu$LO) addresses them [i2Y6,oHHy]. We are particularly pleased that s7Qy thinks the use of $\mu$P in this context is innovative and that $\mu$LOs offer a practical way to dramatically reduce model tuning cost [i2Y6, s7Qy]. Finally, we are happy that umKW believes our method improves overall performance and oHHy finds the empirical performance to be impressive.

**Computational cost of $\mu$LO$_M$and $\mu$VeLO$_M$** We chose not to include any such comparison in our paper as the main focus of our work is to study the meta-generalization of $\mu$LOs applied to larger-width networks (orthogonal to improving learned optimizer computational cost). As mentioned on L538, the benefits of $\mu$LOs come at zero added computational cost relative to learned optimizers from previous work. Nonetheless, we have included results in the table below that compare the computational cost of the learned optimizers in this study to that of AdamW. We refer reviewers to [1] (section B.7) and [3] (sections 4.1,4.2,4.3,4.4) for a more detailed account of the computational cost of $\mu$VeLO$_M$ and $\mu$LO$_M$, respectively.  


*Table 1 Computational cost of the learned optimizers in our study and AdamW* We report the forward and backward pass time of the optimizee, the optimizer step time, and the total step time per iteration when training ResNet18 on ImageNet. Each number is the median step time over the last 250 training iterations. All timings were measured on an L40S Nvidia GPU.


| Optimizer         | Total Step Time | Optimizer Step Time | FwBw Time |
|-------------------|-----------------|---------------------|-----------------|
| AdamW            | 0.0199         | 0.0022              | 0.0177          |
| $\mu$LO$_M$      | 0.0219         | 0.0039              | 0.0180          |
| $\mu$VeLO$_M$    | 0.0253         | 0.0076              | 0.0177          |
| LO$_M$           | 0.0220         | 0.0041              | 0.0179          |
| VeLO$_M$         | 0.0255         | 0.0077              | 0.0178          |


**$\mu$VeLO$_M$ variance and performance relative to $\mu$LO**  $\mu$VeLO$_M$ uses an LSTM to generate the parameters of a small per-parameter MLP for each tensor in the optimizee. This two-level structure substantially increases its capacity relative to $\mu$LO$_M$(only a  single per-parameter MLP). While $\mu$VeLO$_M$ achieves generalization to wider networks, it appears to suffer from greater variance than $\mu$LO$_M$, which is not observed in VeLO-4000 (trained on a large distribution of tasks).  Moreover, $\mu$LO$_M$ outperforms $\mu$VeLO$_M$ on most tasks, despite its lower capacity. We hypothesize $\mu$VeLO$_M$ may require more extensive meta-training beyond the scope of this research to see it outperform $\mu$LO$_M$. Our goal here was simply to demonstrate that our parameterization also works for the VeLO architecture. 


**Figure 4 legend is incorrect** We have fixed this in the updated version. Thank you all for pointing this out.


---

**Local References**

[1] [VeLO: Training Versatile Learned Optimizers by Scaling Up, Metz et al.]

[2] [Tasks, stability, architecture, and compute: Training more effective learned optimizers, and using them to train themselves, Metz et al.]

[3] [Practical Tradeoffs Between Memory, Compute, and Performance in Learned Optimizers, Metz et al., CoLLAs 2022]

[4] [Yang et al., Tensor Programs V: Tuning Large Neural Networks via Zero-Shot Hyperparameter Transfer, Neurips 2021]
</content>
</thread_object>

</thread>
<thread>
creation date:
1732312855948

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Changes to the manuscript

comment:
**Here is a list of our changes to the manuscript based on reviewer feedback.**

- We have revised the text to clarify and tone claims in Sec. 3.3 with respect to i2Y6’s feedback. 
- We have added a brief discussion of Depth-$\mu$P and our choice not to study Depth-$\mu$P in section 2, paragraph 3. This change was made due to i2Y6's feedback.
- In section F of the appendix, we have added a comparison of the performance of our $\mu$LOs to $\mu$Adam tuned for 500 trials on a ViT task. This change was made due to i2Y6's feedback.
- We have added missing details from Figure 2 caption with respect to s7Qy’s feedback.
- We have added an evaluation of ResNets in section G of the appendix concerning reviewer umKW’s feedback.
- We have updated the manuscript by adding a paragraph to section 3.2 that adds more background about $\mu$P to the text in accordance with oHHy’s feedback.
- We have minimally updated the abstract, contributions, and section 4 introductory paragraph to emphasize that the findings of improved generalization to deeper networks and more training steps are purely empirical and not theoretically motivated.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732679085579

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Changes to the manuscript II

comment:
**Here is a list of our newest changes to the manuscript based on reviewer feedback.** 

- We have updated Figure 5a. While responding to reviewer s7Qy, we noticed that the runs for $\mu$LO$_M$ in Fig. 5a (6 months old) were affected by a bug in the ViT’s $\mu$-parameterization that was fixed a while ago. We have updated the plot with corrected runs and verified that no other results are affected in the paper. $\mu$LO$_M$ now consistently decreases the loss throughout training in Fig. 5a.
	
- We have added section H to the appendix of our paper, showing training curves of $\mu$VeLO$_M$ for 40,000 training steps.

- We have made sure that there is a sentence in the captions of every figure stating the number of trials used for each curve and explaining what the error bars designate. 

- We have added a clarifying sentence to section 4.1.2 paragraph 4 to explain the instabilities of $\mu$Adam in figure 5c.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732730092673

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Changes to the manuscript III

comment:
**Here is a list of our latest changes to the manuscript based on reviewer feedback.**
- We have enhanced the presentation of Figures 2,11,12,13,14 by replacing the legends in each figure with a global color bar and removing redundant x and y labels.
</content>
</thread_object>

</thread>


Metareview:
metareview:
The paper introduces µ-learned optimizers (µLOs), leveraging Maximal Update Parametrization (µP) to enhance generalization to larger models, deeper networks, and longer training horizons. µLOs achieve competitive or superior performance to state-of-the-art methods like VeLO, with significantly lower computational costs, effectively addressing challenges in hyperparameter tuning. Strenths of the paper include computational efficiency and performance of the approach in the benchmarks studied. However, concerns were raised about the limited novelty of the contribution, as it primarily applies µP to learned optimizers. Also, some key baselines, such as VeLO-4000 and per-task tuned optimizers, were not fully evaluated. Performance inconsistencies on certain out-of-distribution tasks were also pointed out. The main reason for my recommendation to reject is the incomplete experimental comparison.