无调度学习

PyTorch 中的无调度优化器。

预印本：The Road Less Scheduled

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

简而言之：无需调度即可实现更快的训练——无需提前指定停止时间或步数！

pip install schedulefree 我们提供了几种无调度优化器的实现：

• SGDScheduleFree 和 SGDScheduleFreeReference：SGD 的无调度变体
• AdamWScheduleFree 和 AdamWScheduleFreeReference：AdamW 的无调度变体
• RAdamScheduleFree：RAdam 的无调度变体，既不需要学习率调度，也不需要预热（由社区贡献的实现）
• 实验性的 ScheduleFreeWrapper，可与其他优化器结合使用

ScheduleFreeReference 版本实现了简化版逻辑，但占用更多内存。此外，还有 ScheduleFreeClosure 版本，可用于 PyTorch 的优化器步进闭包。

作为 Optax 的一部分，还提供了一个 Jax 实现。

方法

无调度学习用插值和平均的组合替代了基础优化器的动量。以梯度下降为例，基本的无调度更新公式如下：

$$ \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$ 的更新对应于基础优化器，在此例中即简单的梯度步进。

顾名思义，无调度学习无需递减的学习率调度，但通常性能优于或至少与 SOTA 调度方法（如余弦退火和线性衰减）相当。每次只需存储两个序列（第三个可通过前两个实时计算），因此该方法的内存需求与基础优化器相同（参数缓冲区 + 动量）。

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

使用方法

由于我们的优化器在梯度计算和测试/验证损失计算时使用不同的点，因此在训练过程中需要在这两者之间切换参数缓冲区。这可以通过在调用 model.train() 的地方调用 optimizer.train()，在调用 model.eval() 的地方调用 optimizer.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}
}

发布说明

新增 1.4 版本添加了由 nhamanasu 提供的 RAdam 实现。 1.3 版本调整了学习率预热期间权重衰减的行为，以提高稳定性，并使其与 PyTorch 中标准 AdamW 的行为更加一致。之前的实现仍可作为 AdamWScheduleFreePaper 使用。

示例

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

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

*示例改编自 Pytorch Examples Repo。

注意事项

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

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

这会将 training_mean/training_var 缓存（在 model.train() 模式下每次前向传播时都会更新）替换为基于 $x$ 而非 $y$ 计算的值。使用 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$ 点计算权重衰减。我们提供了在 $y$ 点计算权重衰减的选项，通过 weight_decay_at_y 参数实现，这在我们的实验中似乎效果更好。

我们还提供了一个 ScheduleFreeWrapperReference 版本，该版本占用更多内存，但数值稳定性更高，建议用于早期实验或研究工作。

许可证

请参阅 许可证文件。

相关工作

无调度学习可以被视为原始平均法（$\beta=1$）与 Polyak-Ruppert 平均法（$\beta=0$）之间的插值。这种插值的优势在于它能够兼得两者的优点：我们既可以获得 Polyak-Ruppert 平均法在训练初期的快速收敛特性（因为 $z$ 序列的更新速度比 $x$ 序列更快），又不会让 $x$ 序列偏离 $z$ 序列太远，从而避免不稳定性。

我们的方法还与 Nesterov 的加速方法（Nesterov, 1983）在 AC-SA 形式下（Ghadimi & Lan, 2010）相关：

$$ \begin{align*} y_{t} & =(1-2/(t+1))x_{t} + (2/(t+1))z_{t}\ z_{t+1} & =z_{t}-\frac{t}{2L}\nabla f(y_{t})\ x_{t+1} & =(1-2/(t+1))x_{t}+(2/(t+1))z_{t+1} \end{align*} $$

我们的方法同样包含三个序列，但使用的权重截然不同，且关键在于并未采用随时间递增的学习率，而后者对于 Nesterov 方法实现加速至关重要。此外，我们在插值操作和平均操作中也采用了不同的权重序列。

