Celo2: Towards Learned Optimization Free Lunch
Authors: Abhinav Moudgil
Boris Knyazev
Eugene Belilovsky
Venue: ICLR 2026 Poster (2026)
Keywords: learned optimization
meta-learning
Primary area: transfer learning, meta learning, and lifelong learning

Abstract: Learned optimizers are powerful alternatives to hand-designed update rules like Adam, yet they have seen limited practical adoption since they often fail to meta-generalize beyond their training distribution and incur high meta-training cost. For instance, prior work, VeLO, scaled meta-training to 4,000 TPU months ($\sim$10$\times$ GPT-3 compute) to meta-train a general-purpose optimizer but it failed to generalize beyond 600M parameters tasks. In this work, we present a surprising finding: by crafting a simple normalized optimizer architecture and augmenting meta-training, it becomes feasible to meta-train a performant general-purpose learned update rule on a tiny fraction of VeLO compute, 4.5 GPU hours to be precise. Our learned update rule scales stably to a billion-scale pretraining task (GPT-3 XL 1.3B) which is six orders of magnitude larger than its meta-training distribution. Furthermore, it shows strong performance across diverse out-of-distribution tasks and is compatible with modern optimization harness that includes orthogonalization, distinct update rules for input-output and hidden weights, and decoupled weight decay. In all, this work paves the way for practically applicable _learnable_ optimization algorithms, unlocking exploration of richer meta-training and data curation recipes to further improve performance.

Published as a conference paper at ICLR 2026
CELO2: TOWARDS LEARNED OPTIMIZATION FREE LUNCH
Abhinav Moudgil 1.2.+ Boris Knyazev!.44
Eugene Belilovsky 1.2
Mila
CConorddaa University
'Samsung AI Lab, Montreal *Université de Montréal
ABSTRACT
Learned optimizers are powerful alternatives to hand-designed update rules like Adam, yet they have seen limited practical adoption since they often fail meta- generalize beyond their training distribution and incur high meta-training cost. For instance, prior work, Vel.O, scaled meta-training to 4,000 TPU months (~ 10x GPT-3 compute) to meta-train a general purpose optimizer but it failed to gener- alize beyond 600M parameters tasks. In this work, we present surprising find- ing: by crafting simple normalized optimizer architecture and augmenting meta- training, it becomes feasible meta-train performant general-purposs learned update rule on tiny fraction of VeLO compute, 4.5 GPU hours to be precise. Our learned update rule scales stably to a billion-scale pretraining task (GPT-3 XL 1.3B) which is six orders of magnitude larger than its meta-training distribution. Furthermore, shows strong performance across diverse uut-of-distribution tasks and is compatible with modern optimization harness that includes orthogonaliza- tion, distinct update rules for input-output and hidden weights, and decoupled weight decay. In all, this work paves the way for practically applicable learnable optimization algorithms, unlocking exploration of richer meta-training and data curation recipes further improve performance.
general-purpose
GPT-3 (1.3B) FineWeb
5.0
Adam
4.5
Colo2
VeLO
4.0
Val
3.5
3.0
2.5
1B
2B
3B
4B
5B
6B
7B
8B
9B
10B
Train tokens
LM (30M) FincWeb-Edu
GPT-2 (124M) Fin
5.0
4.50
4.25
4.00
3.75
3.5
3.25
0
120M
240M
360M
480M
600M
0.58
18
1.58
28
2,58
Train
Train tokens
Figure 1: Celo2, our learned update rule scales stably to large-scale pretraining tasks. Our learned update meta-trained on very limited compute budget (4.5 GPU hours) but generalizes to 1,000,000x larger language modeling task (GPT-3) and outperforms strong well-tuned baselines such as Adam (Loshchilov] 2017]] Kingma & Ba (015)) and VeLO (Metz et al.) 2022b). In con- trast, VeLO (Metz et al. 2022b), meta-trained with exorbitant compute (4000 TPU months), fails to generalize. All the language modeling tasks have a modern architecture that includes RMS normal- ization, rotary positional embeddings (RoPE), QK-norm, and GELU (Marek et al. 2025
"Correspondence Abhinav Moudgil <abhinawmoudgil.5ea wwwp::mmooddi199900....mmm

---

Published as a conference paper at ICLR 2026
INTRODUCTION
Pre-training of foundation models has begun to heavily dominate workloads in the past years, pro- ducing progressively more powerful models at the cost of great computational resources. Existing optimization strategies used at scale often rely on the now classic Adam optimizer (Kingma & Ba, 2013). As result, progress made in the direction of improving optimizers for these large-scale tasks thus shows great promise in reducing the computational burden of pre- tre-training. One direction that has yielded strong initial results learned optimization (Metz al.) 2019b): leveraging learning to improve the optimization process itself. The field of learned optimization has progressively seen more practical results since being introduced (Andrychowicz et al., 2016). However, progress in this area has not been as fast-pacs as that in pre-training, primarily due to the inability of learned optimizers to generalize to data that lies outside their training distributions.
The most powerful learned optimizer to date is the Versatile Learned Optimizer (VeLO) (Metz et al. 2022b), which uses a hierarchical architecture for the network representing the optimizer and then performs meta-leaming at immense scale to achieve generalization across tasks. One key obser- vation regarding VeLO is its high computational requirement: VeLO meta-training used 4000 TPU months' worth of compute meta-learn the optimizer for achieving high task coverage. It could be argued that this would be partially or fully offset by the fact that Vel.O outperformed Adam across a wide variety of tasks, becoming a prototype of the vast potential of learned optimization. However, while the performance gains over Adam were significant, they did not scale to the training of large foundation models like LLMs with billions of parameters. We believe that this result presents great opportunity for the learned optimization community to tackle the problem of generalization at scale. More recent work (Moudgil et 2025 Thérien et al., 2024. Knyazev al. 2025 in this direction, notably Celo (Moudgil et al. (025). attempts to address this through architectural modifications and training recipes aimed at improving meta- vetaageneralization While successful in significantly reducing meta-training compute, these approaches show some performance degradation compared to VeLO, which remains the state-of-the-art learned optimizer at scale.
In this work we propose Celo2, a simple recipe for meta-training stable learned optimizers that significantly improves meta-gent gestaaeneralizttin, and is able to outperform VeLO on large-scale tasks while meta-learning on orders of magnitude less. Our proposed recipe aims to tackle the three key challenges: generalization, scalability, and primary objective while developing this simple approach for learned optimization was stability (SA):: we hope that this work eventually evolves to the level that it can be used directly in the training of foundation models. Finally, we highlight a key finding: we do not explicitly scale up meta-training for performance in order to keep meta-training compute-efficiency, however, we obtain it as by- product of following the Celo2 recipe, essentially offering a "free lunch" from limited compute.
We would like to emphasize that our
Our key contributions are as follows:
We present Celo2, a stable learned optimizer that achieves strong meta- generalization to billion-scale tasks (GPT-3 XL 1.3B), outperforming prior sfatc-of-the- learned optimizer VeLO (Metz et al. 2022b). Code is available at: https:://iithb. com/ amoodgl/ceeel
Our proposed meta-training recipe ($3) is compute-efficient and we demonstrate that even with few GPU hours of meta-training, it demonstrates strong performance out-of- distribution standard machine learning tasks in language modeling (GPT-2, GPT-3, LM- 30M), vision (ViT ImageNet), and reinforcement learning (Atari PPO) domains.
We conduct systematic ablations keeping meta veta-generalizztion pose Celo2-base a computationally lighter variant of Celo2 that trades some performance for improved efficiency while maintaining core parts of our generalization recipe.
the key focus and also pro-
2 RELATED WORK
Hand-designed optimizers have long been the cornerstone of neural network training, with meth- ods like SGD with momentum (Sutskever et al.] 2013) and Adam (Kingma & Ba 2013) providing robust baselines through adaptive learning rates and momentum terms. To further enhance conver- gence and stability, second-order optimizers such as Shampoo (Gupta et al.] 2018)) introduce pre- conditioning by maintaining running statistics of gradients in each tensor mode, effectively rescaling
2

---

