File size: 4,085 Bytes
d040338 b1bfe10 d040338 b1bfe10 cc0e792 b1bfe10 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 | ---
license: mit
library_name: optax
tags:
- optimizer
- learned-optimizer
- meta-learning
- jax
---
# Celo2: Towards Learned Optimization Free Lunch
<p>
<a href="https://arxiv.org/abs/2602.19142"><img alt="Paper" src="https://img.shields.io/badge/arXiv-2602.19142-b31b1b.svg"></a>
<a href="https://github.com/amoudgl/celo2"><img alt="Code" src="https://img.shields.io/badge/GitHub-black?logo=github&logoColor=white&labelColor=grey"></a>
<a href="https://opensource.org/licenses/MIT"><img alt="License: MIT" src="https://img.shields.io/badge/License-MIT-yellow.svg"></a>
</p>
Official pretrained weights for **Celo2**: This variant applies (and is meta-trained with) a harness that includes Newton-Schulz orthogonalization on top of the learned update for matrix parameters and uses AdamW for biases/embeddings. For a fully-learned variant without any harness, see [celo2-base](https://huggingface.co/amoudgl/celo2-base).
## Quickstart
Download checkpoint and install:
```bash
pip install git+https://github.com/amoudgl/celo2.git
hf download amoudgl/celo2 --local-dir ./celo2
```
Use `load_checkpoint` method to fetch pretrained params from checkpoint path:
```python
from celo2_optax import load_checkpoint
pretrained_params = load_checkpoint('./celo2/theta.state')
```
Standard optax usage with `scale_by_celo2` method that takes pretrained params as input:
```python
import optax
from celo2_optax import scale_by_celo2
optimizer = optax.multi_transform(
transforms={
'celo2': optax.chain(
scale_by_celo2(pretrained_params, orthogonalize=True),
optax.add_decayed_weights(weight_decay),
optax.scale_by_learning_rate(lr_schedule),
),
'adam': optax.adamw(lr_schedule, 0.9, 0.95, weight_decay=weight_decay),
},
param_labels=lambda params: jax.tree.map_with_path(
lambda path, val: 'adam' if val.ndim <= 1 or 'embed' in jax.tree_util.keystr(path) else 'celo2',
params,
),
)
```
## Loading and inspecting MLP update rule weights
```python
from celo2_optax import load_checkpoint
import jax
pretrained_params = load_checkpoint('./celo2/theta.state') # dictionary containing weights
print(jax.tree.map(lambda x: x.shape, pretrained_params))
```
The checkpoint contains a small MLP stored under the `ff_mod_stack` key with weight matrices (`w0__*`, `w1`, `w2`) and biases (`b0`, `b1`, `b2`). Each `w0__*` key contains weights corresponding to particular input feature such as momentum, gradient, parameter, etc.
## Meta-training config
| Key | Value |
| ----------------------- | ------------------------------------------------------------ |
| **Optimizer architecture** | MLP, 2 hidden layers, 8 units each |
| **Meta-training tasks** | 4 image classification tasks (MNIST, FMNIST, CIFAR-10, SVHN) |
| **Task architecture** | MLP (64-32-10) |
| **Meta-trainer** | Persistent Evolution Strategies (PES) |
| **Outer iterations** | 100K |
| **Truncation length** | 50 |
| **Min unroll length** | 100 |
| **Max unroll length** | 2000 |
For more details, see config JSON included in the repo [here](./config.json).
## Files
| File | Description |
| ------------- | -------------------------------- |
| `theta.state` | Pretrained MLP optimizer weights |
| `config.json` | Meta-training configuration |
## Citation
```bibtex
@misc{moudgil2026celo2,
title={Celo2: Towards Learned Optimization Free Lunch},
author={Abhinav Moudgil and Boris Knyazev and Eugene Belilovsky},
year={2026},
eprint={2602.19142},
archivePrefix={arXiv},
primaryClass={cs.LG},
url={https://arxiv.org/abs/2602.19142},
}
```
|