诸如随机权重平均（Stochastic Weight Averaging, Izmailov 等人, 2018）和最新权重平均（LAtest Weight Averaging, Kaddour, 2022；Sanyal 等人, 2023）等尾部平均方法，将平均化与较大或循环学习率相结合。然而，这些方法仍然需要使用学习率调度策略，引入额外的超参数进行调优，并且相比我们的方法需要更多的内存。值得注意的是，SWA 和 LAWA 也可以叠加在我们的方法之上，从而可能带来进一步的性能提升。

Portes 等人（2022）使用周期性递增的学习率调度策略，以探索训练时间与评估性能之间帕累托前沿上的多个点。每个周期结束时的模型近似于在该时刻终止的调优学习率调度所得到的模型。相比之下，我们的方法能够生成整个帕累托前沿，而不仅仅是路径上的少数几个点。

流行的 Lookahead 优化器（Zhang 等人, 2019）则使用迭代序列的指数移动平均（EMA）。从某种意义上说，Lookahead 方法可以视为原始平均法的 EMA 版本，正如指数权重平均是 Polyak-Ruppert 平均法的 EMA 版本一样。我们额外的插值步骤同样有可能与 Lookahead 优化器结合使用。

Schedule-Free 快速上手指南

Schedule-Free 是一种无需预设学习率调度（Learning Rate Schedule）即可实现高效训练的优化器方法。它通过插值和平均机制替代传统动量，通常能匹配或超越余弦退火等主流调度策略的效果，且无需提前指定训练停止步数。

环境准备

• 系统要求：Linux, macOS, Windows
• Python 版本：建议 Python 3.8+
• 核心依赖：
  • PyTorch >= 1.8
  • 无需其他特殊前置依赖

安装步骤

推荐使用 pip 直接安装官方发布版本：

pip install schedulefree

如需加速下载，可使用国内镜像源（如清华大学源）：

pip install schedulefree -i https://pypi.tuna.tsinghua.edu.cn/simple

基本使用

Schedule-Free 优化器需要在训练模式（计算梯度）和评估模式（计算验证损失）之间切换参数缓冲区。请务必在调用 model.train() 时调用 optimizer.train()，在 model.eval() 时调用 optimizer.eval()。

1. 导入与初始化

以 AdamWScheduleFree 为例（也支持 SGDScheduleFree 等）：

import torch
from schedulefree import AdamWScheduleFree

# 定义模型
model = MyModel()

# 初始化优化器
# 注意：通常需要使用比传统方法更大的学习率 (1x-10x)
optimizer = AdamWScheduleFree(model.parameters(), lr=0.005, weight_decay=0.1)

2. 训练循环示例

for epoch in range(num_epochs):
    # --- 训练阶段 ---
    model.train()
    optimizer.train()  # 关键：切换到训练模式
    
    for batch in train_loader:
        inputs, targets = batch
        
        optimizer.zero_grad()
        outputs = model(inputs)
        loss = criterion(outputs, targets)
        loss.backward()
        optimizer.step()

    # --- 验证阶段 ---
    model.eval()
    optimizer.eval()  # 关键：切换到评估模式，此时参数位于 x 序列
    
    with torch.no_grad():
        for batch in val_loader:
            inputs, targets = batch
            outputs = model(inputs)
            # 计算验证指标...
    
    # 保存检查点时也需保持 optimizer.eval() 状态
    if should_save_checkpoint:
        torch.save({
            'model_state_dict': model.state_dict(),
            'optimizer_state_dict': optimizer.state_dict(),
        }, 'checkpoint.pth')

3. 进阶提示

• 学习率设置：
  • SGD: 尝试比经典值大 10-50 倍的学习率。
  • AdamW: 尝试比经典值大 1-10 倍的学习率。
• Warmup：虽然不需要复杂的调度，但建议使用 warmup_steps 参数进行学习率预热以提升稳定性。
• BatchNorm 注意事项：如果模型包含 BatchNorm 层，在验证前可能需要运行少量训练批次以更新统计信息，确保评估基于正确的参数序列 $x$。