在 Opacus 中启用快速梯度裁剪和 Ghost 裁剪

介绍与背景

差分隐私随机梯度下降 (DP-SGD) 是通过差分隐私训练机器学习模型的规范方法。与非隐私的随机梯度下降 (SGD) 相比，它包含以下两点修改：

样本级梯度裁剪 (Per-sample gradient clipping)：对小批量中的每个样本的梯度进行裁剪，确保其范数在每次迭代中不超过预设的值“裁剪范数” C。
噪声添加 (Noise addition)：在每次迭代中，根据裁剪范数和隐私参数，向平均后的裁剪梯度添加预设方差的高斯噪声。

第一个改动，即样本级梯度裁剪，引入了额外的复杂性，因为它通常需要实例化样本级梯度。

Opacus 是 DP-SGD 的 PyTorch 实现。Opacus 通过使用钩子函数 (hook functions)来解决上述任务，这允许对前向和反向传播等特定事件进行干预。关于 Opacus 的更多详细信息，我们鼓励读者参阅之前的博客文章：DP-SGD 算法详解、《Opacus 中高效的样本级梯度计算》以及《Opacus 中更多层的高效样本级梯度计算》。

虽然与朴素方法相比，Opacus 提供了显著的效率提升，但实例化样本级梯度的内存成本很高。具体而言，内存使用量与批大小 (batch size) 乘以可训练参数的数量成正比。因此，内存限制了 Opacus 只能处理较小的批大小和/或较小的模型，这严重限制了其应用范围。

我们在 Opacus 中引入了快速梯度裁剪 (Fast Gradient Clipping) 和Ghost 裁剪 (Ghost Clipping)，使开发人员和研究人员无需实例化样本级梯度即可执行梯度裁剪。例如，这使得在单张 16GB GPU 上，以 1024 的批大小微调 700 万参数的 BERT 模型时，其内存占用与使用普通 PyTorch（不应用 DP-SGD）相当。相比之下，旧版本的 Opacus 在相同设置下仅支持大约 256 的最大批大小。我们提供了一个教程，以该任务为例说明如何在 Opacus 中使用快速梯度裁剪。

快速梯度裁剪与 Ghost 裁剪

这些技术背后的核心思想基于以下观察：假设已知样本级梯度范数，那么可以通过对重新加权的损失函数 $\bar{L}$ 进行反向传播来实现梯度裁剪。该损失函数定义为 $\bar{L} = \sum_{i} R_{i} L_{i}$，其中 $R_i = \min\left(\frac{C}{C_i}, 1\right)$ 是根据样本级梯度范数 $\{C_i\}$ 计算出的裁剪系数，而 $\{L_i\}$ 是样本级损失。

上述想法初看似乎是循环论证，因为它似乎需要实例化样本级梯度才能计算样本级梯度范数。然而，对于神经网络架构中某些广泛使用的组件（如全连接/线性层），确实可以在一次反向传播中获得样本级梯度范数，而无需显式计算样本级梯度。这提示了一种包含两次反向传播的工作流：第一次用于计算样本级梯度范数，第二次用于计算聚合（非样本级）裁剪后的梯度。第二次反向传播就是标准的批处理反向传播。

图 1：原生 Opacus（左上）、快速梯度裁剪（右上）和 Ghost 裁剪（下）之间的比较。我们用红色标记了成为内存瓶颈的梯度实例化过程。对于原生 Opacus，必须实例化样本级梯度。快速梯度裁剪为每一层实例化样本级梯度以计算其范数，该梯度一旦反向传播进入下一层即被立即释放。Ghost 裁剪直接从样本级激活梯度和样本级激活值工作，从而避免了梯度实例化的需求。

快速梯度裁剪
在快速梯度裁剪中，样本级梯度范数的计算分为三步：

1. 对于每一层，实例化样本级梯度并计算其范数。
2. 随后立即丢弃该样本级梯度。
3. 将每一层的（平方）样本级梯度范数求和，得到整体的（平方）样本级梯度范数。

Ghost 裁剪
作为快速梯度裁剪方法的扩展，Ghost 裁剪利用了以下事实：对于线性层¹，样本级梯度范数仅通过激活梯度和激活值即可计算。特别地，设 backprops 和 activations 分别为维数为 batch_size ✕ output_width 和 batch_size ✕ input_width 的样本级激活梯度和激活值。样本级梯度是这两者的外积，其时间和空间复杂度为 O(batch_size ✕ input_width ✕ output_width)。

Ghost 裁剪技巧则是逐样本计算 backprops 和 activations 的（平方）范数，并取其乘积，从而得到梯度的（平方）范数。其时间复杂度为 O(batch-size ✕ (input_width + output_width))，空间复杂度为 O(batch-size)（用于存储范数）。由于样本级激活值和样本级激活梯度已经存储在内存中，因此仅需额外存储范数。

快速梯度裁剪与 Ghost 裁剪的关系

快速梯度裁剪和 Ghost 裁剪是互补技术。快速梯度裁剪可应用于任何类型的层，而对于受支持的层，Ghost 裁剪在技术上更优。
当某一层不被 Ghost 裁剪支持时，我们的实现会自动切换到快速梯度裁剪。

如何在 Opacus 中使用快速梯度裁剪

训练循环与标准 PyTorch 循环完全相同。和以前的 Opacus 一样，我们使用 PrivacyEngine() 来“净化”模型和优化器。要启用 Ghost 裁剪，需使用参数 grad_sample_mode="ghost"。此外，make_private() 将损失准则作为额外输入并对其进行净化。这使我们能够在 loss.backward() 中隐藏两次反向传播和中间的损失重缩放过程。

