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

引言与背景

差分隐私随机梯度下降 (DP-SGD) 是使用差分隐私训练机器学习模型的典型方法。它对其非隐私对应方法——随机梯度下降法——进行了以下两项修改。

逐样本梯度裁剪：对小批量（mini-batch）中每个样本的梯度进行裁剪，确保其在每次迭代中的范数最多为一个预设值，即“裁剪范数” C。
添加噪声：在每次迭代中，向平均裁剪后的梯度添加高斯噪声，其方差取决于裁剪范数和隐私参数。

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

Opacus 是 DP-SGD 的一个 PyTorch 实现。Opacus 通过使用钩子函数（hook functions）来解决上述任务，这允许在特定事件（如前向和后向传播）上进行干预。关于 Opacus 的更多详情，我们鼓励读者回顾之前的博客文章：DP-SGD 算法详解、在 Opacus 中高效计算逐样本梯度以及在 Opacus 中为更多层高效计算逐样本梯度。

虽然与朴素方法相比，Opacus 提供了显著的效率提升，但实例化逐样本梯度的内存成本很高。具体来说，内存使用量与批次大小乘以可训练参数数量成正比。因此，内存限制了 Opacus 只能使用小批次和/或小模型，这极大地限制了其应用范围。

我们向 Opacus 引入了快速梯度裁剪（Fast Gradient Clipping）和幽灵裁剪（Ghost Clipping），使开发者和研究人员能够在不实例化逐样本梯度的情况下进行梯度裁剪。例如，这使得在一块 16GB 的 GPU 上，以 1024 的批次大小微调 BERT 的 700 万个参数成为可能，其内存消耗与使用 PyTorch（不应用 DP-SGD）相当。相比之下，在相同设置下，先前版本的 Opacus 最多只能支持约 256 的批次大小。我们提供了一个教程，以该任务为例，介绍如何在 Opacus 中使用快速梯度裁剪。

快速梯度裁剪与幽灵裁剪

这些技术背后的关键思想基于以下观察：假设逐样本梯度范数是已知的，那么梯度裁剪可以通过对一个重加权的损失函数 $ \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**（左上）、**快速梯度裁剪**（右上）和**幽灵裁剪**（底部）的比较。我们用红色标记了成为内存瓶颈的梯度实例化。对于原生 Opacus，它必须实例化**逐样本梯度**。**快速梯度裁剪**为每一层实例化逐样本梯度以计算其范数，一旦反向传播移至下一层，该梯度就会被立即释放。幽灵裁剪直接利用**逐样本激活梯度**和**逐样本激活**进行操作，避免了梯度实例化的需要。

快速梯度裁剪
在快速梯度裁剪中，逐样本梯度范数通过三个步骤计算：

对于每一层，实例化逐样本梯度并计算其范数。
然后立即丢弃逐样本梯度。
将每一层的（平方）逐样本梯度范数相加以获得总的（平方）逐样本梯度范数。

幽灵裁剪
在快速梯度裁剪方法的基础上，幽灵裁剪利用了这样一个事实：对于**线性层¹**，逐样本梯度范数仅需通过**激活梯度**和**激活**即可计算。具体来说，设 backprops 和 activations 分别为逐样本激活梯度和激活，维度分别为 batch_size × output_width 和 batch_size × input_width。逐样本梯度是这两者的外积，其时间和空间复杂度为 O(batch_size × input_width × output_width)。

幽灵裁剪技巧则是逐样本地计算 backprops 和 activations 的（平方）范数，然后将它们相乘，从而得到梯度的（平方）范数。这需要 O(batch_size × (input_width + output_width)) 的时间，并且需要 O(batch_size) 的空间来存储范数。由于**逐样本激活**和**逐样本激活梯度**已经存储，额外的内存仅用于存储范数。

快速梯度裁剪与幽灵裁剪的关系

快速梯度裁剪和幽灵裁剪是互补的技术。快速梯度裁剪可以应用于任何类型的层，而对于受支持的层，幽灵裁剪是明显更好的技术。
当层不受幽灵裁剪支持时，我们的实现会自动切换到快速梯度裁剪。

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

训练循环与标准的 PyTorch 循环相同。和之前的 Opacus 一样，我们使用 PrivacyEngine()，它会对模型和优化器进行“净化”。要启用幽灵裁剪，需使用参数 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

在内部，在第一次传播之前，我们启用*钩子*，这使我们能够捕获与前向和后向调用相对应的逐层值。它们用于计算逐样本梯度范数。然后我们计算裁剪系数，重缩放损失函数并禁用钩子，这让我们能够使用标准的 PyTorch 反向传播。

内存复杂度分析

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

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

与普通（PyTorch）SGD 相比，使用幽灵裁剪的 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 和幽灵裁剪 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 的一个内部用例上测试了幽灵裁剪 DP-SGD，该用例包含一个约 100B 参数大小、40M 可训练参数的模型。我们的初步结果显示，幽灵裁剪 SGD 减少了原生 DP-SGD 95% 的内存，并达到了与 PyTorch SGD 相当的内存使用量。

结论

在本文中，我们描述了 Opacus 中快速梯度裁剪和幽灵裁剪的实现，这些实现使得使用差分隐私进行机器学习模型的内存高效训练成为可能。目前，幽灵裁剪的实现仅适用于线性层，但正如本系列第三部分所概述的，它可以扩展到“广义”线性层，如卷积和多头注意力。当前技术需要两个明确的反向传播步骤，这增加了运行时间。我们将探索在幽灵裁剪之上的发展，例如Book-Keeping 算法以缓解此问题。

要了解更多关于 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 提供的宝贵反馈和建议。

有多种方法可以将幽灵裁剪扩展到非线性层。 ↩