Published as a conference paper at ICLR 2026
updates account for curvature in high- dighndimensional spaces. Similarly, SOAP (Vyas et al., 2024 explores optimization in modular norms, allowing for flexible preconditioning that adapts to the ge- ometry of parameter spaces. Recent advancements have incorporated orthogonalization techniques to mitigate issues like ill-conditioned gradients; for instance, Muon (Jordan et al.) 2024) applies a Newton-Schuly iteration approximately orthogonalize momentum- sasrenttu----eddddd updates, projecting them onto the nearest semi semi-orthogonal matrix to amplify rare directions and improve sample efficiency, as demonstrated in speed-running benchmarks like NanoGPT and CIFAR-10. Complementary work by Tuddenham et al. (2022) proposes gradient orthogonalization via SVD, showing speedups in neural network optimization for SGD with momentum. These approaches highlight the value of precondi- tioning and orthogonalization in hand-crafted optimizers, but they often require manual tuning and lack the adaptability of learned methods.
Learned optimizers (LOOs) aim to meta-learn update rules from data, potentially surpassing hand- designed counterparts by discovering task-agnostic strategies. Early efforts, such as those by Andrychowicz et al. (2016), framed optimization as a recurrent process, but scaling remained chal- lenging. State-of-the-art LOs like VeLO (Metz et al. 2022b) represent a breakthrough, training large hierarchical LO on thousands of diverse tasks to achieve broad generalization, though at immense computational cost (4000 TPU-months). Other works have focused on efficiency: [Metz] al. (2022a) analyze trade-offs in L.O design, proposing simple MLP-based optimizer augmented with strong features like Adafac (Adafac MI.) balancing memory and compute while maintaining competitive performance. More recently, pLOs (Thérien et al.. 2024) introduce micro-scale learned optimizers tailored for vision and language models, emphasizing compute-efficient meta-traing training. These works underscore the potential of LOs but reveal persistent limitations in metaageneralizattoon i.e. the ability to apply learned rules to unseen tasks or longer optimization horizons (unrolls). For example, many LOs struggle with extrapolation beyond meta-training distributions, leading in- stability on outrof-distribution problems or when scaling to models with > IB parameters. Celo (Moudgil et al. 2025) addresses some of these by optimizing LO architectures and meta-training protocols for better generalization with minimal compute (24 GPU-hours). outperforming tuned baselines on diverse tasks, yet it stops short of integrating advanced normalization or orthogonal- ization for further stability. Building on these foundations, work integrates insights from both hand-designed and learned optimizers. Unlike prior LOSs that overlook robust normalization for optimization stability, we introduce specialized normalization scheme that enhances generaliza- tion across tasks and extends unroll lengths. Furthermore, by incorporating the Newton-Schulz orthogonalization from Muon into learned framework inspired by Celo, we achieve superior meta- generalization and enable stable and performant training of larger models, bridging the gap between efficient meta-traine veta-training and practical scalability.
APPROACH
We introduce a new compute-efficient learned optimizer, Celo2 (Alg. []. that is just meta-trained for 4.5 GPU hours and generalizes to practical large-scale tasks such GPT-3, ImageNet-Vit and Reinforcement Learning. The key ingredients of our approach are discussed below:
Tunable step-size. Unlike prior work in learned optimization (Metz al., 2022a b Moudgil al., 2025), our approach solely focuses on learning the update rule, thus decoupling step size tuning from update rule learning. This does result in an extra tunable knob in the form of step size. However, we find that with this decoupling, the learned update rule meta- vrct-trriined on small-scale tasks generalizes surprisingly well to large-scale tasks. This provides a scalable approach to improve the optimization update rule used to train deep networks. This effect of decoupling step size tuning from learning the update rule is also noted by prior work such as Celo (Moudgil et al. 2025). Since Celo learns the scheduler as well as the update rule, it fails to generalize large-scale tasks. We keep the step size tunable as per the user and leave the learning of a general scheduler for all tasks to future work.
Simple design that scales. Building on prior work (Metz et al. just learn a small MLP as a drop-in replacement for Adam in large-scale distributed and sharded training setups ().). We ablate its design in Section [3] to improve meta- gentraaenraliioiooon and perfor- mance, and provide full details on its inputs, outputs, and architecture in Appendix [ The MLP can also be meta-trained on specific tasks; notably, we find that meta-training on simple 8x 8 image-
2022a b Moudgil et al. 2025), we
3

---

Published as a conference paper at ICLR 2026
Algorithm Celo2
Input:
0,
optimizee parameters
V0,
gradients
learning rate
8:
optimizer state
(optional) weight decay
Require: learned update rule fuip
accumulators update function face
1: 81+1 (ac( v0 ; Le, 1)
update accumulators state
2: initialize next params 01+1 + Q
3: for each tensor with params pt 0g in parallel do
D parallel scan over tensors
4: prepare per-param features F for pt using updated state 81+1
5: Apt Fmip(F)
D learned update rule forward pass
6: Ap: NewtonSchulzs(App)
Orthogonalization D> RMS normalization update params gather params
7: Apt- - Apt/ ||Ap: ll rms
8:
Pi+1 - - npAPp aap:
9: 0(+1- - 0.41 P1+1
return 81+1
classification tasks ($4) suffices learn a descent rule that generalizes. This provides scalable path for advancing opumization algorithms.
Normalized learned update. Prior work in learned optimizers (Metz et al. 2022a Harrison al. 2022), which learn per-parameter update rule like ours, normalizes inputs F (Alg. ! to the MLP by RMS of each feature across the parameter tensors. However, all of them directly use raw output from the MLP as the step update. We find that simply normalizing the oulputs of the MLP update, thus resulting in fully-normalized learned update rule, generalizes well to large-scale tasks. This not only forces the MLP during meta-trane tetta-trannigg learn task-invariant update rulc (instead of raw MLP outputs) but also results in a similar optimization dynamics as AdamW as demonstrated in Figure 2 The challenges of learning such MLP update rule that operates directly
gradients have been discussed in prior work (Almeida et al., 2 Metz et al.) 2019a). Several concurrent works (Liu al. 2025 Si et 2025) also suggest the benefits of using RMS scaling of the step in order to match the learning rate of AdamW which is used to optimize ID parameters. We extensively ablate different normalization choices in Section A (appendix) including RMS norm clipping used in AdaFactor (Shazeer & Stern] 2018), RMS accumulation, etc. Empirically. we find that simply normalizing by RMS of the current step works the best. Note that this RMS normalization plays key role in Celo2-base that uses the learned MI.P update for all params; in Celo2, we orthogonalize the ML. update (Step 6, Alg. []. making it effectively shape-dependent (Liu et al.) 2025).
Compatible with modern optimization harness. Celo2 is highly compatible with modern tech- niques such as orthogonalization (Jordan al. 2024 Gupta et al., (018)), distinct update rules for 1- D and 2-D params (Jordan et al.) (224)). and decoupled weight decay (Loshchilov] (017), demon- strated by our experiments we describe in the next section. Since Celo2 a learned update rule, this provides a scalable path to improving optimization algorithms for deep neural networks by simply leveraging the data and compute while building upon and extending recent advances in optimization area.
4 GENERALIZATION AT SCALE
Meta-training. Following prior work (Moudgil et al. 2025). we meta-train our update rule on image MLP classification tasks consisting of mixture of datasets such as MNIST (LeCun & Cortes, 1998 Fashion-MNIST (Xiao et al. 2017), CIFAR-10 (Krizhevsky et al. 2009) and SVHN (Netzer al. 2011). The images are resized to 8x 8 in the prior work for faster meta-training, thus yielding a high research throughput. We use Persistent Evolutionary Strategies (PES) (Vicol et al, 2021) to meta-train our optimizer with unroll length logarithmically sampled between 100 and 2000 steps. PES yields unbiased gradients with truncated training. Truncated training which updates the learned optimizer every K steps (we K=50 in experiments) during unroll in inner loop allows
4

---

Published as a conference paper at ICLR 2026
7
Train
0.8
6
0.6
5
Accuracy
Loss
4
3
AdamW
0.2
Cclo2
VeLO
0.0
10K
20K
30K
40K
50K
10K
20K
30K
40K
50K
Train steps
Train steps
Parameter norm
Gradient norm
Update norm
20K
10.0
15K
7.5
10K
5.0
5K
2.5
2
0K
0.0
25K
50K
25K Train steps
50K
25K
50K
Train steps
Train steps
Figure 2: ImageNet classification with ViT. We test our learned update rule, Celo2, on the Im- ageNet classification task with batch size 512 and 50K steps, which is 25x longer than its meta- training unroll length and 30,000 larger than the tasks seen during training. Since VeLO trained with final loss as the meta-objective, it shows non-trivial dynamics during training in order to achieve low final loss (see norm plots). Celo2 achieves VeLO's final loss within ~50 steps. As test ac- curacy reaches 666, all optimizers start overfitting in this task; this is consistent with findings in prior work (Dahl et al. (023). Since our update rule is normalized, shows training norm dynamics consistent with AdamW. Moreover, VeLO is meta-trained with 200K unroll length large num- ber of diverse tasks including ViTs and ImageNet dataset, whereas Celo2 is only meta-trained on small image MLP tasks ($4). which highlights its strong meta-generalizationn capability.
multiple updates in single outer loop. This is much more computationally efficient than ES (used in VeLO) that does only a single optimizer update within an unroll; see Appendix full background on the meta-tran vetaa-training setup. We use average loss from the unroll during meta- vpetattaining as the meta- objective. Despite being meta-trained with average loss as the meta-objective, our optimizer yields strong final performance large-scale tasks as demonstrated by experiments discussed next.
Implementation details. We conduct large-scale evaluations in JAX (Bradbury et al. 2018 on v4 TPU pod with 32 chips and 4 VM hosts, using fully sharded data parallelism (FSDP). For RL experiments, we use JAX (jit, pmap) to vectorize agent and environment computations end-to-end device, avoiding TPU-CPU transfers. Our learned optimizer is meta-trained for 100K iterations with K=50 inner steps and 8 parallel tasks on Nvidia L40S GPU. Unless stated otherwise, optimizer evaluations sweep 7 learning rates logarithmically between le-3 and lc-5, cosine decay schedule with linear warmup fraction 0.05 and report the best performing hyperparameter setting for each optimizer based on final performance (see { more details). VeI.O tested in default self-tuned mode by specifying target duration at initialization along with training loss at each step. Since Celo2 is learned update rule, we implemented it in Optax ($C) and evaluated with a I-line drop-in replacement in all the standard tasks used in this work. We used float32 as the default precision for language modeling experiments to decouple instabilities introduced by the learned optimizer experiments from mixed-precistie instabilities during research. However, testing with bfloat16 on ImageNet ViT (Fig. 2 revealed no instability. We leave targeted study with respect to lower/mixed precisions reduce memory footprint of learned optimizers for future work.
update rule
5

---

Published as a conference paper at ICLR 2026
Freeway
SpaceInvaders
Asterix
60-
15
AdamW
150
Celo2
VeLO
Return
40
10
100
50
20
0-
2M
4M
6M
10M
2M
4M
8M
4M
8M
10M
Frames
Frames
Figure 3: Reinforcement Learning. We directly evaluate our learned optimizer, Celo2, on Atari RL tasks using the PPO algorithm to learn the RL policy. Our results clearly indicate that Celo2 performs par with well-tuned AdamW baseline on these out-of-distributiin tasks, while the VeLO baseline stagnates at a much lower return. The latter result can be corroborated by Figure in Metz et al. (2022b).
4.1 RESULTS
Generalizing billion-scale models (GPT-3). Although learned optimizers have been shown to generalize to small and medium-scale tasks (Metz al.) 2022b) Moudgil et al.) 2025), they have fallen short of generalizing to large-scale pretraining tasks such as GPT (Achiam et 2023), which highly important to the machine learning community. We show that Celo2, proposed learned update rule, generalizes to GPT-3 1.3B, GPT-2 124M, and a 30M transformer pretraining tasks, see Fig. All pretraining tasks follow the Chinchilla (Hoffmann al. 2022) recommended data-to- model scaling ratio 20, except GPT-3 1.3B, which uses a ratio of 10 due to compute constraints; this setup matches prior work (Marek et al. 2025). Notably, GPT-3 1.3B is 1, 0,000,000x, or six orders of magnitude, larger than the tasks Celo2 is meta-trained on, so this is purely out-of-distribution performance. Not only does scale gracefully and remain stable at this scale, it is also competitive with strong, well-tuned baselines such as VeLO and Adam W under the same tuning budget. We also found Celo2 to be competitive with Muon (Jordan et al., (024). recent state-of-the af-tth--rt art optimizer (Figure 7 in Appendix), despite being meta-trained on simple 8x 8 image classification tasks ($4)
Generalizing to longer unrolls. In order to stress-test our learned update rule on long horizon tasks, we choose the ImageNet (Krizhevsky et al.) 2012) classification task with the Vision Transformer (ViT) as our testbed. We perform unrolls for 50,000 steps; this is 25x larger than the unrolls seen during meta-training. As result, this is a strong test of length generalization for Celo2, exceeding previously attempted length generalization evaluations et al. 2024). Our results in Fig. 2 show that Celo2 not only generalizes reliably and stably to large-scale tasks, it also significantly outperforms VeLO and tuned Adam baselines. Note that Vel.O is optimized
literature (Moudgil et al.) 2025 Thérien
achieve the lowest final validation loss; Celo2, on the other hand, achieves Vel.O's target final validation loss ~2x faster than VeLO. addition, Celo2 surpasses 80% training accuracy much faster than the baselines within the given number of steps.
Reinforcement learning. This is another out-of-distribution generalization result from noisy (pol- icy) gradients. Prior work (Metz et al., 2022at Harrison et al. 2022 Moudgil al., 2025 Thérien et al. 2024), with the exception of VeLO in learned optimizers, has not reported results on Rein- forcement Learning tasks from a generalization perspective. However, such tasks are particularly valuable for evaluating generalization, as they involve horizoo credit assignment, high variance training dynamics, and non-stationary objectives, making them challenging and informative benchmark for optimizer robustness. We reuse the hyperparameters for PPO and the environment from prior work (Lu et al., 2022). We further tune the Adam LR using the search space sug- gested by prior work and our AdamW baseline closely matches the reported results. As result, our AdamW V baseline has optimal hyperparameters, for instance e=le-5. Fig. 3 clearly indicates that Celo2, evaluated zero-shot these RL tasks, outperforms VeLO by significant margin, and competitive with Adam W. This is despite it being meta-trained on 8 x image classification tasks with potentially unfavorable hyperparameters for RL (for example, low epsilon value). It is im-

---

Published as a conference paper at ICLR 2026
GPT-2 (124M) FineWeb
LM (30M) FineWeb
7
3.85
oclo2
base
celo2 (test-only)
3.80
celo2 hase
meta- training
Vil loss
adamw
3.75
adamm
3.65
0.5B
1B
1.5B
2B
2.5B
celo2-base
orthogonalize adlmmw @ ID
Train tokens
Figure 4: As shown in the figure the left, Celo2-base that a simple learned MLP rule for all parameters without orthogonalization or AdamW, is able to scale stably on GPT-2 task. We find that both techniques (1) Orthogonalization and (2) Adam for ID params improve performance when applied top of Celo2- base. Applying these two techniques directly at test-time improves performance meta-training with them even better.
hidden size
val loss
RMS decay
val loss
task aug.
val loss
4
4.128
0.999
3.893
4.417
8
3.812
0.95
3,812
3.812
16
3.857
0.99
3.865
32
3.869
0.9
3.922
(a)
(b)
(c)
RMS norn
meta-training
test-time
val loss
update form
val loss
3.961
\rd
3.812
Aid A(d-e(Am)) e(2am)
3.992
3.900
3.812
Aid eldamme)
4.075
(d)
(c)
Table 1: Ablations the ingredients of our proposed recipe. We find that learned optimizer generalization performance is sensitive to architecture and meta-training hyperparameters and they must be ablated thoroughly
achieve strong performance. (a) Hidden size: learned update rule with hidden size 8 emerges as the ideal choice, resulting in a smaller yet more performant learned optimizer; (b) RMS decay: we find that a value of 0.95 for the RMS decay yields the best results; (c) Task augmentation: one of the ingredients in our proposed recipe is task augmentation; this result clearly shows significant improvement validation loss using task augmentation; (d) RMS norm: RMS norm is another component we propose for Celo2-base. Using RMS norm helps both during meta-testing and also meta-training. This trend is consistent with our overall results; (e) Update form: the best learned update rule predicts simply d instead of more complex update rules that need to predict both and m. All ablations are evaluated on a 30M transformer decoder model trained on Fine-WebEdu with 600M tokens ($5). Celo2-base with our default settings are highlighted
portant to note here that the VeLO paper also reported similar trends in the stagnation of reward and the failure of the model learn after point (refer Figure 11 in Metz al. (202bb)
ABLATIONS
We now present ablations of each component in the Celo2 recipe. All experiments are on the LM (30M) task with the FineWeb-Edu dataset using 600M tokens, following the Chinchilla (Hoffmann ct al. 2022) scaling law. The learning rate of the meta-learned rule tuned on this task, and the best result for each ablation is reported in Table T Full hyperparameter details are in Appendix B

---

Published as a conference paper at ICLR 2026
Orthogonalization. Recent progress in optimization (Jordan et al. 2024). has proposed an opti- mizer that relies on orthogonalization to achieve faster convergence. Although this incurs a penalty in the form of wall clock time, the downstream benefits include lower final losses in addition to the already mentioned faster convergence. We find that Celo2 is highly compatible with orthogo- nalization, as demonstrated in Fig. | Specifically, adding orthogonalization to base Celo2 model directly at meta-testing time yields an improvement over Celo2-base (referred to as celo2 (test-only) in Fig. 4.. This improvement is amplitied by adding orthogonalization while meta-training Celo2. We find that similar trends hold by using Adam W for 1-D / bias parameters, which is another re- cently adopted practice (Jordan et al., 2024). Fig. Aclearly shows that adding orthogonalization to Celo2-base, followed by Adam update for I-D gives a compounding benefit in performance, both for meta-testing and meta-training. However, as shown Fig. + our Celo2-base model which uses a fully learned update rule across all parameters is also highly stable on the GPT-2 (124M) task.
Task augmentation. Task augmentation is useful technique (Metz et al. 20226 Moudgil et al.) 2025) to simulate diverse range of tasks from a limited set of tasks during meta-training to en- hance generalization. This is done by randomly sampling a scaling parameter and then rescaling the parameters of the optimizee network during meta-training. This leads to perturbation in the gradients, thereby resulting in a higher coverage of the optimization landscape, analogous to the role that data augmentation plays in supervised learning. We find that task augmentation is critical for good performance (refer to Table (e) for task augmentation ablation). Qualitatively, we see that meta-training with task augmentation allows the use of higher learning rate, thereby speeding up the optimization process and resulting in lower final validation loss.
RMS normalization. In order achieve high performance by using a purely learned rule without any hand-designed update rules for I-D parameters such as Adam, we find that RMS normalization plays a key role. The results of our ablations in Table (d)) indicate that adding RMS normalization during the meta-training phase yields improved performance. This trend did not hold for using it simply during meta-testing with the Celo2-base update rule. In Appendix ş we provide ablations on different normalization variants, showing that the choice in Alg. [ performs best.
Meta-training hyperparameters that matter. In addition to the components discussed above, we find that generalization performance is sensitive to several architectural choices in the learned MLP update and to specific hyperparameters. The first is the magnitude multiplier: the exponential mag- nitude multiplier proposed by prior work (Metz et al. 2022a) does not improve meta--eneralization (Table Te)). Using only the direction coefficient from the MLP update is sufficient. Averaging the exponential magnitude multiplier per tensor, as in Harrison et al. 2022). also reduces performance. The second factor is the RMS decay term with coefficient 33 used similarly to Adam (see (D) for background). For LM pretraining tasks, 0.95, the standard choice in recent LLM works Marek eet al. (2025) gives the best generalization, and this is the default in Celo2. The third factor is the MLP hidden size: prior work (Metz et al. 2022b a) uses 2-layer MLP with 4 units per layer, but we find that using 8 units per layer provides the best balance between performance and efficiency.
Runtime and memory overhead. Celo2-base and Adam have identical wall clock time, since the optimizer update is an inexpensive part of the entire training step. However, in the number of parameters, Adam has a memory overhead of 3x due to additional (momentum and RMS decay) accumulators per-parameter. On the other hand, Celo2 has memory overhead of ~5x since it maintains 3 momentum accumulators with 3 decays (unlike in Adam), along with one RMS decay per-parameter and tensor-level Adafactor features (Metz et al. (022b). Celo2 incurs higher runtime cost (1.3x) due to Newton-Schulz orthogonalization, though this may vary by task.
6 CONCLUSION
In this paper, we proposed practically applicable, inexpensive, yet effective recipe for learned opti- mization. Our approach results in a win-win situation on several axes: stability, compute-efficient and performance. Our learned update rule has several benefits: 1) it is more compute efficient than prior work, taking merely 4.5 GPU hours for meta-training, 2) it shows strong out-of-distribution generalization across datasets, 3) it demonstrates powerful scaling on large models like GPT-3. We provide a detailed set of ablations study the impact of each individual design decision proposed in our recipe. We hope that this work serves as an important milestone in the development and discovery of high-performing scalable learned optimizers.
8

