# PyLO: Towards Accessible Learned Optimizers in PyTorch

arXiv: https://arxiv.org/abs/2506.10315 · HTML: https://mlsanity.com/paper/2506.10315/

## Key facts

- Cited by 1 papers in the map (data through 14 Oct 2026).
- Builds on 17 of its 37 references that are in the map.
- First posted on arXiv 12 Jun 2025.

## Builds on (most cited first, 17 of 17)

- [Learning Transferable Visual Models From Natural Language Supervision](https://mlsanity.com/paper/2103.00020/) (2021)
- [An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale](https://mlsanity.com/paper/2010.11929/) (2020)
- [PyTorch: An Imperative Style, High-Performance Deep Learning Library](https://mlsanity.com/paper/1912.01703/) (2019)
- [Llama 2: Open Foundation and Fine-Tuned Chat Models](https://mlsanity.com/paper/2307.09288/) (2023)
- [Segment Anything](https://mlsanity.com/paper/2304.02643/) (2023)
- [SGDR: Stochastic Gradient Descent with Warm Restarts](https://mlsanity.com/paper/1608.03983/) (2016)
- [FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness](https://mlsanity.com/paper/2205.14135/) (2022)
- [ZeRO: Memory Optimizations Toward Training Trillion Parameter Models](https://mlsanity.com/paper/1910.02054/) (2019)
- [Adafactor: Adaptive Learning Rates with Sublinear Memory Cost](https://mlsanity.com/paper/1804.04235/) (2018)
- [PyTorch FSDP: Experiences on Scaling Fully Sharded Data Parallel](https://mlsanity.com/paper/2304.11277/) (2023)
- [Learning to learn by gradient descent by gradient descent](https://mlsanity.com/paper/1606.04474/) (2016)
- [Learning to Optimize: A Primer and A Benchmark](https://mlsanity.com/paper/2103.12828/) (2021)
- [Real-time Neural Radiance Caching for Path Tracing](https://mlsanity.com/paper/2106.12372/) (2021)
- [VeLO: Training Versatile Learned Optimizers by Scaling Up](https://mlsanity.com/paper/2211.09760/) (2022)
- [Practical tradeoffs between memory, compute, and performance in learned optimizers](https://mlsanity.com/paper/2203.11860/) (2022)
- [μLO: Compute-Efficient Meta-Generalization of Learned Optimizers](https://mlsanity.com/paper/2406.00153/) (2024)
- [Is Scaling Learned Optimizers Worth It? Evaluating The Value of VeLO's 4000 TPU Months](https://mlsanity.com/paper/2310.18191/) (2023)

## Built on by (most cited first, 1 of 1)

- [Learning to Optimize Radiotherapy Plans via Fluence Maps Diffusion Model Generation and LSTM-based Optimization](https://mlsanity.com/paper/2605.13713/) (2026)

## Abstract (arXiv)

Learned optimizers have been an active research topic over the past decade, with increasing progress toward practical, general-purpose optimizers that can serve as drop-in replacements for widely used methods like Adam. However, recent advances such as VeLO, which was meta-trained for 4000 TPU-months, remain largely inaccessible to the broader community, in part due to their reliance on JAX and the absence of user-friendly packages for independently using the optimizers after meta-training. To address this gap, we introduce PyLO, a PyTorch-based library that brings learned optimizers to the remaining ~70% of machine learning community via the familiar torch.optim.Optimizer interface. Unlike prior work focused on limited-scale academic tasks, our emphasis is on applying learned optimization to real-world large-scale pre-training tasks. Our systems contribution includes CUDA-accelerated implementations of the small fc lopt(Metz et al., 2022a) and VeLO(Metz et al., 2022b) learned optimizers, achieving substantial performance gains, with training throughput on ViT-B/16 (batch size 32) increasing from 39.36 and 49.73 to 205.59 and 191.18 samples per second, respectively. PyLO has the versatility that allows us to easily combine learned optimizers with existing optimization tools such as learning rate schedules and weight decay. When doing so, we discover that learned optimizers can substantially benefit from it. Our code is available at https://github.com/Belilovsky-Lab/pylo