|
|
--- |
|
|
license: apache-2.0 |
|
|
tags: |
|
|
- model_hub_mixin |
|
|
- pytorch_model_hub_mixin |
|
|
- learned-optimizer |
|
|
--- |
|
|
|
|
|
|
|
|
# Description |
|
|
`btherien/mulo` is a learned optimizer meta-trained in μ-parameterization. It corresponds to the μLO_M optimizer from [μLO: Compute-Efficient Meta-Generalization of Learned Optimizers](https://arxiv.org/abs/2406.00153). Due to being meta-trained in μP, μLO_M has strong meta-generalization capabilities (i.e., the ability to optimize unseen tasks), despite its relatively short and inexpensive meta-training distribution. |
|
|
### Learned optimizer meta training and architectural details |
|
|
| **Field** | **Value** | |
|
|
|------------------------------|---------------------------------------------------------------------------| |
|
|
| **Meta-training distribution** | ImageNet classification, 3-layer MLP, width ∈ {128, 512, 1024} | |
|
|
| **Number of meta-training steps** | 5000 | |
|
|
| **Target inner problem length** | 1000 iterations | |
|
|
| **Gradient estimator** | Persistent Evolution Strategies | |
|
|
| **Truncation length** | 50 | |
|
|
| **Architecture** | small_fc_lopt | |
|
|
| **Optimizer Input size** | 39 | |
|
|
| **Optimizer Hidden size** | 32 | |
|
|
| **Optimizer Output size** | 2 | |
|
|
|
|
|
# Usage |
|
|
--- |
|
|
## 1) Install PyLO |
|
|
The following |
|
|
```bash |
|
|
git clone https://github.com/Belilovsky-Lab/pylo |
|
|
cd pylo |
|
|
pip install . --config-settings="--build-option=--cuda" #Optional installation with Cuda |
|
|
``` |
|
|
|
|
|
## (2) Use $\mu$LO as a drop-in replacement for pytorch learned optimizers |
|
|
|
|
|
```python |
|
|
if USE_CUDA_KERNEL: |
|
|
from pylo.optim import MuLO_CUDA |
|
|
optimizer = MuLO_CUDA(model.parameters(), hf_key='btherien/mulo') |
|
|
else: |
|
|
from pylo.optim import MuLO_naive |
|
|
optimizer = MuLO_naive(model.parameters(), hf_key='btherien/mulo') |
|
|
``` |
|
|
|
|
|
## (3) A simple example |
|
|
|
|
|
The following example is for illustration purposes and does not implement the correct parameterizaiton. For a correct implementation see https://github.com/Belilovsky-Lab/pylo/tree/main/examples |
|
|
|
|
|
```python |
|
|
import torch |
|
|
import torch.nn as nn |
|
|
import torch.optim as optim |
|
|
from torchvision import datasets, transforms |
|
|
from torch.utils.data import DataLoader |
|
|
|
|
|
# Model |
|
|
class MLP(nn.Module): |
|
|
def __init__(self): |
|
|
super().__init__() |
|
|
self.net = nn.Sequential( |
|
|
nn.Flatten(), |
|
|
nn.Linear(28 * 28, 128), |
|
|
nn.ReLU(), |
|
|
nn.Linear(128, 10) |
|
|
) |
|
|
def forward(self, x): |
|
|
return self.net(x) |
|
|
|
|
|
model = MLP().to(device) |
|
|
|
|
|
######################### |
|
|
Setup Learned Optimizer |
|
|
######################### |
|
|
#USE_CUDA_KERNEL=True # Uncomment for accelerated kernels |
|
|
if USE_CUDA_KERNEL: |
|
|
from pylo.optim import MuLO_CUDA |
|
|
optimizer = MuLO_CUDA(model.parameters(), hf_key='btherien/mulo') |
|
|
else: |
|
|
from pylo.optim import MuLO_naive |
|
|
optimizer = MuLO_naive(model.parameters(), hf_key='btherien/mulo') |
|
|
|
|
|
# Device |
|
|
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') |
|
|
|
|
|
# Data |
|
|
transform = transforms.ToTensor() |
|
|
train_loader = DataLoader(datasets.MNIST(root='./data', train=True, download=True, transform=transform), |
|
|
batch_size=64, shuffle=True) |
|
|
|
|
|
criterion = nn.CrossEntropyLoss() |
|
|
# Training loop |
|
|
for epoch in range(1): # Just 1 epoch for simplicity |
|
|
for x, y in train_loader: |
|
|
x, y = x.to(device), y.to(device) |
|
|
optimizer.zero_grad() |
|
|
loss = criterion(model(x), y) |
|
|
loss.backward() |
|
|
optimizer.step() |
|
|
|
|
|
print("Done!") |
|
|
``` |
|
|
|
|
|
|
|
|
### Per-Parameter Input Features Used by MuLO |
|
|
|
|
|
| **Type** | **# Features** | **Description** | **Equation** | |
|
|
|---------------------------|----------------|------------------------------------------------------------------------------------------------------|---------------------------------------------------------------------------| |
|
|
| **Accumulators** | 3 | Momentum accumulators with coefficients βᵢ, i ∈ {1, 2, 3}. | mₜ,ᵢ = βᵢ·mₜ₋₁,ᵢ + (1 − βᵢ)·∇ₜ | |
|
|
| | 1 | Second moment accumulator with coefficient β₄. | vₜ = β₄·vₜ₋₁ + (1 − β₄)·∇ₜ² | |
|
|
| | 3 | Adafactor row accumulators with coefficients βᵢ, i ∈ {5, 6, 7}. | rₜ,ᵢ = βᵢ·rₜ₋₁,ᵢ + (1 − βᵢ)·row_mean(∇ₜ²) | |
|
|
| | 3 | Adafactor column accumulators with coefficients βᵢ, i ∈ {5, 6, 7}. | cₜ,ᵢ = βᵢ·cₜ₋₁,ᵢ + (1 − βᵢ)·col_mean(∇ₜ²) | |
|
|
| **Accumulator Features** | 3 | Normalized momentum: momentum divided by sqrt of second moment for i ∈ {5, 6, 7}. | mₜ,ᵢ / √v | |
|
|
| | 1 | Reciprocal sqrt of second moment value. | 1 / √v | |
|
|
| | 6 | Reciprocal sqrt of Adafactor accumulators. | 1 / √(rₜ,ᵢ or cₜ,ᵢ) | |
|
|
| | 3 | Adafactor gradient features for i ∈ {5, 6, 7}. | ∇ₜ × rₜ,ᵢ × cₜ,ᵢ | |
|
|
| | 3 | Adafactor momentum features for (i, j) ∈ {(5,1), (6,2), (7,3)}. | mₜ,ⱼ × rₜ,ᵢ × cₜ,ᵢ | |
|
|
| **Time Features** | 11 | Time features for x ∈ {1, 3, 10, 30, 100, 300, 1000, 3000, 10⁴, 3·10⁴, 10⁵}. | tanh(t / x) | |
|
|
| **Parameters** | 1 | Parameter value. | wₜ | |
|
|
| | 1 | Gradient value. | ∇ₜ | |
|
|
| **Total** | 39 | — | — | |
|
|
|
|
|
|
|
|
# Cite |
|
|
If you found this optimizer useful in your research, please consider citing our work: |
|
|
```bibtex |
|
|
@misc{therien2024mulo, |
|
|
title = {$\mu$LO: Compute-Efficient Meta-Generalization of Learned Optimizers}, |
|
|
author = {Benjamin Thérien and Charles-Étienne Joseph and Boris Knyazev and Edouard Oyallon and Irina Rish and Eugene Belilovsky}, |
|
|
year = {2024}, |
|
|
eprint = {2406.00153}, |
|
|
archivePrefix = {arXiv}, |
|
|
primaryClass = {cs.LG}, |
|
|
url = {https://arxiv.org/abs/2406.00153} |
|
|
} |
|
|
``` |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
