无调度学习

PyTorch中的无调度优化器。

作者：Aaron Defazio、Xingyu (Alice) Yang、Harsh Mehta、Konstantin Mishchenko、Ahmed Khaled、Ashok Cutkosky

简述无需调度的更快训练 - 无需提前指定停止时间/步数！

pip install schedulefree

主要实现包括SGDScheduleFree和AdamWScheduleFree。我们还提供了一个AdamWScheduleFreeReference版本，它具有简化的实现，但使用更多内存。要与其他优化器结合使用，请使用ScheduleFreeWrapper版本。

Jax实现可作为Optax的一部分使用。

方法

无调度学习用插值和平均的组合替代了底层优化器的动量。对于梯度下降，基本的无调度更新为：

$$ \begin{align*} y_{t} & = (1-\beta)z_{t} + \beta x_{t},\ z_{t+1} & =z_{t}-\gamma\nabla f(y_{t}),\ x_{t+1} & =\left(1-\frac{1}{t+1}\right)x_{t}+\frac{1}{t+1}z_{t+1}, \end{align*} $$

这里$x$是应进行测试/验证损失评估的序列，它与主要迭代$z$和梯度评估位置$y$不同。对$z$的更新对应于底层优化器，在这种情况下是简单的梯度步骤。

顾名思义，无调度学习不需要递减的学习率调度，但通常优于或至少匹配最先进的调度，如余弦衰减和线性衰减。只需同时存储两个序列（第三个可以从其他两个实时计算），因此该方法与基础优化器具有相同的内存要求（参数缓冲区+动量）。

我们在这个仓库中提供了AdamW和SGD版本，以及一个可与任何基础优化器一起使用的实验性包装器版本。

如何使用

由于我们的优化器在梯度调用和测试/验证损失计算中使用两个不同的点，因此在训练期间需要在两者之间切换参数缓冲区。这可以通过在调用model.train()的同地方调用optimizer.train()，以及在调用model.eval()的同地方调用optimizer.eval()来完成。在存储检查点时，优化器也应该处于eval模式。

如果您的代码支持PyTorch优化器步骤闭包，您可以使用优化器的闭包形式，这不需要.train()和.eval()调用。

论文

如果您在工作中使用无调度训练，请引用我们的预印本：

@misc{defazio2024road,
      title={The Road Less Scheduled}, 
      author={Aaron Defazio and Xingyu Yang and Harsh Mehta and Konstantin Mishchenko and Ahmed Khaled and Ashok Cutkosky},
      year={2024},
      eprint={2405.15682},
      archivePrefix={arXiv},
      primaryClass={cs.LG}
}

示例

使用schedulefree包的示例可以在examples文件夹中找到。这些包括：

使用卷积网络进行图像分类（MNIST）*
更多示例将陆续添加

*示例修改自Pytorch示例仓库。

注意事项

如果您的模型使用BatchNorm，则需要进行额外修改以使测试/验证评估正常工作。在评估之前，需要类似以下操作：

model.train()
optimizer.eval()
with torch.no_grad():
  for batch in itertools.islice(train_loader, 50):
    model(batch)
model.eval()

这将用在$x$处计算的值替换training_mean/training_var缓存（在model.train()模式下，每次前向传播都会更新）。使用PreciseBN也可以避免这个问题。

许多代码库使用可能需要额外更改才能兼容的其他功能。例如，如果参数以fp16缓存，则需要手动更新缓存版本，以确保使用正确的$x$序列进行评估，而不是$y$序列。一些GradScaler可以做到这一点。
训练对$\beta$的选择比标准动量更敏感。我们默认的0.9在大多数问题上都有效，但对于非常长的训练运行，可能需要将值增加到0.95或0.98。
不需要使用学习率调度器，但代码与之兼容。
建议使用学习率预热。这通过warmup_steps参数支持。
此方法确实需要调整 - 如果不同时调整正则化和学习率参数，它不一定会优于调度方法。
对于SGD，比经典速率大10-50倍的学习率似乎是一个好的起点。
对于AdamW，比基于调度的方法大1-10倍的学习率似乎有效。

包装器版本

我们提供了一个实验性的包装器版本ScheduleFreeWrapper，可以包装任何基础优化器。使用此版本时，您可以禁用基础优化器的动量，因为在使用我们的包装器动量时不再需要它（尽管如果需要，您可以同时使用两种类型的动量）。

示例用法：

base_optimizer = torch.optim.RMSprop(model.parameters(), lr=0.0025)
optimizer = ScheduleFreeWrapper(
  base_optimizer, momentum=0.9, weight_decay_at_y=0.1)

如果您在基础优化器上设置权重衰减，它会在$z$处计算权重衰减。我们通过weight_decay_at_y参数提供在$y$处计算权重衰减的选项，这在我们的实验中似乎能得到更好的结果。

我们还包括一个ScheduleFreeWrapperReference版本，它使用更多内存但数值上更稳定，我们建议将此版本用于早期实验或研究工作。