Description

VeLO is a learned optimizer meta-trained on thousands of diverse machine learning tasks. It corresponds to the VeLO (Versatile Learned Optimizer) from VeLO: Training Versatile Learned Optimizers by Scaling Up.

Learned optimizer meta training and architectural details

Usage

---

1) Install PyLO

The following

git clone https://github.com/Belilovsky-Lab/pylo
cd pylo
pip install .
python setup.py install --cuda

(2) Use VeLO as a drop-in replacement for pytorch optimizers

from pylo.optim import VeLO
optimizer = VeLO(model.parameters(), lr=1.0 , num_steps=150_000)

(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

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
#########################
optimizer = VeLO(model.parameters(), lr=1.0 , num_steps=150_000)

# Device
device = torch.device('cuda') 

# 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(loss)

print("Done!")

                                                    |

Note: VeLO requires the total number of training steps as input to initialize its internal states and compute training progress features.

Official Resources

• Paper: VeLO: Training Versatile Learned Optimizers by Scaling Up
• Adapted from Repository: Google Learned Optimization
• Available in PyLO [PyLO] (https://github.com/Belilovsky-Lab/pylo)

Important Notes

• Step Requirement: VeLO requires knowledge of total training steps for proper initialization
• Computational Cost: Meta-trained using ~4000 TPU-months, representing significant computational investment
• Generalization: Designed for meta-generalization across diverse optimization tasks
• No Hyperparameter Tuning: Automatically adapts to problem specifics without manual tuning

Cite

If you found this optimizer useful in your research, please consider citing the original work:

@article{metz2022velo,
  title={{VeLO}: Training Versatile Learned Optimizers by Scaling Up},
  author={Luke Metz and James Harrison and C. Daniel Freeman and Amil Merchant and Lucas Beyer and James Bradbury and Naman Agrawal and Ben Poole and Igor Mordatch and Adam Roberts and Jascha Sohl-Dickstein},
  journal={arXiv preprint arXiv:2211.09760},
  year={2022},
  url={https://arxiv.org/abs/2211.09760}
}