DLG梯度泄露攻击复现：从MNIST还原训练数据

1. 实验背景与整体思路拆解这篇东西拖延了挺久终于找时间把Deep Leakage from Gradients后面统一叫DLG在MNIST上完整复现了一遍顺手记录一下整个过程和踩过的坑。之所以选MNIST做靶场而不是CIFAR或者ImageNet主要原因是MNIST足够小、训练够快、可视化直观单张图片才28×28像素就算真把梯度泄漏出来用肉眼就能确认有没有成功不需要额外写复杂的评估脚本。对于第一次接触这个方向的人来说用MNIST跑通DLG是性价比最高的入门路径。先讲清楚DLG到底在做什么。它的核心论点很简单在分布式训练场景里我们以为只上传梯度不下传原始数据就能保护隐私但实际上梯度携带的信息量远超想象。攻击者只需要拿到某个batch的梯度完全不知道标签和输入仅通过优化一个随机初始化的虚拟输入和虚拟标签就能还原出训练数据的近似副本。这个过程不依赖任何前向推理阶段的中间激活值也不需要攻击者控制模型结构只要梯度可被访问攻击就成立。整个攻击的数学逻辑并不复杂。模型参数为W真实训练数据是x和标签y假设损失函数是交叉熵那么训练时梯度g ∂L(W, x, y) / ∂W是公开可见的。攻击者这边随便初始化一对假数据x和假标签y然后计算假数据的梯度g ∂L(W, x, y) / ∂W只要让g和g之间的欧氏距离不断减小x就会逐步逼近真实的x。这里有个关键前提真实标签y是one-hot向量它的取值空间有限所以攻击者可以先通过暴力枚举直接确定标签或者干脆把标签也当作可优化变量。MNIST只有10类枚举成本极低这是这个数据集对DLG特别友好的原因之一。复现实验的目标很明确搭建一套最小化的攻击闭环包括训练方产生梯度、攻击方接收梯度并还原数据、可视化对比原图和重建图的差异。选了经典的LeNet变体做目标网络交叉熵作为损失函数Adam做优化器。整个流程跑通下来单张图片的还原只需要大约100~200轮迭代在一张普通消费级显卡上几秒钟就能完成CPU大概在十几秒到半分钟的范围完全在可接受范围内。这篇实验记录会从环境搭建写到代码细节再到调参技巧和常见报错尽量把从零复现一个DLG攻击的完整路径交代清楚。适合正在学习隐私攻击方向的开发者、做联邦学习相关课题的学生以及任何对梯度究竟泄露了多少信息这件事好奇的人。看完之后你不仅能跑通实验也能理解每个参数在攻击中扮演的角色。2. 实验环境构建与数据管线准备2.1 依赖安装与硬件选择DLG攻击本身不需要任何特殊依赖TensorFlow或PyTorch都可以实现。这里选PyTorch因为它在灵活性和调试便利性上更适合这种需要大量自定义forward/backward循环的研究型代码。基础依赖包括pip install torch torchvision matplotlib numpy版本方面PyTorch 1.x或2.x都能跑代码里没有用到任何新特性。CUDA不是必须的MNISTLeNet这个规模用CPU也能在合理时间内完成迭代但如果你想多跑几张图或者尝试更高的分辨率一张显存不低于4GB的显卡会舒服很多。我实测在纯CPU环境八核笔记本下单张图片还原耗时约25秒CUDA环境下大概3秒差距很明显但没有到不可用的程度。硬件之外还有个隐藏需求最好有一个能正常显示matplotlib图像的桌面环境。如果是在无显示器的服务器上跑记得把plt.switch_backend(Agg)加上否则会卡在plt.show()那一步。这个细节在后面可视化小节还会再提。2.2 数据加载与预处理细节MNIST获取方式有两种一种是通过torchvision的datasets模块自动下载另一种是手动从官网下载并组织目录结构。推荐第一种省事且不容易出错transform transforms.Compose([ transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,)) ]) trainset torchvision.datasets.MNIST( root./data, trainTrue, downloadTrue, transformtransform ) trainloader torch.utils.data.DataLoader( trainset, batch_size1, shuffleTrue )这里有几个容易忽略的点。Normalize用的均值和标准差是MNIST官方统计值很多新手直接写(0.5, 0.5)这会导致图像像素值分布被拉到一个不是你预期范围内的区间。DLG攻击的优化过程对像素值绝对大小是敏感的因为梯度匹配用的欧氏距离会受到量纲影响。推荐直接用0.1307和0.3081这两个官方统计参数后续虚拟输入初始化也要用同样的分布去初始化否则分布偏移会让重建收敛变慢。另一个细节是batch_size1。原论文用的是单个样本做攻击演示虽然理论上batch size大于1也能攻击但那是后续Improved DLG、Inverting Gradients等改进工作做的事情。如果你刚开始复现老老实实batch_size1将注意力集中在攻击是否成功上不要一上来就给自己加难度。2.3 目标网络结构选择网络结构是鸡和蛋的问题结构太简单梯度信息太少重建会模糊结构太复杂显存和计算开销上来攻击反而在优化上变得不稳定。实验选了最经典的LeNet变体class LeNet(nn.Module): def __init__(self): super(LeNet, self).__init__() self.conv1 nn.Conv2d(1, 6, kernel_size5, padding2) self.conv2 nn.Conv2d(6, 16, kernel_size5) self.fc1 nn.Linear(16*5*5, 120) self.fc2 nn.Linear(120, 84) self.fc3 nn.Linear(84, 10) def forward(self, x): x F.relu(self.conv1(x)) x F.max_pool2d(x, 2) x F.relu(self.conv2(x)) x F.max_pool2d(x, 2) x x.view(x.size(0), -1) x F.relu(self.fc1(x)) x F.relu(self.fc2(x)) x self.fc3(x) return x选这个结构的原因有三个。第一LeNet是CNN入门标配参数量不大梯度空间维度适中既能体现多层网络梯度携带像素级信息又不至于让优化陷入困难。第二MNIST本身是灰度图第一层卷积核数量设为6足够捕获笔画特征。第三这个网络收敛快训练几轮就能得到一个判别能力不错的模型用高质量模型产生的梯度做攻击实验重建效果更稳定。如果你的显卡很强也可以试试ResNet18但要注意浅层梯度对重建贡献更大层数加深后攻击需要更多迭代才能收敛。2.4 损失函数与梯度构造的关键点训练方使用的损失函数建议直接选cross-entropy。DLG原论文里就是基于分类任务做演示交叉熵的梯度分布对标签信息编码足够清晰。这里有个数学上的细节交叉熵梯度在softmax层会产生抑制项和激励项这种结构让重建优化过程更容易定位到正确类别。代码如下def compute_grad(net, data, label, loss_fnnn.CrossEntropyLoss()): net.zero_grad() output net(data) loss loss_fn(output, label) loss.backward() grads [p.grad.detach().clone() for p in net.parameters()] return grads我的建议是不要在训练方这一侧加入weight decay之类的正则项也不要混合多batch。干净的单样本梯度是攻击成功的最有利条件这样你在调试攻击代码时如果重建失败原因只会出现在攻击侧而不是被训练侧的正则干扰。等完整流程跑通、确认所有环节正常后再去加大batch、加dropout、加weight decay观察防御手段如何影响攻击成功率那就是下一步的研究路线了。3. DLG攻击核心代码实现3.1 攻击者模型与参数初始化策略攻击侧的核心是构造一个和训练目标完全一致的网络副本初始化一个随机输入和随机标签然后用梯度匹配去优化。这里有一个容易踩的坑你不需要区分训练方和攻击方的网络对象直接复用同一个网络结构即可因为攻击场景假定攻击者能拿到网络结构和参数状态。这在联邦学习的默认联邦平均协议下是合理的假设但在某些别的研究场景下可能只拿到部分梯度那就是另一段故事了。攻击代码核心结构如下def dlg_attack(net, real_grads, target_shape(1, 1, 28, 28), num_classes10): # 初始化虚拟数据高斯噪声 dummy_data torch.randn(target_shape, requires_gradTrue) # 初始化虚拟标签用均匀概率分布比one-hot枚举更高效 dummy_label torch.randn((1, num_classes), requires_gradTrue) optimizer torch.optim.Adam([dummy_data, dummy_label], lr0.1) criterion nn.CrossEntropyLoss() for it in range(200): optimizer.zero_grad() net.zero_grad() output net(dummy_data) loss criterion(output, dummy_label) # 计算虚拟梯度 dummy_grads torch.autograd.grad( loss, net.parameters(), create_graphTrue ) # 梯度匹配损失欧氏距离平方和 grad_diff sum( ((d - r) ** 2).sum() for d, r in zip(dummy_grads, real_grads) ) grad_diff.backward() optimizer.step() if it % 20 0: print(fIteration {it}: {grad_diff.item():.6f}) return dummy_data.detach(), dummy_label.detach()关于虚拟标签的初始化原论文建议用均匀分布而非one-hot因为one-hot的离散性决定了它无法通过梯度下降直接优化。如果把标签当作连续变量来优化最终结果虽然不会严格等于某个one-hot向量但softmax层的输出会自然地逼近one-hot分布预测类别也就明确了。这个方法在实际实验中被验证是可行的并且比每轮尝试枚举10个标签效率高得多。3.2 梯度匹配损失与计算图构建这个函数里最难理解的部分是torch.autograd.grad(..., create_graphTrue)。这里必须解释清楚为什么不能省略create_graph参数。如果不用create_graphloss对net.parameters()的梯度会被立即计算并返回这些梯度张量本身不再参与后续自动微分的计算图。但DLG攻击需要的是虚拟梯度与真实梯度的距离对dummy_data的导数也就是说梯度匹配损失grad_diff要二次求导先按普通流程优化网络的权重再把dummy_data当作可优化变量看梯度路径是从grad_diff反向传播到dummy_data。create_graphTrue的含义是保留这一层求导记录让计算图保持连通从而允许二次求导。这个细节初学者特别容易漏掉漏掉之后的典型报错是RuntimeError: element 0 of tensors does not require grad and does not have a grad_fn。如果你遇到这个错误第一反应就是检查create_graph是不是设成了True。另外注意一点real_grads应该使用前面计算好的真实梯度detach().clone()之后的版本避免真实梯度的计算图被反向传播遍历否则会浪费大量显存。你可能还会遇到反向传播时梯度累积的问题所以每次迭代开始时的optimizer.zero_grad()和net.zero_grad()都是必要的。这两个调用缺一不可前者清空dummy_data和dummy_label的梯度缓冲后者是防止模型内部参数由于上一次前向的.backward操作残留梯度。3.3 可视化还原结果的实用技巧还原效果怎么看直接print一个像素数组毫无意义一定要用图像方式对比。标准做法是把原始的real_data和重建的dummy_data都画在同一张图上import matplotlib.pyplot as plt def visualize(real, dummy, label, pred, pathNone): fig, axes plt.subplots(1, 2, figsize(6, 3)) axes[0].imshow(real.squeeze().cpu().numpy(), cmapgray) axes[0].set_title(fReal (Label {int(label)})) axes[1].imshow(dummy.squeeze().cpu().numpy(), cmapgray) axes[1].set_title(fRecovered (Pred {int(pred)})) for ax in axes: ax.axis(off) if path: plt.savefig(path, dpi150, bbox_inchestight) else: plt.show()有个从像素数组到真实图像之前的坑torchvision里MNIST的Normalize变换用的均值0.1307和标准差0.3081所以从网络输出拿到的dummy_data是归一化后的版本。直接用squeeze().cpu().numpy()得到的数值会有一半是负值、最大值也不在0~1范围内。imshow虽然能自动拉伸colorbar但视觉上会略微偏移。更好的做法是手动反归一化到[0,1]区间实际操作中可以用一个简单的线性映射dummy_img dummy_data.squeeze().cpu().numpy() dummy_img (dummy_img - dummy_img.min()) / (dummy_img.max() - dummy_img.min())这样虽然本质上做了一个min-max归一化但还原结果的对比度会明显改善更接近MNIST原始图的实际观感。如果不想做这个映射也可以直接调用dummy_data.squeeze().clamp(0, 1).cpu().numpy()但clamp之后超过1的部分会被截断某些重建图中的强噪声区域会被掩盖我更喜欢用min-max映射。4. 训练流程与攻击流程串联4.1 准备真实梯度的完整训练流程先把攻击脚本的入口和训练侧流程完整串一遍。训练侧的目标很简单加载一个预训练好的LeNet取某个batch的输入和标签向前传播并反向传播拿到梯度。这里并不需要模型完成整个训练过程因为联邦学习场景里每个客户端上传的梯度本来就是基于当前模型权重对本地数据的一轮更新结果权重如何获得并不影响攻击的有效性。但有个关键问题如果模型权重完全随机未训练梯度可能太小或方向太随机攻击效果不稳定。所以建议先快速训练一个基础模型。在MNIST上训练LeNet到测试准确率98%大概只需要3个epoch代码不超过20行net LeNet() optimizer torch.optim.SGD(net.parameters(), lr0.01, momentum0.9) criterion nn.CrossEntropyLoss() for epoch in range(3): for images, labels in trainloader: optimizer.zero_grad() output net(images) loss criterion(output, labels) loss.backward() optimizer.step() torch.save(net.state_dict(), mnist_lenet.pt)这个预训练模型只需要训练一次后面所有攻击实验都复用它。有个细节值得提醒攻击实验里的net.load_state_dict(torch.load(...))之后记得把它设置为eval模式或train模式实际上两种模式都能跑因为LeNet里没有BN层和Dropout模式切换不影响梯度计算。但为了严谨我还是习惯在攻击前调用net.eval()。如果以后换成带BN的网络模式切换的影响就大了BN层在train和eval模式下统计方式完全不同生成的梯度也会差异很大。4.2 攻击主流程的参数配置攻击主流程需要明确的几个超参数迭代次数、学习率、Adam的betas、虚拟数据初始化分布。迭代次数我设为200因为在实际运行中发现第50轮之后重建图像轮廓就清晰可见了第100轮之后数字笔画基本稳定再往后主要是微调像素值。250以上收益很小只会增加时间成本没必要。学习率设为0.1这是原论文推荐的值。这个值唯一要注意的是不要太贪心。设成0.5确实前几轮loss下降更快但会导致后期在最优解附近震荡重建图像出现明显噪点。设成0.01则收敛太慢200轮内可能只能看到模糊轮廓。0.1是个折中点跑完200轮基本稳定。Adam的betas参数保持默认(0.9, 0.999)。这个细节很多人会忽略但实际上DLG对Adam的momentum非常敏感。有改进方法专门针对这个问题做了调整比如把betas调成(0.98, 0.999)来降低历史梯度权重让优化更聚焦当前梯度方向。有兴趣可以试一下效果确实有差异但原实验用默认betas已经能成功。虚拟数据初始化推荐均值为0、方差为1的标准正态分布。为什么不用均匀分布或者全零初始化因为全零初始化会让所有卷积核的输出完全一样梯度为零优化直接卡死均匀分布则可能让像素初始值距离真实分布太远增加迭代次数。标准正态分布的像素值范围和MNIST归一化后的真实像素分布比较接近初始loss不是特别高优化也平稳。4.3 攻击结果评估标准怎么判断攻击成功肉眼判断是最直观的方法毕竟MNIST就10个数字重建得像不像一眼就能看出来。但仅凭肉眼不够严谨建议加两个量化指标第一是预测标签正确性。把重建的dummy_data喂给模型看输出的argmax是否等于真实标签with torch.no_grad(): pred net(dummy_data).argmax(dim1).item()如果预测标签和真实标签一致说明重建数据在语义层面已经被正确还原。第二是像素级相似度PSNR峰值信噪比。虽然DLG原论文用的是视觉对比不追求逐像素恢复但PSNR仍然能提供一个客观数字。MNIST单通道8-bit图像PSNR计算如下import math mse ((real_data - dummy_data) ** 2).mean().item() psnr 10 * math.log10(1.0 / (mse 1e-10))经验值参考正常攻击成功时PSNR在18~25dB之间这个范围内肉眼看非常接近如果PSNR低于12dB重建结果基本不可辨识。注意归一化后的像素值范围不同计算PSNR时要么用反归一化的图要么直接把0作为最大值因为去均值后像素范围大概在-0.4到0.6之间。实际操作我直接用归一化后的值算MSE再用1作为最大像素值得到的PSNR只能做相对比较不能和图像处理论文里的绝对数值直接类比。为了省事更多时候我还是先肉眼看再加个标签预测准确率兜底PSNR作为参考线。4.4 单次完整攻击输出示例把前面各部分组合起来最终的攻击脚本逻辑如下if __name__ __main__: # 加载模型 net LeNet() net.load_state_dict(torch.load(mnist_lenet.pt)) net.eval() # 取一个真实样本 real_data, real_label next(iter(trainloader)) real_label torch.nn.functional.one_hot(real_label, num_classes10).float() # 计算真实梯度 real_grads compute_grad(net, real_data, real_label) # 执行攻击 dummy_data, dummy_label dlg_attack(net, real_grads) # 可视化与评估 with torch.no_grad(): pred net(dummy_data).argmax(dim1).item() visualize(real_data, dummy_data, real_label.argmax().item(), pred)我实测在随机采样下攻击成功率在90%左右。所谓成功指的是重建结果能让人一眼认出对应数字并且预测标签和真实标签一致。也有个别样本会失败通常是数字结构过于简单、笔画纤细的图像比如某些手写体1或7它们在低像素分辨率下包含的梯度特征量较少重建时容易被噪声干扰。5. 调参过程与避坑经验汇总5.1 标签恢复陷阱初学DLG时最容易犯的错误是忽略标签的恢复。如果直接用原始整数标签比如real_label 3然后计算交叉熵时把它传给网络攻击代码会报错或者得到混乱的结果。正确做法是把标签转成one-hot向量再通过softmax层映射成可微的logits。有个替代方案值得提一句既然MNIST只有10类你可以直接把10个可能标签逐个尝试。写法是先用真实标签的梯度做一次匹配若重建失败再换另一个标签重新初始化并攻击。这种方式效果也成立原论文里就提到标签在类别少的情况下可以通过穷举确定。但穷举10次的开销比直接把标签当作优化变量多10倍效率太低了。除非出现一优化标签就陷入局部最优的情况否则还是用梯度法联调。5.2 梯度获取的版本差异GPU和CPU的梯度计算在数值上完全一致但float32在复现时偶尔会出现微小误差这不影响攻击效果。需要注意的反而是torch.autograd.grad返回的梯度顺序和net.parameters()的顺序完全一致这是后续做zip匹配的前提。如果你在compute_grad里用[p.grad for p in net.parameters()]拿梯度中间有任何一个detach或clone顺序乱了梯度匹配就会出问题。另一个常见问题是多次backward导致梯度累积。比如在攻击循环里为了调试打印了中间状态的loss保险起见每次调试完都执行一次net.zero_grad()和optimizer.zero_grad()。我做实验时因为漏了一次net.zero_grad()前60轮loss都异常偏高查了半天才发现是上一个iteration残留的梯度在干扰这个场景虽然存在但只在batch内多step时才会明显这里只要规范操作就不会有问题。5.3 可视化失真的坑文章前面提过Normalize造成的可视化偏移这里再展开讲一个更隐蔽的问题plt.imshow默认的cmap是viridis它会把灰度图渲染成彩色渐变图。很多复现DLG的帖子贴出来的重建图是蓝色到黄色的渐变那就是忘了加cmapgray参数。如果你照着别人的结果对比一定要统一颜色映射否则会误判重建质量。另外matplotlib在无显示环境下运行需要提前设置import matplotlib matplotlib.use(Agg)否则plt.show()会一直挂在GUI事件循环里不出结果。我最初在远程服务器上跑实验就卡在这里十分钟还以为是模型结构写错导致死循环了。5.4 学习率边界与局部最优学习率0.1是基准值但根据batch大小和目标网络的深度可以适当微调。如果你的网络加深梯度空间维度更大0.1可能略激进需要降到0.05如果替换成线性回归类简单模型0.1又太缓慢可以考虑0.3。关键是观察loss曲线如果loss在50轮前快速下降后进入平台期且重建图模糊加大学习率往往没用反而应该调整betas里的β2如果loss震荡剧烈说明学习率偏大直接减半重跑。局部最优是DLG另一个恼人的问题。同一个数字不同随机初始化下重建结果可能时好时坏。解决办法很简单随机初始化跑5次选PSNR最高的一次。这听起来很粗暴但确实是复现实验的真实做法。实践角度更推荐的做法是先用较小的迭代次数快速筛出表现好的初始点再加大迭代精修时间成本能降一半。6. 常见问题速查表与防御延伸思考把实际踩过的各类问题整理成速查表方便读者复现时直接对照症状可能原因解决方式loss一直不下降学习率过低或虚拟数据初始化不当调高学习率至0.1~0.2改用标准正态初始化loss前期下降慢后期震荡学习率偏高或标签初始化分布偏差大降低学习率到0.05标签初始化改用均匀分布重建图是完整但错乱的数字标签恢复失败检查是否用one-hot标签确认标签参与可优化程序报错grad_fn缺失create_graph未设为True检查torch.autograd.grad的create_graph参数重建图模糊无细节模型未预训练或batch过大先训练到高准确率batch_size改为1再试重建结果依赖随机种子陷入局部最优多次初始化重新攻击选取最优结果输出结果是彩色伪像matplotlib未指定gray添加cmapgray参数防御方面DLG攻击的复现结论对设计隐私保护方案也有直接参考价值。即使梯度只包含权重更新量也足以恢复训练数据细节因此在联邦学习场景里单纯加密传输通道并不能防住恶意服务器。可以从几个方向设计防御梯度加噪在梯度上传前加一层满足差分隐私的噪声。加噪强度需要权衡太强会影响模型收敛太弱则会被攻击方通过多次攻击平均来抵消。梯度压缩只上传Top-K个梯度值大幅降低梯度中蕴含的信息量。研究证明Top-K压缩能显著破坏梯度与原始图像之间的映射关系但模型精度也会有一定损失。安全聚合利用多方安全计算技术让服务器只能看到聚合后的梯度无法区分单个客户端的贡献。这个方法防御效果好但通信开销大实现复杂度也高。模型拆分或Split Learning把网络架构切分成两部分一部分在客户端运行另一部分在服务器端运行只在中间层传输特征而不是梯度从而规避梯度直接暴露的问题。实际做防御实验时我的体会是梯度加噪最容易被工程落地因为实现简单、无需改动训练框架安全聚合则在系统级防御中更规范但需要额外建设密码学基础设施。这些方向都可以基于本文的复现代码直接做对抗测试攻击侧脚本不变防御侧在compute_grad之后、传输之前加一个扰动函数即可。7. 实验心得与后续扩展方向DLG实验复现本身不难但要做到理解原理、跑通流程、对比分析、尝试防御四步都走完才算真正把这个攻击吃透。我的个人感受是这个实验最大的价值不在攻击本身而在于逼迫你重新思考模型参数中到底存了什么。每当你觉得梯度只是数值时看看重建出来清晰可辨的数字就会对分布式训练中的数据隔离有更深的敬畏。实际操作中的一个小技巧把攻击迭代过程每隔20轮保存一帧图像合成一个GIF或短视频能直观看到噪声如何一步步变成手写数字。这个动画效果对理解和向别人讲解DLG攻击尤其有帮助。实现方式不复杂每次迭代时保存当前dummy_data的numpy数组最后用imageio或PIL合成GIF即可。后续扩展实验可以从三个方向做。一是尝试CIFAR-10数据集观察彩色图像在相同攻击条件下的重建效果你会发现图像从灰度变彩色后攻击难度上升需要更多迭代次数和更精细的标签处理。二是增加batch_size到8或16观察多样本梯度下攻击是否还能成功、还原的图片是混合体还是逐张清晰。三是把攻击目标从恢复数据改为恢复类别信息只重建一个低分辨率类别原型那么所需的迭代次数更少攻击成本更低。这些扩展实验对于理解隐私攻击的边界很有帮助我也正在整理CIFAR版本的结果后续再单独写一篇。