cs.LGFeb 22, 2026

Celo2: Towards Learned Optimization Free Lunch

Authors: Abhinav Moudgil, Boris Knyazev, Eugene Belilovsky

Organizations: Mila · Concordia University · Samsung AI Lab, Montreal · Université de Montréal

Read PDF
Ask AI
Cite

Copy citation

DOINot available for this paper

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 (∼\sim10×\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.

Figures & tables

Figure 1: Celo2, our learned update rule scales stably to large-scale pretraining tasks. Our learned update is meta-trained on very limited compute budget (4.5 GPU hours) but generalizes to 1,000,000 ×\times larger language modeling task (GPT-3) and outperforms strong well-tuned baselines such as AdamW ( Loshchilov, 2017 ; Kingma and Ba, 2015 ) and VeLO ( Metz et al., 2022b ) . In contrast, 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 normalization, rotary positional embeddings (RoPE), QK-norm, and GELU ( Marek et al., 2025 ) .
Figure 2: ImageNet classification with ViT . We test our learned update rule, Celo2, on the ImageNet classification task with batch size 512 and 50K steps, which is 25 ×\times longer than its meta-training unroll length and 30,000 ×\times larger than the tasks seen during training. Since VeLO is 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 ∼\sim 50% steps. As test accuracy reaches ∼\sim 66%, all optimizers start overfitting in this task; this is consistent with findings in prior work ( Dahl et al., 2023 ) . Since our update rule is normalized, it shows training norm dynamics consistent with AdamW. Moreover, VeLO is meta-trained with 200K unroll length on a large number 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-generalization capability.
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 at par with a well-tuned AdamW baseline on these out-of-distribution tasks, while the VeLO baseline stagnates at a much lower return. The latter result can be corroborated by Figure 11 in Metz et al. (2022b) .
Figure 4: As shown in the figure on the left, Celo2-base that uses 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 1D params improve performance when applied on top of Celo2-base. Applying these two techniques directly at test-time improves performance but meta-training with them is even better.
Table 1: Ablations over 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 to 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 a significant improvement in 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\bm{d} instead of more complex update rules that need to predict both d\bm{d} and m\bm{m} . All ablations are evaluated on a 30M transformer decoder model trained on FineWeb-Edu with 600M tokens (§ 5 ). Celo2-base with our default settings are highlighted .
Appendix figures & tables9 assets

Supplementary material from the paper’s appendix.

Appendix

 functional formval loss ↓\downarrow 
raw outputΔpt\Delta\bm{p}_{t} = MLP(…)3.961 
rolling normΔpt / ∑t∥Δpt∥rms\Delta\bm{p}_{t}\,/\,\sum^{t}\left\lVert\Delta\bm{p}_{t}\right\rVert_{\text{rms}}3.860 
rolling norm clipΔpt⋅min⁡(1,τ⋅(∑t∥Δpt∥rms)/∥Δpt∥rms)\Delta\bm{p}_{t}\cdot\min\bigl(1,\tau\cdot(\sum^{t}\left\lVert\Delta\bm{p}_{t}\right\rVert_{\text{rms}})/\left\lVert\Delta\bm{p}_{t}\right\rVert_{\text{rms}}\bigr)3.999 
norm clipΔpt⋅min⁡(1,τ/∥Δpt∥rms)\Delta\bm{p}_{t}\cdot\min\bigl(1,\tau/\left\lVert\Delta\bm{p}_{t}\right\rVert_{\text{rms}}\bigr)4.084 
normΔpt / ∥Δpt∥rms\Delta\bm{p}_{t}\,/\,\left\lVert\Delta\bm{p}_{t}\right\rVert_{\text{rms}}3.812 
Appendix
Table 2: We evaluate several ways of scaling the predicted update vector Δpt\Delta\mathbf{p}_{t} . “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 a time-dependent threshold τ\tau . “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 a 30M-parameter Transformer trained on the FineWeb-Edu dataset with 600M tokens (§ 5 ).
Figure 5: Validation loss curves for Celo2 with various learning rates uniformly sampled on log-scale between 1e-5 and 1e-3 on LM (30M) FineWeb-Edu dataset.
Figure 6: Parallel coordinate plot showing Celo2 hyperparameter sensitivity on SpaceInvaders RL task. Yellow curves indicate high-return configurations.
Figure 7: Comparison of Celo2 and Muon on GPT-2 (124M) FineWeb pretraining task.
modelDLFHNT
30M (LM-30M)384615361283512
124M (GPT-2)76812307212861024
1.3B (GPT-3)2048248192128162048
D: hidden dimension, L: number of transformer layers, F: feedforward inner dimension, H: attention head dimension, N: number of attention heads, T: context/sequence length
Appendix
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=4DF=4D and N=D/HN=D/H . Vocab size for FineWeb dataset is 50257 and rounded to match device sharding requirements.
hyperparameterLM-30MGPT-2 (124M)GPT-3 (1.3B)
learning rate (celo2)0.0004640.00020.00002
learning rate (adamw)0.019760.00480.0002
warmup fraction0.050.050.05
batch size (global)256512512
weight decay0.10.10.1
adamw epsilon1e-81e-81e-8
Appendix
Table 4: Hyperparameters used for the language modeling experiments in Figure 1 .
hyperparametervalue
learning rate (celo2)0.000046
learning rate (adamw)0.0001
datasetimagenet-256
modelViT
hidden size768
depth (layers)12
Appendix
Table 5: Hyperparameters used for the Vision Transformer (ViT) experiments in Figure 2 .
hyperparameterfreewayspace-invadersasterix
learning rate (celo2)0.0002150.0004640.0001
learning rate (adamw)0.0010.0010.001
total timesteps10M10M10M
discount factor ( γ\gamma )0.990.990.99
GAE parameter ( λ\lambda )0.950.950.95
num steps per update128128128
Appendix
Table 6: Hyperparameters used for the RL experiments in Figure 3 .
Figure 8: Comparing Adam with learned MLP update rule. Given a parameter tensor of shape (N,M)(N,M) , the Adam update (left) maintains two additional accumulators per parameter, leading to a 3 ×\times memory overhead over SGD. In contrast, the learned MLP optimizer proposed by Metz et al. (2022a) (right) incorporates additional momentum, root mean square-normalized features, and Adafactor-style row and column features. While the MLP optimizer maintains extra row and column accumulators, their memory cost scales sublinearly with the total parameter count N×MN\times M . In total, MLP consists of <200 params ( Metz et al., 2022a ) and provides an effective trade-off between runtime and memory overhead, making it practical for large-scale tasks. We adapt the same learned MLP optimizer illustrated on the right in our work for per-parameter updates (with a few changes, see § 5 ). Tensor shape annotations are shown in blue , with tiled dimensions in red .
More in this topic

Learning to Optimize

Learning to Optimize uses data to adapt algorithms, solver decisions, and solution maps for faster or better optimization.

See all 138 papers

Explore similar work

CardsList
Jul 7, 2026

Efficient Long-Horizon Learning for Learned Optimization

Learned optimization aims to improve upon hand-designed optimizers (e.g., Adam and Muon) by meta-learning small neural network optimizers over a distribution of tasks. While recent work has greatly advanced the architectural design and inductive biases of learned optimizers (LOs), their meta-training remains biased toward short-unroll learning on particular tasks, resulting in redundant computation and leaving LOs often unable to compete with hand-designed optimizers. We introduce Efficient Long-hOrizon (ELO) learning, an efficient meta-training algorithm that (1) reallocates wasted meta-training compute to longer failure regimes, achieving efficient long-horizon learning, and (2) enforces decoupled progressive expert supervision, providing stable meta-learning signals that additionally improve the generalization of LOs. Our empirical study evaluates ELO for meta-training both element-wise and matrix-based LOs. Across downstream language modeling (GPT-2 and GPT-3 on FineWeb) and image classification (ViT-B/16, ResNet-50 on ImageNet-1K) tasks, ELO substantially improves the long-unroll performance and out-of-distribution generalization of the base LOs. In particular, ELO-Celo2 outperforms well-tuned AdamW and Muon on GPT-2 (30M and 124M) language pretraining. Notably, all ELO baselines require less than 7 H100 GPU-hours for meta-training.

Xiaolong Huang, Benjamin Thérien, James Harrison +1

Mila - Quebec AI Institute · Concordia University · Universit´e de Montr´eal +1

Meta-LearningLearning to Optimize
AskPDF
Oct 7, 2026

Lightweight and Versatile Learned Optimization by Recombination of Gradient History

This paper presents a lightweight and versatile learned optimizer that dynamically recombines gradient history, represented as averages over disjoint time spans. The optimizer reduces the prediction space to one scalar coefficient per gradient average, shared by multiple parameters. Progressively averaging older gradients minimizes memory cost of long history, while keeping their contributions independently accessible. A 37k-parameter network trained in 0.87 GPU-hours generalizes zero-shot to unseen tasks, lowering validation loss by 9.1% and 0.4% on BERT-Tiny and GPT-Tiny, and improving test accuracy over Adam by 3.5 %p on a Vision Transformer and by 2.7 %p on average across nine graph models, with FLOPs overhead as low as 0.3%.

Minyoung Choi, Dalta Imam Maulana, Wanyeong Jung

KAIST Daejeon, Republic of Korea

Learning to Optimize
AskPDF
Feb 24, 2026

Extending μμP: Spectral Conditions for Feature Learning Across Optimizers

Several variations of adaptive first-order and second-order optimization methods have been proposed to accelerate and scale the training of large language models. The performance of these optimization routines is highly sensitive to the choice of hyperparameters (HPs), which are computationally expensive to tune for large-scale models. Maximal update parameterization (μ(μP)) is a set of scaling rules which aims to make the optimal HPs independent of the model size, thereby allowing the HPs tuned on a smaller (computationally cheaper) model to be transferred to train a larger, target model. Despite promising results for SGD and Adam, deriving μμP for other optimizers is challenging because the underlying tensor programming approach is difficult to grasp. Building on recent work that introduced spectral conditions as an alternative to tensor programs, we propose a novel framework to derive μμP for a broader class of optimizers, including AdamW, ADOPT, LAMB, Sophia, Shampoo and Muon. We implement our μμP derivations on multiple benchmark models and demonstrate zero-shot learning rate transfer across increasing model width for the above optimizers. Further, we provide empirical insights into depth-scaling parameterization for these optimizers.

Akshita Gupta, Marieme Ngom, Sam Foreman +1

Purdue University · Argonne National Laboratory

Hyperparameter TransferAdaptive Gradient Methods
AskPDF