from opacus import PrivacyEngine
criterion = nn.CrossEntropyLoss() # example loss function

privacy_engine = PrivacyEngine()
model_gc, optimizer_gc, criterion_gc, train_loader, = privacy_engine.make_private(
        module=model,
        optimizer=optimizer,
        data_loader=train_loader,
        noise_multiplier=noise_multiplier
        max_grad_norm=max_grad_norm,
	 criterion=criterion,
        grad_sample_mode="ghost",
)

# The training loop below is identical to that of PyTorch

for input_data, target_data in train_loader:
    output_gc = model_gc(input_data) # Forward pass
    optimizer_gc.zero_grad()
    loss = criterion_gc(output_gc, target_data)
    loss.backward()
    optimizer_gc.step()  # Add noise and update the model

在内部，在第一次传递之前，我们会启用钩子 (hooks)，以便捕获对应于前向和后向调用的层级数值。它们被用于计算样本级梯度范数。然后我们计算裁剪系数，重缩放损失函数并禁用钩子，这样就可以使用标准的 PyTorch 反向传播。

内存复杂度分析

考虑具有以下属性的多层神经网络：

L：层数
d：最大层宽度
B：批大小
K：非受支持/非线性层的数量

与普通 (PyTorch) SGD 相比，带有 Ghost 裁剪的 DP-SGD 的额外内存开销为 O(BL)，这是存储所有层的样本级梯度范数所需的。此外，如果存在非受支持的层 (K≥1)，则需要额外的 O(Bd²) 内存来实例化该层的梯度。

内存基准测试

我们提供了各种设置下的内存使用结果。

微调 BERT

我们考虑了对 BERT 最后三层进行隐私微调以处理文本分类任务的问题。基础模型拥有超过 1 亿个参数，我们微调了最后三层：BertEncoder、BertPooler 和 Classifier，共约 760 万个参数。实验在具有 16 GB 内存的 P100 GPU 上运行。

下表报告了各种方法在每次迭代中的最大内存占用和耗时：

	批次大小
	B = 32		B = 128		B = 512		B = 1024		B = 2048
	内存	时间	内存	时间	内存	时间	内存	时间
PyTorch SGD	236 MB	0.15 秒	1.04 GB	0.55 秒	5.27 GB	2.1 秒	12.7 GB	4.2 秒	OOM (内存溢出)
DP-SGD	1,142 MB	0.21 秒	4.55 GB	0.68 秒	OOM (内存溢出)		OOM (内存溢出)		OOM (内存溢出)
FGC DP-SGD	908 MB	0.21 秒	3.6 GB	0.75 秒	OOM (内存溢出)		OOM (内存溢出)		OOM (内存溢出)
GC DP-SGD	362 MB	0.21 秒	1.32 GB	0.67 秒	5.27 GB	2.5 秒	12.7 GB	5 秒	OOM (内存溢出)

就峰值内存占用而言：DP-SGD > FGC DP-SGD ≫ GC DP-SGD ≈ PyTorch SGD。此外，运行时间相似，因为大多数参数已冻结，而前向传播占用了大部分时间。

合成设置：内存分析

我们考虑以下设置来分析 PyTorch SGD、原生 DP-SGD 和 Ghost 裁剪 (GC DP-SGD) 的内存使用情况。

2 层全连接神经网络
- 输入：5120
- 隐藏层：2560
- 输出：1280
- 模型参数总数 = 1560 万
- 模型大小 = 62.5 MB
批大小（不同值，如下表所示）。

下表总结了每种方法训练循环各阶段的最大内存增加量（单位：MB）。

批大小	方法	模型至 GPU	前向传播	第一次反向传播	第二次反向传播	优化器步骤
32	PyTorch SGD	62.5	0.5	62.5	不适用	0
	原生 DP-SGD	62.5	0.47	3,663	不适用	162.5
	GC DP-SGD	62.5	0.47	63.13	50	125
2¹⁷	PyTorch SGD	62.5	1920	1932.5	不适用	0
	原生 DP-SGD	OOM (内存溢出)
	GC DP-SGD	62.5	1920	2625	1932.5	125

行业案例

我们在 Meta 的一个内部用例上测试了 Ghost 裁剪 DP-SGD，该用例的模型大小约为 100B，具有 4000 万个可训练参数。初步结果表明，Ghost 裁剪 SGD 比原生 DP-SGD 减少了 95% 的内存占用，并达到了与 PyTorch SGD 相当的内存使用水平。

结论

在本文中，我们介绍了 Opacus 中快速梯度裁剪和 Ghost 裁剪的实现，它们实现了差分隐私下机器学习模型的内存高效训练。目前，Ghost 裁剪实现仅适用于线性层，但正如《系列文章第 3 部分》中所述，它可以扩展到卷积和多头注意力等“广义”线性层。当前技术需要两个显式的反向传播步骤，这增加了运行时间。我们将探索在 Ghost 裁剪之上的改进，例如记账算法 (Book-Keeping algorithm) 来进行缓解。

要了解关于 Opacus 的更多信息，请访问 opacus.ai 和 github.com/pytorch/opacus。

致谢

感谢 Iden Kalemaj、Darren Liu、Karthik Prasad、Hao Shi、Igor Shilov、Davide Testuggine、Eli Uriegas、Haicheng Wang 和 Richard Zou 提供的宝贵反馈和建议。

有一些方法可以将 Ghost 裁剪扩展到非线性层。↩

在 Opacus 中启用快速梯度裁剪和幽灵裁剪