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

一、研究背景与核心问题

扩散模型(DMs)在文本到图像生成领域取得显著成功,但存在生成有害内容(如NSFW图像)和侵犯版权等安全风险。机器遗忘(概念擦除)技术旨在缓解这些风险,却易受对抗性提示攻击------通过对输入提示进行微小扰动,可使已完成概念擦除的扩散模型重新生成需擦除的内容(如裸体图像)。

核心研究问题:如何高效提升概念擦除后扩散模型对对抗性提示攻击的鲁棒性,同时兼顾图像生成质量(模型效用)?

二、关键原理

2.1 扩散模型基础(潜在扩散模型LDM)

扩散模型通过"逐步去噪"将随机高斯噪声转化为清晰图像,其训练目标是最小化去噪误差。

设:

  • \(x\):清晰图像,\(x_t\):\(t\)时刻含噪声的图像( latent 空间表示)
  • \(c\):文本提示,\(\epsilon_\theta(x_t|c)\):参数为\(\theta\)、条件为\(c\)的噪声估计器
  • \(\mathcal{D}\):训练数据集,\(\epsilon \sim \mathcal{N}(0,1)\):随机噪声

训练目标函数(最小化去噪误差):

\[\underset{\theta}{minimize} \mathbb{E}{(x, c) \sim \mathcal{D}, t, \epsilon \sim \mathcal{N}(0,1)}\left[\left\| \epsilon - \epsilon{\theta}\left(x_{t} | c\right)\right\| _{2}^{2}\right] \tag{1} \]

含义:使模型估计的噪声\(\epsilon_\theta(x_t|c)\)尽可能接近真实噪声\(\epsilon\),保证去噪过程准确性。

2.2 概念擦除基础(ESD方法)

ESD(Erased Stable Diffusion)是主流概念擦除方法,通过调整噪声估计器,引导模型生成远离需擦除概念的图像。

设:

  • \(c_e\):需擦除的概念(如"裸体")
  • \(\theta_o\):原始预训练模型参数,\(\theta\):概念擦除后模型参数
  • \(\epsilon_\theta(x_t|\emptyset)\):空提示(无条件)下的噪声估计
  • \(\eta>0\):擦除引导参数(控制擦除强度)

噪声估计器调整规则

\[\epsilon {\theta }(x{t}|c_{e}) \gets \epsilon _{\theta {o}}(x{t}|\emptyset ) - \eta \left( \epsilon {\theta {o}}(x{t}|c{e}) - \epsilon _{\theta {o}}(x{t}|\emptyset )\right) \tag{2} \]

含义:通过"减去原始模型在\(c_e\)与空提示下的噪声差",降低模型生成\(c_e\)相关图像的概率。

ESD训练目标函数(最小化调整后的噪声误差):

\[\underset{\theta}{minimize} \ell_{ESD}\left(\theta, c_{e}\right) := \mathbb{E}\left[\left\| \epsilon_{\theta}\left(x_{t} | c_{e}\right) - \left( \epsilon_{\theta_{o}}\left(x_{t} | \emptyset\right) - \eta\left( \epsilon_{\theta_{o}}\left(x_{t} | c_{e}\right) - \epsilon_{\theta_{o}}\left(x_{t} | \emptyset\right)\right) \right) \right\| _{2}^{2}\right] \tag{3} \]

简化:省略期望中的\(t\)和\(\epsilon\),专注于\(\theta\)的优化,确保\(\theta\)满足"远离\(c_e\)"的生成约束。

2.3 对抗性提示攻击模型

对抗性提示通过微小扰动(如 token 替换、嵌入空间扰动)生成\(c'\),使概念擦除后的模型仍生成\(c_e\)相关内容。

设:

  • \(c'\):扰动后的提示,\(\|c' - c\|_0 \leq \epsilon\)(\(\ell_0\)范数约束:扰动token数不超过\(\epsilon\))

对抗性提示生成目标(最小化模型差异):

\[\underset{\left\| c'-c\right\| {0} \leq \epsilon}{minimize} \mathbb{E}\left[\left\| \epsilon{\theta}\left(x_{t} | c'\right) - \epsilon_{\theta_{o}}\left(x_{t} | c\right)\right\| _{2}^{2}\right] \tag{4} \]

