Hugging Face Trainer的差分隐私保护：无需重写训练循环

众所周知，用敏感数据训练 LLM 会引发严重的隐私问题。在最近的博客文章中，我们讨论了成员推断攻击以及我们为缓解这类攻击所做的研究。

在 JetBrains Research，我们深切关注用户隐私，并持续开发新的方法和工具来加强隐私保护。在本文中，我们介绍 DPTrainer，这是我们已经开发并开源的新库。DPTrainer 将 Opacus 和 Hugging Face Trainer 无缝集成，让你无需重写训练循环或修改 trainer 源代码，就能训练隐私保护模型。

差分隐私的重要性

人们已广泛观察到，模型质量沿三个维度扩展：规模、算力和数据。更大的模型提供更强的容量，但训练效率更低，推理成本也更高。训练中使用更多算力自然意味着更高成本和更长耗时。而数据这一维度，则主要受限于能否获取到足够质量和数量的数据。

差分隐私是我们应对数据获取障碍的解决方案。本质上，差分隐私是一种保护用于训练的个人数据点的数学框架。其核心保证是：使用差分隐私训练的模型，无论某个样本是否被包含在训练集中，其行为都几乎完全一致。对于已知会记忆训练数据、并可能在被恶意提示时复现这些数据的 LLM 来说，这是目前已知最强的防泄露防御。即使是复杂的成员推断攻击，在能够访问模型权重、置信度分数和基础模型架构的情况下，也无法确定受该方法保护的某个样本是否包含在训练集中。

在实践中，差分隐私通过被称为差分隐私随机梯度下降（DP-SGD）的方法应用于神经网络训练。DP-SGD 不是在整批数据上计算一个梯度，而是为每个样本分别计算梯度、裁剪梯度以限制离群值，再聚合批内梯度，然后注入噪声，使任何单个样本的痕迹都变得无法区分。

通过保证训练方法的隐私性，我们可以利用之前无法使用的渠道，使用通过 IDE 每天生成的数据（参见我们的数据收集政策和最近一篇关于面向 AI 的数据共享的文章）。由于我们的用户规模，这带给我们高数据量；同时由于数据是在编写代码的过程中生成的，而不仅仅是从最终产品中提取的，数据质量也很高。这些优势确保我们即将推出的模型能够超越其体量地发挥（这里的 weight 是双关：既指“体重级别/体量”，也指模型权重）。

它弥补的空白

Opacus 是 PyTorch 中实现 DP-SGD 的首选库。它提供了所需的全部能力：逐样本梯度计算、DPOptimizer、隐私会计和泊松采样数据加载器。问题是它围绕手动 PyTorch 训练循环设计，使用不便，并且与 Hugging Face 平台集成得不够好。

Hugging Face 的 Trainer 和 Transformers 强化学习（TRL）中的对齐训练器（例如 SFTTrainer、DPOTrainer）是 transformers 最高层级的高层训练 API。它们处理分布式训练、检查点、评估、回调，以及许多你不想重新实现的功能。然而，它们对差分隐私毫无感知。

将 Opacus 接入基于 Trainer 的工作流，需要涉及模型包装、优化器创建、数据加载、损失计算、检查点和回调管理。这些环节以微妙的方式相互影响，任何一个出错都可能在无声无息之间破坏你的隐私保证。

为了解决这个问题，我们的研究人员 Evgeny Grigorenko 和 David Stanojevic 创建了 DPTrainer，目前由 Mihajlo Linic 维护。DPTrainer 谨慎而妥善地处理了这些问题。

真正的即插即用替代品

差分隐私中的一个关键概念是隐私预算。它代表我们愿意接受的理论上最大信息泄露风险。换句话说，它是任何单个数据点可能使输出分布产生的最大偏移量。隐私预算的一个重要特性是消耗具有累积性，这迫使我们在隐私和性能之间做权衡，因为更高的隐私要求意味着需要向梯度中注入更多噪声。

DPTrainer 扩展了 transformers.Trainer，并通过新增的 PrivacyArguments 数据类引入隐私预算。所有标准的训练参数、回调、检查点和评估工作流都可以不变地使用，如下代码所示：

$ python

