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

Thank you for taking the time to review our paper. We are pleased to hear oHHy believes our paper is well written, that oHHy thinks it addresses an important problem for learned optimization, and that oHHy finds the empirical performance of our approach to be impressive.

We now address the reviewer's concerns. Each bolded paragraph title refers to one or multiple concerns that have been grouped together. All line numbers refer to the originally submitted manuscript (not the updated version).

Technical contribution beyond applying $\mu$P We agree that the change of meta-training LOs in $\mu$P is simple, however, the simplicity of our approach is one of its greatest strengths. This small but non-trivial change relative to the existing meta-training framework in learned optimization enables substantial improvements to meta-generalization with no added computational cost to the learned optimizer and no restriction on the problems that can be considered (e.g., any optimizee can be considered in the $\mu$LO framework). As such, it represents a net improvement over the existing meta-training framework in learned optimization [4.1,4.2,4.3].

Another contribution of our work is alleviating the limitations of $\mu$-transfer [4.4], itself. As outlined in lines 226-245, $\mu$-transfer increases the complexity of hyperparameter optimization (HPO), while requiring re-tuning for each new task and dataset considered—without the possibility of transferring HPO knowledge from one task to another (e.g. $\mu$Adam underperforms in heldout tasks Fig 4. (c),(d),(e)). In contrast, $\mu$LOs allows amortizing the cost of future HPO during meta-training at low cost (in contrast to [4.1], which also amortizes HPO) and can benefit from knowledge sharing across multiple tasks during meta-training, unlike $\mu$-transfer.

Therefore, learned optimization helps further alleviate the tuning costs of $\mu$-transfer for new problems, and parameterizing optimizees in $\mu$P improves the generalization of LOs from small to large tasks. As such, $\mu$ learned optimization is neither a simple application of learned optimization to $\mu$-transfer, nor a simple application of $\mu$-transfer to LOs. It constitute a new meta-training framework for learned optimization that is of interest for amortizing HPO and improving the meta-generalization of LOs.

The paper could be more self-contained with respect to [4.4] We have updated the manuscript by adding a paragraph to section 3.2 that adds more background about $\mu$P to the text. If oHHy could provide more details on anything they believe is still missing we would be grateful.

Question regarding generalization to depth and longer unrolls Our observations of stronger generalization to longer unrolls and deeper networks are purely empirical. We have amended the manuscript to emphasize that these findings are only empirical. [4.4] also shows similar empirical findings for longer and deeper versions of the tasks they tune on. Note that our findings for $\mu$LOs extend beyond the meta-training tasks to deeper and longer out-of-distribution tasks (e.g. ViT and LM in Figures 5 and 6).

Local References

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

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

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

[4.4] [Yang et al., Tensor Programs V: Tuning Large Neural Networks via Zero-Shot Hyperparameter Transfer, Neurips 2021]

Thank you for taking the time to review our paper. We are pleased to hear the reviewer believes that overall our proposed optimizers show improved performance and stability and that the generalization to longer training horizons is interesting.

We now address the reviewer's concerns. Each bolded paragraph title refers to one or multiple concerns that have been grouped together. All line numbers refer to the originally submitted manuscript (not the updated version).

Results on ResNets and Plain ResNets The loss landscape is independent of the optimizer. It depends on the data, the loss function, and the optimizee’s (network being optimized) functional form/parameters. It is, nonetheless, interesting to assess the performance of our learned optimizers for plain and residual networks. In section G of the appendix, we report the performance of $\mu$LOs and SP LOs on plain and residual networks, finding that our $\mu$LOs perform well for both outperforming SP LOs.

The technical contribution is simple While we agree that the change is simple, the strength of our approach, indeed, lies in its simplicity. Changing the optimizee’s parameterization is a small but non-trivial change relative to the existing meta-training framework in learned optimization: it enables substantial improvements to meta-generalization with no added computational cost to the learned optimizer and no restriction on the problems that can be considered. As such, it represents a net improvement over the existing meta-training framework in learned optimization [4.1,4.2,4.3].

