AdvUnlearn阅读笔记：基于对抗训练的扩散模型鲁棒概念擦除

AdvUnlearn阅读笔记基于对抗训练的扩散模型鲁棒概念擦除引言概念擦除的挑战与对抗训练的必要性扩散模型如Stable Diffusion在图像生成领域取得了巨大成功但其训练数据中可能包含有害概念如暴力、色情、侵权风格导致模型可能生成不安全内容。概念擦除Concept Unlearning旨在从预训练模型中移除特定知识而无需完全重训。然而现有方法如EraseDiff、ESD存在一个致命缺陷对对抗性提示不鲁棒。攻击者可以通过精心构造的提示例如添加噪声或同义词替换轻易绕过擦除机制重新生成被禁止的概念。AdvUnlearnAdversarial Unlearning由Li等人提出核心思想是将对抗训练融入概念擦除过程使模型对对抗性扰动具有鲁棒性。本文从原理出发深入剖析其数学框架与实现细节并提供可运行代码示例。## 核心原理对抗训练与概念擦除的融合### 1. 问题形式化设扩散模型为ϵθ\epsilon_\thetaϵθ​原始概念为ccc目标擦除概念为ctargetc_{\text{target}}ctarget​。标准擦除目标为LeraseEx0,c,t[∥ϵθ(xt,c,t)−ϵref(xt,ctarget,t)∥2]\mathcal{L}_{\text{erase}} \mathbb{E}_{x_0, c, t} \left[ \|\epsilon_\theta(x_t, c, t) - \epsilon_{\text{ref}}(x_t, c_{\text{target}}, t)\|^2 \right]Lerase​Ex0​,c,t​[∥ϵθ​(xt​,c,t)−ϵref​(xt​,ctarget​,t)∥2]其中ϵref\epsilon_{\text{ref}}ϵref​是冻结的参考模型用于提供“擦除后”的预测。但此目标对输入ccc的微小扰动敏感。### 2. 对抗训练框架AdvUnlearn引入对抗性提示cadvc_{\text{adv}}cadv​定义为在原始提示上施加ℓp\ell_pℓp​范数约束的扰动cadvarg⁡max⁡∥c′−c∥p≤ϵLerase(c′)c_{\text{adv}} \arg\max_{\|c - c\|_p \leq \epsilon} \mathcal{L}_{\text{erase}}(c)cadv​arg∥c′−c∥p​≤ϵmax​Lerase​(c′)这本质上是一个最大-最小优化问题min⁡θmax⁡∥c′−c∥p≤ϵLerase(c′)\min_{\theta} \max_{\|c - c\|_p \leq \epsilon} \mathcal{L}_{\text{erase}}(c)θmin​∥c′−c∥p​≤ϵmax​Lerase​(c′)内层最大化通过投影梯度下降PGD近似求解。关键洞察对抗性提示迫使模型在擦除时考虑最坏情况从而提升鲁棒性。### 3. 算法流程1.对抗采样对每个原始提示ccc运行K步PGD生成对抗提示cadvc_{\text{adv}}cadv​。2.联合训练同时优化原始提示和对抗提示上的擦除损失。3.正则化加入原始模型预测的KL散度防止灾难性遗忘。## 关键实现细节PGD生成对抗提示对抗提示的生成在文本嵌入空间而非原始文本中执行。这是因为文本嵌入是连续向量可直接应用梯度优化。以下代码实现PGD攻击pythonimport torchimport torch.nn.functional as Fdef pgd_attack(model, ref_model, text_embeds, epsilon0.1, alpha0.01, steps10): 在文本嵌入空间上执行PGD攻击生成对抗性提示。 Args: model: 待训练的扩散模型UNet ref_model: 冻结的参考模型 text_embeds: 原始提示的CLIP文本嵌入形状为[1, 77, 768] epsilon: 扰动半径L2范数约束 alpha: PGD步长 steps: 迭代步数 Returns: adv_embeds: 对抗性文本嵌入 # 初始化对抗扰动为零 adv_embeds text_embeds.clone().detach().requires_grad_(True) for _ in range(steps): # 前向传播计算当前对抗提示下的擦除损失 noise torch.randn_like(text_embeds) # 模拟噪声输入实际需采样时间步 pred model(noise, adv_embeds, timesteptorch.tensor([500])) ref_pred ref_model(noise, text_embeds, timesteptorch.tensor([500])) loss F.mse_loss(pred, ref_pred.detach()) # 擦除损失 # 反向传播梯度 grad torch.autograd.grad(loss, adv_embeds, retain_graphFalse)[0] # 更新对抗嵌入 adv_embeds adv_embeds alpha * grad.sign() # 投影到epsilon球内L2范数约束 delta adv_embeds - text_embeds.detach() norm torch.norm(delta, p2, dim-1, keepdimTrue) delta torch.where(norm epsilon, delta / norm * epsilon, delta) adv_embeds text_embeds.detach() delta # 确保梯度跟踪 adv_embeds adv_embeds.detach().requires_grad_(True) return adv_embeds.detach()关键点注释- 攻击在连续文本嵌入空间执行而非离散词元。- 使用符号梯度grad.sign()加速收敛这在PGD中常见。- L2范数投影确保扰动不超过ϵ\epsilonϵ维持语义相似性。## 训练循环对抗擦除完整的训练循环需交替进行对抗采样和模型更新。以下代码展示核心训练步骤pythondef train_adv_unlearn(model, ref_model, dataloader, optimizer, epsilon0.1, pgd_steps5): AdvUnlearn单轮训练函数。 Args: model: 待训练的UNet ref_model: 冻结的参考模型 dataloader: 包含原始提示和图像的数据加载器 optimizer: 优化器如AdamW epsilon: 对抗扰动半径 pgd_steps: PGD迭代步数 Returns: avg_loss: 平均损失值 model.train() total_loss 0.0 for batch in dataloader: # batch包含: text_embeds, images, timesteps text_embeds batch[text_embeds].cuda() images batch[images].cuda() timesteps batch[timesteps].cuda() # 1. 生成对抗提示 with torch.enable_grad(): adv_embeds pgd_attack( model, ref_model, text_embeds, epsilonepsilon, alphaepsilon/pgd_steps/2, stepspgd_steps ) # 2. 计算原始提示的擦除损失 noise torch.randn_like(images) noisy_images model.q_sample(images, timesteps, noise) pred_clean model(noisy_images, text_embeds, timesteps) pred_ref ref_model(noisy_images, text_embeds, timesteps) loss_clean F.mse_loss(pred_clean, pred_ref.detach()) # 3. 计算对抗提示的擦除损失 pred_adv model(noisy_images, adv_embeds, timesteps) loss_adv F.mse_loss(pred_adv, pred_ref.detach()) # 4. 总损失可加入KL正则化项 loss loss_clean loss_adv # 5. 反向传播与优化 optimizer.zero_grad() loss.backward() torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm1.0) optimizer.step() total_loss loss.item() return total_loss / len(dataloader)核心设计思想- 使用相同的噪声和参考预测计算两个损失确保公平比较。- 对抗提示的梯度计算需要torch.enable_grad()因为pgd_attack内部可能使用了detach()。- 梯度裁剪防止训练不稳定这在对抗训练中尤为重要。## 实验验证与效果分析### 鲁棒性提升原理AdvUnlearn的鲁棒性来源于对抗样本的多样性。在训练过程中模型被迫处理来自不同方向的对抗扰动PGD迭代产生从而学习到更平滑的损失景观。数学上这等价于在经验风险最小化中引入局部Lipschitz连续性约束Ec[max⁡∥c′−c∥≤ϵL(c′)]≤Ec[L(c)]ϵ⋅Ec[∥∇cL(c)∥]\mathbb{E}_{c} \left[ \max_{\|c-c\|\leq\epsilon} \mathcal{L}(c) \right] \leq \mathbb{E}_{c} [\mathcal{L}(c)] \epsilon \cdot \mathbb{E}_{c} \left[ \|\nabla_c \mathcal{L}(c)\| \right]Ec​[∥c′−c∥≤ϵmax​L(c′)]≤Ec​[L(c)]ϵ⋅Ec​[∥∇c​L(c)∥]对抗训练实际上最小化了右侧上界从而控制梯度范数。### 与标准方法的对比| 方法 | 干净提示FID ↓ | 对抗提示擦除成功率 ↑ | 计算开销 ||------|--------------|---------------------|---------|| ESD | 12.3 | 45% | 1x || AdvUnlearn | 13.1 | 92% | 3x |AdvUnlearn在对抗场景下将擦除成功率从45%提升至92%代价是轻微的图像质量下降FID增加0.8和3倍训练时间。## 总结AdvUnlearn通过将对抗训练引入扩散模型概念擦除有效解决了现有方法对对抗性提示脆弱的根本问题。其核心贡献在于1.在文本嵌入空间执行PGD攻击生成语义保持但具有误导性的对抗提示。2.最大-最小优化框架同时优化原始和对抗样本上的擦除损失。3.理论保证对抗训练等价于梯度正则化提升模型局部平滑性。实践表明该方法在保持图像生成质量的同时显著增强了擦除机制的鲁棒性。未来的方向包括探索更高效的对抗采样策略如单步对抗训练以及将方法扩展到多概念擦除场景。对于开发者而言理解并应用AdvUnlearn的思想是构建安全可控生成模型的重要一步。