Merge pull request #12 from Belilovsky-Lab/xiao_elo

Add learned optimizers meta-trianed with ELO

Author: xiaol827

.gitignore

@@ -8,4 +8,6 @@ wandb
 mup
 _build
 dist
-snippets
+snippets
+# Learned-optimizer checkpoints (hosted on the Hugging Face Hub)
+*.pickle

README.md

@@ -17,6 +17,9 @@ PyLO provides efficient PyTorch implementations of cutting-edge learned optimize
 - **PyTorch-native API** designed for simplicity and familiarity
 - **Hugging Face integration** for sharing and loading meta-models
 
+Learned optimizers:
+* ELO series: [https://arxiv.org/abs/2506.10315](https://arxiv.org/abs/2506.10315)
+
 # Installation
 
 ### Via URL (slow, no Kernels)
@@ -72,22 +75,28 @@ python -m pylo.util.patch_mup
 
 ## Quick Start
 
+Taking `ELO-CELO2` (the strongest LO) for example.
+
 ```python
 import torch
-from pylo.optim import VeLO_CUDA
+from pylo.optim import ELO_CELO2_CUDA
 
-# Initialize a model
 model = torch.nn.Linear(10, 2)
 
-# Create a learned optimizer instance
-optimizer = VeLO_CUDA(model.parameters())
+num_steps = 1000  # total optimization steps
+
+# Meta-learned weights download automatically from the Hugging Face Hub on first use.
+# The optimizer has no built-in LR schedule; drive it with a standard
+# torch.optim.lr_scheduler (warmup, cosine, etc.).
+optimizer = ELO_CELO2_CUDA(model.parameters(), lr=3.16e-4, weight_decay=0.1, adam_lr_mult=20)
+scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=num_steps, eta_min=3.16e-5)
 
-# Use it like any PyTorch optimizer
-for epoch in range(10):
+for step in range(num_steps):
     optimizer.zero_grad()
     loss = loss_fn(model(input), target)
     loss.backward()
-    optimizer.step(loss) # pass the loss 
+    optimizer.step()
+    scheduler.step()
 ```
 
 ## Sharing Learned Optimizers

pylo/__init__.py

@@ -5,9 +5,16 @@
     "MetaMLP",
     "VeLOMLP",
     "VeLORNN",
+    "CELO2MLP",
     "AdafacLO_naive",
     "MuLO_naive",
     "VeLO_naive",
+    "CELO2_naive",
+    "ELO_CELO2_naive",
+    "ELO_naive",
+    "CELO2",
+    "ELO_CELO2",
+    "ELO",
     # Default aliases (re-exported from pylo.optim). These resolve to the
     # CUDA implementations when available, falling back to the naive ones
     # otherwise, so downstream code can simply `from pylo import VeLO`.