Sub-optimal performance relative to VeLO-4000 Fig.4 (d) & (e) As mentioned on L321-323 the tasks in Fig.4 (d) and (e) fall under the meta-training distribution of VeLO-4000. That is, VeLO-4000 was meta-trained on variants of these tasks while $\mu$LO$_M$ and $\mu$VeLO$_M$ were only meta-trained on MLPs. When smaller width versions of the tasks are included in $\mu$LO meta-training (subfigures (a) and (b)), we observe that our optimizers outperform VeLO-4000. Note that due to the simplicity of our approach, there is no reason it would not be compatible with VeLO-scale meta-training.

Why study transformers and MLPs We chose to study transformers and MLPs because MLPs allow reporting performance on the meta-training distribution and transformers are the main architecture used in practice today.

Local References

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

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

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

[3.4] [Yang et al., Tensor Programs V: Tuning Large Neural Networks via Zero-Shot Hyperparameter Transfer, Neurips 2021]

Dear reviewer umKW,

Thank you for taking the time to review our paper.

As the ICLR discussion period comes to an end, we would like to offer our assistance in responding to any remaining questions or concerns you may have about our paper.

Best regards,

The authors

Local References

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

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

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

[2.4] [Yang et al., Tensor Programs V: Tuning Large Neural Networks via Zero-Shot Hyperparameter Transfer, Neurips 2021]

[2.5] [Junjie Yang et al., Learning to Generalize Provably in Learning to Optimize, AISTATS 2023]

[2.6] [Harrison et al. A Closer Look at Learned Optimization: Stability, Robustness, and Inductive Biases, Neurips 2022]

[2.7] [Understanding and correcting pathologies in the training of learned optimizers, ICML 2019]

Thank you for taking the time to review our paper. We are pleased that reviewer s7Qy believes we effectively identify issues with currently learned optimizers, that our use of $\mu$P in this context is innovative, that $\mu$LOs offer a practical way to reduce HPO costs, and that our work can make LOs “more practical for large-scale applications”.

We now address the reviewer's concerns. Each bolded paragraph title refers to one or multiple concerns that have been grouped together. All line numbers refer to the originally submitted manuscript (not the updated version).

Computational Cost of Meta-training is 16 TPU-months False. The computational cost of meta-training our LOs is provided on line 19 of the abstract, line 356 puts it in the context of VeLO’s meta-training, and line 356 discusses it again. The value on line 356 is given as a percentage, therefore, this corresponds to 0.16 TPU months, less than 5 TPU-days.

Computational cost of $\mu$LOs This is addressed in the general reply.

Figure 2 captions missing details Thank you for pointing this out. We have updated the caption accordingly.

Regarding downstream evaluation The main focus of our paper is to study the meta-generalization of $\mu$LOs applied to larger-width networks. Therefore, similar to prior published work [2.3], we choose to exclusively study training loss for simplicity. Studying the generalization of the optimizee to downstream tasks is an important orthogonal direction that has been explored multiple times in previous works, as mentioned in L124-127. Specifically, it is well known that the following strategies can improve optimizee generalization: targeting validation loss during meta-training [2.7], using weight decay [2.6], or using flatness-aware regularization [2.5].

“In Section 4.1.1, the authors mention collecting data for 1,000 training steps” This is incorrect, line 375 (section 4.1.1) reports that these plots show 500 steps of training.

Performance relative to $\mu$Adam in Figure 4a and 4b and computational cost Indeed, $\mu$ Adam (tuned for 500 trials) is a strong baseline. However, our goal is to improve meta-generalization to wider networks, a longstanding problem in learned optimization [2.1]. Therefore, it is most important to compare our approach to previous work in learned optimization, that is to the SP LOs and VeLO, which we handily beat in 4a and 4b. Note that our optimizers outperform$\mu$Adam in out-of-distribution tasks 4c, 4d, and 4e.

Behavior of muAdam in Figure 4c $\mu$Adam is tuned for the ImageNet MLP width=1024 task (see Table 1). The results in Figure 4c demonstrate that $\mu$Adam’s hyperparameters aren't transferable to CIFAR-10 while $\mu$LO$_M$and $\mu$VeLO’s hyperparameters are. This shows the strength of our $\mu$LO framework which can achieve stronger meta-generalization than $\mu$-transfer [2.4].

No mention that $\mu$Adam does not use weight decay Thank you for suggesting this. We have updated section 4 to explicitly specify that no weight decay is used with $\mu$Adam.

Inconsistent best performer among $\mu$LO$_M$and $\mu$VeLO This is addressed in the general reply.