含义:使概念擦除模型(\(\theta\))在\(c'\)下的噪声估计,尽可能接近原始模型(\(\theta_o\))在\(c_e\)下的噪声估计,从而"欺骗"模型生成需擦除内容。

2.4 AdvUnlearn框架核心(双层优化)

AdvUnlearn通过"对抗训练(AT)+ 效用保留正则化"解决鲁棒性与效用的平衡问题,采用双层优化(BLO) 结构:

  • 下层优化:生成对抗性提示\(c^*\)(基于式(4))
  • 上层优化:基于\(c^*\)优化模型\(\theta\),同时保留生成质量

2.4.1 效用保留正则化

直接应用AT会导致生成质量下降,因此引入"保留集"\(\mathcal{C}_{retain}\)(含与\(c_e\)无关的良性提示),通过正则化约束模型在良性提示下的生成质量。

设:

  • \(\overline{c} \sim \mathcal{C}_{retain}\):保留集中的良性提示
  • \(\gamma>0\):正则化权重(平衡擦除与效用)

上层优化目标函数(结合ESD损失与效用正则化):

\[\ell_{u}\left(\theta, c^{*}\right) = \ell_{ESD}\left(\theta, c^{*}\right) + \gamma \mathbb{E}{\overline{c} \sim \mathcal{C}{retain }}\left[\left\| \epsilon_{\theta}\left(x_{t} | \overline{c}\right) - \epsilon_{\theta_{o}}\left(x_{t} | \overline{c}\right)\right\| _{2}^{2}\right] \tag{6} \]

分解:

  1. \(\ell_{ESD}(\theta, c^*)\):对抗性提示\(c^*\)下的概念擦除损失,保证鲁棒性;
  2. 正则化项:约束模型在良性提示\(\overline{c}\)下的噪声估计与原始模型尽可能一致,保留生成质量。

2.4.2 快速对抗生成(FGSM)

为提升效率,采用快速梯度符号法(FGSM) 生成对抗性提示,仅需1步迭代。

设:

  • \(\delta\):提示扰动(如前缀向量),\(c' = c + \delta\)("+"表示前缀拼接)
  • \(\delta_0\):扰动初始值,\(\alpha\):步长,\(sign(\cdot)\):元素-wise符号函数

FGSM扰动更新规则

\[\delta = \delta _{0} - \alpha \cdot sign\left( \nabla _{\delta }\ell _{atk}(\theta ,c+\delta _{0})\right) \tag{7} \]

含义:沿攻击损失\(\ell_{atk}\)(式(4)的损失函数)的负梯度方向更新\(\delta\),快速生成具有攻击性的\(c'\)。

三、关键实验验证(数学指标支撑)

3.1 核心评价指标

  • ASR(攻击成功率):越低表示鲁棒性越强(对抗性提示下生成需擦除内容的概率);
  • FID(Fréchet Inception Distance):越低表示生成质量越高(生成图像与真实图像分布的相似度);
  • CLIP得分:越高表示文本-图像对齐性越好(生成图像与提示的匹配度)。

3.2 关键实验结果(以"裸体擦除"为例)

方法 ASR(%) FID CLIP
原始SD v1.4 100 16.70 0.311
ESD(基线) 73.24 18.18 0.309
AT-ESD(无正则) 43.48 26.48 -
AdvUnlearn 21.13 19.34 0.290

数学意义验证

  1. AdvUnlearn的ASR(21.13%)远低于ESD(73.24%),证明对抗训练有效提升鲁棒性;
  2. AdvUnlearn的FID(19.34)接近ESD(18.18),且远低于AT-ESD(26.48),证明效用保留正则化成功平衡鲁棒性与生成质量。

3.3 模块选择验证(文本编码器vs UNet)

AdvUnlearn选择优化文本编码器(而非UNet),原因是文本编码器参数更少、可迁移性强,且对"文本-图像对齐"的控制更直接。实验结果如下:

方法 优化模块 ASR(%) FID
ESD UNet 73.24 18.18
ESD 文本编码器 3.52 59.10
AdvUnlearn 文本编码器 21.13 19.34

数学意义:AdvUnlearn通过效用正则化,解决了"文本编码器优化导致的FID飙升问题",同时保持低ASR(鲁棒性)。

四、总结与贡献

  1. 数学框架创新:提出双层优化的AdvUnlearn,通过"对抗提示生成(下层)+ 效用正则化优化(上层)",首次将对抗训练系统融入扩散模型概念擦除;
  2. 效用-鲁棒性平衡 :通过保留集\(\mathcal{C}_{retain}\)的正则化项(式6),量化平衡"概念擦除强度"与"生成质量";
  3. 模块优化验证:数学实验证明"文本编码器"是更优的鲁棒化模块,且可作为"即插即用"组件迁移到不同扩散模型(如SD v1.5、DreamShaper)。

代码开源地址:https://github.com/OPTML-Group/AdvUnlearn