---

Published as a conference paper at ICLR 2026
ACKNOWLEDGEMENTS
We acknowledge support from the FRQNT Doctoral (B2X) scholarship [AM], the FRONT New Scholar and FRQNT-NOVA [EB] and resources provided by Compute Canada, Calcul Québec and Mila. We also extend our sincere gratitude to the Google TPU Research Cloud (TRC) program for generously providing TPU resources that made this research possible.
REFERENCES
Josh Achiam, Steven Adler, Sandhini Agarwal, Lama Ahmad, Ilge Akkaya, Florencia Leoni Ale- man, Diogo Almeida, Janko Altenschmidt, Sam Altman, Shyamal Anadkat, et al. Gpt-4 technical report. arXiv preprint arXiv:2303.08774, 2023.
14
Diogo Almeida, Clemens Winter, Jie Tang, and Wojciech Zaremba. A generalizable approach to learning optimizers. arXiv preprint arXiv:2106.00958 2021. 4 22
Marein 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 neural information processing systems, 29, 2016.
23
James Bradbury, Roy Frostig, Peter Hawkins, Matthew James Johnson, Chris Leary, Dougal
Maclaurin, George Necula, Adam Paszke, Jake VanderPlas, Skye Wanderman- and Qiao
Zhang. JAX: composable transformations of Python+NumPy programs, 2018. URL http:
Bttcuub.comm/goole///aa
16
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. arXiv preprint arXiv:2306.07179 2023.
Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, Jakob Uszko- reit, and Neil Houlsby, An image worth 16x16 words: Transformers for image recognition at scale. ICLR, 2021. 15
Vineet Gupta, Tomer Koren, and Yoram Singer. Shampoo: Preconditioned stochastic tensor opti- mization. In International Conference on Machine Learning, pp. 1842-1850. PMLR, 2018.
Nikolaus Hansen and Andreas Ostermeier. Completely derandomized self-adaptation in evolution
strategies. Evolutionary computation, 9(2):159-195, 2001. 24
James Harrison, Luke Metz, and Jascha Sohl-Dickstein. A closer look at learned optimization: Stability, robustness, and inductive biases. In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho (eds.), Advances in Neural Information Processing Systems, 2022. URL.
https://openrevieew httpps///oooorervvvveeeeeerrrrrr
24
Jordan Hoffmann, Sebastian Borgeaud, Arthur Mensch, Flena Buchatskaya, Trevor Cai, Eliza Rutherford, Diego de Las Casas, Lisa Anne Hendricks, Johannes Welbl, Aidan Clark, al. Train- ing compute-optimal large language models. arXiv preprint arXiv:2203. 15556, 2022. 6
Keller Jordan, Yuchen Jin, Vlado Boza, You Jiacheng, Franz Cesista, Laker Newhouse, and Jeremy Bernstein. Muon: An optimizer for hidden layers in neural networks, 2024. URL https:// keller jobln.ggodda thub. io/posts/muon/
13
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 http:
wwarxiv.org/abs//ie abtp//1bs//11..... 1412. 6980.
22
Boris Knyazev, Abhinav Moudgil. Guillaume Lajoie. Eugene Belilovsky, and Simon Lacoste-Julies Julien. Accelerating training with neuron interaction and nowcasting networks. In The Thirteenth Inter- national Conference Learning Representations (ICLR), 2025. 2

---

