[论文解读] Defensive Unlearning with Adversarial Training for Robust Concept Erasure in Diffusion Models

Yimeng Zhang, Xin Chen+7 more|arXiv (Cornell University)|May 24, 2024

Adversarial Robustness in Machine Learning被引用 4

一句话总结

该论文提出 AdvUnlearn，一种新颖的鲁棒遗忘框架，通过将对抗训练整合到扩散模型（DM）概念擦除中，以抵御对抗性提示攻击。通过在保留提示集上对文本编码器应用保持实用性的正则化，AdvUnlearn 在保持高图像生成质量（FID ≈ 19.34）的同时实现了强大的鲁棒性（例如，在 SD v1.4 上攻击成功率仅为 21.13%），在鲁棒性与实用性方面均优于现有基线方法。

ABSTRACT

Diffusion models (DMs) have achieved remarkable success in text-to-image generation, but they also pose safety risks, such as the potential generation of harmful content and copyright violations. The techniques of machine unlearning, also known as concept erasing, have been developed to address these risks. However, these techniques remain vulnerable to adversarial prompt attacks, which can prompt DMs post-unlearning to regenerate undesired images containing concepts (such as nudity) meant to be erased. This work aims to enhance the robustness of concept erasing by integrating the principle of adversarial training (AT) into machine unlearning, resulting in the robust unlearning framework referred to as AdvUnlearn. However, achieving this effectively and efficiently is highly nontrivial. First, we find that a straightforward implementation of AT compromises DMs' image generation quality post-unlearning. To address this, we develop a utility-retaining regularization on an additional retain set, optimizing the trade-off between concept erasure robustness and model utility in AdvUnlearn. Moreover, we identify the text encoder as a more suitable module for robustification compared to UNet, ensuring unlearning effectiveness. And the acquired text encoder can serve as a plug-and-play robust unlearner for various DM types. Empirically, we perform extensive experiments to demonstrate the robustness advantage of AdvUnlearn across various DM unlearning scenarios, including the erasure of nudity, objects, and style concepts. In addition to robustness, AdvUnlearn also achieves a balanced tradeoff with model utility. To our knowledge, this is the first work to systematically explore robust DM unlearning through AT, setting it apart from existing methods that overlook robustness in concept erasing. Codes are available at: https://github.com/OPTML-Group/AdvUnlearn

研究动机与目标

• 解决概念擦除后的扩散模型在面对可绕过遗忘机制的对抗性提示攻击时的脆弱性问题。
• 开发一种鲁棒的遗忘框架，能够在抵抗对抗性提示越狱攻击的同时保持高图像生成质量。
• 识别在遗忘过程中对抗鲁棒化最有效的模型组件（文本编码器 vs. UNet）。
• 实现在多种扩散模型架构间可即插即用的鲁棒遗忘器迁移能力。
• 系统性地研究对抗训练在扩散模型机器遗忘背景下的有效性与效率。

提出的方法

• 将对抗训练（AT）集成到机器遗忘流程中，以提升对对抗性提示攻击的鲁棒性。
• 专门对文本编码器而非 UNet 应用对抗训练，因为其在鲁棒遗忘方面表现更优。
• 在保留提示集上应用保持实用性的正则化，以在对抗微调过程中维持图像生成质量。
• 采用多步对抗提示生成策略，构造能最大化攻击未学习模型成功率的扰动。
• 实现训练后文本编码器在其他扩散模型上的即插即用部署，无需重新训练。
• 优化文本编码器的前 N 层，消融实验表明，对更深层进行优化可提升对复杂概念（如裸露）的鲁棒性。

Figure 1: Visualizations of SD v1.4 (Rombach et al., 2022 ) and its concept-erased version, ESD (Gandikota et al., 2023a ) , conditioned on the prompt ‘Hot sexy girl’ . The generation of ESD is shown against no attack and UnlearnDiffAtk (Zhang et al., 2023b ) .

实验结果

研究问题

• RQ1 对抗训练能否被有效且高效地集成到扩散模型遗忘中，以提升对对抗性提示的鲁棒性？
• RQ2 在遗忘过程中，扩散模型的哪个组件（文本编码器或 UNet）在对抗鲁棒化方面更有效？
• RQ3 在对抗遗忘过程中，如何有效平衡鲁棒性与图像生成实用性的权衡？
• RQ4 能否在不微调的情况下，将在一个扩散模型上训练的鲁棒遗忘器迁移到其他模型？
• RQ5 不同文本编码器层配置对遗忘的鲁棒性与实用性有何影响？

主要发现

• 与原始模型的 100% 攻击成功率相比，AdvUnlearn 在 SD v1.4 上对裸露概念的遗忘将攻击成功率（ASR）降低至 21.13%。
• 该框架保持了强大的图像生成质量，FID 为 19.34，仅略高于原始模型的 16.70。
• 文本编码器在对抗鲁棒化方面显著优于 UNet，前者实现了更低的 ASR 和更优的实用性权衡。
• 通过 AdvUnlearn 学习到的鲁棒文本编码器可迁移到其他模型（如 SD v1.5、DreamShaper、Protogen）中，性能下降极小，在 SD v1.5 上实现 16.20% 的 ASR，在 Protogen 上实现 42.96% 的 ASR。
• 优化更多层的文本编码器可提升鲁棒性，全层优化可提供最强防御，尤其对裸露等复杂概念效果显著。
• 保持实用性的正则化有效缓解了对抗训练通常导致的图像质量下降，成功保留了模型实用性。

Figure 2: Performance overview of our proposal AdvUnlearn and various DM unlearning baselines when unlearning the nudity concept under the SD v1.4 model. The robustness is measured by attack success rate (ASR) against UnlearnDiffAtk (Zhang et al., 2023b ) . The performance of image generation retent