Celo: Training Versatile Learned Optimizers on a Compute Diet
Authors: Abhinav Moudgil
Boris Knyazev
Guillaume Lajoie
Eugene Belilovsky
Venue: Accepted by TMLR ()

Abstract: Learned optimization has emerged as a promising alternative to hand-crafted optimizers, with the potential to discover stronger learned update rules that enable faster, hyperparameter-free training of neural networks. A critical element for practically useful learned optimizers, that can be used off-the-shelf after meta-training, is strong meta-generalization: the ability to apply the optimizers to new tasks. Recent state-of-the-art work in learned optimizers, VeLO (Metz et al., 2022), requires a large number of highly diverse meta-training tasks along with massive computational resources, 4000 TPU months, to achieve meta-generalization. This makes further improvements to such learned optimizers impractical. In this work, we identify several key elements in learned optimizer architectures and meta-training procedures that can lead to strong meta-generalization. We also propose evaluation metrics to reliably assess quantitative performance of an optimizer at scale on a set of evaluation tasks. Our proposed approach, Celo, makes a significant leap in improving the meta-generalization performance of learned optimizers and also outperforms tuned state-of-the-art optimizers on a diverse set of out-of-distribution tasks, despite being meta-trained for just 24 GPU hours.

Published in Transactions on Machine Learning Research (06/2025)
Celo: Training Versatile Learned Optimizers on a Compute Diet
Abhinav Moudgil*1.2 Boris Knyazev Guillaume Lajoie Eugene Belilovskyl.2
Mila Quebec AT Institute, "Concordia University, "Somsung- SATT AT Lab, Montreal, Université de Montréal
Reviewed on OpenReview: http:///ooorwweee..e/////o//oott
Abstract
Learned optimization has emerged as a promising alternative to hand-crafted optimiz- ers, with the potential to discover stronger learned update rules that enable faster, hyperparameter-free training of neural networks. A critical element for practically use- ful learned optimizers, that can be used off-tho-shelf after meta-training, is strong meta- generalization: the ability to apply the optimizers to new tasks. Recent state-of-the-art work in learned optimizers, VeI.O (Metz et al., (022b)), requires a large number of highly diverse meta- testa-trinnig tasks along with massive computational resources, 4000 TPU months, to achieve netta-eneeaalizttinn. This makes further improvements such learned optimizers impractical. In this work, we identify several key elements in learned optimizer architec- tures and meta-training procedures that can lead to strong meta-generalliation. We also propose evaluation metrics to reliably assess quantitative performance of an optimizer at scale on a set of evaluation tasks. Our proposed approach, Celd] makes a significant leap in improving the geteaaeenrllizationn performance of learned optimizers and also perporfoms tuned stato-of-tho-art optimizers on a diverse set of cut-o--distribution tasks, despite being mela-trained for just 24 GPU hours.
1 Introduction
Deep learning has advanced remarkably over the past decade by significant improvements architectures and training techniques, allowing us to train large neural networks on big datasets (Krizhevsky et al.
2012
Achiam et al. 2023 Brooks et al. 2024 Kirillov et al. 2023 Jumper et al. 2021). Much of this
progress can be attributed to advancements in optámizers (Sutskever et al. 2013
Kingma & Ba 2015 Gupta et al. 2018 Rajbhandari et al. 2020) updating the neural networks (optimizee)
Loshchilov 2017
over the course of training. In deep learning, the most popular optimizers are adaptive variants of stochastic
gradient descent such as Adam (Kingma & Ba 2015
Loshchilov 2017), AdaGrad (Duchi et al.) 2011
Adafactor (Shazoer & Storn 2018). etc. These variants, although supported by some thcoretical guaranteos.
are often heuristically developed and lested on large-scale neural network tasks (Goodfellow al. 2016
Schmidt et al.) 2021 Bottou et al., 2018 Reddi et al.) 2019). Tn addition to being hand-designed (thus being potentially suboptimal), these optimizers have hyperparameters that need to be tuned for every task, which is quite compute intensive (Schmidt et al. 2021 Sivaprasad et al. 2020). Additionally, recent works also use schedules (Goyal 2017: You et al.) 2017 Smith 2017 Loshchilov & Hutter 2016) which adjust
optimizer hyperparameters dynamically throughout training. These schedules introduce additional tunable
hyperparameters, making the whole training process even more computationally expensive (Metz et al.,
2020b Dahl et al. 2023).
Deep learning advances allowed us to reduce feature engineering efforts and to greatly improve performance by replacing hand-crafted features (Lowe] 2004) with learned ones. Inspired by the advent of deep lcarning, learned optimizers present a promising avenue to make the optimization process hyperparameter free
* *Correspondence Abhinav Moudgil <abhinav. beaahiinav.noddgieeeiee... .
Code is available htttps://githb. com/ anoudg//celo

---

