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

作者：Enayat Ullah, Huanyu Zhang, Will Bullock, Ilya Mironov

引言与背景

差分隐私随机梯度下降 (DP-SGD) 是用于训练具有差分隐私的机器学习模型的规范方法。它与其非隐私对应方法（随机梯度下降）相比，包含以下两项修改。

1. 逐样本梯度裁剪：裁剪迷你批次中每个样本的梯度，确保其范数在每次迭代中至多为一个预设值“裁剪范数” C。

2. 噪声添加：在每次迭代中，向平均裁剪梯度添加高斯噪声，其方差取决于裁剪范数和隐私参数。

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

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

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

我们在 Opacus 中引入了快速梯度裁剪和幽灵裁剪，使开发者和研究人员能够在不实例化逐样本梯度的情况下执行梯度裁剪。例如，这允许在单张 16GB GPU 上，使用批次大小 1024，对 BERT 的 7M 参数进行微调，其内存使用量与使用 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，它必须实例化逐样本梯度。快速梯度裁剪为每一层实例化逐样本梯度以计算其范数，并在反向传播移动到下一层后立即释放。幽灵裁剪直接使用逐样本激活梯度和逐样本激活，避免了梯度实例化的需要。

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

1. 对于每一层，实例化逐样本梯度并计算其范数。
2. 然后立即丢弃逐样本梯度。
3. (平方) 逐样本梯度范数对各层求和，以获得总体 (平方) 逐样本梯度范数。

幽灵裁剪
幽灵裁剪在快速梯度裁剪方法的基础上进行了扩展，它利用了对于线性层¹，可以仅从激活梯度和激活计算逐样本梯度范数这一事实。特别是，设 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) 的空间来存储。由于逐样本激活和逐样本激活梯度已经被存储，额外的内存只需要用于存储范数。

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

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

如何在 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 的最后三层以进行文本分类任务。基础模型拥有超过 100M 个参数，其中我们微调最后三层：BertEncoder、BertPooler 和 Classifier，包含大约 7.6M 个参数。实验在具有 16 GB 内存的 P100 GPU 上运行。

下表报告了各种方法的每次迭代的最大内存和时间

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

合成设置：内存分析

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

• 2层全连接神经网络
  • 输入: 5120
  • 隐藏层: 2560
  • 输出: 1280
  • 模型参数总数 = 15.6M
  • 模型大小 = 62.5 MB
• 批次大小，不同值，如下表所示。

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

行业用例

我们在 Meta 的一个内部用例上测试了幽灵裁剪 DP-SGD，该用例包含一个约 100B 大小、具有 40M 个可训练参数的模型。我们的初步结果表明，幽灵裁剪 SGD 将原生 DP-SGD 的内存减少了 95%，并实现了与 PyTorch SGD 相当的内存使用量。

结论

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

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