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

题目描述
SGD with Polyak Averaging（也称Polyak-Ruppert平均）是一种优化技术，旨在提升随机梯度下降（SGD）的收敛性和稳定性。其核心思想是：在SGD的迭代过程中，不仅记录当前参数，还维护参数的运行平均值（即历史参数的加权平均），最终使用平均后的参数作为模型输出。这种方法能有效减少SGD的方差，尤其在非凸优化中帮助逼近更平坦的最小值，从而提高泛化能力。本题目将详细解释其数学原理、平均机制的设计动机，以及实际实现中的关键细节。

解题过程

1. 基础SGD的局限性

   • 标准SGD更新规则：\(\theta_{t+1} = \theta_t - \eta \nabla f_t(\theta_t)\)，其中\(\eta\)是学习率，\(\nabla f_t\)是第\(t\)步的随机梯度。
   • 问题：SGD的迭代路径会因随机噪声剧烈波动，导致收敛到尖锐的最小值（泛化能力差）或需要大量迭代才能稳定。

2. Polyak Averaging的核心思想

   • 在SGD迭代的同时，计算参数的指数移动平均（EMA）或简单算术平均：

\[ \bar{\theta}_t = \frac{1}{t} \sum_{i=1}^t \theta_i \quad \text{或} \quad \bar{\theta}_t = \beta \bar{\theta}_{t-1} + (1-\beta)\theta_t \]

 其中$\beta \in [0,1)$是平均系数（常用0.99）。  

• 动机：平均操作平滑了随机梯度引入的噪声，使参数更可能收敛到平坦区域（理论证明其逼近期望风险的最小值）。

3. 数学原理与收敛性分析

   • 假设目标函数\(f\)是强凸且平滑的，SGD的迭代误差满足\(\mathbb{E}[\|\theta_t - \theta^*\|^2] \leq O(1/t)\)，其中\(\theta^*\)是最优点。
   • Polyak平均能加速收敛：\(\mathbb{E}[\|\bar{\theta}_t - \theta^*\|^2] \leq O(1/t)\)，但常数项更小，且实际中对非凸问题也有效。
   • 关键直觉：平均后的参数\(\bar{\theta}_t\)比单步参数\(\theta_t\)更接近最优解，因为噪声在平均中被抵消。

4. 实现细节

   • 初始化：设\(\bar{\theta}_0 = \theta_0\)（初始参数一致）。
   • 在线更新：每完成一步SGD后，更新平均值：

\[ \bar{\theta}_t = \frac{t-1}{t} \bar{\theta}_{t-1} + \frac{1}{t} \theta_t \]

 或使用指数平均（更常用）：  

\[ \bar{\theta}_t = \beta \bar{\theta}_{t-1} + (1-\beta)\theta_t \]

• 偏差校正（针对指数平均）：初期\(t\)较小时，平均值的期望有偏差。可进行校正：

\[ \bar{\theta}_t^{\text{corrected}} = \frac{\bar{\theta}_t}{1 - \beta^t} \]

• 最终输出：训练结束后，使用\(\bar{\theta}_T\)作为模型参数，而非最后的\(\theta_T\)。

5. 超参数选择与注意事项

   • 平均系数\(\beta\)：越大则历史权重越高，平滑效果强但可能延迟适应新梯度。一般取0.99–0.999。
   • 学习率\(\eta\)：需与SGD原设置一致，通常随训练衰减（如\(\eta_t = \eta_0 / \sqrt{t}\)）。
   • 适用场景：尤其适合噪声大的数据集或复杂模型（如深度学习），能提升测试精度。

6. 代码示例（PyTorch伪代码）

   import torch
   model = MyModel()
   optimizer = torch.optim.SGD(model.parameters(), lr=0.01)
   beta = 0.99  # 平均系数
   avg_params = {name: param.clone().detach() for name, param in model.named_parameters()}  # 初始化平均参数
   
   for t in range(epochs):
       for batch in dataloader:
           optimizer.zero_grad()
           loss = compute_loss(model, batch)
           loss.backward()
           optimizer.step()
   
           # 更新Polyak平均参数
           for name, param in model.named_parameters():
               avg_params[name] = beta * avg_params[name] + (1 - beta) * param.data
   
   # 训练完成后，将平均参数加载回模型
   for name, param in model.named_parameters():
       param.data.copy_(avg_params[name])
   

7. 与其他优化器的对比

   • 相比SGD直接输出最终参数，Polyak平均通过历史平均提升稳定性。
   • 与动量法（Momentum）不同：动量是加速梯度方向，而Polyak平均是直接平滑参数路径。
   • 常与自适应优化器（如Adam）结合，进一步改善收敛性。