深度学习中的优化器之SGD with Polyak Averaging 算法原理与实现细节

题目描述

在深度学习的优化过程中，随机梯度下降（SGD）是基础且广泛使用的优化算法。然而，SGD及其变体在训练后期常常在最优解附近振荡，影响收敛的稳定性和最终模型的泛化性能。为此，Polyak Averaging（或Polyak-Ruppert Averaging）作为一种模型参数平滑技术被提出。其核心思想是：在优化过程中，不仅记录当前时刻的参数，还对历史参数进行指数衰减平均或简单平均，从而得到一个更平滑、更稳定的参数估计，有助于模型收敛到更平坦的最小值，并提升泛化能力。本题目将详细讲解SGD with Polyak Averaging的原理、计算步骤和实现细节。

解题过程（原理与实现细节讲解）

第一步：理解基础SGD的局限性

1. 标准SGD更新规则：
   对于模型参数 θ，在第 t 次迭代时，SGD的更新公式为：

\[ θ_{t+1} = θ_t - η ∇L(θ_t) \]

其中，η 是学习率，∇L(θ_t) 是在当前批次数据上计算的损失函数梯度。

2. 振荡问题：
   由于SGD使用小批量数据估计梯度，梯度估计本身具有噪声。在接近最优解时，梯度幅度变小，噪声相对影响增大，导致参数更新路径在最小值附近振荡，难以精确收敛。

3. 泛化与平坦最小值：
   研究表明，收敛到“平坦”的最小值（即损失函数在参数空间中变化平缓的区域）的模型，通常比收敛到“尖锐”最小值的模型具有更好的泛化能力。SGD的振荡可能使其在尖锐最小值附近徘徊，而平均化参数有助于平滑路径，倾向于找到更平坦的区域。

第二步：Polyak Averaging 的核心思想

1. 基本概念：
   Polyak Averaging 并不是一个独立的优化算法，而是一种后处理或在线平滑技术。它在优化迭代过程中，维护一个参数的“平均版本” \(\bar{θ}\)，这个平均版本是通过对历史参数进行平均得到的。

2. 目标：

   • 减少方差：通过对噪声梯度导致的参数振荡进行平均，降低最终参数的方差。
   • 促进收敛：理论上可以证明，在某些条件下，平均参数序列的收敛速度可以比原始参数序列更快（达到最优收敛速率）。
   • 提升泛化：平均参数往往对应于损失函数中更平坦的区域，可能提高模型在测试集上的性能。

第三步：Polyak Averaging 的两种主要实现方式

方式一：简单平均（Simple Average）

1. 做法：
   在训练过程中，从某个时刻开始（例如，训练的后半段），对之后所有迭代得到的参数进行算术平均。

\[ \bar{θ}_T = \frac{1}{T - T_0 + 1} \sum_{t=T_0}^{T} θ_t \]

其中，T 是总迭代次数，T_0 是开始平均的迭代次数（例如 T_0 = T/2）。最终使用 $\bar{θ}_T$ 作为模型参数。

2. 特点：
   • 实现简单。
   • 通常需要在训练的后阶段进行，以避免早期不稳定参数的影响。
   • 需要存储所有历史参数或在线更新平均值。

方式二：指数移动平均（Exponential Moving Average, EMA）

1. 做法：
   这是更常用、更高效的方法。它在线维护一个平均参数 \(\bar{θ}\)，每次迭代后都用当前参数 θ_t 对其进行更新，更新规则为：

\[ \bar{θ}_{t} = β \cdot \bar{θ}_{t-1} + (1 - β) \cdot θ_t \]

其中，β 是衰减率（通常接近1，例如0.99, 0.999），控制历史平均值的权重。初始时，$\bar{θ}_0 = θ_0$。

2. 特点：
   • 只需存储一个额外的平均参数变量，内存开销小。
   • 赋予了近期参数更高的权重，是一种平滑操作。
   • 在实际训练中，我们通常用EMA得到的 \(\bar{θ}\) 来进行最终的模型评估和预测，而用原始参数 θ 进行梯度计算和更新。这被称为 “影子参数”。

第四步：SGD with Polyak Averaging 的完整算法流程

我们以最常用的SGD with EMA为例，详述其迭代步骤：

1. 初始化：

   • 初始化模型参数 \(θ_0\)。
   • 初始化平均参数 \(\bar{θ}_0 = θ_0\)。
   • 设置学习率 η，衰减率 β（例如0.995），总迭代次数 T。

