首页 / 资讯中心 / 文章详情

PyTorch复现DLG梯度泄露攻击:从MNIST看联邦学习隐私风险

PyTorch复现DLG梯度泄露攻击:从MNIST看联邦学习隐私风险 ★ FEATURED ARTICLE
如果说联邦学习最大的卖点是“数据不动模型动”那梯度泄露攻击就像一盆冷水直接浇在头上你只把梯度共享出去服务器端却有能力把参与方的原始训练样本反推回来。这篇文章记录的是我用MNIST完整复现Deep Leakage from Gradients简称DLG梯度攻击实验的全过程从原理解读、实验选型、代码实现到调参踩坑全部摊开来讲。DLG这个攻击思路最早来自一篇名为《Deep Leakage from Gradients》的论文核心结论非常反直觉共享梯度与原始图像之间存在强相关性攻击者在不知道任何训练数据的前提下仅通过梯度和模型结构就能优化出一张与真实样本几乎一致的图片。对做联邦学习、隐私计算、对抗攻防研究的人来说这个实验几乎是必跑一遍的入门项目。我的目标很明确用最少的时间把实验完整跑通验证攻击效果同时把影响恢复效果的关键因素摸清楚。先说结论在MNIST 简单CNN Adam优化器的组合下通常几十轮迭代就能看出数字轮廓100到200轮基本恢复出可辨识的数字图形。整个过程不依赖特殊框架PyTorch就能搞定。这个实验非常适合作为理解“梯度究竟携带了多少隐私信息”的起点也适合用来给后续的防御方案做baseline。1. 先说清楚这个实验到底在还原什么1.1 梯度泄露的直观解释要理解梯度攻击先要理解梯度是什么。训练神经网络时反向传播会计算出每个参数对损失函数的影响程度这个影响程度就是梯度。梯度的维度通常等于模型参数的总数而MNIST图像只有28×28784个像素点。一个卷积神经网络轻轻松松就有几万甚至几十万个参数梯度空间远大于图像空间所以从数学上看从高维的梯度反推低维图像是完全可行的。我习惯用一个生活化类比来解释这个现象模型训练就像学生做错题本每个参数是试卷上的一个空位梯度就是老师在错题旁边标注的“这里偏差了多少、应该往哪个方向改”。如果这本错题本足够详细别人完全可以从批改笔迹推测出原题长什么样。DLG做的事情就是拿到“批改记录”梯度反推出“原题”训练图片。具体到算法上DLG的思路可以拆成四步攻击者拿到模型当前参数和一批梯度梯度记为∇W对应某条真实训练样本(x, y)。攻击者随机初始化一个虚拟输入x它与真实输入形状完全一致。把x输入模型算出虚拟梯度∇W。不断调整x让虚拟梯度与真实梯度之间的差距越来越小。差距用欧式距离衡量即损失函数L ||∇W - ∇W||²。当梯度距离逼近0时虚拟输入x就逼近真实输入x攻击成功。# 攻击核心最小化真实梯度与虚拟梯度的欧式距离 grad_loss sum(((dummy_grad[i] - real_grad[i]) ** 2).sum() for i in range(len(real_grad))) grad_loss.backward()这个优化过程听起来简单但真正跑起来会发现很多细节决定成败后面第3章我会详细展开。1.2 实验选型的逻辑为什么是MNIST 简单CNNMNIST几乎是梯度攻击复现的首选数据集。原因一是维度低、训练快一张图只有784个像素单张图片的梯度计算在普通CPU上都能秒级完成原因二是可视化直观恢复出来的图片是灰度图直接用imshow就能看出效果不像CIFAR-10在早期迭代阶段容易被误判成噪声原因三是MNIST的标签任务非常简单CIFAR-10里的猫、狗、汽车在攻击恢复时经常出现类别混淆而MNIST的数字结构相对独立混淆概率低很多。当然MNIST也有明显的“副作用”。它太简单了恢复效果容易“好过头”导致一些人误以为攻击是件轻松的事。一旦切到CIFAR-10或者ImageNet需要用更大的网络、更多的迭代轮次和更精细的初始化策略难度会显著上升。所以我更愿意把MNIST当作“教学载体”而不是“最终战场”。网络结构选了简单的CNN而不是全连接网络原因在于卷积核天然具有局部感受野梯度中包含的空间信息更丰富恢复出的图像会更平滑自然。但注意不要添加BatchNorm层——BN层依赖batch内统计量在batch size1时会引入严重的统计噪声导致恢复过程极不稳定。如果一定要用BN建议换成InstanceNorm。torch.manual_seed(42) # 固定种子保证攻击实验可复现 class SimpleCNN(nn.Module): def __init__(self): super().__init__() self.conv1 nn.Conv2d(1, 16, 5, padding2) # 输出 16×28×28 self.pool1 nn.MaxPool2d(2) # 输出 16×14×14 self.conv2 nn.Conv2d(16, 32, 5, padding2) # 输出 32×14×14 self.pool2 nn.MaxPool2d(2) # 输出 32×7×7 self.fc nn.Linear(32 * 7 * 7, 10) def forward(self, x): x torch.relu(self.conv1(x)) x self.pool1(x) x torch.relu(self.conv2(x)) x self.pool2(x) x x.view(x.size(0), -1) return self.fc(x)2. 动手前的准备与关键参数设计2.1 框架选择与模拟拓扑复现DLG不需要真正搭建一个联邦学习系统。核心场景只需要三个角色一张真实图片、一个固定参数的模型、一份由该模型计算出的梯度。攻击者知道模型结构也知道模型当前参数这在联邦学习场景下完全成立因为服务端为了聚合模型必然要下发参数和接收梯度。我选择PyTorch主要是因为它对自动微分的控制非常灵活torch.autograd.grad可以精确计算损失对任意模型参数的梯度配合create_graphTrue还能继续对输入求二阶梯度这正是DLG实现需要的能力。TensorFlow也能实现但需要自己处理GradientTape的嵌套代码更绕一些。整个模拟拓扑如下# 取一条真实样本作为攻击目标 transform transforms.Compose([transforms.ToTensor()]) dataset datasets.MNIST(data, trainTrue, downloadTrue, transformtransform) x, y dataset[0] x x.unsqueeze(0) # shape: (1, 1, 28, 28) y torch.tensor([y], dtypetorch.long)这里有个关键点真实样本不需要经过训练甚至不需要模型已经收敛。DLG对任意训练阶段的梯度都有效因为梯度在反向传播过程中携带着输入的信息与模型是否收敛无关。这一点经常被新手误解以为“必须训练一个模型再用它的梯度”完全没必要。2.2 网络结构简单CNN还是全连接我对比过全连接网络和CNN的恢复效果。全连接网络的梯度同样能恢复图像但恢复出的图像会出现明显的“网格状伪影”尤其是数字边缘的过渡不够自然。原因在于全连接层的梯度是对全局像素的组合加权缺乏卷积核的局部平滑先验。CNN之所以效果更好可以这样理解卷积核对输入做局部加权求和它的梯度本质上是“每个位置的输入像素对输出特征的贡献量”。多个卷积层叠加后梯度中同时包含浅层的边缘信息、纹理信息和深层的语义信息这些信息交织在一起给反推图像提供了丰富得多的约束条件。具体到参数规模我试过16通道和64通道两种配置。通道数翻倍后恢复速度确实更快但训练时间和显存占用也同步增加对MNIST来说16通道已经足够。我最终选定的是两层卷积一层全连接的小网络参数总量约5万做一次前向和反向传播非常快CPU上单轮迭代也就几十毫秒。2.3 攻击参数距离度量、优化器、迭代轮数DLG的原始论文比较了两种距离度量L2欧式距离和L1曼哈顿距离。实验结论是L2在图像恢复任务上更平滑、更容易收敛L1则更容易产生边缘锐利但整体结构偏差的结果。所以默认选择L2距离。优化器的选择对攻击成功率影响极大。论文中推荐Adam和LBFGS我实测下来Adam要稳得多。LBFGS的收敛速度更快但经常陷入局部最优恢复出的图像是一团噪声Adam虽然收敛慢一些但几乎每次都能成功。所以普通复现优先Adam学习率设为0.1是一个非常可靠的起点。迭代轮数方面MNIST 简单CNN Adam在100轮左右就能看到数字轮廓200轮基本稳定。继续增加轮数收益很小反而可能出现过度拟合噪声导致图像发暗的情况。# 初始化虚拟输入全0.1灰度图 dummy_data torch.full((1, 1, 28, 28), 0.1, requires_gradTrue) optimizer torch.optim.Adam([dummy_data], lr0.1)有个容易忽视的细节虚拟输入需要设置为requires_gradTrue而且优化器必须在这个声明之后创建否则PyTorch会报“参数不需要梯度”的错误。另外不要直接优化模型参数优化目标是虚拟输入本身。3. 代码实现一步步把攻击跑起来3.1 项目结构与公共代码整个复现项目只需要一个脚本就能跑通大概200行以内。我会分成三部分展示模型定义与真实梯度获取、攻击主循环、可视化。正式开跑前建议用torch.manual_seed(42)固定随机种子。DLG对随机性比较敏感同样的代码不固定种子可能这轮能恢复出清晰的“7”下一轮就恢复成一团糊。固定种子至少能让实验可复现便于后续调整参数时做横向对比。3.2 真实梯度的获取真实梯度是整个攻击的目标。需要注意的是DLG中的梯度是模型参数相对于损失函数的梯度而不是输入图像的梯度。代码写法如下criterion nn.CrossEntropyLoss() net SimpleCNN() net.train() real_loss criterion(net(x), y) real_grad torch.autograd.grad( real_loss, net.parameters(), retain_graphTrue, create_graphTrue, allow_unusedTrue )这里有三点需要解释。第一retain_graphTrue是为了保留计算图方便后续对虚拟输入求梯度时重复使用网络结构第二create_graphTrue表示构建二阶导的计算图这样后面才能把梯度距离的梯度回传到虚拟输入上第三allow_unusedTrue是为了防止某些层参数没有被当前输入用到时抛出异常加了它更保险。我在实际复现时还犯过一个低级错误先对真实损失调用backward()再把梯度值保存下来。这样做的结果是计算图被释放后续无法对虚拟输入求梯度。正确做法是只用torch.autograd.grad()获取梯度不要调用backward()。3.3 攻击主循环的实现攻击循环的核心逻辑是用虚拟输入计算虚拟梯度然后最小化虚拟梯度与真实梯度之间的距离梯度下降只更新虚拟输入。for it in range(200): optimizer.zero_grad() dummy_loss criterion(net(dummy_data), y) # 虚拟输入的前向损失 dummy_grad torch.autograd.grad( dummy_loss, net.parameters(), retain_graphTrue, create_graphTrue, allow_unusedTrue ) grad_loss sum( ((dg - rg) ** 2).sum() for dg, rg in zip(dummy_grad, real_grad) ) grad_loss.backward() optimizer.step() if it % 20 0: print(fiter {it:4d} | grad_loss {grad_loss.item():.6f})这段代码的运行逻辑可以总结为“三步循环”虚拟输入前向传播得到虚拟损失再反向得到虚拟梯度。计算虚拟梯度与真实梯度的L2距离。通过grad_loss.backward()把误差回传到虚拟输入再用Adam更新虚拟输入。重要提醒grad_loss.backward()之前必须调用optimizer.zero_grad()否则梯度会累加。虽然优化器只管理dummy_data这一个参数PyTorch中zero_grad()最好还是保留免得后续扩展代码时出现问题。3.4 恢复效果的可视化与观察迭代过程中需要定时把虚拟输入画出来便于观察恢复进度。可视化代码很简单import matplotlib.pyplot as plt import numpy as np def show_dummy(dummy_data, title): img dummy_data.detach().cpu().squeeze().numpy() img np.clip(img, 0, 1) plt.imshow(img, cmapgray) plt.title(title) plt.axis(off) plt.show()注意三个细节。第一必须用detach()切断梯度否则matplotlib转换numpy时会报错第二需要np.clip(img, 0, 1)把像素值限制在合法范围因为优化过程中虚拟像素可能跑到[−0.3, 1.5]这种越界区间第三squeeze()去掉batch和channel维度得到28×28的二维数组。我跑下来的典型现象是前20轮看到的图像几乎是一张带噪灰图50轮左右开始出现数字的粗略轮廓100轮以后轮廓变得清晰150到200轮基本稳定。如果看到的是数字和噪声的混合体通常是学习率偏大或者迭代轮数不够。4. 复现路上最容易踩的坑4.1 问题速查表我把复现过程中遇到过的典型问题整理成了表格方便你对照排查。现象可能原因解决方法训练时报“element 0 of tensors does not require grad”dummy_data没开requires_grad初始化时加requires_gradTrue恢复出的图像始终是随机噪声学习率太小或初始化方式不对试试lr0.1初始化换全0.1灰度图loss下降很快但图像不清晰梯度距离收敛到局部最优换初始值random seed或降低学习率到0.05图像整体发白或发黑像素值越界后clip导致信息丢失可视化时clip但不要每次迭代都clip到0~1加了BN层后恢复效果奇差BN统计量在batch1时不稳定去掉BN或者改用InstanceNormbatch size2以后无法恢复出独立的图片多张图的信息在梯度中混叠DLG基本只在batch1下有效换batch14.2 针对常见场景的调参建议如果是最初测试直接用全0.1灰度图作为虚拟输入初始化通常最稳。为什么这个看起来“无脑”的初始化效果好因为MNIST数字图像的背景就是接近0的黑色前景笔画接近1的白色0.1这个中间灰度值既不偏向前景也不偏向背景优化器可以从这个中立点出发自由决定每个像素往亮还是往暗走。如果换用torch.rand均匀随机初始化恢复效果时好时坏。原因在于随机的亮暗分布提供了一个偏置如果初始噪声的某些区域恰好接近某个数字的结构优化会被引导到那个数字方向如果初始噪声和真实数字差太远容易陷在局部最优里。关于学习率我实测的规律是lr0.1时前30轮loss快速下降随后进入平稳优化阶段lr0.01时下降非常慢100轮后图像仍然模糊lr0.5时前10轮会剧烈震荡虚拟图像会出现大面积过曝虽然最终也能恢复但稳定性差。所以默认0.1是经过验证的安全选择。4.3 小技巧从局部恢复现象看攻击机制有几次实验恢复出来的数字跟真实标签不是同一个。比如真实样本是“4”恢复出来却是“9”。这不是bug而是DLG优化陷入了局部最优。原因是MNIST手写体的“4”和“9”在部分样本上结构相似梯度也相近优化器沿着梯度下降方向走进了混淆区域。遇到这种情况有两个处理思路。一是换一个样本再试不同手写样本的梯度约束强度不同二是同时叠加多个约束比如把标签也纳入优化目标或者同时使用多个模型的梯度进行联合约束。另外我做了一个很有价值的观察中间层的梯度比输出层的梯度对图像恢复贡献更大。输出层梯度只包含类别判别信息而中间卷积层的梯度保留了更多空间纹理信息。如果你想做更深入的实验可以尝试用“中间层梯度输出层梯度”拼接的方式计算距离往往能获得比单纯用全参数梯度更稳定的恢复结果。5. 攻击之后防御思路与可扩展玩法5.1 常见防御手段与攻防博弈复现攻击的目的不是单纯炫技而是理解防御为什么难做。我简单列举几个主流的防御方向全部可以在DLG框架里做对照实验。梯度裁剪是最常见的手段把梯度限制在一个小范围里比如torch.clamp(grad, -0.1, 0.1)。这样做会破坏梯度中的精细结构使攻击恢复的图像噪点增多。但代价是模型训练的收敛速度也会受影响梯度裁剪太狠会导致模型根本不收敛。梯度扰动是另一条路典型代表是DP-SGD。在梯度上传前加入高斯噪声噪声强度和隐私预算挂钩。噪声越大攻击恢复的难度越高。我的对照实验显示当噪声幅度接近真实梯度幅度的10%时MNIST恢复的图像已经明显退化噪声幅度到30%时攻击基本失效只能恢复出模糊的数字轮廓。还有个更轻量的思路本地多次迭代。客户端在本地多跑几个epoch再把梯度传上去梯度携带的单个样本信息会被多个batch平均稀释攻击难度随之上升。这也是联邦学习框架普遍采用的做法。5.2 还能延伸做的实验方向如果MNIST版本的DLG跑通了我建议按下面几个方向延伸每个方向都能单独成文。第一个方向是切换数据集。把MNIST换成Fashion-MNIST或CIFAR-10观察恢复效果的变化。CIFAR-10是32×32×3的彩色图像素维度比MNIST高得多但网络也更大通常需要增加迭代轮数和调整初始化策略。第二个方向是标签未知的恢复。原始DLG论文提出把虚拟标签也作为可学习参数甚至可以让标签从均匀分布开始通过优化自动得到正确的类别。代码改动很小只需在虚拟输入之外再创建一个dummy_label用torch.softmax平滑后参与损失计算。第三个方向是防御对抗实验。把梯度裁剪、噪声扰动和本地迭代三种防御手段分别叠加到攻击流程里评估恢复图像的PSNR值和结构相似度量化攻击成功率随防御强度的变化曲线。这一类实验非常适合作为安全方向毕业论文的实验章节。我在实际运行中还有一个体会梯度攻击与其说是一个“漏洞”不如说是一种对隐私边界的重新定义。它逼迫我们重新思考什么数据可以放心共享什么信息即使经过“匿名化”处理仍然可能被还原。最后分享一个实用小技巧如果恢复过程中loss降不下去可以先检查真实梯度和虚拟梯度的数量级是否一致。方法很简单打印两者的L2范数做个对比如果差异超过10倍多半是网络初始化或者数据归一化不一致导致的把真实数据和虚拟数据用同样的预处理管道过一遍问题就能解决。这个小检查我每次跑新模型都会先做一遍能省掉大量排查时间。
阅读完成 · 觉得有帮助?
咨询建站