本文第一作者为密歇根州立大学计算机系博士生张益萌，为 OPTML 实验室成员，指导教师为刘思佳助理教授。OPtimization and Trustworthy Machine Learning (OPTML) 实验室的研究兴趣涵盖机器学习 / 深度学习、优化、计算机视觉、安全、信号处理和数据科学领域，重点是开发学习算法和理论，以及鲁棒且可解释的人工智能。
扩散模型（Diffusion Models, DMs）已经成为文本到图像生成领域的核心技术之一。凭借其卓越的性能，这些模型可以生成高质量的图像，广泛应用于各类创作场景，如艺术设计、广告生成等。然而，随着扩散模型的日益普及，其带来的安全问题也逐渐显现。模型在处理开放式互联网数据时，可能会在生成过程中输出有害的、不适当的内容，例如裸露、暴力、侵犯版权的图像等，这为其实际应用带来了道德和法律上的挑战。
为应对这些问题，研究者提出了机器遗忘（Machine Unlearning）技术，也称为概念擦除（Concept Erasing）[1-3]。通过这一技术，我们能够有选择性地 “抹去” 扩散模型中的某些不应生成的概念。然而，尽管这一技术能够在非对抗性环境中取得一定效果，现有的概念擦除方法在面对对抗性提示攻击（Adversarial Prompt Attacks）时仍然存在明显的脆弱性。攻击者可以通过对输入提示的微小修改，诱导模型生成本应被擦除的内容，这使得现有的概念擦除技术不够鲁棒。
对抗性提示攻击：机器遗忘中的隐患
对抗性提示攻击的关键在于通过改变文本输入提示的细微细节，诱使模型生成不当的内容。这类攻击可能会通过调整拼写、加入特定符号、或是改变句子结构来规避模型的防护。例如，即便模型已经接受了擦除 “裸露” 相关内容的训练，但通过输入一些巧妙设计的提示词，攻击者依然可能让模型重新生成被擦除的图像。这种攻击的成功率显示了现有擦除技术在对抗恶意输入时的薄弱性。
为了解决这一问题，我们提出了 AdvUnlearn 框架。这是一个将对抗性训练（Adversarial Training, AT）与概念擦除相结合的框架，旨在提升扩散模型在概念擦除任务中的鲁棒性，防止模型在面对对抗性攻击时再次生成被遗忘的内容。
- 论文题目：Defensive Unlearning with Adversarial Training for Robust Concept Erasure in Diffusion Models
- 论文地址：https://arxiv.org/abs/2405.15234
- 代码地址：https://github.com/OPTML-Group/AdvUnlearn
- Unlearned Diffusion Model Benchmark: https://huggingface.co/spaces/Intel/UnlearnDiffAtk-Benchmark
AdvUnlearn 框架：对抗性训练与概念擦除的融合
AdvUnlearn 框架的独特之处在于，它系统性地结合了对抗性训练与概念擦除方法，从而提升模型在对抗恶意输入时的安全性和鲁棒性。传统的对抗性训练主要应用于图像分类任务，而 AdvUnlearn 首次将其引入到扩散模型的概念擦除中。通过双层优化策略，AdvUnlearn 能够在增强模型鲁棒性的同时保持其图像生成质量。
双层优化（Bi-level Optimization）：AdvUnlearn 框架的优化分为两个主要层次：
- 下层优化（Lower-level Optimization）：这一层的任务是生成对抗性提示，即通过对输入文本提示的细微扰动，生成能够诱使模型生成本应被擦除的内容的提示。这个过程的目标是找到这些对抗性提示，使得模型在面对这些提示时生成不合规内容。
- 上层优化（Upper-level Optimization）：这一层的任务是通过最小化模型对目标概念的响应来强化模型的鲁棒性，同时确保模型生成质量不受影响。我们通过更新模型参数 θ，优化概念擦除的效果，同时保留模型的正常生成能力。
其中
为目标遗忘内容，
为遗忘目标函数，
为生成攻击的目标函数。
保留效用的正则化（Utility-Retaining Regularization）
一个重大挑战是，直接引入对抗性训练可能会导致图像生成质量下降。这是因为模型在面对对抗性提示时，往往会过于专注于抵御不当输入，而忽视了对其他正常内容的生成能力。为此，我们提出了保留效用的正则化方法。
该方法的核心思想是引入一个 “保留集” （
），这是一组与目标概念无关的文本提示。这些提示用于确保模型在训练过程中保留其生成正常内容的能力。例如，在处理 “裸露” 相关概念时，保留集中的提示可能涉及描述风景、动物、建筑等非敏感内容。
保留集的选择尤为重要。我们用 COCO 和 ImageNet 等数据集中的非敏感的物体类别制作成文本提示作为保留集，并通过大模型进行筛选，确保保留集中的提示与待擦除概念无关。这些提示集通过以下正则化项加入优化目标：
其中
为未经过机器遗忘的初始模型参数。通过该机制，AdvUnlearn 框架能够有效解决对抗性训练带来的质量下降问题，确保生成图像的多样性和精度。
最佳优化模块的选择 （Modularity Exploriation）: 文本编码器 vs. UNet
扩散模型通常由多个核心模块组成，其中文本编码器（Text Encoder）和 UNet 是两个主要组件。UNet 通常用于处理图像的逐步去噪过程，而文本编码器则将文本提示转换为适合生成图像的表示。
在传统的概念擦除方法中，UNet 被视为模型生成图像的核心，因此大多数现有技术都集中在 UNet 上进行优化。然而，在我们的研究中，我们发现文本编码器在应对对抗性提示攻击时表现更佳。文本编码器在处理输入提示的过程中，负责将文本信息转化为图像生成的关键特征。因此，对文本编码器进行优化，可以更好地控制模型对输入中细微变化的反应，从而提升模型在面对恶意输入时的鲁棒性。
具体来说，我们发现，优化文本编码器能够显著降低对抗性攻击的成功率，同时保留生成高质量图像的能力。此外，经过对抗性训练优化的文本编码器可以作为一个 “即插即用” 的模块，适用于不同类型的扩散模型，从而提高了框架的扩展性和通用性。
实证结果：AdvUnlearn 在多个任务中的表现
为了验证 AdvUnlearn 框架的有效性，我们进行了广泛的实验证明。实验涵盖了多个概念擦除任务，包括裸露概念擦除、艺术风格擦除和对象擦除。UnlearnDiffAtk [4] 作为默认的对抗性提示攻击来检查遗忘模型的鲁棒性。
裸露概念擦除
在处理敏感的 “裸露” 概念时，AdvUnlearn 展现了卓越的性能。相比传统的概念擦除方法，AdvUnlearn 显著降低了对抗性提示攻击的成功率（ASR）。例如，在常见的对抗性提示攻击场景下，ASR 从原始模型的 100% 下降至 21.13%，表明模型能够有效抵御恶意输入。同时，通过引入保留集，模型在生成无关内容（如风景或物体图像）时保持了高质量，FID 仅比标准模型略微上升。
艺术风格擦除
我们还对 “梵高风格” 的擦除进行了实验。在这一任务中，AdvUnlearn 同样表现出了优异的性能。经过对抗性训练后的模型，能够有效抵御对抗性提示攻击，成功率降至 2%。相比之下，未经优化的模型在攻击下几乎无法抵御，生成的图像仍保留了梵高的风格。此外，AdvUnlearn 在这一任务中成功保持了高质量的图像生成能力。
对象概念擦除
在对象擦除任务中，我们对 “教堂” 这一具体对象进行了擦除实验。AdvUnlearn 不仅能够有效抹去生成 “教堂” 图像的能力，还能在对抗性提示攻击下维持图像生成的质量。与其他擦除方法相比，AdvUnlearn 框架能够在保证生成质量的前提下，显著提升鲁棒性，使得模型在面对恶意输入时更加可靠。
总结与展望
AdvUnlearn 框架是首个结合对抗性训练和概念擦除技术的系统性方法，旨在提升扩散模型在概念擦除任务中的鲁棒性。通过对文本编码器进行优化，以及引入保留效用的正则化机制，AdvUnlearn 成功平衡了鲁棒性和模型效用。我们的实验结果显示，AdvUnlearn 能够在多种场景下有效抵御对抗性提示攻击，并保持高质量的图像生成。
随着生成式 AI 的应用范围不断扩大，如何确保模型的安全性和可靠性成为一个关键问题。AdvUnlearn 框架为解决这一问题提供了新的路径，并为未来的生成模型技术提供了重要的参考。如果你想深入了解 AdvUnlearn 框架的技术细节或实验结果，欢迎访问 GitHub 项目页面（https://github.com/OPTML-Group/AdvUnlearn）。