from dptrainer import DPTrainer, PrivacyArguments

privacy_args = PrivacyArguments(
    target_epsilon=8.0,
    per_sample_max_grad_norm=1.0,
)

trainer = DPTrainer(
    model=model,
    args=training_args,
    train_dataset=train_dataset,
    privacy_args=privacy_args,
    data_collator=data_collator,
)

trainer.train()

设置一个与你的隐私预算匹配的 target_epsilon，剩下的交给 DPTrainer 处理。内部会计会跟踪预算消耗，剩余预算会在保存检查点时一并保存，因此可以轻松恢复训练。

让 TRL 和其他专用 Trainer 也支持差分隐私

真正的威力来自 privatize_trainer。许多工作流使用 Trainer 的子类：例如用于偏好学习的 DPOTrainer、用于指令微调的 SFTTrainer 和用于生成的 Seq2SeqTrainer。它们都在基类之上添加了特定任务的损失函数和生成逻辑。若将这些类重写为继承自 DPTrainer，将是侵入性和脆弱的改动。

privatize_trainer 会在运行时对任何基于 Trainer 的类打补丁，将 DPTrainer 注入其继承链，而无需改动该类自身的逻辑：

$ python

from trl import DPOTrainer
from dptrainer import PrivacyArguments, privatize_trainer

privatize_trainer(DPOTrainer)  # one line

trainer = DPOTrainer(
    model=model,
    args=training_args,
    train_dataset=train_dataset,
    processing_class=tokenizer,
    privacy_args=PrivacyArguments(target_epsilon=8.0, per_sample_max_grad_norm=1.0),
)
trainer.train()

被修补后的 trainer 保留其所有原始行为（例如奖励计算、DPO 损失、生成），同时获得 DP-SGD 能力。

DPTrainer 替你处理的事情

DPTrainer 自动管理以下事项：

噪声添加。 通过 DPOptimizer 向聚合梯度添加校准后的高斯噪声。
梯度裁剪。 在聚合之前对每个样本的梯度单独裁剪，而不是裁剪批梯度。通过 clipping 和 per_sample_max_grad_norm 支持扁平（flat）、自适应（AdaClip）和逐层（per-layer）策略。
梯度计算。 将模型包装在 Opacus 的 GradSampleModule 中，以计算逐样本梯度，这是 DP-SGD 正确性的必要条件。
优化器创建。 拦截 create_optimizer，用 DPOptimizer 包装 Hugging Face 创建的优化器。
数据加载。 重写 get_train_dataloader，返回带泊松子采样的 DPDataLoader，这正是通过采样实现隐私放大的机制。
噪声校准。 给定 target_epsilon 和你的训练配置，DPTrainer 会自动计算正确的 noise_multiplier，无需手动二分搜索。
隐私会计。 DPCallback 挂钩优化器步骤，在每次更新后跟踪正在消耗的隐私预算。
检查点。 将会计状态与模型权重一起保存和恢复，确保恢复训练后隐私预算跟踪依旧正确。
早停。 一个感知隐私预算的停止机制会在预算完全耗尽时自动停止训练。

灵活的配置

PrivacyArguments 提供了你期望的各类调节旋钮：

target_epsilon / noise_multiplier：设置其中一个即可，二者互斥。
clipping：选择 "flat"（标准）、"adaptive"（AdaClip）或 "per_layer"。
poisson_sampling：切换泊松子采样，以启用隐私放大。
grad_sample_mode："hooks"（默认）。
accountant：隐私会计类型（默认 RDP）。
epsilon_log_mode：在训练步骤、评估时、两者都记录，或完全不记录预算消耗。

试试我们的 DPTrainer

差分隐私正日益成为合规要求，而不仅仅是一种研究上的“锦上添花”。围绕个人数据训练的法规日益严格，人们对针对 LLM 的成员推断攻击的认识也在不断提高，这意味着团队需要实用、可审计的差分隐私训练方案。困难从来不在数学本身，而在于工程实现。DPTrainer 正是移除了这一障碍。

如果你正在敏感数据上训练 transformers，并且使用 Hugging Face 生态中的任何部分，这都值得一试。
立即查看

分类导航