Published as a conference paper at ICLR 2026
Alex Krizhevsky, Vinod Nair, and Geoffrey Hinton. Cifar- 10 and cifar-100 datasets. https://wwww cs. toronto. edu/krizccifar html,, 6(1):1, 2009. 4
URI:
Alex Krizhevsky, Ilya Sutskever, and Geoffrey E Hinton. Imagenet classification with deep convo- lutional neural networks. Advances in neural information processing systems, 25, 2012.
Yann LeCun and Corinna Cortes, The mnist database of handwritten digits. 1998.
URL http:
Jayann,lecuum www.enncoo/eeeeeeeeeeeeeeeeeeeee commexxb/mnist
Jingyuan Liu, Jianlin Su, Xingcheng Yao, Zhejun Jiang, Guokun Lai, Yulun Du, Yidao Qin, Weixin Xu, Enzhe Lu, Junjie Yan, et al. Muon is scalable for Ilm training. arXiv preprint arXiv:2502. 16982, 2025.
I Loshchilov. Decoupled weight decay regularization. arXiv preprint arXiv: 171105101, 2017.
Chris Lu, Jakub Kuba, Alistair Letcher, Luke Metz, Christian Schroeder de Witt, and Jakob Foerster. Discovered policy optimisation. Advances in Neural Information Processing Systems, 35:16455- 16468, 2022.
15
Martin Marek, Sanae Lotfi, Aditya Somasundaram, Andrew Gordon Wilson, and Micah Goldblum. Small batch size training for language models: When vanilla sgd works, and why gradient accu- mulation is wasteful, 2025. I
14
Luke Metz, Niru Maheswaranathan, Brian Cheung, and Jascha Sohl-Dickstein. Learning unsuper- vised learning rules. 2019a.
23 24
Luke Metz, Niru Maheswaranathan, Jeremy Nixon, Daniel Freeman, and Jascha Sohl-Dickstein. Understanding and correcting pathologies in the training of Iearned optimizers. In International Conference on Machine Learning. pp. 4565- PMLR, 2019b.
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. arXiv preprint arXiv:200. 2020. 22 23 24
Luke Metz, C Daniel Freeman, James Harrison, Niru Maheswaranathan, and Jascha Sohl-Dickstein. Practical tradeoffs between memory, compute, and performance in learned optimizers. In Confer- ence on Lifelong Learning Agents (CoLLAs), 2022a. optiimization
24
Luke Metz, James Harrison, Daniel Freeman, Amil Merchant, Lucas Beyer, James Bradbury, Naman Agrawal, Ben Poole, Igor Mordatch, Adam Roberts, et al. Velo: Training versatile learned optimizers by scaling up. arXiv preprint arXiv:2211.09776, 2022b.
I
Abhinav Moudgil, Boris Knyazev, Guillaume Lajoie, and Eugene Belilovsky. Celo: Training ver- satile learned optimizers on a compute diet. Transactions on Machine Learning Research, 2025.
ISSN 2835-8856. URL httpt:://oppeeeeiieeee wttp:/wwwwww...ooom/////sataaaa.
23
24 25
Yuval Netzer, Tao Wang, Adam Coates, Alessandro Bissacco, Baolin Wu, Andrew Y Ng. et al. Reading digits natural images with unsupervised feature learning. In NIPS workshop on deep learning and unsupervised feature learning, volume 2011, pp. 4. Granada, 2011.
Ingo Rechenberg. Exolutionsstrategie:: Optimierung technischer systeme nach prinzipien der biolo- gischen evolution. 1973. 24
Tim Salimans, Jonathan Ho, Xi Chen, Szymon Sidor, and Ilya Sutskever. Evolution strategies as
scalable alternative to reinforcement learning. arXiv preprint arXiv: 2017. 24
Silvia Sapora, Gokul Swamy, Chris Lu, Yee Whye Tch, and Jakob Nicolaus Foerster. Evil: Evolution strategies for generalisable imitation learning. arXiv preprint arXiv:2406.11905. 2024. 15
Noam Shazeer and Mitchell Stern. Adafactor: Adaptive learning rates with sublinear memory cost.
arXiv preprint arXiv: 1884004235, 2018.
10

---

Published as a conference paper at ICLR 2026
Chongjie Si, Debing Zhang, and Wei Shen. Adamuon: Adaptive muon optimizer, 2025. URL wttpar/www...rgr//////.../ 11005
Andreas Steiner, Alexander Kolesnikov, Xiaohua Zhai, Ross Wightman, Jakob Uszkoreit, and Lucas Beyer. How to train your vit? data, augmentation, and regularization in vision transformers, 2022.
URL
http://...gg/bb//////// 10270 15
Ilya Sutskever, James Martens, George Dahl, and Geoffrey Hinton. On the importance of initial- ization and momentum in deep learning. In International conference on machine learning, pp-
1139-1147. PMLR, 2013. 2
Benjamin Thérien, Charles-Etienne Joseph, Boris Knyazev, Edouard Oyallon, Irina Rish, and Eu- gene Belilovsky. μLO:: Compute-efficien meta-generalization of learned optimizers. arXiv preprint arXiv:2406.0011 00153, 2024. 2 6
Hugo Touvron, Matthieu Cord, Alexandre Sablayrolles, Gabriel Synnaeve. and Hervé Jégou. Train- ing data-efficient image transformers & distillation through attention. In ICML, 2021. 15
Mark Tuddenham, Adam Prügel-Bennett, and Jonathan Hare. Orthogonalising gradients to speed up neural network optimisation. arXiv preprint arXiv:2202.07052, 2022.
Paul Vicol, Luke Metz, and Jascha Sohl-Dickstein. Unbiased gradient estimation in unrolled com- putation graphs with persistent evolution strategies. In International Conference Machine Learning, pp. 10553-10563. PMLR, 2021.
Nikhil Vyas, Depen Morwani, Rosie Zhao, Mujin Kwun, Itai Shapira, David Brandfonbrener, Lucas Janson, and Sham Kakade, Soap: Improving and stabilizing shampoo using adam. arXiv preprint 11XXv::4099.1322.0 2024. 3
Olga Wichrowska, Niru Maheswaranathan, Matthew W Hoffman, Sergio Gomez Colmenarejo, Misha Denil, Nando Freitas, and Jascha Sohl-Dickstein. Learned optimizers that scale and gen- eralize. In International conference machine learning, pp. 3751-3760. PMLR, 2017. 23
Daan Wierstra, Tom Schaul, Tobias Glasmachers, Yi Sun, Jan Peters, and Jürgen Schmidhuber. Natural evolution strategies. Journal of Machine Learning Research, 15())994--980 2014. 24
Han Xiao, Kashif Rasul, and Roland Vollgraf. fashion-mniste a novel image dataset for benchmark- ing machine learning algorithms. arXiv preprint 1rXivv:/08.007444 2017.
Xiaohua Zhai, Alexander Kolesnikov, Neil Houlsby, and Lucas Beyer. Scaling vision transformers. In CVPR, 2022.
15

---

