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

## Abstract

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 italic_μ P) for two popular learned optimizer architectures and propose a simple meta-training recipe for μ 𝜇 \mu italic_μ -parameterized LOs ( μ 𝜇 \mu italic_μ 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 italic_μ 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 observe that learned optimizers trained with our μ 𝜇 \mu italic_μ LO recipe also exhibit substantially improved meta-generalization to deeper networks ( 5 × 5\times 5 × meta-training) and remarkable generalization to much longer training horizons ( 25 × 25\times 25 × meta-training).

## 1 Introduction

Deep learning (Goodfellow et al., ) has enabled a great number of breakthroughs (Brown et al., ; Brooks et al., ; Radford et al., ; Alayrac et al., ; Kirillov et al., ; Rombach et al., ; Oquab et al., ). Its success can, in part, be attributed to its ability to learn effective representations for downstream tasks. Notably, this resulted in the abandonment of a number of heuristics (e.g., hand-designed features in computer vision (Dalal and Triggs, ; Lowe, )) in favor of end-to-end learned features. However, one aspect of the modern deep-learning pipeline remains hand-designed: gradient-based optimizers. While popular optimizers such as Adam or SGD provably converge to a local minimum in non-convex settings (Kingma and Ba, ; Li et al., ; Robbins, ), there is no reason to expect these hand-designed optimizers reach the global optimum at the optimal rate for a given problem. Given the lack of guaranteed optimality and the clear strength of data-driven methods, it is natural to turn towards data-driven solutions for improving the optimization of neural networks.

To improve hand-designed optimizers, Andrychowicz et al. (); Wichrowska et al. (); Metz et al. (; ) replaced them with small neural networks called learned optimizers (LOs). Metz et al. () showed that scaling up the training of such optimizers can significantly improve wall-clock training speeds and supersede existing hand-designed optimizers. However, LOs have limitations in meta-generalization – optimizing new problems. For example, despite training for $4000$ TPU months, VeLO (Metz et al., ) is known to (1) generalize poorly to longer optimization problems (e.g., more steps) than those seen during meta-training and (2) have difficulty optimizing models much larger than those seen during meta-training. Given the high cost of meta-training LOs (e.g., when meta-training, a *single training example* is analogous to training a neural network for many steps), it is essential to be able to train learned optimizers on small tasks and generalize to larger ones. Harrison et al. () explore preconditioning methods to improve the generalization from shorter to longer optimization problems (e.g., ones with more steps). However, no works have tackled the meta-generalization of LOs to wider models in a principled way.

To address the meta-generalization problem of LOs, we recognize that this problem can be reformulated as zero-shot hyperparameter transfer (Yang et al., ).
The latter involves selecting optimal hyperparameters of hand-designed optimizers for training very large networks (that one cannot afford to tune directly) by transferring those tuned on smaller versions of the model. Under the standard parametrization (SP)(^1^11When we refer to SP, we follow the same meaning as Yang et al. (). That is, we consider SP to designate a parameterization that does not admit HP transfer. However, we note that recent work (Everett et al., ) shows hyperparameter transfer is possible in SP under certain alignment assumptions.), the optimal hyperparameters of an optimizer used for a small model do not generalize well to larger versions of the model. However, when a small model is tuned using the Maximal Update Parametrization ($\mu$P), and its larger counterparts are also initialized with $\mu$P, the small and large models share optimal hyperparameters (Yang et al., ).
Given the appealing connection between zero-shot hyperparameter transfer in hand-crafted optimizers and meta-generalization in LOs, we ask the following questions: *Can learned optimizers be meta-trained under $\mu$P? How would the resulting optimizers perform on wider unseen tasks?* We seek to answer these questions in the following study. Specifically, we consider the recent LO architectures (Metz et al., ; ) and demonstrate that $\mu$P can be adapted to these optimizers leading to our $\mu$LO optimizers. We subsequently conduct an empirical evaluation which reveals the power of our $\mu$LOs and their advantages for scaling learned optimizers.

Our contributions can be summarized as follows:

- •
We derive $\mu$-parameterization for two popular learned optimizer architectures (VeLO and small_fc_lopt) and propose a training recipe for $\mu$LOs.
- •
We demonstrate that $\mu$LOs meta-trained with our recipe significantly improve generalization to wider networks when compared to their SP counterparts and several strong baselines and that, for wider counterparts of the meta-training tasks, they outperform VeLO (meta-trained with 4000 TPU-months of compute).
- •
We demonstrate that $\mu$LOs meta-trained with our recipe show improved generalization to deeper networks ($5\times$ meta-training) when compared to their SP counterparts.
- •
We demonstrate that $\mu$LOs meta-trained with our recipe significantly improve generalization to longer training horizons ($25\times$ meta-training) when compared to their SP counterparts.

Our results show that $\mu$LOs significantly improve learned optimizer generalization without increasing meta-training costs. This constitutes a noteworthy improvement in the scalability of meta-training LOs.

Figure: ((a))
Refer to caption: x1.png

## 2 Related Work

##### Learned optimization.

While research on learned optimizers (LOs) spans several decades (Schmidhuber, ; Thrun and Pratt, ; Chen et al., ; Amos, ), our work is primarily related to the recent meta-learning approaches utilizing efficient per-parameter optimizer architectures of Metz et al. (). Unlike prior work (Andrychowicz et al., ; Wichrowska et al., ; Chen et al., ), which computes meta-gradients (the gradients of the learned optimizer) using backpropagation, Metz et al. () use Persistent Evolutionary Strategies (PES) (Vicol et al., ), a truncated variant of evolutionary strategies (ES) (Buckman et al., ; Nesterov and Spokoiny, ; Parmas et al., ). ES improves meta-training of LOs by having more stable meta-gradient estimates compared to backpropagation through time, especially for longer sequences (i.e. long parameter update unrolls inherent in meta-training) (Metz et al., ). PES and most recently ES-Single (Vicol, ) are more efficient and accurate variants of ES, among which PES is more well-established in practice making it a favourable approach to meta-training.

##### Generalization in LOs.

One of the critical issues in LOs is generalization in the three main aspects (Chen et al., ; Amos, ): (1) optimize novel tasks (often referred to as meta-generalization); (2) optimize for more iterations than the maximum unroll length used in meta-training; (3) avoid overfitting on the training set.
Among these, (3) has been extensively addressed using different approaches, such as meta-training on the validation set objective (Metz et al., ), adding extra-regularization terms (Harrison et al., ), parameterizing LOs as hyperparameter controllers (Almeida et al., ) and introducing flatness-aware regularizations (Yang et al., ).
The regularization terms (Harrison et al., ; Yang et al., ) often alleviate issue (2) as a byproduct.
However, meta-generalization (1) has remained a more difficult problem.
One approach to tackle this problem is to meta-train LOs on thousands of tasks (Metz et al., ). However, this approach is extremely expensive and does not address the issue in a principled way leading to poor meta-generalization in some cases, e.g. when the optimization task includes much larger networks.
Alternatively, Premont-Schwarz et al. () introduced Loss-Guarded L2O (LGL2O) that switches to Adam/SGD if the LO starts to diverge improving meta-generalization. However, this approach needs tuning Adam/SGD and requires additional computation (e.g. for loss check) limiting (or completely diminishing in some cases) the benefits of the LO. In this work, we study aspects (1) and (2) of LO generalization, demonstrating how existing SP LOs generalize poorly across these dimensions and showing how one can apply $\mu$P to learned optimizers to substantially improve generalization in both these aspects.

##### Maximal Update Parametrization.

First proposed by Yang and Hu (), the Maximal Update Parametrization is the unique stable abc-Parametrization where every layer learns features. The parametrization was derived for adaptive optimizers by Yang and Littwin () and was applied by Yang et al. () to enable zero-shot hyperparameter transfer, constituting the first practical application of the tensor programs series of papers. Earlier works in the *tensor programs series* build the mathematical foundation that led to the discovery of $\mu$P. Yang () shows that many modern neural networks with randomly initialized weights and biases are Gaussian Processes, providing a language, called Netsor, to formalize neural network computations. Yang () focuses on neural tangent kernels (NTK), proving that as randomly initialized network’s width tends to infinity, its NTK converges to a deterministic limit. Yang () shows that randomly initialized network’s pre-activations become independent of its weights when its width tends to infinity. Building on the latest works in the tensor program series (Yang and Littwin, ; Yang et al., ), in this work, we show that $\mu$P can be extended to the case of learned optimizers and empirically evaluate its benefits in this setting.

## 3 Method

### 3.1 Background

A standard approach to learning optimizers (Metz et al., ) is to solve the following meta-learning problem:

$$ $\min_{\phi}\;\mathbb{E}_{({\mathcal{D}},{\bm{w}}_{0})\sim{\mathcal{T}}}\left[ \mathbb{E}_{(X,Y)\sim{\mathcal{D}}}\left[\frac{1}{T}\sum_{t=0}^{T-1}{\mathcal{ L}}(X,Y;f_{\phi}({\bm{u}}_{t}),{\bm{w}}_{t})\right]\right],$ (1) $$

where ${\mathcal{T}}$ is a distribution over optimization tasks defined as pairs of dataset ${\mathcal{D}}$ and initial weights ${\bm{w}}_{0}$ associated with a particular neural architecture (we refer to this network as the *optimizee*), $\phi$ represents the weights of the learned optimizer, $f_{\phi}$, that takes gradient-based features ${\bm{u}}_{t}$ as input. Finally, ${\mathcal{L}}$ is the loss function used to train the optimizee. $T$ is the length of the unroll which we write as a fixed quantity for simplicity. In our experiments, during meta-optimization, $T$ is varied according to a truncation schedule (Metz et al., ).
A clear goal of the learned optimization community is not only learning to solve optimization problems over ${\mathcal{T}}$, but also to apply the learned optimizer, $f_{\phi}$, more generally to unobserved datasets and architectures. This *transfer* to new tasks is referred to as meta-generalization. This problem can be seen as a generalization of the zero-shot hyperparameter transfer problem considered in Yang et al. (); for instance, when the optimizer is a hand-designed method such as SGD or Adam and $\phi$ represents optimization hyper-parameters such as the learning rate.

Gradient descent is a standard approach to solving equation 1. However, estimating the meta-gradients via backpropagation for very long unrolls is known to be noisy (Metz et al., ). Instead, gradients are estimated using evolutionary strategies (Buckman et al., ; Nesterov and Spokoiny, ; Parmas et al., ). Evolutionary strategies work by sampling perturbations to the LO’s weights (similar to SPSA (Spall, )), unrolling an optimization trajectory for each perturbation, and estimating gradients with respect to evaluations of the meta-objective (usually the loss of the network being optimized, see eq. 1). In contrast to ES, which estimates one gradient per full unroll, PES (Vicol et al., ) allows estimating unbiased gradients at many points (called truncations) during the full unroll. This allows updating the optimizer’s parameters more often during meta-training. We use PES to estimate meta-gradients in our experiments.

Learned optimizer features ${\bm{u}}_{t}$ are constructed based on momentum, second-order momentum, and adafactor values as in (Metz et al., ), with the full list of features described in the (Table 6 of the Appendix). In our experiments, the architectures of our $f_{\phi}$ are similar to small_fc_lopt of Metz et al. () and VeLO of Metz et al. () except that their dimensions differ slightly (see sec. C for details). $f_{\phi}$ has three outputs $m$, $d$, and $\alpha$, the magnitude, scale, and learning rate of the update respectively. For small_fc_lopt, $\alpha_{w}=1$ always, $\alpha_{w}$ is produced by the tensor-level LSTM for VeLO. The standard LO update is given as

$$ $w_{t}=w_{t-1}-\alpha_{w}\lambda_{1}d_{\phi}\exp{\left(\lambda_{2}m_{\phi} \right)},$ (2) $$

where $\lambda_{1}$ and $\lambda_{2}$ are constant values of $0.001$ to bias initial step sizes towards being small.

### 3.2 μ 𝜇 \mu italic_μ -parametrization for Learned Optimizers

Consider a model to be optimized $g_{w}$ with weights in layers $l$ denoted $w_{l}$. We apply and construct $\mu$LOs as follows.

##### Initialization- μ 𝜇 \mu italic_μ .

$w_{l}$ which are hidden and input layers (aka the output sizes that can increase) have their weights initialized as $\mathcal{N}(0,\frac{1}{\sqrt{\textsc{fan_in}}})$. While output layers (those mapping to fixed sizes) have their weights initialized as $\mathcal{N}(0,1)$.

##### Multipliers- μ 𝜇 \mu italic_μ .

Output layers pre-activations are multiplied by $\frac{1}{\textsc{fan_in}}$ during the forward pass.

##### Updates- μ 𝜇 \mu italic_μ .

The update by $f_{\phi}$ on the parameters of $g_{w}$, at both meta-training and evaluation is modified as follows:

$$ $w_{t}=\begin{cases}w_{t-1}^{i}-\frac{1}{\textsc{fan_in}}\cdot\left(\alpha_{w} \lambda_{1}d_{\phi}^{i}\exp{\left(\lambda_{2}m_{\phi}^{i}\right)}\right)&\text {if $w^{i}$ is part of a hidden layer}\\ w_{t-1}^{i}-\alpha_{w}\lambda_{1}d_{\phi}^{i}\exp{\left(\lambda_{2}m_{\phi}^{i }\right)}&\text{otherwise.}\end{cases}$ (3) $$

We now show that this can lead to a maximal update Parametrization, following the analysis of (Yang et al., , Appendix J.2.1) which studies the initial optimization step. For our analysis, we consider a simplified input set for $f_{\phi}$ which takes as input only the gradient
while producing an update for each layer. Note that this analysis extends naturally to other first-order quantities.

###### Proposition 1 .

Assume that the LO $f_{\phi}$ is continuous around 0. Then, if $f_{\phi}(0)\neq 0$, the update, initialization, and pre-activation multiplier above is necessary to obtain a Maximal Update Parametrization.

###### Proposition 1 .

### 3.3 μ 𝜇 \mu italic_μ LO Meta-training Recipe

$\mu$P for hand-designed optimizers involves tuning the optimizer on a small width version of the target architecture and transferring the hyperparameters to the larger width target model (Yang et al., ). While $\mu$-transfer makes hyperparameter search for large models tractable, it has the following limitations: (1) the smaller scale hyperparameter search suffers from increased complexity as it requires sweeping various multipliers in addition to the standard hyperparameters, (2) tuning the hyperparameters on too small of a model may results in sub-optimal hyperparameters for the largest models, (3) Yang et al. () recommend repeating the procedure for every new task/dataset.

Figure: ((a))
Refer to caption: x4.png

Meta-training flexible $\mu$-parametrized learned optimizers can address these limitations. Due to their flexible functional forms (as opposed to just a learning rate hyperparameter), $\mu$LOs can learn to optimize networks in $\mu$P without tuning multipliers (we set all multipliers to $1$ in our experiments). Therefore, by training our $\mu$LOs with fixed multipliers on multiple tasks that are large enough to admit strong transfer but still tractable and reusing it on new tasks, we address (1), (2), and (3) by amortizing the tuning cost during the optimizer pre-training stage avoiding costly hyperparameter tuning on each new architecture and task.

To verify the effectiveness of this multi-task strategy for learned optimizers, we compare $\mu\text{LO}_{S}$, trained on a single small task (see Tab. 1), to $\mu\text{LO}_{M}$, trained on 3 small tasks of the different width (see Tab. 1), in figure 2. When training for $1000$ steps (meta-training unroll length), we observe that $\mu\text{LO}_{M}$ outperforms $\mu\text{LO}_{S}$ as the width of the model is increased (Fig. 2 (a)). Moreover, we observe that there is a discrepancy in performance between both models after $5000$ steps (Fig. 2 (b)), showing that meta-training with multiple tasks of different widths has benefits for generalization to longer unrolls in addition to improved generalization to larger optimizees. Given the improved generalization of $\mu\text{LO}_{M}$ compared to $\mu\text{LO}_{S}$, we adopt the multiple-width single-task meta-training recipe as part of our method. Subsequent experiments (e.g., figures 1 and 4(c)) will show that it is also effective for meta-training $\mu$VeLO.

Figure: ((a))
Refer to caption: x6.png

## 4 Empirical evaluation

We use a suite of optimization tasks of varying width to evaluate meta-generalization properties of our $\mu$LOs vs tuned $\mu$Adam (Yang et al., ), SP AdamW Loshchilov and Hutter (), and baseline SP LOs. We also include pre-trained VeLO (Metz et al., ) as an oracle which we denote as VeLO-4000. Meta-trained for $4000$ TPUv4 months, it is the strongest publicly available pre-trained learned optimizer. We focus on evaluating generalization to wider networks, however, we also verify whether $\mu$LOs can improve generalization to longer training horizons and deeper networks.

Table: Table 1: Meta-training configurations of LOs and baselines in our empirical evaluation.

Baseline LOs and $\mu$LOs. The meta-training configuration of each learned optimizer is summarized in Table 1. Each learned optimizer (ours and the baselines) in our empirical evaluation is meta-trained using the multiple-width single-task meta-training recipe proposed in section 3.3. The baseline sheds light on whether simply varying the SP optimizee width during meta-training is enough to achieve generalization of the LO to wider networks in SP. During meta-training, we set the inner problem length to be $1000$ iterations. Therefore, any optimization beyond this length is considered out-of-distribution. For all meta-training and hyperparameter tuning details, including ablation experiments, see section C of the appendix.

$\mu$Adam $\mu$Adam is a strong hand-designed $\mu$P baseline. It follows the Yang et al. () Adam $\mu$-parametrization and is tuned on the largest meta-training task seen by our learned optimizers (Table 1). We tune the learning rate and three multipliers: input multiplier, output multiplier, and the hidden learning rate multiplier. These multipliers correspond to adding a tunable constant to the pre-activation multiplier for input weights, the pre-activation multiplier for output weights, and the Adam LR for hidden weights. More details about the grid search over $500$ configurations are provided in Section B.1 of the appendix.

AdamW AdamW (Loshchilov and Hutter, ) is a strong hand-designed SP baseline. It is tuned on the largest meta-training task seen by our learned optimizers (Table 1). We tune the learning rate, $\beta_{1}$,$\beta_{1}$, and the weight decay. More details about the grid search over $500$ configurations are provided in Section B.1 of the appendix.

Pre-trained VeLO (VeLO-4000). VeLO (Metz et al., ) is a learned optimizer that was meta-trained on a curriculum of progressively more expensive meta-training tasks for a total of $4000$ TPU months. These tasks include but are not limited to image classification with MLPs, ViTs, ConvNets, and ResNets; compression with MLP auto-encoders; generative modeling with VAEs; and language modeling with transformers and recurrent neural networks. During meta-training, VeLO-4000 unrolls inner problems for up to 20k steps ($20\times$ ours); the final model was then fine-tuned on tasks with up to 200k steps of optimization. VeLO-4000, therefore, represents the strongest baseline in our empirical evaluation and we consider it to be an oracle.

Is VeLO-4000 a fair baseline?
While we believe the comparison is important given the relevance of our results to scaling up LOs, we highlight that the comparison will unfairly advantage VeLO-4000 as all tasks in our suite fall within its meta-training distribution and VeLO-4000 was meta-trained on inner unroll horizons well beyond those we evaluate. Thus, when comparing our LOs to VeLO-4000, it is important to keep in mind that ours are meta-trained with only $0.004\%$ of VeLO-4000’s compute budget.

Evaluation tasks. Our evaluation suite includes $35$ tasks spanning image classification (CIFAR-10, ImageNet) using MLPs and Vision Transformers (ViTs) (Dosovitskiy et al., ) and autoregressive language modeling with a decoder-only transformer on LM1B (Chelba et al., ). To create the tasks, we further vary image size (for image classification), width, and depth of the optimizee network, and the number of optimization steps.
See Table 7 of the appendix for an extended description of all the tasks.

Figure: ((a))
Refer to caption: x12.png

### 4.1 Results

In the following sections, we first (Sec. 4.1.1) present results empirically verifying the pre-activation stability of our $\mu$LOs. Subsequently, we present the results of our main empirical evaluation of meta-generalization to wider networks (Sec. 4.1.1), a study of $\mu$LOs generalization to deeper networks (Sec. 4.1.3), and a study of $\mu$LOs generalization to longer training horizons (Sec. 4.1.4).

#### 4.1.1 Evaluating pre-activation stability

We now verify that desiderata J.1 of Yang et al. () is satisfied empirically. In Figure 3, we report the evolution of the coordinate-wise standard deviation of the difference between initial (t=0) and current (t) second-layer pre-activations of an MLP during the first $500$ steps of training. We observe that all models parameterized in $\mu$P enjoy stable coordinates across widths, while the pre-activations of the larger models in SP blow up after a number of training steps. Notably, SP Adam’s pre-activations blow up immediately, while $\text{LO}_{S}$ and $\text{LO}_{M}$ take longer to blow up and have a more erratic pattern; we hypothesize that this is a side effect of meta-training where the optimizers may learn to keep pre-activations small by rescaling updates. Section F of the appendix contains similar plots for the remaining layers of the MLP which show similar trends.

In summary, we find, empirically, that pre-activations of $\mu$LOs and $\mu$Adam are similarly stable across widths, while the activations of SP Adam and SP LOs both blow up but behave qualitatively differently.

Table: Table 2: In-distribution and out-of-distribution average performance of optimizers. We report the average rank of different optimizers across the five tasks in our suite. We evaluate in-distribution at a Base width of $1024$ as this is the width used to tune the hand-designed baselines. We also evaluate out-of-distribution widths: Large (2048) and XL (largest size for each task see Tab.7 of the appendix). We bold the strongest, underline the second strongest, and italicize the third strongest average rank in each column. We do not bold entries of VeLO-4000 as it is reported only for reference since it is not a fair comparison. We observe that, across all iterations, when compared to fair baselines, $\mu\text{LO}_{M}$ obtains the best rank for all settings except for the XL task at $5000$ iterations, where it is only bested by $\mu$VeLO.

#### 4.1.2 Meta-generalization to wider networks

Given our goal of improving LO generalization to unseen wider tasks, the bulk of our empirical evaluation is presented in this section. Specifically, we evaluate the behavior of $\mu$LOs as the width of tasks increases well beyond what was seen during meta-training. To accomplish this, we fix the depth of each task and vary the width (see Table 7 for a full list of tasks), leading to a testbed of $32$ different tasks. We then train each task using the baselines and $\mu$-optimizers outlined in section 4 for $5000$ steps for $5$ different random seeds. This involves training $1120$ different neural networks. To make the results easily digestible, we summarize them by width and final performance in Figure 4(c) and by average optimizer rank in Table 2. We also highlight the smooth training dynamics of our optimizers at the largest widths in figure  4(c).

##### Performance measured by final loss as a function of width

Figure 1 compares the training loss after $1000$ steps of SP learned optimizers to $\mu$-parameterized learned optimizers for different widths. This is shown in three subfigures for three MLP image classification tasks: (a) Imagenet $32\times 32\times 3$ (IN32), (b) Imagenet $64\times 64\times 3$ (IN64), and (c) Cifar-10 $32\times 32\times 3$ (C10). Subfigure (a) shows the performance of learned optimizers on larger versions of the meta-training tasks. We observe that the $\mu$LOs achieve lower final training loss as the width of the task is increased. In contrast, $\text{LO}_{M}$ diverges for widths larger than $2048$ and VeLOM fails to substantially decrease the loss at larger widths, falling behind the $\mu$LOs. Subfigure (b) evaluates our $\mu$LOs of larger ImageNet images (e.g., when the input width is larger). Similarly, we observe smooth improvements in the loss as the optimizee width increases for $\mu$LOs, while their SP counterparts either diverge at width $512$ (LOM) or fail to substantially improve the loss beyond width $1024$ (VeLOM). Finally, Subfigure (c) shows the performance of our $\mu$LOs on Cifar-10 (smaller output width) as the width of the model is increased. Similarly, we observe smooth improvements in the loss as the width increases for $\mu$LOs, while their SP counterparts either diverge immediately at small widths (VeLOM) or diverge by width $1024$ (LOM).

##### Performance measured by average optimizer rank

Table 2 reports the average rank of different optimizers on in-distribution width tasks (Base, width 1024) and out-of-distribution width tasks (Large (width 2048) and XL (maximum width)). Each entry of the table corresponds to the optimizer’s average rank (within the $7$ optimizers evaluated) over the $5$ tasks in our suite: Cifar 10 MLP image classification, ImageNet 32 MLP image classification, ImageNet 64 MLP image classification, ImageNet 32 ViT image classification, and LM1B transformer language modelling. The optimizers are ranked by their training loss at the given iteration. We report average ranks for $1000$ iterations (inner-problem length), $3000$ iterations, and $5000$ iterations. We bold the strongest, underline the second strongest, and italicize the third strongest average rank in each column. We do not bold entries of VeLO-4000 as it is reported only for reference since it is not a fair comparison. We observe that, across all iterations, when compared to fair baselines, $\mu\text{LO}_{M}$ obtains the best rank for all settings except for the XL task at $5000$ iterations, where it is only bested by $\mu$VeLO. When only looking at the out-of-distribution Large and XL tasks, we observe that $\mu\text{LO}_{M}$ and $\mu$VeLOMdominate the first two spots of the optimizer podium in all cases except one. For the Large task at $1000$ steps, $\mu$VeLOMis bested by AdamW. When comparing our $\mu$LOs to VeLO-4000, we observe that at least one of $\mu\text{LO}_{M}$ and $\mu$VeLOMbests VeLO-4000 on all tasks except for the large task at $5000$ iterations. This is remarkable as our $\mu$LOs are trained on many orders of magnitude less compute than VeLO-4000. These results demonstrate that meta-training LOs using our recipe yields substantial improvements in meta-generalization (across various tasks and widths) over LOs from previous work and strong hand-designed baselines.

##### Training dynamics at the largest widths

Figure 4(c) reports the training curves of different optimizers on the largest width tasks in our suite. Despite training for $5\times$ longer than the maximum meta-training unroll length, our $\mu$LOs are capable of smoothly decreasing the loss for the largest out-of-distribution tasks in our suite. In contrast, the strong SP LO baselines diverge by $1000$ steps (subfigures (a),(b),(c),(d)), or fail to decrease the training loss (subfigure (e)). Our $\mu$LOs also substantially best the well-tuned hand-designed baselines on LM and ViT tasks (subfigures (d) and (e)) and best or match the best performing hand-designed optimizer in subfigures (a),(b), and (c). When comparing with VeLO-4000, we observe that our $\mu$LOs substantially outperform VeLO in subfigures (a),(b), and $\mu\text{LO}_{M}$ outperforms VeLO-4000 in subfigure (c). In contrast, VeLO-4000 outperforms our $\mu$LOs on transformer language modeling and ViT image classification, the most out-of-distribution tasks for them. These findings show that $\mu$LOs can outperform VeLO-4000 on larger in-distribution tasks, suggesting that scaling meta-training in SP (e.g., as done for VeLO) may not be sufficient to achieve strong meta-generaliztion to the largest tasks, but that meta-training in $\mu$P could be.

In summary, the results in Fig. 1,Tab. 2 and Fig.4(c) demonstrate that our $\mu$LO meta-training recipe represents a considerable advancement to low-cost meta-generalization for learned optimizers. The technique is shown to be a substantial improvement over previous work.

Figure: ((a))
Refer to caption: x18.png

#### 4.1.3 Meta-generalization to deeper networks

In this section, we evaluate LO meta-generalization to deeper networks. Specifically, we increase the number of layers used in MLP, ViT, and LM tasks from $3$ to $16$, while being sure to select models that have widths within the meta-training range ($128-1024$) to avoid confounding the results. Figure 5 reports the performance of our multi-task learned optimizers on deeper networks. We observe that both $\mu\text{LO}_{M}$ and $\mu\text{VeLO}_{M}$ optimize stably throughout and generally outperform their counterparts, $\text{LO}_{M}$ and $\text{VeLO}_{M}$, by the end of training on each task, despite being meta-trained on MLPs of exactly the same depth. Moreover, $\text{LO}_{M}$ immediately diverges when optimizing the deep MLP while $\mu\text{LO}_{M}$ experience no instability. Similarly, $\text{VeLO}_{M}$ diverges on ViTs and Transformers, while $\mu\text{VeLO}_{M}$ performs well, especially on ViTs. This is remarkable as, unlike width, there is no theoretical justification for $\mu$P’s benefit to deeper networks. We hypothesize that $\mu$P’s stabilizing effect on the optimizee’s activations leads to this improvement generalization.

In summary, we find, empirically, that using $\mu$P during meta-training benefits the generalization of learned optimizers, including VeLO, to deeper networks.

Figure: ((a))
Refer to caption: x21.png

#### 4.1.4 Meta-generalization to longer training horizons

In this subsection, we evaluate the capability of $\mu$LOs to generalize to much longer training horizons than those seen during meta-training. Specifically, we use $\mu\text{LO}_{M}$ and $\text{LO}_{M}$ as well as $\mu\text{VeLO}_{M}$ and $\text{VeLO}_{M}$ to train three networks with width $w=1024$: a 3-layer MLP, ViT on $32\times 32\times 3$ ImageNet and a 3-layer Transformer for autoregressive language modeling on LM1B. Each model is trained for $25,000$ steps (25$\times$ the longest unroll seen at meta-training time). Figure 6 reports the training loss averaged over $5$ random seeds. We observe that $\mu\text{LO}_{M}$ and $\mu\text{VeLO}_{M}$ stably decrease training loss over time for each task, while $\text{LO}_{M}$ and $\text{VeLO}_{M}$ fail to decrease training loss (a), decreases it but suffers instabilities (b), or diverges after $8000$ steps (c). These results suggest that generalization to longer training horizons is another benefit of using $\mu$P with learned optimizers.

In summary, we find, empirically, that using $\mu$P during meta-training significantly benefits the generalization of learned optimizers to longer training horizons.

## 5 Limitations

While we have conducted a systematic empirical study and shown strong results within the scope of our study, there are some of limitations of our work. Specifically, (1) we do not meta-train on tasks other than MLPs for image classification and we do not provide evaluation of models wider than $8192$ for MLPs and $3072$/$12288$ (hidden/FFN size) for transformers due to computational constraints in our academic environment.

## 6 Conclusion

We have demonstrated that applying or $\mu$LO meta-training recipe produces optimizers with substantially improved meta-generalization properties when compared to strong baselines from previous work. Remarkably, our $\mu$LOs even surpass VeLO-4000 (meta-trained for $4000$ TPU months) on wider versions of in-distribution tasks. Moreover, our experiments also show that $\mu$LOs meta-trained with our recipe generalize better to wider and deeper out-of-distribution tasks than their SP counterparts. Moreover, when evaluated on much longer training tasks, we observe that $\mu$LOs have a stabilizing effect, enabling meta-generalization to much longer unrolls ($25\times$ maximum meta-training unroll length). All of the aforementioned benefits of $\mu$LOs come at *zero* extra computational cost compared to SP LOs. Our results outline a promising path forward for low-cost meta-training of learned optimizers that can generalize to large unseen tasks.

## Acknowledgements

We acknowledge support from Mila-Samsung Research Grant, FRQNT New Scholar [*E.B.*], the FRQNT Doctoral (B2X) scholarship [*B.T.*], the Canada CIFAR AI Chair Program [*I.R.*], the French Project ANR-21-CE23-0030 ADONIS [*E.O.*], and the Canada Excellence Research Chairs Program in Autonomous AI [*I.R.*]. We also acknowledge resources provided by Compute Canada, Calcul Québec, and Mila.

## Appendix A Proof of proposition 1

###### Proposition 1 .

Assume that the LO $f_{\phi}$ is continuous around 0. Then, if $f_{\phi}(0)\neq 0$, the update, initialization, and pre-activation multiplier above is necessary to obtain a Maximal Update Parametrization.

###### Proof.

The update produced by $f_{\phi}$ is denoted $\Delta W$ and we write $\nabla W$ the corresponding gradient, so that $\Delta W=f_{\phi}(\nabla W)$. For the sake of simplicity, $n$ will be the output size and $d$ the feature input size of our neural network. Our goal is to satisfy the desiderata of [Yang et al., , Appendix J.2]. We assume our initialization follows Initialization-$\mu$ in Sec 3. Overall, our goal is to study strategy so that if $x_{i}=\Theta(1)$, then one needs to renormalize/initialize so that $(Wx)_{i}=\Theta(1)$ while $((W+\Delta W)x)_{i}=\Theta(1)$ so that the update is as large as possible. Note that given the assumptions on $f$, if $x=\Theta(\frac{1}{n})$, then $f(x)=\Theta(1)$.

###### Proposition 1 .

###### Proof.

##### Output weights.

Here, if input $x$ has some $\Theta(1)$ coordinates, we initialize $W=(w_{i})_{i\leq n}$ with weights of variance 1 (which is necessary) and rescale the preactivations with $\frac{1}{n}$. For the update, we thus have that $\nabla W$ scales (coordinate wise) as $\Theta(\frac{1}{n})$ and we do not rescale the LR, given that $f_{\phi}(\nabla W)$ will also have coordinates in $\Theta(1)$.

##### Hidden weights.

Now, for the update, we observe that the gradient $\nabla W$ has some coordinates which scale as $\Theta(\frac{1}{n})$, due to the output renormalization choice. Thus, the LO $f_{\phi}(\nabla W)$ satisfies that $f(\nabla W)=\Theta(1)$, given that $f_{\phi}$ is continuous in 0 and satisfies $f_{\phi}(0)\neq 0$. Thus for the update, we need to use $\Delta W=\frac{1}{n}f_{\phi}(\nabla W)$ so that $\Delta Wx$ is coordinate wise bounded.

##### Input weights.

In this case, the gradient has coordinates which already scale in $\Theta(\frac{1}{n})$ (due to the output renormalization) and there is no need to rescale the LR.
∎

## Appendix B Hand Designed Optimizer Hyperparameter Tuning

### B.1 Tuning μ 𝜇 \mu italic_μ Adam

We tune the $\mu$Adam baseline on the largest meta-training seen by our learned optimizers.$\mu$AdamM was, therefore, tuned using a $1024$ width MLP for $32\times 32\times 3$ ImageNet classification. As mentioned in section 4, we tune the learning rate and three multipliers: input multiplier, output multiplier, and the hidden learning rate multiplier. These multipliers correspond to adding a tunable constant to the pre-activation multiplier for input weights, the pre-activation multiplier for output weights, and the Adam LR for hidden weights (e.g., in Table 8 of Yang et al. []). Specifically, we search for the learning rate in $\{0.1,0.01,0.001,0.0001\}$ and for each multiplier in $\{2^{-4},2^{-2},1,2^{2},2^{4}\}$. This results in a grid search of $500$ configurations, whose optimal values are reported in table 3.

**Table 3: Best hyperparameters values for $\mu$Adam baseline. $\mu$Adam is tuned to optimize 3-layer W$=1024$ MLP for $32\times 32\times 3$ ImageNet classification.**
| Optimizer | LR | Input Multiplier | Output Multiplier | Hidden LR Multiplier |
| --- | --- | --- | --- | --- |
| $\mu$AdamM | $0.1$ | $0.25$ | $0.25$ | $4$ |

### B.2 Tuning AdamW

We tune the AdamW baseline on the largest meta-training seen by our learned optimizers. AdamW was, therefore, tuned using a $1024$ width MLP for $32\times 32\times 3$ ImageNet classification. As mentioned in section 4, we tune the learning rate, betas, and weight decay: LR, $\beta_{1}$, $\beta_{2}$, and the weights decay. Specifically, we search over the values of each hyperparameter reported in Table 4. This results in a grid search of $500$ configurations, whose optimal values are reported in table 5.

Table: Table 4: Optimal Hyperparameters Found AdamW. Similar to the $\mu$Adam baseline, we tune all optimizers on a 3-layer W$=1024$ MLP ImageNet classification task and use a budget of approximately $500$ total runs.

Table: Table 5: Optimal Hyperparameters Found AdamW. Similar to the $\mu$Adam baseline, we tune all optimizers on a 3-layer W$=1024$ MLP ImageNet classification task and use a budget of approximately $500$ total runs.

## Appendix C Meta-training with μ 𝜇 \mu italic_μ LOs

##### General meta-training setup for small_fc_lopt

Each small_fc_lopt [Metz et al., ] learned optimizer is meta-trained for $5000$ steps of gradient descent using AdamW [Loshchilov and Hutter, ] and a linear warmup and cosine annealing schedule. We using PES [Vicol et al., ] to estimate meta-gradients with a truncation length of $50$ steps and sampling $8$ perturbations per task at each step with standard deviation $0.01$. For the inner optimization task, we used a maximum unroll length of $1000$ iterations; that is, all our learned optimizers see at most $1000$ steps of the inner optimization problem during meta-training. Unlike with $\mu$Adam, we do not tune the $\mu$P multipliers when meta-training $\mu\text{LO}_{S}$ and $\mu\text{LO}_{M}$, instead, we set the all to $1$. All optimizers are meta-trained on a single A6000 GPU. $\mu\text{LO}_{S}$ and $\text{LO}_{S}$ take $8$ hours each to meta-train, while $\mu\text{LO}_{M}$ and $\text{LO}_{M}$ take $103$ hours.

##### General meta-training setup for VeLO

Each VeLO [Metz et al., ] learned optimizer is meta-trained for $45000$ steps of gradient descent using AdamW [Loshchilov and Hutter, ] and a linear warmup and cosine annealing schedule. We using PES [Vicol et al., ] to estimate meta-gradients with a truncation length of $20$ steps and sampling $8$ perturbations per task at each step with standard deviation $0.01$. For the inner optimization task, we used a maximum unroll length of $1000$ iterations; that is, all our learned optimizers see at most $1000$ steps of the inner optimization problem during meta-training. Unlike with $\mu$Adam, we do not tune the $\mu$P multipliers when meta-training $\mu\text{LO}_{S}$ and $\mu\text{LO}_{M}$, instead, we set the all to $1$. All optimizers are meta-trained on a single A6000 GPU. $\mu$VeLOMand VeLOM take $250$ hours to meta-train.

##### Meta-training hyperparameters for small_fc_lopt in μ 𝜇 \mu italic_μ P

While there are very few differences between $\mu$LOs and SP LOs, the effective step size for hidden layers is changed (see eq. 3) which could alter the optimal meta-training hyperparameters. Consequently, we conduct an ablation study on hyper-parameters choices for $\mu\text{LO}_{S}$. Specifically, using AdamW and gradient clipping with a linear warmup and cosine annealing LR schedule, we meta-train $\mu\text{LO}_{S}$ to train 3-layer width 128 MLPs to classify $32\times 32\times 3$ ImageNet Images. By default, we warmup linearly for 100 steps to a maximum learning rate of $3e-3$ and anneal the learning rate for $4900$ steps to a value of $1e-3$ with $\lambda_{1}=0.001$ (from equation 3) and sampling $8$ perturbations per step in PESVicol et al. []. The above ablation varies the maximum learning rate $\in\{1e-2,3e-3,1e-3\}$ (always using $100$ steps of warmup and decaying to $0.3\times$MaxLR), $\lambda_{1}\in\{0.001,0.01,0.1\}$, the number of steps (5k or 10k), and the number of perturbations (8 or 16). We observe that using all default values except for $\lambda_{1}=0.01$ yields one of the best solutions while being fast to train and stable during meta-training. We, therefore, select these hyperparameters to meta-train $\mu\text{LO}_{S}$ and $\mu\text{LO}_{M}$.

##### Meta-training hyperparameters for VeLO in μ 𝜇 \mu italic_μ P

Unlike small_fc_lopt, we do not find it necessary to $\lambda_{1}$ from its default value. However, we do remove the multiplication by the current parameter norm used in the update equation to VeLO as it causes meta-training problems when initializing tensors to zero.

##### μ 𝜇 \mu italic_μ P at Meta-training time

While we use the same $\mu$P at meta-training and testing time, it is important to consider meta-training tasks that have similar training trajectories to their infinite width counterparts. In [Yang et al., ], authors provide discussions of these points for zero-shot hyperparameter transfer. Two notable guidelines are to initialize the output weight matrix to zero (as it will approach zero in the limit) and to use a relatively large key size when meta-training transformers. For all our tasks, we initialize the network’s final layer to zeros. While we do not meta-train on transformers, we suspect that the aforementioned transformer-specific guidelines may be useful.

Figure: Figure 7: Ablating Meta-training Hyperparameter for $\mu\text{LO}_{S}$. Using AdamW with a linear warmup and cosine annealing schedule, we meta-train $\mu\text{LO}_{S}$ to train 3-layer width 128 MLPs for classifying $32\times 32\times 3$ ImageNet Images. By default, we warmup linearly for 100 steps to a maximum learning rate of $3e-3$ and anneal the learning rate for $4900$ steps to a value of $1e-3$ with $\lambda_{1}=0.001$ (from equation 3) and sampling $8$ perturbations per step in PESVicol et al. []. The above ablation varies the maximum learning rate $\in\{1e-2,3e-3,1e-3\}$ (always using $100$ steps of warmup and decaying to $0.3\times$MaxLR), $\lambda_{1}\in\{0.001,0.01,0.1\}$, the number of steps (5k or 10k), and the number of perturbations (8 or 16). We observe that using all default values except for $\lambda_{1}=0.01$ yields one of the best solutions while being fast to train and stable during meta-training.
Refer to caption: x24.png

## Appendix D Features of the learned optimizer

**Table 6: Per-parameter features used as input to our learned optimizers. All the coefficients, $\beta_{i}$, are learnable parameters adjusted during meta-optimization. We replicate the table of [Joseph et al., ] for convenience.**
| Description | value |
| --- | --- |
| parameter value | $w_{t}$ |
| 3 momentum values with coefficients $\beta_{1},\beta_{2},\beta_{3}$ | $m_{t,i}=\beta_{i}m_{t-1,i}+(1-\beta_{i})g_{t}$ |
| second moment value computed from $g_{t}$ with decay $\beta_{4}$ | $v_{t}=\beta_{4}v_{t-1}+(1-\beta_{4})g_{t}^{2}$ |
| 3 values consisting of the three momentum values normalized by the square root of the second moment | $\frac{m_{t,i}}{\sqrt{v}}$ |
| the reciprocal square root of the second moment value | $\frac{1}{\sqrt{v}}$ |
| 3 $\Delta_{t}$ Adafactor normalized values | $g_{t}\times\textsc{row factor}\times\textsc{column factor}$ |
| 3 tiled Adafactor row features with coefficients $\beta_{5},\beta_{6},\beta_{7}$, computed from $g_{t}$ | $r_{t,i}=\beta_{i}r_{t-1,i}+(1-\beta_{i})\textsc{row_mean}(\Delta_{t}^{2})$ |
| 3 tiled Adafactor column feature with coefficients $\beta_{5},\beta_{6},\beta_{7}$ computed from $g_{t}$ | $c_{t,i}=\beta_{i}c_{t-1,i}+(1-\beta_{i})\textsc{col_mean}(\Delta_{t}^{2})$ |
| the reciprocal square root of the previous 6 features | $\frac{1}{\sqrt{r_{t,i}\textsc{ or }c_{t,i}}}$ |
| 3 $m$ Adafactor normalized values | $m_{t,i}\times\textsc{row factor}\times\textsc{column factor}$ |

## Appendix E List of Meta-testing Tasks

Table 7 reports the configuration of different testing tasks used to evaluate our optimizers. We note that we do not augment the ImageNet datasets we use in any way except for normalizing the images. We tokenize LM1B using a sentence piece tokenizer[Kudo and Richardson, ] with 32k vocabulary size. All evaluation tasks are run on A6000 $48$BG or A100 $80$GB GPUs for $5$ random seeds.

Table: Table 7: Meta-testing settings. We report the optimization tasks we will use to evaluate the LOs of Table 1.

## Appendix F Coordinate evolution of MLP layers in μ 𝜇 \mu italic_μ P for different optimizers

The following section presents the continuation of our experiments comparing pre-activation growth during training for SP LOs and $\mu$LOs with different meta-trainnig recipes, SP adam, and $\mu$Adam.

Figure: ((a))
Refer to caption: x25.png

Figure: ((a))
Refer to caption: x31.png

Figure: ((a))
Refer to caption: x37.png

Figure: ((a))
Refer to caption: x43.png

## Peer Reviews

confidence: 4
contribution: 3
initial_rating: 5
presentation: 2
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?
soundness: 2
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.
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.
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.

---

confidence: 3
contribution: 3
initial_rating: 6
presentation: 2
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?
soundness: 3
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.
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.
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.

---

confidence: 3
contribution: 3
initial_rating: 6
presentation: 3
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.
soundness: 3
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.
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.
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.

---

confidence: 4
contribution: 2
initial_rating: 5
presentation: 3
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.
soundness: 3
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 …).
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.
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

Meta-review: 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.

Decision: Reject