$\mu$LO$_M$ diverges in Figure 5a We disagree that $\mu$LO diverges in Figure 5a. We acknowledge that $\mu$LO’s loss slightly increases after iteration 2000 but it is inconclusive whether it will diverge later on in training. It is important to compare $\mu$LO$_M$ to the LO$_M$ baseline in this figure. We note that $\mu$LO$_M$ reaches the final loss of value LO$_M$ after ~ 300 training iterations (16x speedup). Note that the SP baseline VeLO$_M$ (purple dashed line) does diverge in Figure 5a after the meta-training horizon (red dashed line). This demonstrates the improvements of $\mu$LOs over previous work with respect to generalization to deeper networks.

$\mu$VeLO$_M$ begins to diverge in Figure 6a $\mu$VeLO$_M$ does NOT diverge in Figure 6a. Perhaps you are referring to the SP baseline VeLO$_M$ (purple dashed line) which immediately diverges after the meta-training horizon (red dashed line) in this figure, demonstrating the improvements of $\mu$LOs over previous work with respect to generalization to longer training horizons. We stand by our claim that "$\mu$VeLO$_M$ stably decrease training loss over time for each task" in Figure 6.

Variance of $\mu$VeLO Answered in the general reply.

Possible divergence in Figure 5a After re-training the models in Figure 5, 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 accordingly and verified that no other results are affected in the paper. The training curve now shows no sign of divergence. Thank you for asking about this.

$\mu$VeLO small loss increase in Figure 6a The VeLO architecture takes as input the number of training steps remaining, thus, requiring the user to specify the total number of training steps (total_steps) a-priori (e.g. as is done for many LR schedules in practice). Therefore, at each step, VeLO’s LSTM is conditioned on an embedding that provides the number of training steps remaining, allowing it to learn a schedule. Previous work analyzing VeLO-4000’s behavior has noted that changing the value of the total_steps hyperparameter leads to variable performance [2.8]. Specifically, they found that increasing the value of total_steps does not always lead to better performance [2.8]. To answer the reviewer’s question about extending training in Figure 6a, we have included training curves for $\mu$VeLO$_M$ on the width $1024$ ViT ImageNet task (same task as in 6a) in section H of the appendix. We show curves for $40,000$ steps of training and $25,000$ steps of training. We observe that $\mu$VeLO$_M$ can successfully optimize for $40,000$ training steps. However, we note this training curve underperforms $\mu$VeLO$_M$ using total_steps=25,000, similar to what was found in previous work for VeLO-4000 [2.8]. We note that this issue is specifically related to the VeLO architecture and not to $\mu$LOs

$\mu$Adam in Figure 5c $\mu$Adam was tuned for 1000 training steps on a width 1024 MLP ImageNet task. Please note that it does not suffer from instability during the tuning window (first 1000 steps) in Figure 5c. While [2.4] shows that, in the cases they study, there is empirical generalization of $\mu$Adam’s hyperparameters to longer training unrolls, there is no reason to expect these findings hold for all tasks. We observe that $\mu$Adam’s hyperparameters are not transferable beyond the tuning horizon (1000 iterations) to the Cifar10 task in Figure 5c. We have added a clarifying sentence in section 4.1.2 paragraph 4 stating that $\mu$Adam suffers from instability outside of its tuning window in figure 5c.

Random seeds used and performance variation across multiple runs All of our plots reporting training loss show the average loss across 5 random seeds. The error bars in these plots report the standard error. Each seed corresponds to a different ordering of training data and a different initialization of the optimizee. We have added a sentence clarifying this in the introductory paragraph of section 4.1 and have made sure the captions of every figure throughout the entire paper information about seeds and error bars if there are any. For the performance variability, we assume the reviewer is referring to the variance in some $\mu$VeLO$_M$ training curves and its performance relative $\mu$LO$_M$. Please see the general comment for our answer to such concerns.

Thank you for taking the time to thoroughly review our paper. Please let us know if you have any other questions or reservations.

Local references

[2.4] [Yang et al., Tensor Programs V: Tuning Large Neural Networks via Zero-Shot Hyperparameter Transfer, Neurips 2021]

[2.8] [Is Scaling Learned Optimizers Worth It? Evaluating The Value of VeLO’s 4000 TPU Months]

