PyLO：面向PyTorch的可访问学习优化器
机器学习
2026-04-20 v3
摘要
学习优化器在过去十年中一直是活跃的研究领域，不断取得进展 toward 实用的、通用的优化器，可作为广泛使用的方法（如Adam）的直接替代品。然而，最近的进展如VeLO（经过4000 TPU-month的元训练）在很大程度上仍难以被更广泛的社区获取，部分原因是它们依赖JAX且缺乏用户友好的包来在元训练后独立使用优化器。为填补这一空白，我们引入了PyLO，一个基于PyTorch的库，通过熟悉的torch.optim.Optimizer接口将学习优化器带给剩余约70%的机器学习社区。与先前专注于有限规模学术任务的工作不同，我们强调将学习优化应用于真实世界的大规模预训练任务。我们的系统贡献包括对小型fc lopt(Metz et al., 2022a)和VeLO(Metz et al., 2022b)学习优化器的CUDA加速实现，实现了显著的性能提升，在ViT-B/16（批量大小32）上的训练吞吐量从39.36和49.73分别提升至205.59和191.18样本每秒。PyLO具有灵活性，允许我们轻松地将学习优化器与现有优化工具（如学习率调度和权重衰减）结合使用。通过这样做，我们发现学习优化器可以从中显著受益。我们的代码可在https://github.com/Belilovsky-Lab/pylo获取。
引用
@article{arxiv.2506.10315,
title = {PyLO: Towards Accessible Learned Optimizers in PyTorch},
author = {Paul Janson and Benjamin Therien and Quentin Anthony and Xiaolong Huang and Abhinav Moudgil and Eugene Belilovsky},
journal= {arXiv preprint arXiv:2506.10315},
year = {2026}
}
备注
Accepted at MLSys Conference 2026