Celo2: Towards Learned Optimization Free Lunch

Abhinav Moudgil; Boris Knyazev; Eugene Belilovsky

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

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.

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

\bm{d}

instead of more complex update rules that need to predict both

\bm{d}

and

\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 form	val loss $\downarrow$
raw output	$\Delta\bm{p}_{t}$ = MLP(…)	3.961
rolling norm	$\Delta\bm{p}_{t}\,/\,\sum^{t}\left\lVert\Delta\bm{p}_{t}\right\rVert_{\text{rms}}$	3.860
rolling norm clip	$\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	$\Delta\bm{p}_{t}\cdot\min\bigl(1,\tau/\left\lVert\Delta\bm{p}_{t}\right\rVert_{\text{rms}}\bigr)$	4.084
norm	$\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

\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.

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
model	D	L	F	H	N	T
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

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=4D

and

N=D/H

. Vocab size for FineWeb 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	0.1	0.1	0.1
adamw epsilon	1e-8	1e-8	1e-8

Appendix

Table 4: Hyperparameters used for the language modeling experiments in Figure 1 .

hyperparameter	value
learning rate (celo2)	0.000046
learning rate (adamw)	0.0001
dataset	imagenet-256
model	ViT
hidden size	768
depth (layers)	12

Appendix

Table 5: Hyperparameters used for the Vision Transformer (ViT) experiments in Figure 2 .

hyperparameter	freeway	space-invaders	asterix
learning rate (celo2)	0.000215	0.000464	0.0001
learning rate (adamw)	0.001	0.001	0.001
total timesteps	10M	10M	10M
discount factor ( $\gamma$ )	0.99	0.99	0.99
GAE parameter ( $\lambda$ )	0.95	0.95	0.95
num steps per update	128	128	128

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)

, 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\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 .

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

Efficient Long-Horizon Learning for Learned Optimization
Jul 7, 2026Xiaolong Huang, Benjamin Thérien, James Harrison +1Meta-Learning Learning to Optimize
Mila - Quebec AI Institute · Concordia University · Universit´e de Montr´eal +1
Lightweight and Versatile Learned Optimization by Recombination of Gradient History
Oct 7, 2026Minyoung Choi, Dalta Imam Maulana, Wanyeong JungLearning to Optimize
KAIST Daejeon, Republic of Korea
Extending $μ$ P: Spectral Conditions for Feature Learning Across Optimizers
Feb 24, 2026Akshita Gupta, Marieme Ngom, Sam Foreman +1Hyperparameter Transfer Adaptive Gradient Methods
Purdue University · Argonne National Laboratory

Jul 7, 2026J/K move · Enter open · S save

Efficient Long-Horizon Learning for Learned Optimization

Xiaolong Huang, Benjamin Thérien, James Harrison, Eugene Belilovsky

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

Meta-Learning Learning to Optimize

Open paper AskPDF arXiv

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.

Abstract

Learning to Optimize

Explore similar work

Efficient Long-Horizon Learning for Learned Optimization

Lightweight and Versatile Learned Optimization by Recombination of Gradient History

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

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