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

题目描述
在深度学习中，随机梯度下降（SGD）是基础优化算法，但存在收敛不稳定和震荡问题。SGD with Polyak Averaging（或Polyak-Ruppert Averaging）通过计算参数的历史平均值来提升收敛稳定性和最终性能。该算法在优化后期平滑参数轨迹，尤其适用于非光滑或强凸问题。我们将深入探讨其数学基础、实现步骤及在深度学习中的实际应用。

解题过程循序渐进讲解

1. 算法核心思想

• 问题背景：标准SGD的迭代参数在最优解附近震荡，即使学习率适当，单个迭代点可能偏离最优解。
• Polyak Averaging解决方案：不直接使用最终迭代点，而是计算所有历史参数的平均值作为最终模型。数学上，对于迭代次数 \(t\)，平均参数 \(\bar{\theta}_t\) 定义为：

\[ \bar{\theta}_t = \frac{1}{t} \sum_{i=1}^t \theta_i \]

其中 \(\theta_i\) 是第 \(i\) 次SGD迭代的参数。

• 直观理解：震荡的轨迹在平均后更接近最优解，类似于信号处理中通过平均过滤噪声。

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

• 理论保证：在强凸且光滑的损失函数下，Polyak Averaging可将SGD的收敛率从 \(O(1/\sqrt{t})\) 提升至 \(O(1/t)\)，达到最优统计效率。
• 关键假设：
  • 损失函数 \(L(\theta)\) 是 \(\mu\)-强凸（即 \(\nabla^2 L(\theta) \succeq \mu I\)）。
  • 随机梯度 \(g_t\) 是无偏估计且方差有界（\(\mathbb{E}[\|g_t\|^2] \leq \sigma^2\)）。
• 收敛证明概要：
  1. 标准SGD迭代：\(\theta_{t+1} = \theta_t - \eta g_t\)，其中 \(\eta\) 为固定学习率。
  2. 通过递归展开，证明 \(\mathbb{E}[\|\bar{\theta}_t - \theta^*\|^2] \leq \frac{C}{t}\)，其中 \(\theta^*\) 是最优点，\(C\) 为常数。

3. 算法实现步骤

• 输入：初始化参数 \(\theta_1\)，学习率 \(\eta\)，总迭代次数 \(T\)。
• 迭代过程：
  1. 初始化累加变量 \(\theta_{\text{sum}} = \mathbf{0}\)（与 \(\theta_1\) 同维度的零向量）。
  2. 对于 \(t = 1\) 到 \(T\)：
     a. 计算当前随机梯度 \(g_t = \nabla L(\theta_t; \xi_t)\)，其中 \(\xi_t\) 是随机样本。
     b. 更新参数：\(\theta_{t+1} = \theta_t - \eta g_t\)。
     c. 累加参数：\(\theta_{\text{sum}} = \theta_{\text{sum}} + \theta_{t+1}\)。
  3. 计算平均参数：\(\bar{\theta}_T = \frac{\theta_{\text{sum}}}{T}\)。
• 输出：平均参数 \(\bar{\theta}_T\) 作为最终模型。

4. 深度学习中的实际应用技巧

• 学习率选择：使用固定学习率或递减策略（如 \(\eta_t = \eta_0 / \sqrt{t}\)），但需注意强凸假设在深度网络中常不成立。
• 计算效率优化：
  • 无需存储全部历史参数，通过累加和动态计算平均值，内存开销为 \(O(1)\)。
  • 可在训练后期（如后50%迭代）开始平均，以减少早期不稳定性的影响。
• 与动量结合：Polyak Averaging可与SGD with Momentum结合，先通过动量加速收敛，再通过平均平滑参数。

5. 代码实现示例（PyTorch）

import torch

def sgd_polyak_averaging(model, train_loader, lr=0.01, epochs=100):
    optimizer = torch.optim.SGD(model.parameters(), lr=lr)
    theta_sum = None  # 用于累加参数
    for epoch in range(epochs):
        for batch in train_loader:
            optimizer.zero_grad()
            loss = model.compute_loss(batch)
            loss.backward()
            optimizer.step()
            
            # 累加参数（Polyak Averaging）
            if theta_sum is None:
                theta_sum = {name: param.data.clone() for name, param in model.named_parameters()}
            else:
                for name, param in model.named_parameters():
                    theta_sum[name] += param.data
    
    # 计算平均参数
    for name in theta_sum:
        theta_sum[name] /= epochs
    model.load_state_dict(theta_sum)  # 加载平均后的参数

6. 适用场景与局限性

• 优势场景：
  • 损失函数近似强凸时（如逻辑回归、线性模型）。
  • 训练数据噪声大或批量较小时。
• 局限性：
  • 深度网络中强凸性不成立时，效果可能不显著。
  • 增加额外计算（尽管可优化），但不影响训练时间。

总结
Polyak Averaging通过简单的历史平均，在不增加计算复杂度的情况下提升SGD的收敛稳定性和统计效率。在深度学习实践中，可作为标准SGD的直接改进，尤其适用于对模型鲁棒性要求高的任务。