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

鬼剪裁技巧实际上计算的是 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 与 Ghost Clipping 的内存开销是可加的 O(BL)，需要存储所有层的样本梯度范数。此外，如果存在不支持层（如果 K≥1），则需要额外的 O(Bd ² )内存来实例化该层的梯度。

内存基准测试

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

微调 BERT

我们考虑了在文本分类任务中私有地微调 BERT 的最后三层的问题。基础模型有超过 1000 万个参数，其中我们微调了最后三层，即 BertEncoder, BertPooler, 和 Classifier ，大约包含 760 万个参数。实验在一个 P100 GPU 上运行，内存为 16GB。

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

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

合成设置：内存分析

我们考虑以下设置来分析 PyTorch SGD、Vanilla DP-SGD 和 Ghost Clipping，GC DP-SGD 使用的内存。

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

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

行业应用案例

我们在 Meta 内部的一个用例上测试了 Ghost Clipping DP-SGD，该用例包含一个大约 100B 大小的模型，具有 4000 万个可训练参数。我们的初步结果表明，Ghost Clipping SGD 将 vanilla DP-SGD 的内存减少了 95%，并且与 PyTorch SGD 具有可比的内存使用量。

结论

在本文中，我们描述了 Opacus 中 Fast Gradient Clipping 和 Ghost Clipping 的实现，这些实现使得机器学习模型在差分隐私的情况下能够进行内存高效的训练。目前，Ghost Clipping 的实现仅适用于线性层，但正如系列的第 3 部分所述，它可以扩展到“广义”线性层，如卷积和多头注意力。当前的这些技术需要两个显式的反向传播步骤，这会增加运行时间。我们将探索 Ghost Clipping 之上的发展，例如用于缓解的 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 提供的宝贵反馈和建议。