让深度学习新手也能用上高效自学习优化器。
PyLO: Towards Accessible Learned Optimizers in PyTorch
- 基于PyTorch实现可直接替换Adam的自学习优化器
- 在ViT-B模型上训练速度提升4倍以上
- 支持与学习率调度等工具自由组合使用
过去十年，自学习优化器成为研究热点，近期如VeLO（Meta训练耗时4000个TPU月）已展现出实用潜力，但因依赖JAX且缺乏易用包，难以普及。为此，我们提出PyLO，一个基于PyTorch的库，通过torch.optim.Optimizer接口将自学习优化器带给约70%的非JAX用户。重点面向真实大规模预训练任务，提供fc lopt与VeLO的CUDA加速实现，在ViT-B/16（batch size 32）上训练吞吐量从39.36和49.73提升至205.59和191.18样本/秒。该库支持与学习率调度、权重衰减等工具灵活结合，实验发现能显著提升优化器性能。代码已开源。
原文摘要 · Abstract (English)
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
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。