Robust and Fast Training via Per-Sample Clipping

技术摘要：通过逐样本裁剪实现鲁棒且快速的训练

问题陈述
本文研究的是随机优化问题 $\min_{x \in \mathbb{R}^d} F(x)$，其中 $F(x) = \mathbb{E}_{\xi \sim \mathcal{D}}[f(x, \xi)]$。虽然随机梯度下降（SGD）是标准求解器，但其收敛保证通常依赖于梯度噪声具有有界方差（轻尾分布）的假设。然而，最近的经验和理论工作表明，在许多机器学习应用中，梯度噪声可能是重尾的，即仅具有 $p \in (1, 2]$ 阶矩，但缺乏有界方差。在这种条件下，标准 SGD 的表现可能很差。现有的鲁棒替代方案，如标准梯度裁剪（Clip-SGD）和归一化 SGD（Normalized SGD），要么无法达到最优样本复杂度下界，要么需要特定的、通常是不切实际的参数调优（例如，依赖于初始化间隙 $\Delta_1$ 或失败概率 $\delta$ 的特定对数因子）。

方法论
作者提出了 逐样本裁剪 SGD（PS-Clip-SGD），这是一种鲁棒的梯度估计器，它在聚合之前对小批量（mini-batch）内的单个样本应用裁剪，而不是对聚合后的批量梯度进行裁剪。

对于规模为 $n$、样本为 $\xi^{(1)}_t, \dots, \xi^{(n)}_t$ 的批量，梯度估计器 $G(x_t, \alpha, \beta)$ 定义为：
$$G(x_t, \alpha, \beta) = \frac{1}{n} \sum_{k=1}^n \gamma^{(k)}_t \nabla f(x_t, \xi^{(k)}_t)$$
其中逐样本裁剪因子为 $\gamma^{(k)}_t := \min\left\{1, \frac{\alpha k^{1/\beta}}{|\nabla f(x_t, \xi^{(k)}_t)|}\right\}$。参数 $\alpha$ 和 $\beta$ 控制裁剪阈值，其中 $\alpha$ 与矩界 $\sigma$ 相关，$\beta$ 与尾部指数 $p$ 相关。更新规则遵循标准 SGD：$x_{t+1} = x_t - \eta_t G(x_t, \alpha, \beta)$。

核心贡献

1. 最优期望收敛性： 作者证明了 PS-Clip-SGD 在非凸优化且存在重尾噪声的情况下，能够实现最优的样本复杂度。具体而言，该方法在平均平方梯度范数上的收敛速率达到 $O(1/T)$，这与 Zhang 等人 (2020) 建立的下界相匹配。由于该界限适用于平方梯度范数的和（比某些先前工作中使用的范数之和更强的准则），且不需要根据初始化间隙 $\Delta_1$ 来调整步长，因此相比 Clip-SGD 和 Normalized SGD 取得了改进。
2. 高概率收敛性： 本文建立了高概率收敛保证，其结果在失败概率 $\delta$ 和步数 $T$ 的多项式对数因子范围内与期望收敛速率相匹配。值得注意的是，与最先进的 Clip-SGD 相比，其对 $\log(1/\delta)$ 的依赖关系得到了改进，并且该方法不需要将步长或裁剪阈值显式地调节为 $\delta$ 的函数。
3. 梯度累积洞察： 作者挑战了仅在所有梯度累积步骤完成后才应用梯度裁剪的常见做法。他们提出并从经验上验证了，在累积过程中在小批量层面应用裁剪（小批量裁剪，Mini-Batch Clipping）可以提高训练性能，且几乎不会增加额外的计算成本，这与认为中间裁剪会引入偏差的普遍共识形成对比。

实验结果

• 理论方面： 理论分析表明，在 $L$-光滑性和有限 $p$ 阶矩的假设下，PS-Clip-SGD 提供了强大的收敛保证。推导出的界限在对 $\delta$ 的依赖以及收敛度量强度（平方范数 vs 范数）方面，均优于 Normalized SGD 和 Clip-SGD。
• 二次优化： 在使用帕累托分布噪声最小化二次函数的实验中，即使在未调优超参数的情况下，PS-Clip-SGD 在各种噪声机制（$p \in \{1.2, 1.5, 1.8\}$）下也显著优于 Clip-SGD 和 Normalized SGD。高概率实验证实了其对 $\log(1/\delta)$ 的亚线性依赖关系。
• 图像分类（CIFAR-100 上的 AlexNet）： 在使用带动量的 AlexNet 进行训练时，带动量的 PS-Clip-SGD 实现了更高的验证准确率（64.84%），高于带动量的 vanilla SGD（58.90%）和带动量的 Clip-SGD（58.94%）。即使考虑到逐样本梯度范数计算带来的约 $1.3\times$ 计算开销，这种提升依然存在。作者将其归功于对具有极大范数的少量离群梯度的有效控制，而标准的批量级裁剪无法检测到这些离群值。
• 语言建模（GPT-2）： 在训练 124M 参数的 GPT-2 模型时的梯度累积设置中，尽管存在中间裁剪是有害的流行观点，但在多种学习率下，小批量裁剪（在每个累积步骤进行裁剪）的表现优于标准 Clip-SGD（仅在完整累积后进行裁剪）。

意义与主张
本文声称是第一项专门提出将逐样本裁剪作为提高重尾噪声存在下训练性能的方法的研究，这与其在差分隐私中的既定用途不同。其意义在于提供了一种满足以下条件的方案：

1. 在非凸重尾问题下实现最优的理论收敛速率。
2. 对超参数选择具有鲁棒性，即使在未经精细调整的设置下也能表现良好。
3. 证明了逐样本裁剪可以为深度学习任务（如 AlexNet）带来实际的性能增益，且其收益超过了计算成本。
4. 提出了梯度累积策略的范式转变，表明中间裁剪是有益的而非有害的。

作者承认了局限性，指出其理论分析需要已知噪声矩界 $\sigma$ 和尾部指数 $p$ 来定义最优参数，且“有限 $p$ 阶矩”的假设比“有限中心 $p$ 阶矩”更强。他们还指出，虽然逐样本裁剪对于像 AlexNet 这样的小型模型很有效，但其计算开销可能使其在没有进一步优化的情况下难以应用于超大规模架构，不过在累积场景中使用小批量裁剪提供了一种更具成本效益的替代方案。