Thank you for taking the time to review our paper. We are pleased that reviewer i2Y6 believes that $\mu$LO addresses generalization issues of previous work and that our approach could reduce the cost of model tuning dramatically.

We now address the reviewer's concerns. Each bolded paragraph title refers to one or multiple concerns that have been grouped together. All line numbers refer to the originally submitted manuscript (not the updated version).

Overclaims in Section 3.3 We believe that the reviewer is referring to Section 3.3 Paragraph 1 and, in particular, to our suggestion that $\mu$LOs allow amortizing the tuning cost of multiple architectures, unlike $\mu$-transfer [1.4]. We understand your concern: while the experiments currently demonstrate strong generalization to unseen tasks and architecture families (e.g., as is the case for ViT, LM in Figure 4 (d) and (e)), the performance is not necessarily optimal in the sense that training an LO on small versions of unseen tasks/architectures may perform better than a model attempting to meta-generalize. Our goal for Section 3.3 Paragraph 1 was to state that, unlike $\mu$-transfer, the $\mu$LO framework allows meta-generalizing across unseen tasks/architecture families (a significant weakness of $\mu$-transfer). We have revised the text to clarify and tone down those claims in Sec. 3.3. Thank you for bringing this to our attention.

Regarding Figure 5 and 6 baselines We emphasize that these figures study additional advantages of $\mu$LO for learned optimization beyond those it was designed to address (width scaling under a certain unroll length). Therefore, these figures only compare with the existing meta-training paradigm in learned optimization. In Figure 6, our focus is to establish the effect of $\mu$LO on generalization to longer unrolls. Thus, we directly compare to learned optimizers meta-trained identically (with the same unroll length) but using Standard Parameterization, which is the current paradigm in the learned optimization literature [1.1,1.2,1.3]. VeLO-4000 is meta-trained on a much longer unroll length (20k-200k steps) and as discussed in L321 is not a fair baseline, thus it is not only an unfair comparison but distracts from the question asked by Fig 6. For Figure 5, similarly, our goal is to directly show the unexpected advantage of $\mu$LO on generalization to deeper networks over learned optimizers meta-trained under the existing paradigm.

Regarding per-task tuned baselines We did not include the per-tasked re-turned $\mu$Adam baseline in our main results because our $\mu$LOs do not benefit from meta-training on these tasks and the comparison would not be fair. We do agree that this is an interesting baseline to consider, however. In section F of the appendix, we include results on the ViT w=3072 task where $\mu$Adam was tuned for 500 trials in the resource-constrained setting (e.g. on a smaller width version of the task) as you suggested. The results show that the re-tuned $\mu$Adam improves over its predecessor, but is outperformed by $\mu$LO$_M$ and $\mu$VeLO.

$\mu$ Depth Thank you for bringing this up. Depth $\mu$P only works for residual networks with a block depth of 1. Since most networks used in practice today (e.g., the transformer) have a block depth > 1 they are incompatible with $\mu$Depth. We, therefore, chose not to study meta-generalization under $\mu$Depth. The main focus of our paper is to study the meta-generalization of $\mu$LOs applied to larger-width networks. However, we noticed that this choice of parameterization is beneficial for meta-generalization to deeper networks and longer unrolls which is very interesting to the learned optimization community, so we also included results for these cases. We have added a brief discussion of \muDepth and our choice not to study \mu depth in section 2, paragraph 3.

Question regarding the meta-training of $\mu$LO$_M$ and $\mu$VeLO$_M$ Yes, the learned optimizers were only meta-trained on the tasks shown in Table 4 for $1000$ steps. This means the tasks seen in Figure 4 (d) and (e) were NOT seen during meta-training and that these subfigures show meta-generalization to new tasks. Moreover, all figures in the paper showing training beyond $1000$ iterations are out-of-distribution with respect to training duration for our $\mu$LO$_M$ and $\mu$VeLO$_M$.

Why is LO better than VeLO on average despite VeLO being a more powerful model? This question is addressed in the general reply.

Overhead of $\mu$LO$_M$ and $\mu$VeLO$_M$ This is addressed in the general reply.

Local References

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

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

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

[1.4] [Yang et al., Tensor Programs V: Tuning Large Neural Networks via Zero-Shot Hyperparameter Transfer, Neurips 2021]