Published as a conference paper at ICLR 2026
A ADDITIONAL EXPERIMENTS AND ANALYSIS
Normalization variants. To achieve high meta-generalization performance, we explore several normalization strategies for the predicted updates. Table 2 compares different approaches: us- ing raw MLP outputs, per-step RMS normalization, rolling RMS normalization (aggregated over time), and variants with hard clipping. Per-step RMS normalization (norm) achieves the best val- idation loss of 3.81268, substantially outperforming the raw output baseline (3.96168). Rolling normalization (rolling norm) also improves over the baseline but falls short of per-step normaliza- tion. Incorporating hard clipping, either with rolling normalization (rolling norm clip) or per-step normalization (norm clip), degrades performance compared to their non-clipped counterparts, sug- gesting that overly constraining update magnitudes can be detrimental. These results demonstrate that simple per-step RMS normalization is the most effective strategy for stabilizing meta-learned updates without requiring hand-designed rules.
functional form
val loss
raw output rolling norm
Ap. MLP(...)
3.961
Apt/ Σ APelll llms
3.860
rolling norm clip norm clip norm
Ap, min (I, (Σ' Apel mmss
/ /pp ml))
3.999
Ap: min 1, T/ | rms
4.084
Ap: |
3.812
Table 2: We evaluate several ways of scaling the predicted update vector Ape. "Raw output" applies the MLP output directly. "Rolling norm" normalizes by the cumulative RMS magnitude of past updates, while "rolling norm clip" additionally limits step size using time-dependent threshold т. "Norm clip" clips each update based on its RMS norm. "Norm" simply normalizes each update by its own RMS magnitude and performs the best, achieving the lowest validation loss among all variants. All variants are evaluated on a language modeling task using 30M-parameter Transformer trained on the FineWee-Edu dataset with 600M tokens ($5).
Stability and hyperparameter sensitivity. As evident from the LM (30M) Fib-WeeEEdu plot on the right (Figure 55. Celo2 demonstrates sta- ble optimization behavior across range of learn- ing rates, with performance characteristics that closely match standard optimizers like A and VeLO. The best tuned Adam W and VeLOO con- figurations are also plotted for reference from Fig- T Learning rates toward the higher end of our sweep generally perform well, though the highest learning rate (0.001) underperforms compared to the best setting with slightly lower learning rate (0.000464). This pattern, where moderately high leamning rates optimize effectively while very high rates cause instability, is consistent across all our supervised learning experiments. We did not find Colo2's stability to be sensitive to any specific hy- perparameter values within reasonable ranges. For reinforcement learning experiments, we present parallel coordinate plot in Figure O showing sen- sitivity to ablated parameters in sweep for the SpaceInvaders environment as a representative example. The highlighted yellow curves indicate high-return trajectories. Here too, Celo2's behavior aligns well with standard optimizers. The opti- mal learning rate value is close to the higher end of the range but not the highest, as the maximum learning rate leads to unstable optimization and consequently lower returns. The optimizer shows no particular sensitivity to random seed initialization, demonstrating robust performance across dif- ferent experimental runs. We also ablate over the choice of linear learning rate decay (a common practice in these Atari RL tasks) and find that in this environment, Celo2 benefits from incorporating linear decay, though it maintains stable performance across both configurations.
LM (30M) Feb--eb-Edu
Adam'
4.8
Vel.D
Colo2 LR=000001
4.6
Colo2 LR-00000444
Coloe2 R--0000215
4.4
Colo2 LR=0.0001
4.2
Colo2 LRR00.00006
4.0
3.6
0
120M
240M
360M
480M
600M
Train tokens
Figure 5: Validation loss curves for Celo2 with various learning rates uniformly sampled on log-scale between le-5 and 1e-3 on LM (30M) Fine Web-Eeb dataset.
12

---

Published as a conference paper at ICLR 2026
Decay LR
Socd
Leaming Rate
Return
200
3.5
450
36-4
23
120
15-
80
40
26-5
Tuse
88
169
Figure 6: Parallel coordinate plot showing Celo2 hyperparameter sensitivity on SpaceInvaders RL
lask. Yellow curves indicate high-return configurations.
Comparison with Muon. We compare Celo2
GPT-2 (124M) FineWeb
2024 Liu al. 2025) the GPT-2 (124M) Wib pretrain- ing task in Figure 7 Muon is a hand-crafted optimizer for hidden layers in neural networks that orthogonalizes momentum with Newton- Schulz iterative process. For non-hidden pa- rameters (embeddings, scalars), uses Adam update rule. We the Kimi variant (Liu et al, 2025). which applies RMS normalization the orthogonalized update (as in Celo2) to match the RMS of Adam updates for ID pa- rameters. This enables a shared learning rate across embeddings and hidden layers and re- duces tuning effort. Kimi-Muon uses an RMS multiplier of 0.2 for hidden layers to match the empirical RMS Adam; we found this also works well for Cclo2 and include it in the plot (Celo2 RMS x 0.2 in Figure 7). Note that Celo2 RMSx00.2 variant and Muon only differ in the update rule which Orthogonalization is applied: Muon uses momentum whereas Celo2 uses learned MI.P update and rest of the optimizer update identical since they both use RMS normalization with 0.2 multiplier and Adam for ID parameters.
with Muon (Jordan al.
4.4
AdamW
Muon
Celo2
4.2
Celo2 (RMSx0.2)
4.0
Val loss
3.6
0.5B
1.5B
2.5B
Train tokens
Figure 7: Comparison of Celo2 and Muon on GPT-2 (124M) FineWeb pretraining task.
Overall, both Celo2 and Celo2 RMS 0.2 demonstrate competitive performance compared to Muon, achieving lower validation loss for majority of the run (Figure 7). In terms of final loss, Celo2 RMS x.2 peforms better, achieving slightly lower validation loss (3.35588) than Muon (3.35636) while default Celo2 (without 0.2 RMS multiplier) achieves 3.36785. For this comparison, we tuned the Muon baseline by log-uniformly sampling 10 peak learning rates between 1c-4 and 1.0: (0.0001, 0.000215, 0.000464, 0.001, 0.00215. 0.00464, 0.01, 0.215, 0.464, 1.0). All optimizers use cosine decay schedule with warmup fraction 0.05, decay to zero and weight decay 0.1. Muon performs best at learning rate 0.01. For Celo2 variants, we use the same log-uniform sampling 7 learning rates between le-5 and le-3 as in our experiments in the main paper ( (0.00001, 0.0000215, 0.0000464, 0.0001, 0.000215, 0.000464, 0.001). Celo2 and Celo2 RMS x0.2 variant perform best 0.000215 and 0.000464 learning rates, respectively. The learned MLP update rules tested in this study preta-traines on simple 8x image classification tasks (44): further meta-traing trainitraiinng or fine- tuning language modeling tasks could improve performance, which we leave for future work.
13

---

Published as a conference paper at ICLR 2026
B EXPERIMENTAL DETAILS
Language Modeling, For our large-scale language-modeling evaluations, we evaluate Celo2 on three autoregressive Transformer architectures ranging from 30M 1.3B parameters, all of which were trained on TPU pod. All models use the same decoder-only Transformer implementation provided in the public codebase] by Marek al. (2025). Each model consists of token embed- dings (input and output), L pre-norm Transformer blocks, and a 2-layer MLP with inner dimension F 4D; full model configs are presented in T below. Multi-head causal self-attention is im- plemented with query/key RMSNorm, RoPE positional embeddings, and either standard dot-product attention or TPU FlashAttention via aatassssatteeetionno kernell (automatically enabled when train- ing on TPU). Weights are initialized using Xavier-uniform except embedding parameters, which initialized using variance scaling with normal distribution, fan-in mode, and scaling factor of 1.0 along the output axis.
model
D
L
H
N
30M (LM-30M)
384
6
1536
128
3
512
124M (GPT-2)
768
12
3072
128
6
1024
1.3B (GPT-3)
2048
24
8192
128
16
2048
D: hidden dimension, L: number of transformer layers,
F: feodforward inner dimension, H: attention head dimension,
N: number of attention heads, T: context/sequence length
Table 3: Language model architectures. We evaluate Celo2 on three autoregressive Transformer architectures. All models share the same architecture template but differ in size. F = 4D and N D/H. Vocab size for FineWeh dataset is 50257 and rounded to match device sharding requirements.
hyperparameter
LM-30M
GPT-2 (124M)
GPT-3 (1.3B)
learning rate (celo2)
0.000464
0.0002
0.00002
learning rate (adamw)
0.01976
0.0048
0,0002
warmup fraction
0.05
0.05
0.05
batch size (global)
256
512
512
weight decay adamw epsilon
0.1
0.1
0.1
1e-8
1e-8
1e-8
adamw 31
0.9
0.9
0.9
adamw $2
0.999
0.999
0.999
learning rate schedule
linear warmup cosine decay
Table 4: Hyperparameters used for the language modeling experiments in Figure
Sharded parameters are placed across a 2D device mesh (data x model). Training follows standard next-token prediction using teacher forcing with cross-entropy loss. Large global batch sizes supported through gradient accumulation. We evaluate all the optimizers by sweeping learning rate and weight decay hyperparameters. All training hyperparameters for tuned LM experiments from Figure I are described in Table [ We search over learning rates and weight decay values for both Celo2 and AdamW. For Celo2, we sample seven learning rates logarithmically between 0.00001 and 0.001 (specifically: 1e-5, 2. 25e5,, 4.6e-5, 1e-4, 2. 1.e5-4, 4.64e-4, 1e-3) and test them with linear warmup and cosine decay to zero schedule. For weight decay, we test three values: 0.0, 0.1, and 10.0. We include the high value of 10.0 because initial experiments with Celo2-base showed it can handle higher weight decay than typical optimizers. For Adam W, we consult tuned values search space from prior work (Marek et al. 2025 Achiam et al., 2023). which were found using larger compute budgets than our Celo2 search (for example, 13 learning rates between 0.0002 and 0.03 the LM- 30M task for AdamW) and we further tune them with additional weight decay values discussed above. This makes AdamW a stronger baseline. Empirically, our search ranges appear well-chosent for both optimizers, mid-range values generally perform best, with larger models preferring slightly lower learning rates.
https: httpas//ggtttta....//eertmrrrrra www..///ttiinaaeeeeeaaaaaaaa
14

---

Published as a conference paper at ICLR 2026
Image Classification. Vision Transformer (ViT) experiments are conducted using the jaxtransformerfi a minimal but performant Transformer repository implemented in JAX. Full hy- perparameter config is presented in the beloo. Models are trained the ImageNet dataset with a base-size configuration (hidden size 768, 12 Transformer layers, 12 attention heads, MLP ratio 4). Images are preprocessed by resizing padding the smaller side to make them square, then resized to 256x256 and passed as input to ViT model. During training, additional data augmenta- tion applied: images are randomly flipped horizontally, and a random resized crop (area range 5 5-100% aspect ratio 0.75-1.333 is performed, followed by resizing back to 256x256. Pixel values are normalized to [-1, 1]. Training uses batch size of 512 for total of 50,000 steps. Weight decay is fixed to 0.1. ViTs are known to be data-hungry (Dosovitskiy et al. 2021 Zhai et al. [022 Steiner et al, 2022 Touvron et al., 202T) and on this task begin to overfit around ~66 validation aceuracy, making stronger regularization crucial. We also tried lower weight values (0.01 and 0.001), but they consistently underperform compared to 0.1. The learning rate is swept log-uniformil from x 10-3 to 1 10 for both Cclo2 and AdamW optimizers. We found this to be an effective search space, as the optimal values for both optimizers lie within this range, with Adam W favoring comparatively higher learning rates. Fully Sharded Data Parallel (FSDP) is leveraged to efficiently distribute model parameters and optimizer states across multiple devices, enabling large-batch training on TPUs.
hyperparameter
value
learning rate (celo2)
0.000046
learning rate (adamw)
0.0001
dataset
imagenet-256
model
VIT
hidden size
768
depth (layers)
12
number of heads
12
MLP expand ratio
4
weight decay batch size (global) training steps warmup fraction learning rate schedule
0.1
512
50,000
0.02
linear warmup cosine decay
Table 5: Hyperparameters used for the Vision Transformer (ViT) experiments in Figure
Reinforcement L Learning. We conduct reinforcement learning experiments using a distributed PPO- style framework iimplemented in JAX. Agents are trained on multiple Atari environments, including Asterix, Freeway, and Space Invaders. The hyperparameter configurations used to generate the sults in Figure are listed in Table 6 Each environment runs in parallel across multiple actors, and observations are flattened and logged for stability. Policies are parameterized as actor-critic networks with two hidden layers of size 64 and tanh activations, producing action distributions and value estimates. Training proceeds in mini-batcher collected over multiple steps per environment. Following common practices and PPO hyperparameters from the purejaxrl repository, we perform multiple PPO update epochs per batch of trajectories, subdivide batches into minibatches for gra- dient updates, and use generalized advantage estimation (GAE) with discount factor y 0.99 and GAE parameter À = 0.95. The PPO objective includes clipped surrogate losses, with clip epsilon defining the maximum allowed deviation. Entropy regularization encourages exploration, weighted by the entropy coefficient. Similar to other experiments, we sweep learning rates log-uniformly from 10-6 1 10-3 for AdamW W and Celo2. Gradients are clipped to maximum norm, and learning rates are either kept constant or linearly decayed to zero (we found linear decay cou- pled with high learning rate to be slightly better) which is a standard practice in the aforementioned
purejaxrl repository that multiple works build on (Lu et al. 2022 Sapora et al. 2024
htttps://github.com /tt.a/////ggttto..o//aa/aa/a///a http:///gggguuucooo///aaasaaaaas
15

---

Published as a conference paper at ICLR 2026
hyperparameter
freeway
space-invaders
asterix
learning rate (cclo2)
0,000215
0.000464
0.0001
learning rate (adamw)
0.001
0.001
0.001
total timesteps
10M
10M
10M
discount factor (7)
0.99
0.99
0.99
GAE parameter (A)
0.95
0.95
0.95
num steps per update parallel actors
128
128
128
64
64
64
PPO clip epsilon
0.2
0.2
0.2
PPO entropy coefficient
0.01
0.01
0.01
gradient norm clipping
0.5
0.5
0.5
value loss coefficient
0.5
0.5
0.5
random seeds weight decay
5
5
0.0
0.0
0.0
adamw epsilon
le-5
le-5
le-5
adamw $1
0.9
0.9
0.9
adamw 32
0.999
0.999
0.999
learning rate schedule
linear decay
linear decay
linear decay
Table 6: Hyperparameters used for the RI experiments in Figure
C JAX IMPLEMENTATION
The complete JAX implementation of Celo2 is available at: https://github com/ amoudg1/cell In this section, we provide high-level overview the implementation and illustrate key design choices using simplified code snippets and pseudocode. The goal is make the core ideas acces- sible to readers who may be new to JAX and Optax, while keeping the presentation focused on the conceptual structure rather than low-level implementation details.
Optax overview. Celo2 is implemented in JAX with Optax (Bradbury et al. 2018 syntax. Specifi- cally, in Optax, each optimizer is expressed transformation that maps parameters and gradients to updated parameters together with any auxiliary state. For example, the standard Adam opti- mizer constructed via optax. adam adam(1earningg_rate) which returns transformation consisting of an init function that initializes the optimizer state and an update function that applies the Adam update rule. Concretely, python, an optimization step can written as:
for example, adam transformation is initialized as: tx optax. adam(learning_rate)
# initialize optimizer state state tx (paraa))
# update state
updates, state = tx. update state, params) params optax. wpd.utesepapaaaaae updates)
Celo2 follows this same abstraction: we implement our method as custom Optax transformation with its own state definition, initialization logic, and update rule, ensuring compatibility with JAX's functional style and seamless integration with existing Optax training pipelines. One of the powerful features of this transformation view is that transformations Multiple transformations, such as gradient clipping, weight decay, and learning rate scaling, can be appended in modular way. For example:
be composed using optax. chain.
# Compose multiple transformations using optax, chain tx optax. chain(
celo2_optaa (oto__ottx.(...)
celo2 transform function
optax add add www.deeyyeeeeiieeeeeeee
weight decay
optax.
scale by learning rate (-IrHupdate)
state tx. (narm(assss)))) updates, state tx.update(grads, state, params)
16

---

Published as a conference paper at ICLR 2026
params optax. optax.apply. upplyy_ updates)
This modular design allows us to implement learned update rules such as Celo2 cleanly, while keeping each component of the update logic separated and reusable within standard Optax pipelines. In addition to modular transformations, Optax also supports flexible learning rate schedules that can be combined seamlessly with any optimizer, For instance, one can define a schedule such as cosine decay, exponential decay and then scale updates accordingly:
from optax. schedules import warmup_cosine__eee harmmp__osineedeccy_ssthhhi _cchuuueececeellee
Define cosine learning rate schedule
schedule warmup_cosine_ http://cossse...eeceeeeeeeeeeeee end_1r)
peak_1r, warmup_steps, num_opt__steps,
Apply schedule to Celo2 transformation
optax. chain
celo2_optax optax. www..callle..lll.mi.aeareeeeeeie
dynamically scales updates with schedule
utility to apply different update rules to different subsets of parameters. This particularly useful when one wants treat parameters (e.g., embeddings. biases, layer norms) differently from 2D+ parameters (e.g., weight matrices, convolution kernels). With optax. tatttttsssoorrm, separate transformations can be defined for each parameter group and combined into single optimizer:
Furthermore, Optax provides the optax.
# Define transformations for 1D and 2D+ parameters tx_1d optax. chaiin adammiooppa...a........ adam(...)) ...) tx_2d optax www.chiinnceo__pptaaaooo.
...)
Create param_labels as function that returns label pytree param_labele 1ambda params: jax. tree_map( lambda p: '1d" if p.ndim else '2d',
params
Combine using multi_ trati__rrasform tx optax. multi_trannsform(
transforms=("'dd: tx_1d, '2d': tx_
Baram_Llbbllspperan_abbllll
# Initialize and use state tx. ix.init updates, state tx. apdate(gradss state, params) params optax.apply boopllyupdates((e
updates)
Celo2 optax transformation. The Celo2 transformation (CeloptaxTTansformmtiioo imple- mented in JAX custom Tradiientrrnnformatt internal state (CeloOptaxState) consists of rolling averages for momentum, RMS, and Adafactor- style scaling, along with step counter. At each update, the rolling statistics are updated, and per-parameter updates are computed by applying a small MLP defined in the Celo2 class.
following the standard Optax interface. Its
import functools
from typing import Optional import flax
import jax
inport jax.numpy jnp
import chex
import optax
dflasasttucttdattcllass
class CeloOptaxState:
"*"Internal state of the Celo2 optimizer.
rms_rolling: AreayTree
mon_rolling: Arex.Arra
fac_rolling: Arra.Arra
step: Imp..darray
17

---

Published as a conference paper at ICLR 2026
class Cellopptxtrennsfoortiiioooooooo http://owtt..orooeeeeerrrrrroooa
Optax-compatible gradient transformation for Celo2.
Computes per-parameter updates using rolling statistics and small MLP.
def iseft celo2, theta: dict):
Args:
celo2: Instance of the Celo2 optimizer theta: Meta-parameters or pretrained state for the MLP
self. celo2
self.theta theta
def init(self, params: Ahex.ArrayTree) CeloOptaxState:
"""Initialize rolling statistics and step counter.
mom_ace, rms_acc, fac_acc = wtt.c//o..ccmmmlaaasssrooroooooo return Celo@ptaxState(
mon_rolling=mom.cc senee.inittpapapt ipit(par)m)), rms_rolling=rms. ini
wae..iiiipppaaanttt... (params))
fac_rolling=fac_ace init(params),
step=jnp.asarrayy dtype=jjp.int333
def update( self,
grads: chex. AreayTree,
state: CeloOptaxState,
params: Optional(chex.Arraaa
None, tuple[chex.. ArrayTree, CelolptaxState]:
"* Computtte par-paraeeter updates and new optimizer state."" ""*
Increment step and clip gradients
Celo/VeLO
step wtwasssafeeiceeeeeetteeeeee step)
grads jax. www.ttreeulll... maree_aappaamldll g: jnp.clip(g, -1000.0, 1000.0), grads)
Update rolling statistics
mom_ace, rms_acc, fac_acc = self.celo2 accumllttors_for_decayy
next_mom mom_acc, nom_acccupdate((S) wwtpeettttttttotlllllleeeoee mommrolling, grads)
next_rms rms_acc, wwt..cuudateststttt roadattestateemmm_rodoo rms_rolling, grads)
next_fac, fac_g = fac_acc. wrc.acc..uddtesstatee. fac_rolling, grads)
Prepare MLP function apply_mlp = functools. partial (werrialllslfll.eeiaiee www.ee...ell...oo melo2.mlp_.pply, self, theta)
Compute per-parameter updates
updates jax. tree_util. mree_map(
apply_mlp,
params,
grads,
next_mom.n
next_rms.rms
fac_g,
vext_fac.v_col,
next_fac.v_row
next_fac.v_diog
# Update optimizer state
new_state Gelo@ptaxState(
mom_roolingenext__ rms_rooling=next__mm fac_rollingsnett_fat step=step,
)
return updates, new_state
18

---

Published as a conference paper at ICLR 2026
The above Celo2 transformation can be directly plugged in standard Optax chain. Following example constructs a simple Celo2 optimizer with weight decay and learning rate scaling as in celo2-base variant:
def create_ www.oeeoooliiieeeeeeeeeeeeeeeeee (params, raarnnng_rate, decayt_ checkpoint.ppach:oonn)) "*"Create Celo2 optimizer with weight decay.
celo2 = Celo2(....)
theta = celo2. init_paaaams())
pretrained_stabe = load_ voad_staaeecheccpoint..ttt theta) if checkpoint_path else None
return optax.
optax.chain( chain(
pretrained_state)) optax. ppttaxscaae_.b_leaaee bay..leaminngggaaeaeariinno
optax. aww_aadd_dde dew.yeed_weighteeeeeeeeeeeesesee
For advanced usage, different update rules can be applied to different parameter groups using optax. multi_ transform multi_transform beddings and Celo2 for other parameters (2D+ tensors) in language modeling experiments ($4))
as described optax overview section above; for example Adam for em-
implementation relies on an instance of the Celo2 Notice that above Celnopttaxranssorraatt optimizer class. This object responsible for maintaining and updating the internal rolling statistics (momentum, RMS, and Adafactor-style scaling), providing the accumulators for these decays, and defining the MLP function that computes per-parameter updates. In essence, the Optax transforma- tion acts as a wrapper that translates standard gradient inputs into updates generated by the Celo2 optimizer logic which described next.
Celo2 optimizer. The Celo2 class below implements learned optimizer that tracks rolling statis- ties for momentum, RMS, and Adafactor-style accumulators, and computes per-parameter updates through a small MIP with normalized inputs and outputs. Optionally, updates can be orthogonal- ized using the Newton-Schul method. The following pseudocode shows the adapted Celo2 class, based on the original implementation (available here), along with an example of its usage.
import jax import naxpy as jnp
import flax import haiku hk import optax
import functools
from learned LammedLoott//zttiinnnnneetttttt
import common
def
ns_coeffs, ns_steps=5, eps=1e-8):
if x.ndim 2:
raise ValueError ("Input must have >= 2 dims")
transposed = False
if x.shape[-2] > shape[-1]:
x = swppsxeaax,s -2, -1)
x /= transposed True
axis=(-2, -1), keepdims=True) eps)
def ns_iter(_, x):
x_mT = swpp.saapxxssxx, -2, -1)
= x_mT
ns_coeffs[[]] ns_coeffs[2]
ns_coeffs[0]
jax lax. fax.ioop.ori ns_steps, ns_iter, x)
if transposed:
jnp. Sxapaxes -2,
return
class Celo2:
def init (self,
ff ff_hidden_siz==., ff ff_hidden_llaerse
comentum_deccys=((0 0.99, 0.999),
rms ros_decays=(0.95.) adafactor (eaacttor_decass(0.__9 orthogomalizesTrue,
0.99, 0.999),
ns_ (t..cfffff334444:0 -4, 7750, 0315), ns ns_iters=5,
19

---

Published as a conference paper at ICLR 2026
ns_eps=le-8):
self. htt.ffffdddd_si2222 Sizee ff_ ff_hiddenn_izee
helffffhide___dd_eee_ ff ff_hidden_layers
self.momentumm_eecay = jnp.assarraymmmmmtttuu.ededaaa
res_dcyyyy www.aaaarrrrssdderea.
self.
jnp aasaarra(adaaaactrrrdeeee.
self orthogonalize
self.ns_coeffi = sinp.asarraynns_coffff
self.ns_iters ns_iters
inleens_epps ns_eps
self.act_fn jax. nn.
Haiku transform for per-parameter MLP self.apply_mie
def hat.iiatoor__ooo_decsssssseeee
Returns rolling statistics for momentum, RMS, and Adafactor-style scaling.
mom_acc common. wec_roolinggmommeeeee memmmsssl..mmmnuumm.ee.... deccayt
rms_acc = common www._erooooonnnrmseeeeeeeeoe
(self. ff.. rms_decays)
fac_acc common.
htt.r//ooonnnnss.aaarrrrrooorear addafactor__ecays)
return mon_acc, rms_acc, fac_acc
def dxecond_moment.e. mot.mmmmnnnooommliiee(ellll x, axis, eps=1e-9): (Ju).suuree)))))) rms jnp.mman (jnp. square(x), axis=axis, keepdims=True) return x jax. waa...ssrr(rmm... eps)
def _apely_ param, grad, mom, rms, fac_g, fac_vec_col, rac_vec_row, vec_vec_,, summary_prefix):
Per-parameter MLP that computes updates.
inps [grad, param, mom, rms, fac_g, fac_vec_col, fac_vec_row, fac_vec_v] axis (aaaaaee(((e(ppaaaaasaappeeeeee)
Input normalization
inps =
axis=axis) for in inps]
# Concatenate features and apply MLP
x = jnp. popp.concteeatee(inng rishapee,, (-1, i.shape(-11)) for in inps], axis=-1) for in range(self. ff_hi dden layers):::
hk. Linear (self. ff wwwddhensizeeeeeee (x) self.act_fn(x)
out hk. Linear(1) (x) step
param. saape)
Newton Newton-Schulz orthogonalization
if self self.orthogen thogoooollleeeeee and step. reep.ndim
step sechogonalize_via_eeeee. sstettonschullsste.. self.ns_coeffs. self. iters, self
_n__pps)
Output normalization wwt.eeeoenmmmmetttoomraaaiieeeee axis=axis)
step = seef._secood_mmnneeeeeee
return step
# --- Example initialization --- celo2 Celo2(
ff ff_hidden.size-f
ff_hidden__aayersh)
0.99, 0.999), rms rms_decaa (ss_deeaay=(00.55))
gomenttu__decass=(0.1
adafaccorr_decyss-CC 0.99, 0.9) or thogonalize=frue,
)
Ir sc_schedeee warmup_ cosiine_dcay.sscheulleeeiii peak_Ir, warmup_steps, numLopt_steps) celo2_opt optax. optax.chain(
optax. add www__adddeddddddceeeeeeeeeaa scrayeewweighhttswegghh_.oee
scale_bc berr___eeeemmmnaareeeeeeeeeeeeee
20

---

Published as a conference paper at ICLR 2026
scaled_lr_sciee sr_scced Lambda step: enbedding_Ir la__i_m__ll_____-- 1r_ (ssscheeulesstep)
adam_opt optax. dddm(ssseeed__r__cchhllll b1=0.9, b2=0.95. .9,, heahhhddecayeeighh_deeaaaa
# Combine using multi_transform labeling embedding parameters "adam' and others
celo2'
tx = optax, multi_.rannsormc transforms=(
'celo2": celo2_opt,
'adam": : adam_opt
},
Larrm_labes=lanmbd params: jax. jax.tree, www.tteaeee............e map__il pi_phhhh(
lambda path, val: 'adam" 'embed' jax. tree_util .ey.tt.rtttttt else 'celo2', params
Use just like standard optax optimizer! updates, state wx.update(grad: state, params)
21

---

Published as a conference paper at ICLR 2026
D ADDITIONAL BACKGROUND
This section provides background on learned optimizers for readers new to the field. We focus on two key components: the architectural design of MLP-based learned optimizers and the meta- training procedure used to train them. For the architecture, we describe the Adafactor MLP learned optimizer introduced by Metz et al. (2022a), which serves as the foundation for several state-of- the-art learned optimizers including VeLO (Metz et al. (2022b) and Celo2, proposed in this work. For meta- pueta-training we explain the bi-level optimization framework and describe key techniques that enable practical meta-training of learned optimizers: Persistent Evolution Strategies (PES) ( (Vicol al. 202T), an efficient gradient estimation method, and task augmentation (Metz et al.) 2022b] Moudgil et al. 2025), which improves generalization to unseen tasks. Finally, we present the com- plete outer-training loop that integrates all these components.
Learned MLP update. Several learned optimizer architectures have been proposed in previous works (Andrychowicz et al.) 2016; Metz et al., 2020 Almeida et al. 2021 Metz et al.) 2022a b), each varying in the types of imputs they process. the elements they maintain within their state, and the functional form of their outputs. In this section, we review a learned optimizer parameterized using a multi-layer perceptron (MLP) (Metz et al.) 2022a) (Adafactor MLP LOpt) that work and other state-of-the-an works like Vel.O (Metz et al. 2022b) build upon. It is a simple MLP-basee learned optimizer that can serve as a drop-in replacement for hand-designed update rules such as SGD or Adam (Kingma & Ba 2015).
shspe
R (N.M
grad
(N.M.3)
with
(K.M.13
momenrum in with
rata iy with 05
jaram
gnd V.0
(4.
(M.4 1)
1)
Adam
Adl, (V (1.
normallieed momemum R,,T, with Pi.PP..
adulactor row festues with Pe,
reciprocal 1ny"
(..M..3
CN.M
(6.10
Icarned MLP
A0 (h,9,13 param
Inerztion
param
adafacnoe recipmical now featurss 1yy
update
koming rate
(N.
apdiae
adaiaacue colunn features with
update
adadaactos reciprocal column
(N.M.3)
adaicit gradicmn:V x 1i My.
(N.M.33
adafector momertum 1n x, IN MAF
(K.M..3
loannng rave d. (N.M.13
Nel (N.0.13
ON M.13
Platt (N. (N))
Swi (N.7)
Sey (Mi.1
example
shape
Figure 8: Comparing Adam with learned MLP update rule. Given parameter tensor of shape (N, M), the Adam update (left) maintains two additional accumulators per parameter, leading to 3x memory overhead over SGD. In contrast, the learned MLP optimizer proposed by [Mtz ct al. (2022a) (right) incorporates additional momentum, root mean square-normalized features, and Adafactor-styl row and column features. While the MLP optimizer maintains extra row and col- umn accumulators, their memory cost scales sublinearly with the total parameter count N x M. In total, MLP consists of <200 params (Metz al., 2022a) and provides effective trade-off between runtime and memory overhead, making it practical for large-scale tasks. We adapt the same learned MLP optimizer illustrated the right in our work for per-parameter updates (with few changes, Tensor shape annotations are shown in blue, with tiled dimensions in red.
see $ 5
Concretely, at a given iteration f, the Adafactor MLP learned optimizer takes as input the parameter vector 0, and gradient VO,. 'hese values are used to build per-parameter input features that are passed to the learned MIP,, as illustrated in Figure g The learned MLP takes these input features and updates different accumulators in its state s,, which we describe in detail next.
The Adafactor MLP optimizer maintains four types of accumulators for each tensor (i.c., layer weight or bias) in its state, as illustrated in Figure g) (1) gradient momentum m, with 82, Ss. (2) root mean square gradient accumulator 24 with 34. (3) Adafactor row feature accumulator r. with Ss. Se, 37, and (4) Adafactor column feature accumulator c, with 8s, Bs, 87. The momentum m, and RMS U, accumulators are "perrpaarmmter," meaning they correspond to each parameter in given tensor. In contrast, the Adafactor accumulators rt and e are "per-tensor," these features
22

---

Published as a conference paper at ICLR 2026
extracted from the entire parameter tensor. Formally, given parameter 0 and its gradient V from parameter tensor 0, E RN* x momentum my and RMS accumulators is for each parameter updated follows:
M4+1 B.m + (1- v
V (1,2,3)
(I)
U++1 B40 +11 (1- 34) V
where ,1,..., 84 are scalar constants. In Metz et al. (2022a b). three momentum accumulators are used with B1, 00, 883 (0.9, 0.99, 0.999) and one RMS accumulator with 8n 0.999. Celo2,
we found 34 0.95 to be optimal for LM pretraining tasks, see $5 Adafactor accumulators (T),
maintain row and column mean of squared gradient tensor in each iteration:
T++1 = Byre + (1 83) rew- V0²)
Vj E & (5,
(2)
C++1 = 3,ce + (1 B3 /ol mol-mean(( V0²)
where $3. P6, 97 are scalar constants set to (0.9, 0.99, 0.999) in Metz et al. (2022b a). These row and column accumulators are tiled (replicated) corresponding to each parameter position and are then used to compute Adafactor momentum m, x 1/vTi x 1/ and Adafactor gradient VO x
1/VT x 1/v/ features, as illustrated in Figure g
These accumulated values are then used to build input features for each parameter, which are passed to the learned MLP rule. Metz et al.] (2022a) also passes global training progress features wP as input to the MLP update rule, which we omit since we defer step-size tun- ing to the practitioner ($3). Concretely, Metz et al. (2022a) constructs training progress fea- tures w computed by taking the current iteration and computing tanh(t/z) where [1,3, 10,30, 100, 300, 1000, 3000, 10k, 30k, 100k]. In total, the learned MI.P optimizer by Metz input for each parameter. These features are processed through a 2-layer MI.P network with 4 hidden units each and ReI.U activations (we ablate over hidden size in S which has total 197 parameters (Metz et al. 2022a). The forward pass through the learned returns two outputs of the same dimensionality as Úr, corresponding to direction d and magnitude m, which are used to compute parameter updates A0; and then parameters 0(+1 as follows:
jet al. (2022a) takes 39 input features
A0: \ A (amm))
0,+1 = 0, 20.,
where a and 22 are fixed scalars set to low values (0.001) to keep meta-training stable. The learned optimizers discussed in this work, including VeLO (Metz al. 2022a), Celo (Moudgil al. 2025). and our proposed Celo2, roughly maintain the same Adafactor MLP architecture but differ in minor details such as accumulator constants (8 values), progress features and update functional form. We refer the reader to the respective works for more details.
Meta-training problem. Unlike hand-designed optimizers, learned optimizers' parameterized up- date function Φ is trained through a meta-training process. A standard approach to meta-training involves solving a bi-level optimization problem that consists of an inner problem, which opti- mizes network parameters 0 using the learned optimizer update & on a sampled task, and outer problem, which optimizes the learned optimizer parameters & based on feedback from the inner loop (Andrychowicz et al.) 2016; Wichrowska et al.] 2017: Metz et al.) [019a] 2020 Vicol et al.) 2021). Formally, given a set of optimization tasks T. the learned optimizer parameters & are ob- tained by sampling an optimization task consisting of cata distribution D, initial network parameters Og. and training objective L, and solving the bi-level problem below:
t = arg min Ep
1
T
Σ
((x::01)))
(4)
T
where the inner loop is recursively defined for L= [0,T 1] as:
(0,, 81) = (0o. 0)
if t 0,
(5)
(0,, s1) fo(01-1, V0-1.
if t 0;
(6)
V0,- cC(Xi: 01, ②)
X,~D D
31.
(7)
20,
23

---

Published as a conference paper at ICLR 2026
Here T 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 commonly, the final loss of the inner loop (Metz et al., 2019a 2020: 2022b). Computing meta- gradients for Φ in Eq. is challenging, as we explain next.
Evolution strategies for gradient estimation. Directly prop-propgg through the inner loop in Eq. to compute meta-graan eeta-gradieets can lead to noisy gradient estimates, especially when the unroll length T' is large, due to the accumulation of errors through the computational graph (Metz et al.) 2019a). An alternative approach is to use Evolution Strategies (ES) (Rechenberg, 1973, Hansen & Ostermeier 200T] Wierstra et al., 2014; Salimans et al. 2017). class of black- box optimization methods that estimate gradients by evaluating the objective function at perturbed parameter val- ues. ES estimates the gradient of a Gaussian-smoothee objective with respect to parameters & by sampling perturbations e ~ (0, I² and evaluating:
n T Σ J( E.)E
(8)
where n is the number of perturbation samples and a is the perturbation scale. This gradient estimate is obtained by evaluating the objective function J different perturbed versions of Φb (specifically at & + e, for each i), and then computing a weighted average of the perturbations based on how much each perturbation improved or worsened the objective. ES has the advantage of not requiring backpropagation through the inner loop, making it memory-effciiet and applicable even when the inner optimization process is non-differentiable or extremely long. However, standard ES applied to the full unroll of length T in Eq. 4 can be computationally expensive, it requires completing entire training runs for each perturbation sample before computing any meta-gradients. This is where Persistent Evolution Strategies (PES) provides practical solution.
Evolution Strategies (PES). PES (Vicol et al. 202I extends ES to work efficiently with truncated unrolls while maintaining unbiased gradient estimates. Instead of unrolling the entire inner loop of length T' before computing meta-gradee veta-gradients, PES divides the inner loop into shorter segments of length K (where K T). After each segment, meta- peeta-gradients are computed using ES on that truncated segment, and the learned optimizer parameters Φ are updated without resetting the inner optimization. The key insight of PES is that it maintains persistence across truncations. Specifically, the optimizer state 8, and inner problem parameters 0, are carried forward after each meta-gradient update, and perturbations than reset. This persistence is crucial: without it, truncated gradient estimates would be biased because they would not account for how parameter updates affect future optimization trajectories. PES provides unbiased gradient estimates in the truncated setting by carefully accounting for how perturbations c to the learned optimizer parameters affect both the immediate truncated objective and the carry-over effects through the persistent state.
Persistent
the learned optimizer parameters are accumulated over time rather
Formally, for a truncated segment from step to L. + K, PES estimates:
n
k+K-1
ss Amuncatdd
C 0(X;00%.000
1
where -,00 = et denotes the accumulated perturbation for particle i, and g denotes the inner parameters at step j when using the perturbed optimizer, starting from the persistent state at step f. This makes meta-training significantly more efficient, as the learned optimizer can be updated frequently (every K steps) rather than only after full unrolls of length T. The frequent updates also lead to better training dynamics, as the learned optimizer can adapt more quickly to improve performance. We use PES for computing meta- meta-gradiee gradients in our work, following prior work (Metz et al.
2022a Harrison et al. 2022 Moudgil et al. 2025).
Task augmentation. To improve the generalization of learned optimizers, meta-training can incor- porate task augmentation (Metz et al.) 2022b] Moudgil et al., (025). Task augmentation modifies the sampled optimization tasks during meta-training to create a more diverse training distribution, helping the learned optimizer generalize better to unseen tasks at test time. One effective form of augmentation is parameter reparameterization, where the initial parameters Oy of each sampled task are scaled by random factors. This reparameterization can be done at three levels of granularity: (1) g global, where all parameters are scaled by the same factor, (2) per-tensor, where each weight ma- trix or bias vector is scaled independently, or (3) per-parameter, where each individual parameter is
24

---

Published as a conference paper at ICLR 2026
scaled independently. Formally, given sampled task with initial parameters uu the reparameterized initialization is O,- a O Oo. where is scaling factor (or tree of scaling factors for per-tensor and per-parameter cases) sampled from a log-uniform distribution:
= cxp(uu)
Uniform(In (In Omin In amax))
(10)
with the range typically set to [Qiinn maax [0.001, 1000]. During training, the learned opti- mizer sees parameters at vastly different scales within each task, which encourages to learn scale- invariant update rules that generalize better across different initialization schemes and architectures. In our work, we follow the task augmentation approach used in Celo (Moudgil et al., 2025), applying reparameterization global-level during meta-training to improve generalization to unseen tasks.
Outer-training loop. Putting these components together, the complete meta-training procedure for learned optimizers proceeds follows. each meta-training iteration, we sample task (D, C. 0o) from the task distribution T and apply task augmentation to obtain reparameterized initial parame- ters 0g. We then perform K steps of inner optimization using the current learned optimizer fo on this task, accumulating a truncated meta-objective over these steps. PES is used compute unbiased gradient estimates V for the learned optimizer parameters by evaluating perturbations Φ + 6; on the same truncated segment. The learned optimizer parameters are then updated using these meta-gradient (typically with Adam another outer optimizer). The inner problem state (0;, 8;) is maintained across truncations, and this process repeats for multiple truncations until the total unroll length T is reached or the task is considered complete. This entire process is repeated across many sampled tasks to train the learned optimizer to generalize across the task distribution T. The combi- nation of efficient gradient estimation via PES, task augmentation for improved generalization, and the expressive MLP architecture enables the practical meta-training of learned optimizers that can match exceed the performance of hand-designed optimizers like Adam variety of tasks.
25


Reviews:
[{"review_id": "eUd2ZEZxEX", "reviewer": "Reviewer_mn2u", "summary": "This paper addresses a major limitation of existing learned optimizers (LOs), which is that they do not meta-generalize well out of distribution, particularly when the optimizee is large or the optimization is long. To mitigate this issue, the authors propose Celo2 with a simple yet effective LO architecture that adopts orthogonalization and RMS normalization, which has been shown to be successful in hand-crafted optimizers. The experiments show that Celo2 meta-trained on relatively small datasets and optimizees in a supervised learning paradigm outperforms strong baselines in three out-of-distribution settings: larger optimizee, longer optimization horizon, and application to reinforcement learning. Detailed ablations are provided to analyze each component of Celo2.", "questions": "The ablations are conducted from a generalization-first perspective in what sense? How does the ablation analysis focus on generalization?\n\nWhat is the purpose of using both TPUs and GPUs in the experiments?\n\nWhat is the optimizer used for meta-training?\n\nHow many steps of optimization were used for the results shown in Fig. 1?\n\nWhy do Celo2-base and Adam have identical wall clock time? Does the MLP application on line 5 of algorithm 1 take significantly higher computation?\n\nWhat are the batch sizes used in the experiments?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 4, "strengths": "The method proposed is simple yet effective, indicating a great potential for generalization and further improvement.\n\nThe experimental results provide strong support for the claims.\n\nThe contribution of this paper is significant towards the practical application of LOs.", "weaknesses": "I suggest avoiding the term \"free lunch\" in the paper, as the no-free-lunch theorem indicates that the inductive bias found by Celo2 will likely downgrade performance on some tasks, although these tasks may be unlikely in realistic applications.\n\nThe paper states, \"We would like to emphasize that our primary objective while developing this simple approach for learned optimization was stability.\" However, no sensitivity analysis is provided.\n\nAlgorithm 1 is not referred to in the text.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "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": "", "extra_scores": {}}, {"review_id": "4v5amkzrFf", "reviewer": "Reviewer_TU4w", "summary": "This paper presents a provocative and potentially groundbreaking finding in learned optimization (LO). It directly challenges the prevailing \"scaling hypothesis\" (e.g., VeLO), which posits that massive meta-training compute is necessary for generalization. The authors introduce Celo2, a learned optimizer meta-trained on a \"toy\" distribution (8x8 image classification) for a mere 4.5 GPU hours.\n\nThe central, surprising claim is that this \"cheap\" LO, built on a simple normalized MLP architecture and task augmentation, generalizes six orders of magnitude beyond its training data. It successfully and stably trains billion-scale models (GPT-3 1.3B) and vision transformers (ViT), outperforming not only the strong AdamW baseline but also the exorbitantly expensive VeLO, which famously fails on such large-scale tasks.", "questions": "1. The decoupling of the learning rate (Sec 3) is a major simplification. How much of Celo2's superior performance and stability can be attributed to relying on a hand-tuned cosine schedule, which VeLO was explicitly designed to avoid?\n\n2. Figure 4 shows that Celo2's best performance comes from hybridization (adding Orthogonalization and AdamW@1D). Does this suggest that the future of LOs is not \"pure\" end-to-end learning, but rather learning \"plugins\" that work with hand-designed rules?\n\n3. Could you please expand on the choice of \"average loss\" vs. \"final loss\" as the meta-objective? Does Celo2's stability come from optimizing this easier, myopic objective? What happens if Celo2 is meta-trained to optimize the final loss, as VeLO does?", "limitations": "", "rating": 4, "confidence": 3, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "1. The paper's core finding—that a 4.5 GPU-hour \"toy\" meta-training can produce an LO that scales to 1.3B parameter models—is a \"free lunch\" that could fundamentally realign research in this field. It moves LOs from a \"computationally impossible\" (VeLO) to a \"highly practical\" domain.\n\n2. The paper rightly centers its comparison against VeLO. The results are stark: VeLO, meta-trained with 4000 TPU-months, is unstable and fails on large tasks, while Celo2, trained for 4.5 hours, is stable and superior. This is the paper's strongest point. It demonstrates robust generalization across the three most critical axes:\n\nModel Scale: From tiny 8x8 MLPs to 1.3B GPT-3 models.\n\nUnroll Length: From short 2k-step unrolls to 10B+ tokens (GPT-3) and 50k steps (ViT).\n\nTask Domain: From 2D Image Classification to 1D Language Modeling and Reinforcement Learning (Atari).\n\n3. aThe proposed LO (\"Celo2-base\") is a small, 8-hidden-unit MLP (Table 1a). The design is explicitly intended as a 1-line \"drop-in replacement\" for Adam (Sec 4.1), which is a massive win for practical adoption.", "weaknesses": "The paper makes extraordinary claims (4.5 GPU-hour meta-training, 6-orders-of-magnitude generalization) with a simple recipe. Such \"too good to be true\" results demand exceptional evidence. However, the paper is missing an appendix and supplementary material, providing no code, implementation details, or full hyperparameter lists beyond what is in the main text. This makes it impossible to verify the claims or assess reproducibility, which is paramount for such a shocking result.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "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": "", "extra_scores": {}}, {"review_id": "EOB2St1Y1U", "reviewer": "Reviewer_9J2f", "summary": "This paper proposes a new learnable optimizer Celo2. It is intended to outperforms Celo and other alternatives in terms of generalization, scalability and compute efficiency. This work is built on top of prior Celo optimizer  (Moudgil et al., 2025). This paper introduces the following changes:\na. It retains a user-tunable step size, not touching the learned scheduler. This helps the scalability.\nb. It learns a small MLP to serve as a drop-in replacement for the standard Adam update rule. This makes the proposed Celo more computing efficient.\nc. The proposed algorithm normalizes the MLP's output, which improves the generalization.\n\nCelo2 shows strong results on ImageNet image classification with ViT. In exp with GPT-2 and GPT-3, it also achieves consistent lower validation loss. On Atari benchmark, it is shown effective in a PPO algorithm to learn an RL policy.", "questions": "1. Could the authors clarify how Celo2 is expected to generalize to even larger models? Large models (LLMs/VLMs) have a massive number of parameters, and their optimizer states (like Adam's moments) consume significant GPU memory. How does the proposed MLP-based approach compare in terms of GPU memory efficiency?\n\n2. What is the architecture of the small MLP used for ImageNet classification versus the one used for the LLM experiments? It is not clear how much this optimizer component needs to be adapted when moving between different problem domains.\n\n3. The paper claims the RMS-normalized learned update leads to better generalization, but the mechanism is unclear. Is there a more detailed ablation study or theoretical justification to support this specific design choice?", "limitations": "", "rating": 2, "confidence": 3, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "The proposed Celo2 does seem to simpler, scalable and generalizable based on the experiments.\n\nThere are a wide variety of experiments to back up the claims, e.g. vision, language, and RL.", "weaknesses": "The proposed method does not provide sound theoretical guarantee, making it feel ad-hoc. This raises significant concerns about its practical application. In particularly, I would be very concerning on how would a practitioner use the proposed Celo2 in real-world application. For example, it is not clear how to replace the Adam with small MLP? It is unclear if this technique was found to work only on specific, \"cherry-picked\" problems and whether the MLP architecture must be significantly re-tuned when switching to a new problem domain\n\nThe experimental results feel incomplete. While reporting validation loss is informative as a preliminary measure, it is not sufficiently convincing on its own. The paper would be much stronger if it included results from more established, concrete benchmarks (e.g., perplexity or BLEU scores for LLMs, not just training/validation loss).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "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": "", "extra_scores": {}}, {"review_id": "RQqfBEGH7g", "reviewer": "Reviewer_D6de", "summary": "The paper proposes a new learned optimizer that improves upon the state-of-the-art VeLO optimizer. It achieves higher performance while requiring less compute and data during training, and it scales effectively to large tasks. The optimizer generates an update direction, after which classical update components (such as orthogonalization, normalization, and step size adjustment) are applied to produce the final parameter update.", "questions": "- How sensitive is the proposed optimizer to the choice of step size and learning rate schedule compared to standard optimizers like AdamW? Could you provide results or plots showing performance across the range of tested hyperparameters?\n\n- During meta-training, which step size or learning rate schedule is actually used? Is it fixed, or does it follow the same cosine decay schedule described for evaluation?\n\n- Are the orthogonalization and normalization steps used during the optimizer’s training phase, or only at test time?", "limitations": "", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "- The proposed approach simplifies the training process and scales effectively to larger problems.\n\n- The usage and implementation of the optimizer (in Algorithm 1) are well described.\n\n- The paper presents comprehensive comparisons across multiple tasks and datasets of varying sizes, demonstrating consistent improvements.", "weaknesses": "- The paper depends on manually tuned step sizes and learning rate schedules. For instance, the authors note that they “test with 7 learning rate values sampled logarithmically between 1e-3 and 1e-5, warmup fraction 0.05 with cosine decay schedule, and plot the best performing hyperparameter setting for each optimizer.” However, only the best-performing configuration is reported, without clarifying the selection criteria. This limits the understanding of tuning sensitivity and comparability—for example how easily the method can be tuned relative to AdamW.\n\n- The training procedure is described briefly and relies on citations rather than detailed explanation. For instance, “We use Persistent Evolutionary Strategies (PES) (Vicol et al., 2021) to meta-train our optimizer with unroll length logarithmically sampled between 100 and 2000 steps.” Re-explaining the main aspects of PES and its role in this work would make the paper more accessible and reproducible.\n\n- It is unclear how specific training hyperparameters are chosen and applied:\n  - Which step size is used during meta-training for each unrolled iteration?\n  - Is the cosine decay schedule also used during training?\n  - Are orthogonalization and normalization applied at training time, or only at evaluation?\n\nThese details are essential for reproducibility and for understanding how much the method’s performance depends on them.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "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": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)