Published in Transactions on Machine Learning Research (06/2025)
and also more performant (ndrychhoicz et al. 2016 Wichrowska et al., 2017 Almeida al. 2021 Metz et al. 2022a b). Formally, given an optimizee neural network with n parameters 0, E IR a iteration t, an optimizer for parameterized by & (e.g. hyperparameters), takes as input the gradient vO,, current state St (c.g. momentum) and returns updated parameters 0t+1 along with updated state $t+1: (1+1,8141) = [o(O,, 8t). While in a hand-designed optimizer an updale rule and Φ are predefined by the practitioner, a lcarned optimizer fs learns the entire update rule parameterized by Φ through a meta-training process. Learning the update rule allows learned optimizers to leverage additional meta-data Mr, such as training loss (Metz et al.] 2020a] 2022b), training progress (Almeida et al. 2021 Metz et al., 2022a b neural network graph (Peebles et al. 2022 Kofinas et al. 2024]) and other relevant information.
Thus, f can learn a more powerful update rule formalized as the following:
(0++1,8111) = (0, VO,, 8t);
(1)
1
where th represents all the "features" of the optimizce network at itcration t. Hence, learned optimizers can automatically change the update function or implement schedules by monitoring the input meta-dala over the course of training, thus removing the expensive manual tuning of optimizer hyperparameters.
Despite the potential of learned optimizers, one of their major challenges is meta-genermlization: learned optimizers once mela-trained on a distribution of ion lasks should be able to generalize to unseen or putt-f-diisrrbuttoo tasks. Prior work (Metz et al.] 2020a 2022b) has shown that akin to scaling laws in deep learning (Kallan et al. 2020 Achiam et al. 2023). the metaa-enneralization of these learned optimizers improves as meta-training is scaled up. However, the meta-training of these learned optimizers as in prior work, VeT.O (Metz et al.; 2022b), is prohibitively expensive (4000 TPU months) which poses a serious challenge in further improvements of these learned optimizers and highlights the need for more compute- efficient approaches (Rezk et al., 2023 Metz et al. 2022b). Therefore, in this work, keep small fixed mela-training set and instead identily several elements in the design and meta-training of these learned optimizers which could lead to strong meta-generalization while using only fraction of VeLO's compute budget (24 GPU hours).
Another challenge in learned optimization is quonlilative evaluation, which can be used reliable measure to track progress in this field. Most prior work in learned optimizers (Schneider et al. 2019 Schmidt et al. 2021 Metz et al. 2022b a) focus on loss curves or final performance of the trained models on selected tasks as the primary mode of evaluation. However, it remains unclear how to obtaín a reliable aggregated performance of an optimizer across a diverse, large set of tasks, each with potentially different evaluation metrics and loss scales, Furthermore, learned optimizers often also exhibit high variance and occasional instability across tasks (Andrychowicz al.) 2016 Metz et al.] 2022a b) which can also significantly bias the evaluation results due to outliers. In this work, we seek inspiration from Reinforcement Learning (Agarwal et al. 2021,, , which also faces similar evaluation challenges and proposes reliable evaluation metries which we employ to quantitatively evaluate learned optimizers. Our contributions in this work are the following:
1. We propose simple recipe to train meta-generalizabble (versatile) learned optimizers on limited com- pute budget; our recipe's main ingredients are task augmentation, simple design consisting of per-param Icarned update rulc with a high-level scheduler and decoupled training of the update rule and schoduler.
2. We provide reliable evaluation metrics to track aggregate performance of optimizers on a set of evaluation tasks which are robust to outliers.
3. Our proposed approach, Celo, outperforms 15 hand-designed optimizers, including Adam (Kingma & Ba 2015 and Shampoo (Anil et all 2020), as well as
learned optimizers such VeLO (Metz)
el al. 2022b)
a diverse set of out-of-dlistribution tasks.
4. We conduct an exhaustive ablation study over each component in our approach with mela-generallizttion the principal focus.
5. We analyze schedules predicted by Celo and find that they are adaptive to unseen tasks and training horizons, further highlighting the meta-generalizationn capabilities of Celo.
2

---

Published in Transactions on Machine Learning Research (06/2025)
2 Background
Learned MLP update rule. Several learned optimizer architectures have been proposed in previous works (Andyychowicz et al. 2016; Metz et al. 2020a] [Almidda et al.) 2021 Metz et al. 2022a b), each varying in the types of inputs they digest, the elements they maintain within their state, and the functional form of their outputs. In this section, briefly review learned optimizer parameterized using a multi-layer perceptron (MLP) (Metz et al. 2022a) (Adafac MLP LOpt) that our work and recent statc-of-tho-art works like VeLO MMet et a (022b) also build upon. It is a simple MLP-based learned oplimizer that can serve a drop-in replacement for hand-designed update rule such as SGD or Adam. Concretely, at a given iteration i, Adafac MLP LOpt takes as input paramcter vector 6,, gradient V0, along with current itcration t. These values are then used to update momentum accumulator buffers in state st. These accumulated values are then used to build input features F for each parameter which are passed to the learned ML. rule along with the global training progress features w (see (Metz al.) 2022aa) for all the input features). The learned MLP returns two outputs of the same dimensionality 0,, corresponding to direction d and magnitude m. which are used to get parameter updates A0, and then parameters 0.+1 as follows:
A0: >jjd c(amm)
(2)
0.+1 02 -00
where Λ and X2 are fixed scalars set low values (0.001) in order to keep meta-training stable. VeLO maintains multiple such learned MLPs (of the same architecture) in its bank which are combined to give a weight update at each step. In addition to learned MLPs that are applied to each parameter (per-parameter MLP), VeLO has per-tensor LSTM network applied to each tensor (e.g. layer weight bias). The per- tensor LSTM takes global features such as training progress, loss and per-tensor features such as variance of momentum, mean of gradient RMS, etc. Furthermore, VeT.O scales per-parameter updates with a tensor norm of the corresponding layer. VeLO's more sophisticated hierarchical design makes it more performant than Adafac MLP LOpt while keeping it efficient at run-time due to its hierarchical design. Specifically, the hierarchical design allows adding capacity at higher levels, such as the layer (or tensor) and global levels, whose added computational costs scale sublinearly with the number of parameters. In this work, also propose a hierarchical architecture similar to VeLO but is much simpler in design and meta-generalizes significantly better than VeLO and Adafac MLP LOpt.
Meta-training. Unlike hand-designed optimizers, learned optimizers' update function is parametrized by & learned through a meta-trainie training process. A standard approach to meta-training involves solving a bi- level optimization problem that involves an inner problem which optimizes network parameters O using the learned optimizer update & on sampled task and an outer problem which optimizes the learned optimizer parameters & based on the feedback from the inner loop (Andrychowicz et al. 2016 Wichrowska el al. 2017 Metz et al.) 2019a 2020a] Vicol et al.) (021). Formally, given a set of optimization tasks T, the learned optimizer parameters & are obtained by sampling an optimization task which consists of data distribution D, initial network parameters On, and a training objective L and solving the bi-level problem below:
1 T-1
* = arg min Bin.
Σ
((x::01
(3)
T
where the inner loop is recursively defined for t [0.TT-- -0 1] as:
(0,, 8t) = (6o.0)
if = 0,
(4)
(0,, 81) fo(0--1. V0,1,
if O:
(5)
V0,- dC oC(Xx: or, ④)
X, D
A1.
(6)
20:
Here denotes the unroll length in the inner loop and X denotes sampled data from D. The outer training objective or meta-objective is based on the mean loss as formalized above or, less common, the final loss of the inner loop (Metz et al. 2019a 2020al 2022b). We compute meta-gradients for & in eq. 3 using Persistent
3

---

Published in Transactions on Machine Learning Research (06/2025)
Evolution Strategies (PES) (Vicol et al 2021 used in prior work (Metz et al.
2022a
Harrison et al. 2022
Gärtner et al. 2023) since directly back-propagatingg through the inner loop in eq. 3 can lead to noisy gradient estimates, especially when the unroll length T' is large (Metz et al. 2019a Vicol et al. 2021). Moreover, PES also gives unbiased gradient estimates the truncated unroll setting in which the learned optimizer is updated in truncations every lew steps in the inner loop instead of updating only once after the full unroll of the inner loop. This allows frequent updates to the optimizer during meta-training, thus making it less
computationally expensive than full unrolls (Vicol et al. 2021 Metz al. 2022bb.
3 Celo: Compute-efficient learned optimizer
Algorithm
Celo update
Input:
0,
optimizee parameters
V0,
gradients
L
current loss
St
optimizer state
Require: learned scheduler fistm, learned update rule fmip-
accumulators update function Sace, fixed scalars a, Å1, X2 E IR
8t+1 - (s,, V0,, Lt,
update accumulators in state
compute progress features wP, loss features wf using updated 8
xx + concat(w" w²)
D scheduler input
4: < fism((")
learned scheduler forward pass
5: ac
D scheduler output
6: initialize next params 0c+1 - ()
7: for each tensor with params P: E 0, in parallel do D parallel scan over all layer weights biases
prepare per-param features F for P: using updated state 8t+1
d, Smip(F) (F)
learned update rule forward pass compute per-param updates >> scale by scheduler output t updated params >> gather params
Ap: \jd e(dam)) prl2
10:
11:
Ap: Ap:
12:
Pi+1 Pt Apt
13: 0.+1= < 0c+1 U PE+1
return 8141
In this work, we treat mela-generalliztiion as the principal component for designing and training learned opti- mizers. We show that proposed approach, Celo (short for "Compute-efficient learned optimizer")), incor- porating recipe of key algorithmic meta-training and architectural improvements yiclds strong generaliza- tion to unseen tasks. Our meta-tran training recipe described below can be followed on small-scale compute budget (more details in 5 and experiments in $ show that each element. is critical for strong metaagenerallizttion:
Task augmentation. Building on the success of data augmentation strategies in deep learning (Krizhevsky et al. 2012 Cubuk et al. 2020 Chen al. 2020b Laskin et al.) 2020), we apply task augmentation during meta- training enhance the robustness of learned optimizers (Metz et al. 2022b). This approach emulates lcarning dynamics across a potentially large number of tasks using a limited set of meta-training tasks (Metz) el al. 2022b). Task augmentation is done by re-scaling parameters (re-paramelrixingg in the inner loop. Recall that each meta-training iteration consists of applying an inner training loop on a randomly initialized, optimizee network with parameters, Og. A random scalar parameter is sampled each meta-training iteration by which the optimizee network parameters each inner training step are scaled by 7 before the forward pass as 70, rrepparametrization) which yields inner task loss L. At initialization, the initial weights On are also scaled by factor of 1/7 in order to keep the underlying optimizee function output same. Modified Eq. 3 affter task augmentation is given below with task augmentation specific changes highlighted in red:
1
C c(x;; (XXTT0p 0,, 01. 0))
arg min EX
(7)
T
4

---

Published in Transactions on Machine Learning Research (06/2025)
Outer metartraining loop
Celo
Inner training loop
learned update rule
0++1
1/1
updated params
every steps
Σ
7|1
task sampler
7
L Joss
update
learned schecloler
learned optimizer
optimizee
2
0
0
re parumetrize
optimizee features
Figure 1: Celo meta-training and architecture.
where in the inner loop for t= = [0, TO,T"- 1]:
(0,, 8t) (00/T, 0)
if 1 0,
(8)
(0,, 8() fo(01-1, V0.-1, M1-1,
if t 0:
(9)
V0. OL Tx,
X1 ~
11.
(10)
00,
Overall, this task augmentation does not change the underlying optimizee network output at initialization, but it does alfect the learning dynamics of the optimizee network in the inner loop when is updated with adaptive or Icarned optimizors, thus of- fectively "simulating" more lasks from a fixed amount of given meta-training tasks (see Fig. 2 for a representative example).
Varying task augmentation param
Tmii bat
Simple learned optimizer design: learned scheduler with learned update rule. Currently, the most performant and dominant paradigm in the optimization of modern deep neural networks is to use scheduler coupled with a hand-designed op- timizer [Goyal 2017; You et al., 2017; Smith 2017 Loshchilov schoduler changes globally the learning rate of the underlying hand-designed optimizer as training pro- gresses. Learning rate schedules have a long history in non- adaptive stochastic gradient descent methods to address the high variance of the gradient estimator close to the optimum (Boyd & Vandenberghe 2004). Deep Learning practitioners have shown that sophisticated schedules such as warmn-us ups and cyclic learning rate (Smith 2017 Loshchilov & Hutter 2016: Touvron et al., 2023 help significantly in the optimization of deep neural networks and empirically, it has been shown that schedulers are beneficial even in the case of adaptive optimizers such as Adam [Loshchilov] 2017 TTouvron et al.) 2023). Inspired by this, focus on separating the global learning rate scheduler from the per-parameter update rule.
500
1000
1500
2000
& Hutter 2016
Iorations
0-0.00
0-1
10.0
-
1.0
100.0
Figure 2: Task augmentation. Learn- ing curves for Adam with fixed (Ic- 4) learning rate with different augmen- tation parameters (r) for CIFAR-10 ConvNet task.
Several designs for learned optimizers have been proposed prior work which extend this notion of schedulers beyond a scalar lcarning rate (Andrychowicz et al. 2016 Wichrowska et al.] 2017 Almeida et al. 2021 Peebles al.) 2022 Metz al.) (2022b), see Fig. g in Appendix. However, in this work, we streamline the learned optimizer design and show that learned recurrent scheduler (LSTM in case) which outputs
5

---

Published in Transactions on Machine Learning Research (06/2025)
scalar schedule parameter coupled with learned per-parameter MLP update rule generalizes well to large- number of unseen tasks, when even meta-trained on limited number of small-scale tasks. Our learned LSTM-based scheduler takes global loss and progress features (Metz et al.
2022b a as input and outputs a positive scalar d/ which is obtained by lincarly projecting the LSTM output and raising it to the exponential
We show that this exponential form in scheduler is crucial for performance in S The final parameter updates are obtained by scaling the per-parameter MLP (Metz et al. 2022a b) outputs with this scalar d) predicted
by the scheduler. The pseudocode for forward pass of our optimizer expressed concretely in Algorithm [
Decoupled training of scheduler and up- date rule. We propose to view optimiza- tion problem in two parts: how to control step- size and what rule to use for parameter updates.
Stage I: learn MI. update rule
Stage 2: learn scheduler with frozen update rule
hand-designed optimizers such as Adam (Kingma & Ba 2015) and Sham- poo (Gupta et al. 2018) propose a hand-designed update rule, leaving the task of controlling or modulating global learning rate and other hyper- parameters up to the practitioners. Following this view, propose to mela-train learned optimiz- ers in a two-stage process: (1) Learn a strong pa- rameter update rule (2) Freeze the learned update rule and learn a step-size scheduler. Our experi- ments show that two-stage training crucial for meta-generalization and avoids meta-overiitiing when training on compute budget with fixed amount of meta-training tasks, hence outperform- ing single-stage training. Moreover, our simple two-stage design allows us to contrast and com- pare with state-of-the-art hand-designed optimiz- ers. Our experiments show that our learned update rule outperforms hand-designed update rules like Adam and when combined with a learned scheduler, the mota-generalization performance further improves.
Current state-of-the-an
perparmm leatures
learned update rule
learned update rule
0++1 updated params
Oc+1 updared parums
state
learned schechiler
V/ optimizee features
loss progress features
frozen
0,
optimizee params at ineration f
V/
Iparams, grads, loss, progresil
V/
optimizee features
1/1
schedule param (scalar)
Figure 3: Two-stage training of Celo.
4 Reliably scoring optimizers at scale
Prior work in learned optimizers often resorts to showing training or validation curves as the primary means of evaluation (Schneider et al. 2019 Schmidt et al.) 2021 Metz et al. 2022bla). This is a common practice, because ovaluating the performance of learned optimizers quantitatively is challenging. For example, different evaluation tasks can have different ranges and semantics of loss values. Moreover, learned optimizers also often show unstable behaviour and evaluating them multiple times to reduce the variance can be infeasible, since each evaluation run involves training a neural network. Therefore, metrics such as mean median (of loss) can be noisy. Alternatively, 2023) However, this approach requires carefully designing target scores for each task and meaningfully aggregating speedups, which becomes cumbersome for a large and diverse set of potentially changing tasks. Morcover, if the targets are set too high, the optimizers that fail to meet these targets reccive a zero score, disregarding improvements between the optimizers. Conversely, setting targets too low introduces a large bias from outliers, compromising the evaluation's 's reliability.
aggregate specdup can be reported, as done in prior work (Dahl et al.,
Hence, in this work, we address the following problem: oplimizer al scale without hand-designing any Uhresholds? We find that this evaluation problem is analogous to the one in Reinforcement Learning (RL) in which an RL agent is evaluated on several tasks, e.g. Atari games. Agarwal et al. (2021) ) propose a widely adopted solution to this evaluation problem (Schwarzer et al., 2023 Ceron al.) Micheli et al. (022). Specifically, they show that the inter- quartile mean (IQM) accurately aggregates results with only few evaluation runs and is less sensitive to outliers than other metrics. For
How do we reliably compute aggregated score for each

---

Published in Transactions on Machine Learning Research (06/2025)
evaluating optimizers, we focus two aspects: final loss performance and speedup with respect to tuned baseline.
Final loss IQM. Given a fixed set of M tasks and N trials for cach task, IQM Agarwal et al. 2021 expects M x N normalized score matrix S for each algorithm. In the RL setup, Agarwal al. (2021) compute scalar normalized score 8 S for the m-th game and n-th trial by scaling the game score in that trial with respect to the human score on the corresponding game. This implics that, in any trial, a normalized score 8m.n >1 1 represents "Super-Human" performance. Equivalently, in optimization, we treat the best Adam run from a fixed number of trials as the "human baseline". The normalized score for each run sm,n is computed as the ratio between the final loss of the best Adam run Lydam after optimizing for T steps, and the final loss obtained by the optimizer being evaluated (LTT))
L Ladam
(11)
8m,,7
opt
where 8m,n > 1 represents "Super-Adam" performance. The M x N normalized scores are then used to compute the aggregate IQM metric. Specifically, IQM is computed by discarding top & bottom 25% scores and calculating the (trimmed) mean of the remaining 50% of the normalized scores.
Similar to the setup described above, we first compute normalized speedup M x N normalized score matrix S which is then used to compute the aggregate speedup IQM metrie. We compute normalized scores as follows: For a given task, we first pick the best Adam baseline out of K trials which achieves the lowest. loss at the end of training, which consists of say, T steps. Next, we check the number of steps for the optimizer being evaluated to achieve the same loss, say, Topt and compute normalized speedup score m as below:
Speedup IQM.
T
8 mm,
(12)
Topt
If an optimizer is unable reach the final loss achieved by the best tuned Adam baseline in any run, we set its corresponding speedup score to zero assuming Topt
5
Experiments
5.1 Baselines
We compare our proposed optimizer, Celo, with 15 state-of-the-arl hand-crafted optimizers and 4 learned optimizers. Among the learned optimizers, we compare with: (1) NNAdam LOpt (Metz et al. 2020b), a hybrid-optimizer using Adam as the update rulc whose hyperparameters are controlled by a learned LSTM- based module; (2) AdaFac MLP LOpt (Metz et al. 2022a) described in Section 2 (3) RNN MLP LOpt (Metz) et al. 2020a) using an T.STM-based per-tensor network that communicates with a learned ML. update rule through embeddings;
(4) VeLO (Metz et al., (022b) with an LSTM- per-tensor network that generates weights of an MLP update rule through hyper-networks Celo and other RNN-based optimizers use a hidden size of 64 in their recurrent networks. VeI.O's 'ss default hidden size is 512, therefore to focus exclusively on differences in architecture, as a baseline we use "VeLO-S" with hidden size reduced to 64 and accordingly the number of MLPs in its bank reduced from 256 to 32. This scaling down does not severely impact meta- generalization (see Appendix A.2). Among the hand-crafted optimizers, we compare with Adam (Kingma & Ba.
2015), Shampoo (Anil et al. 2020 Gupta et al., 2018) with SGD and Adagrad grafting (default block
size
128), Nesterov accelerated AdamW (NAdamW) (Dozat 2016 Loshchilov 2017), RAdam (Liu et al.
2019
Yogi (Zaheer et al.) 2018). AdaBelief (Zhung et al.
2020 LAMB (You et al. 2019). LARS (You
et al.
2017), RMSProp (Tieleman & Hinton 2012), Adafactor (Shazeer & Stern 2018). AdaGrad (Duchi
ct al.
2011). SM3 (Anil et 2019), SGD with momentum (Sutskever et al. 2013)) and Fromage (Bernstein
eel al.
2020) for which there is released baseline data from VeLOO with 15 trials per lask tuning learning rate
logarithmically with half powers of 10. Morever, NAdamW, which serves as a superset of many adaptive optimizers (Choi et al. 2019 Dahl et al.) 2023), is aggressively tuned with 1000 trials scarching over different configurations of learning rate, 31, A3g e and cosine learning rate schedule (Metz et al., 2020b). We pick the best trial based on final loss for evaluation and visualization.
7

---

Published in Transactions on Machine Learning Research (06/2025)
Final loss IQM
Speedup IQM T
2.0
1.2
Adam normallzed
1.5
1.1
1.0
0.9
0.5
Fromage
AdaGrad SCD M
LARS
AdaBeliet
NAdamW ShampooSGD
Celo (ours)
SM3
RMSProp
LAMB
Yogi
Shampookdagrad
Figure 4: Comparing Celo with state-of-the-art hand-crafted optimizers. Celo is our proposed Iearned optimizer meta-trained on a limited budget 24 GPU hours. Optimizers are evaluated with final loss (left.) and speedup (right) criteria with respect to Adam on a diverse set of 17 tasks which are out-of- distribution for Celo and include image classification, language modeling, autoencoders, learned optimizer training, IQM score above 1.0 indicates "Super-Adam" performance, read more about our evaluation
methodology in 4 and meta-training/ ceaeriiiinn/valaaatiinn evaluation tasks in 5
5.2 Meta-training
All the learned optimizers are meta-trained with the same setup on fixed compute budget i.e. given fixed set of meta-training tasks, we study how far can we push the meta- generaaceraiinoiinnn performance of learned optimizers. We take set of four meta-training tasks proposed by Metz et al. (2022b) for all our experiments which contains four image-clasiffiction datasets including MNIST (LeCun & Cortes 1998)
Mashion-MNIST (Xiao et al. 2017), SVHN (Netzer et al.) 2011), and CIFAR-10 (Krizhevsky et al., 2009) paired with a one-layer MLP network with 32 hidden units and ReLU activations. resized to 8x8 to keep the meta-training computationally efficient and fast like in Metz et al. (2022b). All the learned optimizers are meta-trained with truncated PES (Vicol et al.) 2021 with maximum 2K inner unroll length for 100K mela-iterations using mean loss of the inner loop the mela-objective following prior work (Metz et al. 2022a Vicol et al.) 2021 Harrison et al.) 2022). With this setup, all our meta-tran meta-training experiments finish in a day (<24 hours) on a single Nvidia RTX8000 GPU. Note that we intentionally did not scale up meta-training here in order to do an exhaustive controlled study within our compute budget and solely focus on improving the mntaa-eenerllizaiion performance. For task augmentation ($3), we sample T uniformly on log scale between 0.001 and 1000 in each reta-iterationn See further implementation details in Appendix A.3
Additionally, images are
5.3 Evaluating on a broad spectrum of tasks
We test all the optimizers on 17 diverse unseen tasks with seeds per task. These tasks span different machine learning problems (including classification, language modeling, learned optimizer training, auto- encoders) with different architecture types like MLPs, ConyNets, Transformers, RNNss etc., and varying model sizes. We pick this evaluation set from larger VeLOdrome task set (Metz al., 2022b) of 83 tasks in order to fit. evaluation of all the experiments within our compute budget. The tasks are described below:
Image MLPs. Total 3 image MLP tasks. Two CIFAR- 10 tasks: 3-layer MLP with 128 hidden units, layer-norm and ReLU activations after each layer; 3-layer MLP with 128 hidden units and Tanh activations. One NNSS task with 2-layer MLP of size 128 and ReLU activations.
CNNs. Total 4 convolutional neural network (CNN) tasks, One CIFAR-100 task with three convolution layers having 32, 64 and channels respectively. Three CIFAR-10 tasks with the same (32-64-64) layers as in CIFAR-100, but different normalization schemes: (1) layer-norm (2) batch-norm (3) without any normalization. All these tasks ReLU activations.
8

---

Published in Transactions on Machine Learning Research (06/2025)
final loss
speedup
w/ task augmentation nnadam lopt (Metz et al.
mediant
OG!
IQM1
mediant
OGL
IQM1
2020b
0.69
0.39
0.67
0.00
0.82
0.00
rmn mlp lopt MMetz et al. 2020a
1.00
0.08
1.00
1.07
0.45
0.69
adafac mlp lopt (Metz et al.
2022a
1.04
0.02
1.05
1.32
0.24
1.30
velo-s (Metz et al. 2022b
0.93
0.17
1.00
0.00
0.53
0.81
celo (ours)
1.14
0.01
1.20
1.92
0.14
1.86
Table 1: Comparing Celo with learned optimizers. All the learned optimizers are meta-trained on
the same set of 4 tasks with task augmentation and evaluated on 17 diverse out-of-distribution tasks (listed in $5) IQM score above 1.0 indicates "Super-Adam" performance (read more in $4). Our learned optimizer, Celo, incorporating the recipe proposed in $ significantly improves neta-generalizatio..
VIT. Two vision transformer (ViT) (Dosovitskiy et al. 2021] tasks CIFAR-100: (1) "wide shallow" with 6 layers, 6 heads and hidden size 384; (2) "skinny deep" with 10 layers, 4 heads and hidden size 128. Both
VIT variants represent images as 16 x 16 patches.
Transformer LM. Three transformer decoder language modeling tasks (Radford et al. 2019 on
LMIB (Chelba et al. 2013) of an increasing scale: (1) hidden size 20, 5 heads, 1 layer, sequence length 8; (2)
hidden size 32, heads, 2 layer, sequence length 8; (3) hidden size 32, 4 heads, 2 layer, sequence length 32.
RNN LM. Two recurrent neural network (RNN) language modeling tasks on the LM1B32K (Brants et al. 2007 Chelba et al. 2013) and wikipedia32k (Merity al., 2016) datasets. Both tasks have the same RNN: LSTM with hidden size 256, input embedding size 128 and sequence length 32.
Auto-encoders. Two image MLP autoencoder tasks: (1) 3-layer 128-32-128 MLP on CIFARIO and batch size 256 (2) 3-layer 128-32- 128 MLP on MNIST and batch size 128.
Learning Optimizers. One learned opt omiimization lask to test the ability of our proposed optimizer to train new learned optimizers. This task meta- meta-trains a learned optimizer with an MLP-based optimizer architecture on PashionMMIST with the maximum inner unroll length 50 and PES as the meta-gradient estimator. The optimizee network a 1-layer MLP with 32 hidden units.
5.4 Evaluation metrics
All the 17 tasks from our suite are first optimized using Celo and the baselines for 2K iterations. We then evaluate them using the proposed final loss and speedup 1QM (Section 4. Additionally, we report aggregate Median and Optimality Gap (OG) metrics from Agarwal et al. [2021) for both criteria, final loss and speedup, using normalized score matrix S defined in $ The optimality gap metric (Agarwal et al. 2021) is computed by first clipping all the normalized scores above 1.0 (optimal baseline score) to 1.0 and then subtracting the mean of all the clipped scores from the optimal baseline performance score 1.0.
6 Results
Celo generalizes to unseen tasks. Figure and Table I show the evaluation performance our Celo on the 17 unseen tasks along two criteria: final loss and speedup. For both criteria, Celo outperforms all the previous learned optimizer approaches well standard hand-crafted optimizers such as Shampoo and NAdamW. This result is remarkable for two reasons. First, these evaluation tasks are much larger in scale in terms of flops and parameter count and far out-of-distributiin w.r.t what Colo has seen during meta-training. Second, all the hand-crafted optimizers use 10-1000s of tuning trials per lask, out of which the best trial is picked for evaluation. Notably, for speedup, Celo achieves IQM 1.86 and outperforms all the basclines by a significant margin. Morcover, even after applying task augmentation in all the learned optimizer baselines, making them stronger (Table 2a). Celo leads by large margin (Table D which further highlights the generalization strength of Celo to out-of-distribution tasks from limited meta-training.
9

---

Published in Transactions on Machine Learning Research (06/2025)
Without task augmentation
With task augmentation
1..0
Fraction of runs with score >
-1.00
Frackion of runs with score
0.75
0.75
0.50
0.50
0.25
0.25
0.00
1.0
1.2
1.4
1.8
1.2
2.0
Adam speedup score (T)
Adam speedup score (t)
L.pt
LOpt
Adafac L.pt
Velo-S
Colo (ours)
Figure Meta-generalization speedup profiles. Task augmentation improves meta-genera lizatioo performance in terms of speedup for all the learned optimizers. Celo outperforms all the learned optimizer baselines with and without task augmentation.
IQM1
optimizer
w/o aug
w/ aug
task aug
nnadam lopt
0.65
0.67
update rule
scheduler
IQM1
OGL
rnn mlp lopt
0.73
1.00
celo w/o aug
1.07
0.04
adafac mlp lopt
0.89
1.05
1.16
0.04
velo-s
0.90
1.01
1.10
0.08
celo
1.07
1.20
celo
1.20
0.01
(a)
(b)
Table 2: Task augmentation. Task augmentation improves generalization across the board for all learned optimizer architectures (a) and helps the most when it is used at both stages of Celo meta-training (b).
Task augmentation is key for generalization. Table 2a and Figure show the impact of using task augmentation for all the learned optimizer baselines and Celo. It is important to note that we use the same meta- vettatraiiing dataset and algorithm for all the learned optimizers and only ablate over task augmentation. As evident from the results, task augmentation improves generalization performance for all the approaches. We also ablate task augmentation different stages of training Celo, namely scheduler and parameter update rule in Table 2b As the results suggest, task augmentation helps the most when it is added at both stages of meta-training and optimality gap (0.01) is close to the optimum (0.0). Moreover, most of the performance gain comes from adding task augmentation at the parameter update rule stage (Table 2b This suggests that in learned optimizers, learning a robust update rule is more crucial for generalization performance than a robust scheduler.
Learned per-marre update rule outperforms Adam. We experiment with replacing the learned MLP update rule in Celo with hand-crafted rule of Adam and report its performance in Table 3a This is equivalent to learning a scheduler over a fixed parameter update. The results clearly indicate that Celo with the lcarned update rule significantly outperforms the Cclo with the Adam basclino, as evidenced by the IQM metric (1.20 vs. 1.06). Moreover, while Celo with Adam baseline surpasses the best-tuned Adam baseline and achieves a lower final loss (TQM > 1.0), still lags behind the best-tuned Adam baseline on few tasks, judging by non-zero optimality gap (0.05). This indicates that there is potential for further improving the scheduler, which could also enhance the performance of fully learned optimizers like Celo.
Scheduler is critical for performance. As we describe above, Celo mainly consists of two components: (1) learned update rule (2) Icarned schoduler. To assess the impact of scheduler on meta-genoralization performance for unseen tasks, we evaluate our learned update rule from Stage 1 ($3) without the scheduler and report its performance in Table 3b The results show that our learned update rule nearly matches the
10

---

Published in Transactions on Machine Learning Research (06/2025)
Learned schedules across horizons
Learned schedules across tasks
0.4%
10"
10
Scheduil param
0.8%
VI7 CIFAR-100
TransfomarLL LM1B
AutoEncader MNIST
10-1
AdatacMLL' Lüpt
CIFAR-10
10-1
0
500
1000
1500
2000
500
1000
1500
2000
Iterations
Iterations
Figure 6: Visualizing schedules predicted by Cclo on unseen tasks. Cclo adapts its schodules basod the evaluation target length as shown for VIT CIFAR-100 task (left. and also as per the given evaluation task by varying warmup and decay phases (right). Notably, Celo discovers schedules with warmup phase followed by cyclic ramp-up and cosine decay-like phases throughout the training.
IQMI
OGL
celo
IQM1
OGL
IQMI
OG]
celo w/ adam celo
1.06
0.05
w/o scheduler
0.99
0.12
1-stage training
0.89
0.16
1.20
0.01
w/ scheduler
1.20
0.01
2-stage training
1.20
0.01
(a) Adam learned update rule. Learned parameter update rule outperforms Adam.
(b) Impact of scheduler. Using a scheduler significantly improves meta-generalization.
(c) Two-stage training. Learn- ing the update rule and scheduler separately in two stages is better.
scheduler
IQMI
OGL
functional form
IQMI
OGL
w/ tensor feats
1.19
0.02
linear
Q.Ot
1.07
0.12
W /o tensor feats
1.20
0.01
lincar + clip
clip(a.0:)
1.05
0.06
exp
a.exp(o.)
1.20
0.01
(d) Tensor features. Scheduler with- out per-tensor features performs well.
(e) Functional form. Exponential form in scheduler
works the best.
Table 3: Ablations. Celo default settings are highlighted
final loss of the best Adam baseline across 14 trials, achieving an IQM score close to 1.0 (0.99). However, the learned scheduler significantly improves the meta-generalization performance of Celo, resulting in TQM score of 1.20. This improvement to the unseen task (see Figure G), enabling Celo to achieve final losses that are lower than those of Adam.
be attributed to the learned scheduler's ability to respond dynamically
Decoupling training of step-size scheduler and param update rule helps. We propose to meta- train learned optimizer in two stages, which consists of first learning a parameter update rule followed a learning scheduler which controls the step-size in the update rule keeping the learned update rule fixed (non-trainable). We compare results of this two-stage mela-training procedure with single-stage which all the prior works follow (Metz et al.) 2022b a 2019b) (Table 3c). Our experiments find that single-stage training quickly overfits on the meta-trriiing tasks, specially in compute-limit setting we considered, without learning a general parameter update rule or scheduler. Two-stage training shows better meta- getta-generalization unseen tasks, as indicated by the TQM Table 3c Two-stage training also has a superior worst-case performance (optimality gap) compared the single-stage learned optimizer.
Task-agnostic scheduler generalizes better. Prior work in learned optimizers (Metz el al. 2022b 2020b) with hierarchy uses tensor-level statistics in the scheduler to control the step size and other hyperparameters of the per-parameter update rule. However, since tensor statistics (c.g. gradient mean, momentum, etc) are highly specific for each task, learned optimizers that rely on such features often fail to generalize newer tasks with completely different statistical properties (Almeida et. al. 2021 To address this, use
11

---

Published in Transactions on Machine Learning Research (06/2025)
LiME Mnist, bs128
VI7 widesballow
Conv 32x81a64
3.04
16-
3.03
12
3.02
100
0.8
250
500
750
100n
1250
1500
1750
2000
250
500
750
1000
1250
1500
17t0
2000
250
500
750
1800
1259
1500
1759
2003
RNNLM Im1b32k PatchS2
LSTM256
Embed120
TransformerLM LM1B Multiuustime
ImageMMPPE Cifar10 120x32x120
bs256
6.0
0,045
6.8
0,040
0.035
5.4
6.2-
0.030
0.020
250
1900
1250
1800
2000
250
500
250
1000
1250
1500
1750
2000
250
100
750
1000
1250
3500
1750
2003
Iterations
Shampoo
NNAndam 1.Opt
1.pt
MIP 10pt.
Callo
Figure 7: Training curves for selected tasks. See Appendix A.8
for detailed description of each task.
task-agnostic features such as normalized loss features and training progress which remain similar between small meta-training lasks and large-scale evaluation lasks. Table 3d compares a learned scheduler with per- tensor statistical features with Celo, which does not use any per-tensor features. The results show that our task-agnostic schoduler without any per-tensor statistics gives comparable (in fact, slightly better) results to the one which uses per-tensor statistics, the same time being computationally cheaper. Therefore, we make it our default choice in Celo.
Functional form of scheduler matters for genoralization. Learned scheduler plays an important during training by modulating the step-size of parameter update rule. We test different functional form for the learned scheduler in Table 3e which guides how the step-size is predicted and find that it crucial for generalization performance. Our experiments show that exponential form performs the best and allows the scheduler to rapidly alter step-size exponential scale which is also in line with our modern band-designed schedulers which use
decays.
7 Related work
Learned optimizer architectures. Learned optimizer architectures can be broadly classified into three categories: (1) Flat learned optimizers (2) Hierarchical learned optimizers (3) Hybrid optimizers. Flat optimizers (Andrychowicz al., 2016 Chen et al.) 2020a Metz et al. 2022a; Harrison al.) (2022) use the same update rule for every parameter in the optimizee neural network without considering any additional context such as structure of the underlying optimizee network. Hence, their parameter update computation is non- non-hierarchical or in other words "flat". These flat learned optimizers can be thought of as drop-in replacement for hand-crafted update rules such as SGD or Adam. Hierarchical optimizers (Wichrowskal et al. 2017 Metz et al. 2020a. 2022b Peebles et al. 2022 Moudgil et al. 2023 unlike flat optimizers, have hierarchical structure their update function. They take additional meta-data of the optimizee neural network such as tensor shapes, neural graph (Kofinas al. 2024). etc input and utilize it in their parameter update function. These hierarchical optimizers not only capture inter- paee-pprrmmtee dependencies for faster optimization but also be scaled in memory-effciintt manner without significantly increasing the runtime overhead since their runtime cost scale sub-linearly with the number of parameters due to hicrarchy (Metz et al. 2022b at. Hybrid optimizers (Almeida et al. 2021 Jang et al.) 2023 Knyazev et. al. 2024 Krist iansen et al. 2024) combine hand-crafted and learned modules in their update function. Using statistics of the underlying optimization task such as loss, gradient norms, momentum, etc, learned modules in hybrid optimizers control hyperparameters of the hand-crafted update rules such as SGD (Kristiansen et al.
2024
Adam Metx et al.) 2020b Kristiansen et. al. 2024h LAMB (Almeida al. 2021). Jang et. al.
2023
12

---

Published in Transactions on Machine Learning Research (06/2025)
Knyazev et al. (024) alternate between Adam updates and updates performed by learned function taking only the past parameters as input showing promising results, but limited scalability. In this work, we focus on enhancing the meta- generalization capabilities of fully learned optimizers while maintaining memory and compute efficiency for practical use. Hence, our proposed approach, Celo, has a hicrarchical architecture.
Mettagenerallzationn of learned optimizers. Several learned optimizers have been proposed for specific settings such as physics-based motion reconstruction [Gärtner et al. 2023), reinforcement learning (Goldie ct al. 2024), ; matcrial design (Merchant et al. 2021), etc. Although these approaches are quite performant, their applicability to broad range of Machine Learning (ML) tasks that practitioners care about is limited. Since mueaa-traiiing a learned optimizer for each setting is computationally expensive, optimizers that meta- gencralize to multiple tasks after being meta-trained are more desirable. A few recent works have explored this direction. Almeida et al. (202T) proposed a hybrid oplimizer with LAMB updale rule that is trained on small-scale tasks and generalize to 3 orders of magnitude higher tasks. VeL.O (short. for "Versatile Learned Optimizer") is fully learned optimizer, meta-trained using an extensive computational budget of 4000 TPU- months across diverse ML tasks. VeLO has demonstrated the ability to generalize effectively to new, unseen tasks and successfully optimize neural networks with parameters up to 600M. In this work, we propose a recipe to meta-train fully learned optimizers on a limited compute budget that meta-generalize to wide range of unseen tasks. Thérien et al. (2024) recently proposed another approach to tackle meta mneaa-eenealliztion by developing learned optimizers within the Maximal Update Parametrization (Yang al.) 2021) framework. This enables meta-generallzattin to wider and deeper models well as to longer unrolls of a meta-training task. Their approach is complementary to ours, so combining it with our recipe is a promising direction to further improve mela-generalliatiin.
Benchmarking optimizers. A few works have attempted to benchmark and compare optimizers (Schneider et al. 2019 Schmidt et al., 2021: Dahl al. 2023) on standard ML. [Schniderr al. (2019) proposed the DeepOBS benchmark that contains set of realistic 8 ML tasks and compares standard (hand-crafted) optimizers using training curves, test curves and learning rate sensitivity analysis. [Schmidt et al. (2021) benchmarked 15 hand-crafted optimizers on 8 ML tasks from DeepOBS by varying tuning budget and learning rale schedules, leading to more than 50K optimization runs from all the combinations. The oplimizers are compared relative to each other by using final loss/test accuracy. Dahl et al. (2023) proposed another benchmark that evaluates standard optimizers on large-scale MI.-Perf tasks and a methodology to score optimizers by using speedup. Speedups for each task are computed using fixed task-specific metric targets that are hand-designed. Tn this work, we propose evaluation metrics that can capture aggregate performance of an optimizer using both speedup and final performance criteria. Unlike prior work, our evaluation metrics are specifically designed to be robust to outlicrs, especially when the number of evaluation tasks is large and doesn't 't require hand-designing any cut-offs/ Hil-offs//hresholle /hresholds. Hence, our proposed evaluation methodology is scalable with respect to number of evaluation tasks.
8 Limitations and future work
Our work focuses on recipe to train learned optimizers that mela-generalize well from a limited training set. However, this work is leaving out lot of "free" performance gains which could be obtained by trivially mod- ifying or expanding the meta-training set, using better/faster architectures like state-space models instead of LSTMs. We intentionally kept our architecture and meta- veeta-training setup simple without any bells and whistles to do a fair head-to-head comparison with prior work and focus solely on the improvements from our recipe. Moreover, although our work significantly improves generalization performance of fully learned optimizers and we firmly believe this work is the right step towards making learned optimizers more practical, it is not focused on delivering the most-performant optimizer for all kinds of large-scale workloads. Additional work needs to be donc to make these learned optimizers ready for production use cases like analyzing the impact of meta-training data distribution on generalization properties, which is out of scope for this work. Further- more, given the amount of compute resources we had, we tried our best to include as diverse and large number of tasks as possible for evaluating all the lcarned optimizers and ablations with multiple scods and hyperpa- rameters. This resulted in hundreds of evaluation runs for each optimizer ablation. However, there is a scope for larger scale study just focused on evaluation using our proposed metrics which we leave for future work.
13

---

Published in Transactions on Machine Learning Research (06/2025)
9 Conclusion
We propose a recipe to train versatile fully learned optimizers on a limited computed budget by identifying three key elements mcta-training and architecture: (1) task augmentation (2) a simple design consisting of learned scheduler with a learned per-parameter update rule and (3) decoupled training of scheduler and per- na-paaameter update rule. We show that our proposed learned optimizer, Celo, outperforms state-of-the art hand-crafted and learned optimizers on unseen tasks, despite being meta-trained on a limited compute budget of 24 GPU hours, demonstrating its strong generaaencil capabilities. Moreover, we propose evaluation metrics that can reliably capture aggregate performance of an optimizer on an evaluation set inspired from RL literature. Furthermore, we perform exhaustive ablations with geteraaeneiliittionn as the principal axis of evaluation. Finally, we analyze the schedules proposed by Cclo on unseen tasks and find that they are adaplive to unseen training tasks and targel length horizons. Overall, believe that this work serves a significant stepping stone towards the development of fully learned optimizers that are efficient, hyperparameterr-fres and practically applicable to large-scale unseen tasks.
Acknowledgements
We acknowledge support from Mila-Samsung Research Grant, FRQNT New Scholar (EB), the FRQNT Doc- toral (B2X) scholarship [AM] and the Canada-CIFAR AI Chair program [GL]. We also acknowledge resources provided by Compule Canada, Calcul Québec, and Mila. We thank NVIDIA for providing the GPUs that made this work possible. We are also grateful to the developers of open-source libraries such as JAX (Brad- bury et al. 2018), NumPy (Harris et al. 2020), google/learneddooptimiza which were instrumental in this research.
and gogge--ssearcch/rllaab.
References
Josh Achiam, Steven Adler, Sandhini Agarwal, Lama Ahmad, Ilge Akkaya, Florencia Leoni Aleman, Diogo Almeida, Janko Altenschmidt, Sam Altman, Shyamal Anadkat, el al. Gpt-4 technical report. arXiv preprint arXim:2303.087 087744 2023.
Rishabh Agarwal, Max Schwarzer, Pablo Samuel Castro, Aaron Courville, and Marc G. Bellemare. Deep reinforcement learning at the edge of the statistical precipice, 2021.
2 6
9 22
Diogo Almeida, Clemens Winter, Jie Tang, and Wojciech Zaremba. optimizers. arXiv preprint Axix:2106.00955 2021. 3
A generalizable approach to learning
Marcin Andrychowicz, Misha Denil, Sergio Gomez, Matthew W Hoffman, David Pfau, Tom Schaul, Brendan Shillingford, and Nando De Freitas. Learning to learn by gradient descent by gradient descent. Advances in neural information processing systems, 29, 2016.
2 3
5 12
Rohan Anil, Vincet Gupta, Tomer Koren, and Yoram Singer. Memory efficient adaptive optimization. Newral Information Processing Systems, 32, 2019.
Rohan Anil, Vineet Gupta, Tomer Koren, Kevin Regan, and Yoram Singer. Scalable second order optimiza- tion for deep learning. arXiv preprint arXi::2002.0900 2020.
2
Jeremy Bernstein, Arash Vahdat, Yisong Yue, and Ming- Yu Liu. On the distance between two neural networks and the stability of learning. In Neural Information Processing Systems, 2020. 7
Léon Bottou, Frank E Curtis, and Jorge Nocedal. Optimization methods for large-scale machine learning.
SIAM review, 60(2):223 311, 2018.
Stephen Boyd and Lieven Vandenberghe. Conver optimization. Cambridge university press, 2004. 5
James Bradbury, Roy Frostig, Peter Hawkins, Matthew James Johnson, Chris Leary, Dougal Maclaurin, George Necula, Adam Paszke, Jake VanderPlas, Skye Wanderman-Milne, and Qiao Zhang. JAX: com- posable transformations of Nython+NumPy programs, 2018. URL http://gitti http::/ggicuub.com/gggge///// com com/google/jax
22
14

---

Published in Transactions on Machine Learning Research (06/2025)
Thorsten Brants, Ashok Popat, Peng Xu, Franz Josef Och, and Jeffrey Dean. Large language models in machine translation. In Proceedings of the 2007 Joind Conference on Empirical Methods in Natural Language Processing and Computational Natural Language Learning (EMNLP-CONLL), pp. 858 867, 2007.
Tim Brooks, Bill Peebles, Connor Holmes, Will DePue, Yufci Guo, Li Jing, David Schnurr, Joc Taylor, Troy Luhman, Eric Luhman, et al. Video generation models as world simulators. 2024. URL htlps://opende
com/research/ sepporch/vvieoo-eee hatperrttoonmmadee------or---sss
3, 2024.
Johan Samir Obando Ceron, Ghada Sokar, Timon Willi, Clare Lyle, Jesse Farebrother, Jakob Nicolaus Foerster, Gintare Karolina Dziugaite, Doina Precup, and Pablo Samuel Castro. Mixtures of experts unlock parameter scaling for deep rl. In Forty-first International Conference on Machine Learning. G
Ciprian Chelba, Tomas Mikolov, Mike Schuster, Qi Ge, Thorsten Brants, Phillipp Kochn, and Tony Robin- son. One billion word benchmark for measuring progress in statistical language modeling, arXiv preprint orXiv:1812.3005, 2013. 9
Tianlong Chen, Weiyi Zhang, Zhou Jingyang, Shiyu Chang, Sijia Liu, Lisa Amini, and Zhangyang Wang. Training stronger baselines for learning to optimize. Advances in Newral Information Processing Systems,
33:7332- 03333-3333.2.. 2020a. 12
Ting Chen, Simon Kornblith, Mohammad Norouzi, and Geoffrey Hinton. A simple framework for contrastivv learning of visual representations. In International conference on machine learning, pp. 1597-1607. PMLR, 2020b.
Dami Choi, Christopher J Shallue, Zachary Nado, Jachoon Lee, Chris J Maddison, and George E Dahl. On empirical comparisons of optimizers for deep learning. arXiv preprint Xrx:::1900.0555 05446, 2019. 7
Ekin D Cubuk, Barret Zoph, Jonathon Shlens, and Quoe V Le. Randaugment: Practical automated data augmentation with a reduced search space. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition workshops, pp. 702-703, 2020.
George E Dahl, Frank Schneider, Zachary Nado, Naman Agarwal, Chandramouli Shama Sastry, Philipp Hennig, Sourabh Medapati, Runa Eschenhagen, Priya Kasimbeg, Daniel Suo, et al. Benchmarking neural network training algorithms. aivi preprint 0rXiw:2306.07177 07179, 2023.
6
13
23
25
Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Un- terthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, Jakob Uszkorcit, and Neil Houlsby. An image is worth 16x16 words: Transformers for image recognition at scale. In International Conference on Learning Representations, 2021. URL https://openrevieet net http:///weeerreeee.eeeeeeeooooaa
Timothy Dozat. Incorporating Nesterov momentum into Adam. In International Conference on Learning Representations, 2016. 7
John Duchi, Elad Hazan, and Yoram Singer. Adaptive subgradient methods for online learning and stochastic optimization. Journal of Machine Learning Research, 2011.
Erik Gärtner, Luke Metz, Mykhaylo Andriluka, C Daniel Freeman, and Cristian Sminchisescu. Transformer- based learned optimization. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 11970-11979, 2023.
4 13
Alexander David Goldie, Chris Lu, Matthew Thomas Jackson, Shimon Whiteson, and Jakob Nico- laus Foerster. Can learned optimization make reinforcement learning less difficult? ar preprint ar 07700:247.007082, 2024. 13
Ian Goodfellow, Yoshua Bengio, Aaron Courville, and Yoshua Bengio. Deep learning, volune 1. MIT Press, 2016. M
P Goyal. Accurate, large minibatch sgd: training imagenet in 1 hour. arXiv preprint arxiv 1rx6:1706.000 026777 2017. 5
15

---

Published in Transactions on Machine Learning Research (06/2025)
Vineet Gupta, Tomer Koren, and Yoram Singer. Shampoo: Preconditioned stochastic tensor optimization. In International Conference on Machine Learning, pp. 1842-1850. PMLR, 2018.
6 T
Charles R. Harris, K. Jarrod Millman, Stéfan J. van der Walt, Ralf Gommers, Pauli Virtanen, David Cournapeau, Eric Wieser, Julian Taylor, Sebastian Berg, Nathaniel J. Smith, Robert Kern, Matti Pi- cus, Stephan Hoyer, Marten H. van Kerkwijk, Matthew Brett, Allan Haldane, Jaime Fernández del Río, Mark Wiebe, Pearu Peterson, Pierre Gérard-Marchant, Kevin Sheppard, Tyler Reddy, Warren Weckesser, Hameer Abbasi, Christoph Gohlke, and Travis E. Oliphant. Array programming with NumPy. Nature, 585(7825):357-300 September 2020. doi: 10. 100088/441586-020-2000 7/doi. org/ 10. 1038//41586--22-60
URL https:
14
James Harrison, Luke Metz, and Jascha Sohl-Dickstein. A closer look learned optimization: Stability, robustness, and inductive biases. In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho (ods.), Advances in Neural Information Processing Systems, 2022. URL htt::///ppenreiieenneeeefoooo
id=cxZEBOFDoFK
20
22
Jinhyeok Jang, Woo-han Yun, Won Hwa Kim, Youngwoo Yoon, Jachong Kim, Jaeyeon Lee, and ByungOk Han. Learning to boost training by periodic nowcasting near future weights. Tn Indernational Conference on Machine Learning (ICML), 2023. 12
John Jumper, Richard Evans, Alexander Pritzel, Tim Green, Michael Figurnov, Olaf Ronneberger, Kathryn Tunyasuvunakool, Russ Bates, Augustin Zídck, Anna Potapenko, et al. Highly accurate protein structure prediction with alphafold. nature, 596(7873):583 -89958933855885588 2021.
Jared Kaplan, Sam McCandlish, Tom Henighan, Tom Brown, Benjamin Chess, Rewon Child, Scott Gray, Alec Radford, Jeffrey Wu, and Dario Amodei. Scaling laws for neural language models. arXie preprint
arXiv:2001.0880 0000:, 2020. 2
Priya Kasimbeg, Frank Schneider, Runa Eschenhagen, Juhan Bae, Chandramouli Shama Sastry, Mark Saroufim, Boyuan Feng, Less Wright, Edward Z. Yang, Zachary Nado, Sourabh Medapati, Philipp Hennig, Michael Rabbat, and George F. Dahl. Accelerating neural network training: An analysis of the algoperf competition, 2025. URL
15015
23 25
Diederik P. Kingma and Jimmy Ba. Adam: A method for stochastic optimization. In Yoshua Bengio and Yann LeCun (eds.), 3rd International Conference on Learning Representations, ICLR 2015, San Diego, CA, USA, May 7-9, 2015, Conference Track Proceedings, 2015. URL htttp://arxiv..rr/bbb////// 6980
2
7 22
Alexander Kirillov, Eric Mintun, Nikhila Ravi, Hanzi Mao, Chloe Rolland, Laura Gustafson, Tote Xiao, Spencer Whitehead, Alexander C Berg, Wan-Yen Lo, et al. Segment anything. In Proceedings of the
IRFE/CVF International Conference Computer Vision, pp. 4015-4026, 2023. ]
Boris Knyazev, Abhinav Moudgil, Guillaume Lajoie, Eugene Belilovsky, and Simon Lacoste-Julien. Acceler- ating training with neuron interaction and nowcasting networks. arXie preprint arXiv:2409. N000:2409004444, 2024.
12
Miltiadis Kofinas, Boris Knyazev, Yan Zhang, Yunlu Chen, Gertjan J. Burghouts, Efstratios Gavves, Coes G. M. Snoek, and David W. Zhang. Graph Neural Networks for Learning Equivariant Representations of Neural Networks. In 12th International Conference on Learning Representations (ICLR), 2024. 12
Gus Kristiansen, Mark Sandler, Andrey Zhmoginov, Nolan Miller, Anirudh Goyal, Jihwan Lee, and Max Vladymyrov. Narrowing the focus: Learned optimizers for pretraincd models. arXiv preprint arXiv:2408.0920 2024. 12
Alex Krizhevsky, Geoffrey Hinton, et al. Learning multiple layers of features from tiny images. 2009. S
Alex Krizhevsky, Ilya Sutskever, and Geoffrey Hinton. Imagenet classification with deep convolutional neural networks. Advances in neural information processing systems, 25, 2012. T
16

---

Published in Transactions on Machine Learning Research (06/2025)
Misha Laskin, Kimin Lee, Adam Stooke, Lerrel Pinto, Picter Abbeel, and Aravind Srinivas. Reinforcement learning with augmented data. Advonces in neural information processing systems,
2020.
Yann LeCun and Corinna Cortes. The mnist databasc handwritten digits. 1998. URL http://yann. con//edb/mnist
Liyuan Liu, Haoming Jiang, Pengcheng He, Weizhu Chen, Xiaodong Liu, Jianfeng Gao, and Jiawei Han. On the variance of the adaptíve learning rate and beyond. arXiv preprint orXi::108 03265, 2019.
Loshchilov. Decoupled weight decay regularization. arXie preprint orXiv:1711.051000 2017.
Nya Loshchilov and Frank Hutter. Sgdr: Stochastic gradient descent with warm restarts. arXiv preprint
arXiv:1608.0300 2016.
David Lowe. Distinctive image features from scale-invariant keypoints. International journal of computer vision, 60:91-110, 2004.
Amil Merchant, Luke Metz, Samuel Schoenholz, and Ekin D Cubuk. Learn2hop: Learned optimization on
rough landscapes. In International Conference Machine Learning, pp. 7643-7653. PMLR, 2021. 13
Stephen Merity, Caiming Xiong, James Bradbury, and Richard Socher. Pointer sentinel mixture models. arXiv preprint arXio:1609.078.4 2016. g
Luke Metz, Niru Maheswaranathan, Brian Cheung, and Jascha Sohl-Dickstein. Learning unsupervised learning rules. 2019a. 3 4
Luke Metz, Niru Maheswaranathan, Jeremy Nixon, Daniel Freeman, and Jascha Sohl-Dickstein. Under- standing and correcting pathologies the training of learned optimizers. In International Conference on Machine Learning, pp.
PMTR, 2019b. 11
Luke Metz, Niru Maheswaranathan, C Daniel Freeman, Ben Poole, and Jascha Sohl-Dickstein. Tasks, stability, architecture, and compute: Training more effective learned optimizers, and using them to train themselves. Xive preprint or arXiv:2009. 1000..1220 11243, 2020a. 3
20
Luke Metz, Niru Maheswaranathan, Ruoxi Sun, Daniel Freeman, Ben Poole, and Jascha Sohl- Dickstein. Using a thousand optimization tasks to learn hyperparameter search strategies. arxio preprint arXiv:2002.1189) 2020b. I
9
12 21
Luke Metz, Daniel Freeman, James Harrison, Nirn Maheswaranathan, and Jascha Sohl-Dickstein. Practical tradeoffs between memory, compute, and performance in learned optimizers. In Conference on Lifelong Learning Agents (CoL.As), 2022a. URL http://tttbbbooo/oooo///////eeee wwwparriatoooiiiiiiiiiiioooo
22 23
Luke Metz, James Harrison, Danicl Freeman, Amil Merchant, Lucas Beyer, James Bradbury, Naman Agrawal, Ben Poole, Igor Mordatch, Adam Roberts, el al. Velo: Training versatile learned optimizers by scaling up. arXiv preprint 17X1:2221..99760 2022b. I
20
21
22
23
24
Vincent Micheli, Eloi Alonso, and François Fleuret. Transformers are sample-efficient world models. arXiv preprint ar Xiv: 00588, 2022.
Abhinav Moudgil, Boris Knyazev, Guillaume Lajoie, and Eugene Belilovsky. Learning to optimize with recurrent hierarchical transformers. In ICML Workshop on New Frontiers in Learning, Control, and Dynamical Systems, 2023. 12
Yuval Netzer, Tao Wang, Adam Coates, Alessandro Bissacco, Baolin Wu, Andrew Y Ng, et al. Roading digits natural images with unsupervised feature learning. In NIPS workshop on deep learning and umsupervised feature learning, volume 2011, pp. 4. Granada, 2011.
17

---

Published in Transactions on Machine Learning Research (06/2025)
William Peebles, Lija Radosavovic, Tim Brooks, Alexei Efros, and Jitendra Malik. Learning to learn with generative models of neural network checkpoints. orXiv preprint arXiv:2209. 12892, 2022. 2 12
Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, Ilya Sutskever, et al. Language models are unsupervisod multitask learners. OpenAT blog, 1(8):9, 2019.
Samyam Rajbhandari, Jeff Rasley, Olatunji Ruwase, and Yuxiong He. Zero: Memory optimizations toward training trillion paramcter models. In SC20: International Conference for High Performance Computing, Networking, Slorage and Analysis, pp. 1-16. IEEE, 2020.
Sashank J Reddi, Satyen Kale, and Sanjiv Kumar. On the convergence of adam and beyond. arXiv preprint arXiv:19440090 092977, 2019. 1
Fady Rezk, Antreas Antoniou, Henry Gouk, and Timothy Hospedales. Is scaling learned optimizers worth it? evaluating the value of velo's 4000 tpu months. In Proceedings on, pp. 65-83 83. PMLR, 2023. 2
Robin M Schmidt, Frank Schneider, and Philipp Hennig. Descending through crowded sellly-beucharking deep learning optimizers. In International Conference on Machine Learning, pp. 9367-9376. PMLR, 2021.
6
13
Frank Schneider, Lukas Balles, and Philipp Hennig. Deepobs: A deep learning optimizer benchmark suite.
arXiv preprint 1tt:1900.053 2019. 2
Max Schwarzer, Johan Samir Obando Ceron, Aaron Courville, Marc G Bellemare, Rishabh Agarwal, and Pablo Samuel Castro. Bigger, better, faster: Human-level atari with human-level efficiency. In Interna- tional Conference on Machine Learning, pp. 30365-30380 PMLR, 2023. 6
Noam Shazeer and Mitchell Stern. Adafactor: Adaptive learning rates with sublinear memory cost. arXiv preprint (8fe::8044042335 2018
Prabhu Teja Sivaprasad, Florian Mai, Thijs Vogels, Martin Jaggi, and François Fleuret. Optimizer bench- marking needs to account for hyperparameter tuning. In International conference on machine learning, pp. 1086-9045. PMLR, 2020.
Leslie N Smith. Cyclical learning rates for training neural networks. In 2017 IEEE winter conference on applications of computer vision (WACV), pp. 464-472. IEEE, 2017.
Ilya Sutskever, James Martens, George Dahl, and GeofTrey Hinton. On the importance of initialization and momentum in deep learning. In International conference on machine learning, pp. 1139 1147, PMLR, 2013.
Benjamin Thérien, Charles-Etienne Joseph, Boris Knyazev, Edouard Oyallon, Irina Rish, and Eu- gene Belilovsky. μLO: Compute-efficient goctaaennralizttion of learned optimizers. arXiv preprint N0.0:2406.0015.. 2024.
13
T. Tieleman and G. Hinton. Lecture 6.5 RmsProp: Divide the gradient by a running average of its recent magnitude. COURSERA: Neural Networks for Machine Learning. 2012. 7
Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie Anne Lachaux, Timothée Lacroix, Baptiste Rozière, Naman Goyal, Eric Hambro, Faisal Azhar, ot al. Llama: Open and efficiont foundation
language models. orXiv preprinl urXiv:2302. 13971, 2023. 5
Paul Vicol, Luke Metz, and Jascha Sohl-Dickstein. Unbiased gradient estimation in unrolled computation graphs with persistent evolution strategies. Tn International Conference on Machine Learning, pp- 10553- 10563. PMLR, 2021. 3
g
Olga Wichrowska, Niru Maheswaranathan, Matthew W Hoffman, Sergio Gomez Colmenarejo, Misha Denil, Nando Freitas, and Jascha Sohl-Dickstein Learned optimizers that scale and generalize. In International conference machine learning, pp. 3751-3760 PMI.R, 2017.
3
12
18

---

Published in Transactions on Machine Learning Research (06/2025)
Han Xiao, Kashif Rasul, and Roland Vollgraf. Fashion-mnist: novel image dataset for benchmarking machine learning algorithms. aiv preprint arXiv: nnn:7088007777 2017.
Ge Yang, Edward Hu, Igor Babuschkin, Szymon Sidor, Xiaodong Liu, David Farhi, Nick Ryder, Jakub Pachocki, Weizhu Chen, and Jianfeng Gao. Tuning large neural networks via zero-shot hy- perparameter transfer. In M. Ranzato, A. Beygelzimer, Y. Dauphin, P.S. Liang, and J. Wortman Vaughan (eds.), Advances in Newral Information Processing Systems, volume 34, pp. 17084 17097. Curran Associates, Inc., 2021. URL Bdf1 Papp Paper-pdf 13
cc/paper httpr//ss...pppppp....se/////p/p filess/papr/22221fiile
Yang You, Igor Gitman, and Boris Ginsburg. Large batch training of convolutional networks. arXiv preprint
wrXiv:1708. arXiv:1708.03888 03888, 2017.
Yang You, Jing Li, Sashank Roddi, Jonathan Hseu, Sanjiv Kumar, Srinadh Bhojanapalli, Xiaodan Song, James Demmel, Kurt Keutzer, and Cho-Jui Hsieh. Large batch optimization for deep learning: Training BERT in 76 minutes. arXiv preprint 19X::1904400062, 2019. 7
Manzil Zaheer, Sashank Reddi, Devendra Sachan, Satyen Kale, and Sanjiv Kumar. Adaptive methods for nonconvex optimization. Newral Information Processing Systems, 2018.
Juntang Zhuang, Tommy Tang, Yifan Ding, Sekhar C Tatikonda, Nicha Dvornek, Xenophon Papademetris, and James Duncan. Adabelief optimizer: Adapting stepsizes by the belief in observed gradients. Newral Information Processing Systems, 2020.
19

---

Published in Transactions on Machine Learning Research (06/2025)
A
Appendix
A.1 Notation table
Optimizee network parameters
Φ
Learned optimizer parameters
V0
Gradient of optimizee network parameters
0.
Optimizee network parameters at iteration / Optimizer state at iteration
8t
M
Meta-data input to optimizer such as current loss, training progress, etc.
tb
Optimizee features
fo
Function parameterized by ø
F
Per-param features
S
Normalized score matrix Sct of meta-training tasks
T
D
Inner task data
C
Inner task objective
L
Loss value at iteration /.
1]
Scheduler scale parameter at iteration f
Task augmentation parameter
A.2
Additional ablations
Results without task augmentation. Table shows results for Celo and all the learned optimizer baselines without task angmentation in prior work (Metz et al. 2022a Harrison et al., 2022 Metz et al., 2020a; 2022b). As evident from the results, all the learned optimizers perform much worse for both speedup and final loss criteria when task augmentation is not used (refer to Table for baseline results with task augmentation). Interestingly, VeLO-S performs better than Adafac MT.P LOpt when task angmentation is not used. Moreover, VeLO-S median speedup score is 0.0 with and without task augmentation. However, it improves when task augmentation used with respect IQM and optimality gap metrics further highlighting the importance of our proposed metrics for tracking improvements which are not obviously captured by standard metrics like median. Except Celo, none of the learned optimizer baselines from prior work are able to achieve a final loss lower than Adam when evaluated on the 17 tasks evident from their IQM results (final loss IQM- < 1 and speedup = 0).
VeLO ablations. In order be fair in comparison with all the RNN-based learned optimizer baselines which use, scale the hidden size of VeLO from 512 to 64 and accordingly the MLP bank size from 256 to 32 as mentioned in $5 In Table architecture (denoted by "velo") meta-trained with the same setup other basclines ($5). Morcover, we also show results for the released pre-trained VeLO model Metz el al. (2022b) (denoted by "pretrained--elo"" in Table 5) which is meta-trainet using 4000 months on large collection of tasks. Note that for pre-traines VeLO, our evaluation task set actually in-distribution (Metz et al. (022b) but for VeLO, VeLO-S and also our proposed Celo optimizer, they are aut-of-distribution. The meta-generallzzttoon performance of VeLO is slightly better with the bigger hidden size (512) than VeLO-S (IQM 0.96 0.90) but still lags behind
show mntaagenerallzztion results with original VeLO's original
20

---

Published in Transactions on Machine Learning Research (06/2025)
final loss
speedup
w/o task augmentation nnadam lopt (Metz et al.
mediant
OG!
IQM1
mediant
OGL
IQM1
2020b
0.62
0.40
0.65
0.00
0.90
0.00
rmn mlp lopt MMetz et al. 2020a
0.77
0.32
0.77
0.00
0.88
0.00
adafac mlp lopt (Metz et al.
2022a
0.93
0.18
0.89
0.00
0.78
0.00
velo-s (Metz et al. 2022b
0.90
0.14
0.90
0.00
0.71
0.13
celo (ours)
1.14
0.01
1.20
1.92
0.14
1.86
Table 4: Additional results without task augmentation. All the learned optimizer baselines are meta-
trained on & fixed set of tasks without task augmentation ($3) as in prior work and evaluated on 17 diverse out-of-distribution lasks (refer to $5). Our proposed approach, Celo, consisting of all the three ingredients proposed in 3 significantly outperforms all the learned optimizer baselines.
our proposed approach Celo (IQM 1.20). This result shows that returns from architecture and algorithmic improvements in learned optimizers can be much higher than simply scaling optimizer sizes.
compute budget
TQ
OG]
velo-s
24 GPU hours
0.90
0.14
velo
24 GPU hours
0.96
0.12
velo-4000
4000 TPU months
1.41
0.00
celo
24 GPU hours
1.20
0.01
Table 5: VeLO ablations. 1QM and optimality gap (OG) metrics for final loss criteria 17 evaluation task set (55) are reported along with meta-training compute budget. Note that our evaluation task set is in-distributton for pretrained-velo but out-of-distribution for celo and all the other velo baselines.
Task augmentation level ablations. Task augmentation i.e. re-paramettrizatio can be applied at differ- ent hierarchical levels of neural network during meta-training (1) Global level, where single augmentation parameter is sampled for the entire network; (2) Tensor level, where an augmentation parameter sampled for cach tensor or layer in the network; (3) Per-parameter level, where an augmentation parameter is sampled individually for each parameter the network; and (4) Mixed stralegy, where the augmentation parameter is sampled uniformly from a combination of the previous three approaches. In our experiments, for mixed augmentation strategy, we first randomly sample strategy from the list ['none", "none', "global', "global', "tensor", "tensor", "parameter"] and then randomly select an augmentation range from the list [[0.001, 1000.0), (0.01, 100.0), [0.1, 10.0]]. Additionally, different augmentation strategies can used in the two stages for our approach presented in $3 We present results for all the ablations in Table 6 Overall, we find that global augmentation performs the best closely followed by tensor augmentation.
task aug level
stage 1
stage 2
IQM
OGL
tensor
global
1.13
0.04
global
tensor
1.13
0.04
mix mix
1.14
0.03
tensor
tensor
1.18
0.04
celo
global global
1.20
0.01
Table 6: Task augmentation level ablations. 1QM and optimality gap (OG) metrics for final loss criteria
on our 17 evaluation task set ($5) are reported. Task augmentation applied at different levels during two- stage meta-training proposed in 3 "mix" refers to randomly sampling from a mixture of level list, see A.2 for more details.
21

---

Published in Transactions on Machine Learning Research (06/2025)
RNN MLP LOpt
VeLO Metz et al. 2022b
Celo (Ours)
Metz et al. 2020a
global ageregrtion
global agaregation
teatures
per MLP
RESt parazss
per frature
per MLP
Bext
perpera
per param MLP
porams
REAT
passen cubbookings
per nenerr conditoning emberdiing
Aggrrggaee
hypernet
?
linear
lincar
per LSTM
per temsor LSTM
global LSTM
Scmoor x global progress features
Temsor global progress Teatures
global prograxx teatures
Figure 8: Comparing Celo architecture with VeLO and RNN MLP LOpt. Our proposed approach, Celo, not only simplifies learned optimizer architecture from prior work as shown above but also vastly outperforms in mota-genoralizztion performance (see $
A.3 Implementation details
We meta-train and evaluate all the learned optimizers in this work in JAX (Bradbury et al. 2018 using the open-sourced learned optimization ior library. For fair comparison, all the learned optimizers are meta- (raine with exactly same PES setup: maximum inner unroll length 2000 with unroll lengths sampled logarithmically sampled between 100 and 2000, standard deviation 0.01, truncation length 50 and mcan inner training loss as the meta-objective. We use Adam (Kingma Ba 2015) as the meta optimizer and meta-train for 100K iterations single Nvidía. RTX8000 GPU. Following prior work (Metz et al. 2022a Harrison et al. 2022) meta-train each learned optimizer using Adamm the meta-optimizer by sweeping over 3 seeds and 5 learning rates [3e-5, 5e-5, 1e-4, 3e-4, 1e-3]. We find that higher and lower learning rates often result in unstable training with PFS, hence learning rates such as 1e-4 and 3e-4 work the best for all learned optimizers (except NNAdamW LOpt for which lower lcarning rates, 30-5 and 56-5, work better). Mcta-training jobs hours.
for all the learned optimizers in this work finish within hyperparameter-fne free and do fair comparison, we do not use weight decay in learned optimizers although it has been shown to help for some tasks (Metz et al.) 2022b Harrison et al.) 2022) but requires additional tuning. We evaluate all the learned optimizers in this work our 17 task evaluation set with 3 seeds per task. For lask augmentation in mela-training, 0.001 and 1000 in each meta-iteration and use it to re-parametrize optimizee network weights at the global level inner-training (Metz et al., 2022b). As noted in prior work (Metz et al. 2022b), this augmentation can simulate even more lasks when applied at different levels in the optimizee neural network: global level, tensor level and per-parameter level by sampling global (one for all parameters), per-tensor or per- parameter a augmentations (rs) respectively. In our experiments, we find out that global-level augmentation performs the best with respect to meta-generan geeta-generalization and hence use it as the default. Moreover, in order to decrease the impact of noise and loss fluctuations during training (e.g. lucky batches with low loss values or loss spike at the end of training), we smooth the loss curves for Adam baseline and optimizers being evaluated using exponential moving average with coefficient 0.9 before computing the normalized scores. We compute final loss by taking average over multiple batches (10) as in prior work (Metz et al. 2022b a). We use the official code rliable open-soureed by (Agarwal al.] 2021 p for computing IQM and other metrics.
order to keep learned optimizers
sample random parameter logarithmically bet ween
hctps: http://wwww...ogggggg////eeeeede
pttps://giii com/ goog resaearch//liablle
22

---

Published in Transactions on Machine Learning Research (06/2025)
A.4 Scaling study
In our main paper, learned optimizers are meta-trained with a 2K unroll length. To study scaling behavior of our proposed approach, we meta-train Celo and VeLO-S with unroll lengths of 1K, 2K, and 5K, then meta-test them on all tasks in our evaluation set using a 10K unroll length, assessing generalization to longer unrolls. We ro- port the IQM metric with the final loss criterion and show scaling trends in the right Figure. Celo performance scales predictably with the meta-training budget and gets IQM 1.85 on 10K unroll length with 5K meta-train unroll length, hence demonstrating strong gener- alization performance with respect to unroll length. VeLO-S, on the other hand, gets 1QM 0.34 with 5K meta-train length. These results clearly demonstrate that Celo scales significantly better than VeLO, achieving much better performance under fixed veta-training budget, thereby highlighting its compute efficiency.
2.0
celo
velo-s
1.5
10M (10K)
1.0
0.5
0.0
1K
2K
5K
meta-train unroll length
A.5 Runtime analysis
Celo updates the parameters of given optimizee network using per-parameter MLP. Since the number of FLOPs for a forward pass of a linear layer with weights
is 2nmm, the cost of a 2-layer MLP update rule in Celo with 30 input features, 2 hidden layers with 4 units and 2 outputs units is 2 x (30 x x 2) = 288 FLOPs per parameter. For an optimizee linear layer with n x m parameters and batch size B, the total amount of FLOPs to perform a forward and backward pass is 6nmB. Hence, the ratio of FLOPs of our 288mm//6mml = 48/B. This shows
learned optimizer's update w.r.t. to fawward-backwarr that the overhead of our optimizer decreases as the tasks get larger with more inputs flowing through the optimizee network with larger batch size B which also makes sense intuitively because as the tasks become computationally expensive, the cost of forward and backward pass dominates optimizer update cost. In this analysis, we omitted the cost of LSTM scheduler but as noted by prior work (Metz et al. 2022b a). the overhead of any tensor or global-level computation in the learned optimizer decreases as the number of parameters increase in the optimizee network since the cost of per-parameter updates dominate and tensor monsor/glo globall computation overhead asymptotically reduces to a fixed constant (Metz et al.) 2022b).. We compare wall-clock time of Celo with Adam and other learned optimizer baselines in Table runtime NVIDIA V100 GPU with 10 seeds per optimizer.
FLOPs
We benchmark
A.6 MLCommons AlgoPerf evaluation
We additionally evaluate on a few standard MT. Commons tasks from AlgoPerf benchmark (Dahl et al. 2023) which are much larger and completely out-of-distribution for meta-trained optimizers in this work. Figures and 10 show training and validation curves respectively, comparing Celo-Adam (Celo scheduler with Adam update rule discussed in Table 3) with Schedule-Free AdamW, the winning entry in the self-tuning track of AlgoPerf 2024 competition as well as tuned NAdamW baseline provided by the AlgoPerf team (Kasimbeg et al. 2025). For baselines, we directly use the results data from the AlgoPerf 2024 competition. AlgoPerf is a tereult- benchmark, where runs are halted either upon reaching the task's validation target or upon timing out. As evident from the results, Celo-Adam manages to optimize these relatively largo-scale unseen workloads despite being veta-trained on small image MLP tasks ($A.7). However, the final validation performance of Celo-Adam still lags behind the tuned baselines and it fails to hit validation targets earlier evident from the validation curves (Figure Lop. Morever, we find that Celo, pre-trained on simple image MIP tasks considered in this work, is unstable for ML.ommons tasks, hence we omit it from Figures g and 10. This highlights that learning a robust update rule is critical to scale fully learned optimizers lo large-scale workloads since learned schedulers can generalize, as evident from Celo-Adam results, even from limited meta-training. We defer the scaling up of meta-training for Celo to future work, which we believe is the right direction to make ready for standard use-cases especially given that our proposed recipe is compute-efficient and scalable A.4).
23

---

Published in Transactions on Machine Learning Research (06/2025)
runtime per step (ms)
evaluation tasks
adam
velo-s
velo
celo
0.20
1.69
1.76
1.31
22t.12212222122___oerrroom LayerNorm__eeeee
0.31
3.41
3.83
2.69
ImageMLP_Ccfar11_1188.
0.24
2.74
3.05
2.20
Conv_Ciffar10332x6xx66
1.12
2.80
3.02
2.44
CowvcCifar10..2264x66.bbottt
1.29
3.68
4.10
3.22
Cow._Cifarr1032266644__leeeee
1.41
3.78
4.11
3.21
Conv_Cifar100__32x66666
1.11
2.70
3.09
2.55
VIT_Cifar1000_wiceshhalll
6.73
43.89
46.00
39.47
VTT_Cifar100_sskimyyeerr
4.52
29.09
32.69
23.70
I
0.44
4.29
4.28
4.16
LmageMLPAE_Mnnist_IPPe
0.19
2.54
2.70
2.22
DNILX_imbss2k_PPtcchh___555 18
23.96
53.42
53.38
52.86
RNNL
Embed128
24.29
53.47
53.81
53.38
TransfoomerLLLLMMMM.MML.
0.71
8.30
8.71
7.10
TTrasformmerLKLLMMMMMMLLiiiiBnnd
0.99
14.50
15.46
12.94
TransformeeLLL_LMMBMMM
1.26
10.93
11.42
9.88
65.12
67.31
68.16
67.06
Table 7: Runtime per step for each task (ms). Cclo is faster in terms wallclock time per step than previous stale-of-the-arl learned optimizers such as VeLO (Metz el al.) 2022b) as well as VeLO with smaller hidden size (VeLO-S from $5). However, compared to Adam, Celo has additional overhead depending on the evaluation task. Notably, as tasks become more computationally expensive, this overhead diminishes, discussed in A.5
A.7 Meta-training tasks
We meta-trail all the learned optimizers this work on a set of small image MLP tasks from prior work (Metz) et al. 2022b order to keep the meta-training fast and perform all the ablation experiments within our compute budget. The same task set is also used by VeLO compare different learned optimizer architectures (Metz al. 2022b). It consists of the 4 tasks with different datascts namcly, Fashion MNIST, CIFAR-10, MNIST and SVHN. All the tasks share the same configuration which consists of 1-layer MLP network with 32 hidden units, 8 x 8 single channel image input, batch size 64 and ReLU activations. Link to code snippets here
A.8 Evaluation tasks
We evaluate Celo and all the optimizer baselines the following 17 tasks. Each task name contains link to its associated code snippet.
ImageMLP_Fastinn htwpp//_aaaaiionnnsseeeluu
Image classification task with Fashion MNIST dataset, MLP network with 2 hidden layers each containing 128 hidden units and ReLU activations, batch size 128, cross-entropy loss.
2. ImageMLP_Ciffar1__1 1t_0_11xx1222000 20xxx22xx222.... httpx//22_aaaeoorooorR..... Image classification task with CIFAR-10 dataset, MLP network with 3 hidden layers each containing 128 hidden units with layer-norm followed by ReLU in each hidden layer, batch size 128, cross-entropy loss.
3. CangeMLP_Ciffr110 1tt_rr00_11x1111________
Image classification task with CIFAR-10 dataset, MLP network with hidden layers each containing 128 hiddon units with Tanh activations, batch size 128, cross-entropy loss.
Conv_Cifar 1c_ffar1__3264x64 Image classification task with CIFAR-10 dataset, convolutional neural net- work with hidden layers containing 32, 64 and 64 channels with ReLU activations, stride 2 in first layer and stride 1 in rest. of the layers, batch size 128, cross-entropy loss.
24

---

Published in Transactions on Machine Learning Research (06/2025)
Fast. MRI
ImageNet ResNet 50
ImageNet ResNet-50 Large BN Init
0.40
25
15.0
200
Train lose
0.35
Train loss
10.0
15
Train
7.5
20
0.30
5.0
2.5
0.25
0.0
4000
8000
12000
150000
300000
450060
150000
300000
450000
Iterations
Berations
Iterations
ImageNet VIT
Librispeech Conformer Layernorm
WMT
3.0
4.0
2.5
Train cTC Joss
101
3.5
Train less
3.0
1.5
Train
2.5
1.0
2.0
0.5
10-1
1.5
1.0
150000
300000
450000
30000
60000 Iterations
90000
120000
80000
160000
240000
Iterations
Iterations
celo-adam
schedule-free adamw
nadamw baseline
Figure 9: Mommons AlgoPerf train curves.
Fast MRI
ImageNet ResNet-50
ImageNet ResNet-50 Large BN Init
0.75
0.70
Top-1 val accurasy
0.6
Top-1 Souuz5y
0.6
SSIM
0.
0.60
0.2
0.2
0.30
0.0
0.0
4000
8000
12000
150000
300000
450000
150000
300000
450000
nerations
Rerations
Itarations
ImageNet ViT
Librispeech Conformer Layernorm
WMT
30
Vil socurayy
0.6
10a
Val BLE
20-
0.4
Val
10
0.2
10-l
150000
300000
450000
30000
E0000
90000
120000
80000
160000
240000
lterations
Iterations
celo-adam
schedule-free adamw
nadamw baseline
Figure 10: MLCommons AlgoPerf validation curves. We compare Celo-Adan with winner in the self- tuning track of AlgoPerf 2024 competition Schedule-Free AdamW and prize qualification baseline NAdamW which is heavily tuned (Kasimbeg al. 2025 Dahl et al., 2023). Although Celo-Adan is just meta-trained small image MLP tasks ($A.7). it manages to optimize much larger tasks, in some cases, even better than Schedule- Free AdamW (ImageNet ResNet tasks). However, it still lags behind baselines in terms of final performance, a gap that can be potentially covered by improving and scaling up meta-training.
5. Conv_Cifar10 1Ffffrr10_22xxx64 44064x64__bbtchormm Image classification task with CIFAR-10 dataset, convolutional neural network with 3 hidden layers containing 32, 64 and 64 channels with ReLU activations and batch-norm, stride 2 in first layer and stride 1 in rest of the layers, batch size 128, cross-entropy loss.
[onv_Cifarie 1e1eeee_xxxxx4__aeeee_
Image classification task with CIFAR-10 datasot, convolutional neural network with 3 hidden layers containing 32, 64 and 64 channels with ReLU activations and layer-norm, stride 2 in first. layer and stride in rest of the layers, batch size 128, cross-entropy loss.
25

---

Published in Transactions on Machine Learning Research (06/2025)
7. Sonv_Cifar100 33x66666_44 Image classification task with CIFAR-100 dataset, convolutional neural network with 3 hidden layers containing 32, 64 and 64 channels with ReLU activations and layer- norm, stride 2 in first layer and stride in rest of the layers, batch size 128, cross-entropy loss.
8. VIT_Cifar100 hw..iffartoowwidehhallllooo Image classification task with CIFAR-100 dataset, vision transformer model with 6 transformer layers cach containing attention heads, 16x16 patch input, hidden size 384, MLP expand factor 4, batch size 128, cross-entropy loss.
VIT_Cifar1 Sarraarr@e__sinnydeeeeee Image classification task with CIFAR-100 dataset, vision transformer model with 10 transformer layers each containing 4 attention heads, 16x 16 patch input, hidden size 128, MLP expand factor 4, batch size 128, cross-entropy loss.
10. LransfoormerLM_MMB Lww.BB/Mliiuuunime....mmm Language modeling task with LMIB dataset, transformer de- coder model with layer containing 5 attention heads, hidden size 20, sequence length 8, vocabulary size 32K, batch size 4, cross-entropy loss.
11. TransformerLM_ LMMBr__erLMMMMMBBMMMMMM__ MaMMMMBBMMMiiuuiiieee. tiRuntimeeme22 Language modeling task with LMIB dataset, transformer de- coder model with 2 layers each containing 4 attention heads, hidden size 32, sequence length 8, vocabulary size 32K, batch size 8, cross-entropy loss.
12.
TransformerLM__MBB Bttp:/meerrMMMMBBMMMLMLAEAIIIECE tttituumutiuunimme555 Language modeling task with LMIB dataset, transformer de- coder model with layers each containing 4 attention heads, hidden size 32, sequence length 32, vocabulary size 32K, batch size 8, cross-entropy loss.
13.
ImageMLPAE__iife SwAPPPE_C_aaaaaee 1al.ffar10112xxx2240 12xx1xxxx2xxxx 28 bs256 Auto-encoder task with CIFAR-10 dataset, MLP network with shape 128-32-128 128 and ReLU activations, batch size 256, MSE loss.
14.
ImageeMPPA__MMiss__12222222 b_1188 Auto-encoder task with MNIST dataset, MLP network with shape 128-32-1222 and ReLU activations, batch size 128, MSE loss.
15.
Embeedd Language modeling task with LMIB dataset, LSTM re- current network with embedding size 128, hidden size 256, sequence length 32, vocabulary size 32K, batch size 128, cross-cntropy loss.
16.
251 Language modeling task with wikipedia dataset, LSTM recurrent network with embedding size 128, hidden size 256, sequence length 32, vocabulary size 32K, batch size 128, cross-entrrop loss.
RNNLM_
17. BaaaraccPPPOOp__Fsshhoionnt
Learned optimizer meta-training task using Fashion MNIST dataset, MLP-based learned optimizer (Adafac MLP LOpt), 1-layer optimizee network with 32 hidden units, 8 image input, maximum inner unroll length 50, outer batch size 2, inner loss cross-entropy, mean loss of inner loop as the mela-objective.
26

---

Published in Transactions on Machine Learning Research (06/2025)
A.9 Evaluation plots
ImageMLP Mushionniiss Reiu 28x126
ImageMLP Cifarl0 LayerNorm
ImageMMP 128x128x128 bs128
24
22x
10
18x
16x 10
14
1000
2000
300
1500
Conv Clar10 32x64x64
Conv Cifar10 32x64x64 betcheorm
Coev Cifar10 32x64x64 layernorm
3
2x100
1.4
101
2000
Cony, Cifar100 32x64xx64
VIT Cifar1 00 wideshallow
VIT Cifar100 skimnydeep
1.3
10a
30
10
10
10-1
10
10
500
1000
1500
2000
303
1000
1300
2000
0
500
1000
1500
2000
ImageMLPAE Cifar10 128x32x126 bs256
ImsgeMLPME Mnist 128x32x128 bs128
RNNLM im:1332k Patch32 LST Embed120
30.1
xx
30.2
10-2
500
1000
1500
2000
300
1000
1500
2000
302
1000
1500
2000
wisipediane.32k
Transfocme:1.M
TransTormrrL.
2 10!
387
6x
207
Transformeri.M LM1B Multiuuntime 5
LOpt AdsfacMLPPOpt PrshionMraitt 50
4x
3x.10"
8x.07
2x
x.20
6x107
16a
3000
303
Iterations
NNAadam
L0pl
LOpt.
Velo-S
Figure 11: All train curves.
27

---

Published in Transactions on Machine Learning Research (06/2025)
ImAgeMGP MoshionMnist Relu128x128
ImageMPP Cifar10 128x128x128 LayerNorm Rola
ImageMLP Clar10 128x 12xx28 Tanh bs128
2.6
10"
24
12x
101
107
18x10
10"
6x30
4x30
L6K10
1011
3
500
1000
1500
2000
500
1000
1500
2000
500
1000
1500
2000
Cifar10 32x644664
Cifar10 32x644664 batchnorm
Cifar10 32x64x64 layernorm
101
2x 10
100
10
500
1000
1:00
2005
500
1000
1500
1500
VIT Cl wideshallow
skinnydeop
107
12a
10
ImageMLLPAE Cifar10 128x32x128 bs256
ImageMLPAE Mrist: 128x322128 bs128
RNNLM Jm1b32kk Patch32 LSUM256 Embed120
101
9x 109
10-4
8x 109
10-9
10"
10
5x 10
1500
2000
1000
1500
1500 TransformerLMM LM1B MultiRuatime 2
RNN Willpediaen322 Pateh32 LSTM256 Embed128
TransformerLM LM1B MaltiRuntime 0
18
101
107
500
1000
1500
2000
300
1000
1500
2000
500
1000
1500
2000
Transfomerr.. MultRnntime
LOpt AdafacMM.PL.pp. PashienMaist
4x 10
201
3x 10"
107
107
12
100
1000
1500
2005
500
1000
15:0
2000
NNAadam LOpt
RNN MLP LOgt
Adafac MLP LOpt
Velo-S
Celo
Figure All test curves.
28

---

Published in Transactions on Machine Learning Research (06/2025)
A.10
Evaluation plots with task augmentation
ImageMLPP
ImegeMLP Cifar10 128x128x126 LayerNorm
ImageMLLP be128
23
2x100.
22x
2.1 K12
1.5
500 Cony
1500
5.00
1000
1500
2050
5:0
1000
1500
2010
Cifari 32x014x64 batchnorm
Cony
32x654x64 layermrrm
x 10"
8.70
x 10
x1
1x
1.x
14x10 10
100
12
Conyy 3dfar100332xt4xx66 32xb44x64
VIT Cifar100 widesballow
VII Cifar1 00 skincydeep
101g
107
10
109
PMA. Cifar1 128x212128 28 ha256
ImageM1.PAF Mhist 128x32X1 128 bs128
RNNLM Im h32k Patch32 LSTM2356 Embed1 28
9x 10
10"1
13-1
Bx
7x
8x
Ex 12.
4x
3x.0
5x
1563
2000
1000
1500
2000
RNNLM wckkpedinen332 Potch32 LSTM256, Embed128
Transformert. LMIB Multitantime 0
Transformert. M. LM1B Multi Runtime 2
19
9x100
100
7x10'
500
1000
1500
2000
500
1000
1302
2000
500
1000
1500
2000
TransformecI.M LMBB MulLRuntime
LOpt AdstecM1PLOppt RasitionMnist 50
4x10
107
3x100
8x.00
1000
x 100
10P
500
1000
1500
2000
500
1000
1509
2030
Iterations
RNN MLP LOpt wlo Aug
RNN MLP LOpt Aug
Figure 13: RNN MLP LOpt augmentation training curves.
29

---

Published in Transactions on Machine Learning Research (06/2025)
ImageMIP FashionKnist Relal 28x128
ImageMLP Cifari 128x128x128 LaverNorm. Relu
ImageMLP Cifar10 Saarr1011222211 126x12xx128 Tanh bs128
1
24
24x
12x12
10a
10
2x30"
16x12
15K10
14x
14x
1.7
103
1000
15.00
2005
503
1000
20:00
2000
2000
15000
2000
Cont Cifar10 32x61xx61
Coay Cifar10 32x64x664
Cony
32x64x66
14
10"
18x 10"
10
16x
14
100
12x10 10
1000
15.00
2000
503
1070
1000
2000
1000
3550
2000
Cifar100 32x014xt64
VIT 00 stdeshallow
CifaCttarr00 00 skmnyydoop
107
100
10
10t4
10
10
10
1500
2000
Train
ImageMi. BAEL Cifar1 128x32x128 bs258
ImageMLPAL PAE Maist, 128x32xd 28, bs128
RNNLM Jm1b32k Patch32 LSTM255 Embed120
10'
8x
3500
2000
1000
15:00
2005
RNNLM wekkpediaen32k Patch32 LSTM256 Embed128
PransformetL. LM1B Multiluatime
Transformerl..M LM1B MultRuattmee 2
101
10
9x10"
00
$* 107
1x
3x 107
503
1000
15:0
2000
303
1000
1520
2000
303
1000
1500
2000
Transforme:LM LMTB MultiRuatime 5
LOpt AdelscMLPLOpt Mashion 50
103
3x 10
2x 10?
80"
7x.00
6x.00
503
1000
3500
3000
50b
1000
15:00
2000
Adafac MLP LOpt w/o Aug
Adafac MLP LOpt w/ Aug
Figure 14: Adafac MLP LOpt augmentation training curves.
30

---

Published in Transactions on Machine Learning Research (06/2025)
ImAgeMGP PoshionMnist Relu128x128
ImageMLP Cifar10 128x128x128 LayerNorm Rola
ImageMLP Ciar10 128x 18x128 Tanh bs128
26611
24
107
10
137
10
14
300
1000
1500
2000
500
1000
1500
2000
500
1000
1509
2000
Cifar10 32x644m64
Cifar10
32x644x644 batchnorm
Cifar10 32x54x64
2x 30
2x 10
309
500
1000
1500
2005
500
1000
15:0
1000
1500
Cl wideshallow
C1far100 skinnydeop
10"
10-1
10-1
13"4
ImageMLPPEE Cifar10 128x32x128 bs258
ImageMLPAE Miistt 128x32x128 bs128
RNNLM Jm1b32k Patch32 LSU M256 Embed120
5x.00
104
4x10-
10t4
1500
2003
1000
1500
1500
RNNLM wilipediae322 Pateh32 LSTM256 Embed128
TransformrLL LM1B MultiRuntime D
TransformerLM LM1B MultiRuatime 2
3x
10
10
9x 10
10"
201
10
107
7x 307
500
1000
3500
2000
300
1000
1500
2000
500
1000
1500
2000
Transfomerr.M 1MIB MultRnntimee5
AdafacMI PashienMnist
101
107
10
1
12a
100
1000
1500
2005
500
1000
15:00
2000
Volo-S wlo Aug
Velo-S w/ Aug
Figure 15: Velo-S augmentation training curves.
31

---

Published in Transactions on Machine Learning Research (06/2025)
Ima PashionMaist Relu126x128
ImageMLi Cifar10, 28xx120x128 LayerNorm Rela
Image MLLP Cifar10, 128x128x128 Tanh bs128
2.6
2.4x 10*
23 100
x2x x0"
2x
18x 10"
1.6x 10
500
1040
1500
2000
500
1000
1500
2000
569
1000
1509
2000
Cifar10 32x64x66
Cons Cifar10 32x644664 batchnorm
Cifor10 32x644 layamorn
10
101
500
1000
150n
2000
500
1000
1500
500
1000
1500
2000
Cony
32x644x644
Cifar100 wideshallow
00 skmnyydopp
10a
10a
10-
10.
10-4
10-2
1044
10
Train
ImageMLPAE Ctfarro_128x311221_ 128x32x128
b bs236
ImsgeMLPAR Mrost_ 128x322128 bs128
RNNLM Jm1b32k Patch32 1ST2256 Embed120
101
10-1
80
1x
6x
10-2
10P
1000
1300
3300
2000
RNNLM wieppeiaenn32 Patch32 LSTX256 Embed1 28
Transformer_. LM1B MultRuntime
Transformerl..M LM1B Multituatimee 2
10
101
1x"
10
10
10"
107
500
1000
1.500
2000
500
1009
1300
2000
300
1000
3300
2000
L LMIB MultRantime
LOpt AdelacMLPLOpp Fashionnaist 50
101
3 10"
10"
10
6x
500
1000
1500
2000
500
1005
1500
2000
Iteeations
Colo w/o Aug
Colo w/ Aug
Figure 16: Cclo augmentation training curves.
32

---

Published in Transactions on Machine Learning Research (06/2025)
A.11
Celo schedules
ImageMLPP FashionMaist Mels128x128
ImageMLPP Cifario 128x128x120 LayerNorm Relu
ImageNLPP Cifar10 320x128x128 Tanh bs128
107
10-1
10
250
500
750
1003
1250
1500
1750
2000
1000
1250
1500
1750
2000
250
520
750
1000
1250
1500
1750
2000
Cifar10
Oifar10
layecnorm
10-1
10
10
250
300
750
1000
1250
1500
1750
2000
150
500
750
1000
1250
1500
1750
2000
250
500
750
1000
1250
1500
1750
2000
Codw Cifar100, 32x64x64
VIT
CI wideshallow
VIT Clear 10 skinaydeep
107
101
10
1000
1250
250
500
1230
ImageMLPPAE Cifar10,
128x32s128 ba256
ImaguMMl MILPAE, Mnist, 128:3221128 bal18
RNNLM Im 1a332 Patch332
1STT2256.
Embed122
10*
10
260
500
750
1000
1250
1500
1750
2000
1250 1600
2000
250
500
750
1000
1250
1500
1750
2000
wikipeddsee332
TransformerLLM
TransformerLM
107
10*
10
10
255
500
750
1003
1250
1500
3750
2000
250
500
750
1000
1250
1500
1750
2003
0
250
300
720
1000
1250
1500
1750
2000
TransformerLM
PLopt RashionMaist
10a
10-4
250
1003
1250
1500
3750
2000
1750
Optimtoor Langet length
0.2x
0.4%
0.6x
0.8x
1x
Figure 17: Celo learned schedules for all tasks.
33


Reviews:
[{"review_id": "PJCZoA2CES", "reviewer": "Reviewer_2iuG", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "S1 - Clear Methodological Innovation: The paper proposes an elegant and efficient L2O architecture named \"Celo,\" carefully designed to be computationally efficient while achieving strong generalization. The methodology clearly identifies and effectively addresses critical issues like truncation bias by employing a two-stage training strategy, clearly separating update rule learning from scheduler learning. This allows the optimizer to achieve state-of-the-art performance on diverse out-of-distribution tasks, demonstrating robustness and efficiency in practical scenarios.\n\nS2- The authors provide extensive experimental results, including comparisons with both learned and hand-designed optimizers across various optimization scenarios. The introduction of robust evaluation metrics derived from reinforcement learning literature, such as meta-generalization measures, provides a more reliable assessment of optimizer quality beyond traditional benchmarks. These careful, thorough experiments add substantial credibility to the claims.\n\nS3 - The methodological innovations are supported by theoretical justifications, particularly regarding the convergence properties and generalization behaviors of the learned optimizer. By clearly identifying critical aspects such as truncation bias, and proposing targeted strategies like curriculum-based training, the paper significantly advances theoretical understanding in the Learning to Optimize (L2O) community.\n\nW1 -  The proposed two-stage training and task augmentation methods bear close resemblance to established curriculum and imitation learning approaches introduced in \"Training Stronger Baselines for Learning to Optimize\" (Chen et al., NeurIPS 2020). Despite this similarity, the submission did not discuss, cite, or directly compare against this prior work.\n\nW2 - While task augmentation is central to achieving strong generalization, the submission does not rigorously investigate how variations in the task distributions (e.g., complexity or diversity of augmented tasks) directly influence performance. A targeted ablation study explicitly varying augmentation settings, along with analysis of sensitivity to task sampling strategies, would be essential to validate claims and strengthen methodological rigor.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The paper introduces Celo, a compute-efficient learned optimizer designed to achieve strong meta-generalization using minimal computational resources. By identifying key ingredients such as task augmentation, a decoupled two-stage training procedure (first learning a parameter update rule, then a scheduler), and simplified architecture, the proposed optimizer significantly outperforms prior learned and hand-designed optimizers on diverse out-of-distribution tasks. Additionally, the authors introduce robust evaluation metrics derived from reinforcement learning literature to reliably measure performance and generalization of optimizers.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "Please address W1 and W2. \n\nBesides, authors are recommended to check more benchmark settings in \"Learning to Optimize: A Primer and A Benchmark\" (JMLR 2022)", "broader_impact_concerns": "NA", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "aPHXiJeOt1", "reviewer": "Reviewer_j6jn", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n- The major selling point of the paper is that it is able to meta-generalize from a few datasets and toy models to unseen datasets and large models.\n- The experiments are thorough. The paper ablated settings such as learned update, scheduler, two stage training, etc. The paper evaluates MLP, conv net, and LM transformers on different datasets and benchmarks.\n- Task augmentation is a novel idea, and is empirically effective.\n- The experiments contain many unseen test datasets ran with multiple random seeds.\n\nWeaknesses and Questions:\n- The metrics are based on the number of iterations, but the learned optimizer could incur more computation. It would be good to clarify on the wall clock time implication.\n- It is surprising that the paper is able to generalize so well on 8x8 images with a one-layer MLP network. It could also be interesting to test what is the limit on such meta-training and meta-generalization. Can any datasets work? Can the datasets be synthetic?\n- Lastly, I would love to see if there are some standard benchmarks with a standard architecture, where we can reach a known accuracy (e.g. ImageNet) with less number of iterations. The CNN models evaluated in the paper are fairly small.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The paper proposed a new meta-optimizer that can be meta-trained on a small computation budget. They proposed a decoupled design of optimizer and scheduler, and a unique way for performing task augmentation for better meta-generalization. The proposed algorithm is able to meta-optimize from a small number of datasets and on a toy model architecture using a very small number of GPU hours. The resulting meta-optimizer outperforms existing optimizers on many unseen datasets and models.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "- Add wallclock time\n- Add more standard benchmarks (e.g. ImageNet) with standard architecture (e.g. ResNet 50)", "broader_impact_concerns": "N/A", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "WieMewUqiq", "reviewer": "Reviewer_Yef6", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n\n* 4000h TPU -> 24h GPU is an impressive achievement, but I think the results need more comparison about the actual cost of the competitor optimizers.\n\n* The discovery that \"task augmentation\" is key for generalization.\n\n* Outperforming a strong suite of learned and handcrafted optimizers.\n\n* Nice set of ablations in Table 3.\n\nWeaknesses:\n\n* The experiments are not scaled, so the gains from the current scales may not transfer to larger scales.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "* Data: Data augmentation to effectively simulate new tasks by doing a reparameterization trick on the optimizee's parameters.\n\n* Optimization: Decoupled, multi-stage training of per-parameter (local) updates and scheduler scalar (global) learning rates.\n\n* Metrics: Final & speedup IQM scores for the task: Inspired from reinforcement learning (Agarwal et al., 2021), the authors replace the human baseline with an Adam baseline, normalize the metric of interest by the baseline, and sum over the meta-learning tasks and the meta-learning steps.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "We need to see costs of the optimization runs in Figure 4 (and corresponding result artifacts).\n\nminor: \n\nlinear projecting <- linearly projecting\n\nbest Adam baseline <- the best Adam baseline\n\nlowest loss <- the lowest loss\n\nEtc. (please check typos and grammar)", "broader_impact_concerns": "None", "additional_comments": "", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: This paper introduces a new meta-optimizer called Celo that appears to train much faster than competing alternatives due to various ideas explored in the paper including task augmentation, and use of learnt schedulers and update rules. It outperforms 15 hand-designed optimizers, including Adam as well as other meta-learning baselines in comprehensive experiments.
Decision: Accept (poster)