2. 迭代训练（对于 t = 0 到 T-1）：
   a. 前向传播与梯度计算：
   使用当前参数 \(θ_t\) 在 mini-batch 数据上前向传播，计算损失 \(L(θ_t)\) 和梯度 \(g_t = ∇L(θ_t)\)。

   b. 参数更新（SGD步骤）：
   使用标准SGD规则更新参数：

\[ θ_{t+1} = θ_t - η \cdot g_t \]

c. **平均参数更新（EMA步骤）**：
   利用刚刚更新得到的 $θ_{t+1}$ 来更新平均参数 $\bar{θ}_{t+1}$：

\[ \bar{θ}_{t+1} = β \cdot \bar{θ}_t + (1 - β) \cdot θ_{t+1} \]

   注意，这里用 $θ_{t+1}$ 还是 $θ_t$ 取决于实现细节，但用更新后的 $θ_{t+1}$ 更为常见。

3. 训练完成：
   训练结束后，我们丢弃原始参数序列 \(\{θ_t\}\)，而将最终的平均参数 \(\bar{θ}_T\) 作为训练好的模型参数用于推理和测试。

第五步：关键细节与实现注意事项

1. 偏差校正：
   在训练初期，由于 \(\bar{θ}_0\) 被初始化为 \(θ_0\)，会导致指数移动平均值存在偏差（偏向初始值）。一种常见的做法是进行偏差校正：

\[ \hat{\bar{θ}}_t = \frac{\bar{θ}_t}{1 - β^t} \]

其中 t 是迭代次数。随着 t 增大，$β^t$趋近于0，偏差校正的作用减弱。很多现代框架（如PyTorch的`torch.optim.swa_utils`）在实现中会处理或提供相关选项。

2. 何时开始平均：
   有时不会从一开始就进行平均，而是等待优化器“预热”（warm-up）一段时间，让参数进入一个相对稳定的区域后再开始EMA。这可以通过动态调整β或控制平均的起始步来实现。

3. 与优化器的结合：
   Polyak Averaging 可以与任何基于梯度的优化器结合，如SGD、Adam等。其逻辑是独立的：优化器负责更新“原始参数”，而另一个EMA模块负责维护“平均参数”。

4. 代码实现示例（PyTorch风格伪代码）：

   import torch
   import torch.nn as nn
   
   # 模型和优化器定义
   model = MyModel()
   optimizer = torch.optim.SGD(model.parameters(), lr=0.01)
   
   # 初始化平均参数影子字典
   ema_weights = {}
   for name, param in model.named_parameters():
       ema_weights[name] = param.data.clone()  # 初始化为当前参数值
   
   beta = 0.999  # EMA衰减率
   
   # 训练循环
   for epoch in range(num_epochs):
       for batch_data, batch_labels in dataloader:
           # 1. 清零梯度
           optimizer.zero_grad()
   
           # 2. 前向传播与损失计算（使用原始参数）
           outputs = model(batch_data)
           loss = criterion(outputs, batch_labels)
   
           # 3. 反向传播
           loss.backward()
   
           # 4. 优化器更新原始参数
           optimizer.step()
   
           # 5. 更新平均参数（EMA）
           for name, param in model.named_parameters():
               ema_weights[name] = beta * ema_weights[name] + (1 - beta) * param.data
   
   # 训练结束后，将平均参数加载回模型用于推理
   for name, param in model.named_parameters():
       param.data = ema_weights[name]
   

第六步：相关变体与扩展

• 随机权重平均（Stochastic Weight Averaging, SWA）：可以看作是Polyak Averaging的一个特定实例。SWA通常以固定周期（而非每一步）保存参数的快照，并对这些快照进行等权平均。它在学习率周期性调整时特别有效，能收敛到更广阔的最小值区域。
• 指数移动平均的衰减率调度：可以动态调整β，例如在训练初期使用较小的β以快速跟踪参数变化，后期使用较大的β以获得更平滑的平均。

总结

SGD with Polyak Averaging 通过维护模型参数的移动平均，有效平滑了优化过程中的噪声和振荡，倾向于使模型收敛到更平坦、泛化能力更好的最小值区域。其实现简单，计算和存储开销小，是一种即插即用且能稳定提升模型性能的实用技术。在具体实践中，指数移动平均（EMA） 因其高效和在线性而最为常用，常作为深度模型训练时的标准技巧之一。