Disentangled Sparse Representations for Concept-Separated Diffusion Unlearning¶
会议: NeurIPS 2026
arXiv: 2605.12122
代码: https://github.com/khyeonjin/SAEParate
领域: 图像生成
关键词: 扩散模型遗忘、稀疏自编码器、对比学习、解纠缠表征、概念消除
一句话总结¶
针对传统稀疏自编码器仅以重构为目标导致概念表征纠缠并引发附带遗忘的问题，本文提出 SAEParate，通过交叉注意力引导的空间聚合与隐空间对比学习显式分离概念表征，配合非线性编码器实现高精度、低干扰的扩散模型概念遗忘。
研究背景与动机¶
文本到图像（T2I）扩散模型在合成高保真、多样化视觉内容方面取得了显著成功，但其潜在的生成风险（如生成暴力、色情内容以及侵犯版权和隐私的图像）也引发了广泛关注。为了规避这些风险，机器遗忘（Machine Unlearning）技术致力于在不重新训练整个模型的前提下精确消除特定概念或数据的影响。现有扩散遗忘方法主要分为两类：直接修改模型参数的方法（如 ESD、SalUn 等），以及不改动基座权重、通过抑制目标概念生成特征的即插即用方法。在后者中，基于稀疏自编码器（SAE）的方法（以 SAeUron 为代表）受到极高关注，因为其通过在中间激活层训练辅助稀疏字典，能够在完全冻结扩散底座的前提下，通过抑制与目标概念关联的离散潜在神经元完成轻量化消除，不仅具备高度模块化特性，还拥有极强的可解释性。
然而，传统 SAE 的训练范式存在根本性的目标错位：其目标函数主要建立在稀疏重构误差之上，并不显式约束不同概念在隐空间的几何解纠缠。精确的稀疏重构往往倾向于促使模型学习共享或组合特征——通过让语义相关的概念共享部分活跃神经元，以极少的稀疏激活换取整体重构质量。当对目标概念关联的特征集合进行推理时消融时，这种跨概念特征共享会导致不可避免的连带干扰，直接破坏那些依赖相同潜变量的保留概念。这一矛盾在细粒度的联合“风格-物体”（style-object）遗忘任务中尤为致命，例如在消除“卡通风格的狗”时，SAeUron 会连带损害“卡通风格的花”或“水墨画风格的狗”，导致保留概念生成质量发生灾难性退化。
面对特征重叠与精准遗忘之间的结构性冲突，本文的核心切入点在于重构 SAE 隐空间的几何分布：与其被动依赖无监督重构带来的隐空间纠缠，不如主动将概念身份作为显式监督信号引入表征学习过程。核心 idea：提出 SAEParate 框架，利用扩散过程相邻时间步构建多视角样本，通过交叉注意力引导的空间加权聚合提取概念特征，并在隐空间引入概念感知对比学习约束与非线性编码器，从表征根源上解耦概念重叠并消除特征共享。
方法详解¶
整体框架¶
SAEParate 的整体流程分为离线概念解耦 SAE 训练和推理时特征消融两个阶段。在训练阶段，模型从扩散模型指定 U-Net 块的交叉注意力层提取中间激活，利用文本提示中概念词元的交叉注意力图计算空间相关性权重，对每个样本的逐像素潜变量进行加权池化，得到紧凑的样本级隐表征；接着，结合扩散去噪轨迹构建的相邻时间步多视角样本批次，计算隐空间对比损失（Latent-Space Contrastive Loss, LSC），拉近同概念表征并推开异概念（尤其是共享部分属性的困难负样本）表征；在推理阶段，计算特征重要性分数筛选出目标概念专属的少量 Top 神经元集合，在生成过程中以负乘子对其进行轻量抑制并解码，实现外科手术式概念清除。
%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
A["扩散模型中间激活 x(t)<br/>U-Net 块激活与交叉注意力图"] --> B["非线性 SAE 编码器<br/>两层 MLP + GeLU 非线性变换"]
B --> C["交叉注意力引导空间聚合<br/>Prompt 概念 Token 权重软池化"]
C --> D["隐空间对比学习 LSC<br/>相邻时间步多视角 + 困难负样本重权"]
D --> E["联合优化目标<br/>稀疏重构损失 + 辅助损失 + LSC 正则"]
E --> F["推理时特征评分与消融<br/>定位目标概念专有神经元并负乘子抑制"]关键设计¶
1. 交叉注意力引导的空间加权聚合：过滤背景噪声以提炼纯净概念表征 传统特征池化往往对整张特征图所有空间位置进行均匀平均，但图像中大量背景区域或共存的不相关实体与目标概念毫无关联，均匀池化会引入严重的背景噪声并稀释概念独有特征。针对这一痛点，SAEParate 直接复用扩散模型 U-Net 交叉注意力层本身提供的空间定位信号。对于样本 \(i\) 的每个空间位置 \(p\)，模型提取所有注意力头在目标概念提示词元集合 \(\mathcal{C}_i\) 上的交叉注意力概率并做归一化，得到软池化权重：
通过计算加权和 \(r_i=\sum_{p=1}^S \alpha_{i,p} z_{i,p}\)，模型能够自适应抑制背景区域潜变量的干扰，仅将概念高度相关的视觉激活送入后续对比学习，显著提升对比信号的纯度。
2. 隐空间对比学习与困难负样本重权：显式塑形概念解纠缠的潜变量簇 为了打破传统 SAE 仅依赖重构所导致的跨概念特征复用，SAEParate 将监督对比学习迁移至 SAE 潜在空间。针对扩散模型天然的时序结构，模型将相同概念在相邻去噪步 \((t_{2m}, t_{2m+1})\) 的激活配对作为正样本视角，避免了额外图像数据增强带来的伪影。在联合风格-物体任务中，语义部分重叠的样本（如目标为“卡通狗”，而负样本为“水墨狗”或“卡通花”）处于表征模糊带，极易共用神经元。SAEParate 显式定义困难负样本集合 \(\mathcal{H}(i)=\{j\neq i \mid (y_j^{cls}=y_i^{cls}) \oplus (y_j^{thm}=y_i^{thm})\}\)，并在对比 logit 中引入重权因子 \(\omega_{\mathrm{HN}}\)：
通过增大困难负样本在配分函数中的权重，模型强行推开仅共享物体或风格的部分重叠概念，迫使 SAE 为复合概念分配真正解耦的专用潜变量。
3. 非线性增强型 SAE 编码器：提供隐空间重构与簇重排所需的表征容量 传统 TopK SAE 编码器仅包含单个仿射变换后接激活函数，网络结构极浅。当同时施加重构约束与严格的概念排斥对比损失时，单层线性投影容量受限，难以同时兼顾稀疏重构保真度与高度非线性的拓扑流形折叠。SAEParate 在编码器中引入额外的隐藏层和 GeLU 激活函数：
利用 GeLU 光滑且局部非单调的几何映射能力，编码器获得了对潜变量局部区域进行精细折叠与聚类的表征弹性，而解码器仍保持单层线性映射，既保障了推理时高效的离散神经元消融机制，又大幅提升了隐空间的类间分离裕度。
4. 推理时高特异性神经元消融：零参数改动的纯净概念抹除 在离线训练完成后，SAEParate 依据特征重要性评分在目标集 \(\mathcal{D}_c\) 与非目标集 \(\mathcal{D}_{\neg c}\) 之间的平均激活差值，筛选出对目标概念响应极高而在其他概念中几乎静默的 Top-\(\tau_c\) 个专有神经元集合 \(\mathcal{F}_c\)。在扩散生成前向过程中，当特征值超过基线阈值时，以负乘子 \(\gamma_c\) 进行乘法抑制：
由于前端训练阶段实现了严格的概念解耦，选出的神经元重叠数降至为零，消融操作得以实现“只杀目标、不伤无辜”的精确剪除。
损失函数 / 训练策略¶
SAEParate 采用联合端到端优化目标： $$ \mathcal{L}{\mathrm{total}} = \mathcal{L} $$ 其中 }} + \lambda \mathcal{L}_{\mathrm{LSC}\(\mathcal{L}_{\mathrm{SAE}} = \|x - \hat{x}\|_2^2 + \alpha \mathcal{L}_{\mathrm{aux}}\) 为标准 BatchTopK SAE 重构损失与死神经元缓解辅助损失，系数 \(\alpha=1/32\)。隐空间对比损失权重 \(\lambda\) 在物体任务中设为 0.1，风格任务中设为 0.01，联合风格-物体任务中设为 1.0（搭配困难负样本加权 \(\omega_{\mathrm{HN}}=4\)）。优化器选用 Adam，学习率设为 \(4\times 10^{-4}\)，批大小为 131,072，每步对解码器权重进行单位范数归一化。
实验关键数据¶
主实验¶
在 UnlearnCanvas 基准（包含 20 个物体类别与 50 种艺术风格）上，基于微调版 Stable Diffusion v1.5 进行全面评测。评估指标包括遗忘准确率（UA）、域内保留准确率（IRA）、跨域保留准确率（CRA）、联合任务中的风格一致性（SC）、物体一致性（OC）和非相关提示质量（UP）。
| 方法 | 单属性: 风格 UA (↑) | 单属性: 物体 UA (↑) | 物体 IRA (↑) | 联合任务: UA (↑) | 联合任务: SC (↑) | 联合任务: OC (↑) | 联合任务: UP (↑) | 联合 Avg. (↑) | 总平均 Avg. (↑) |
|---|---|---|---|---|---|---|---|---|---|
| ESD | 98.58% | 92.15% | 55.78% | 91.42% | 4.88% | 14.72% | 84.38% | 48.85% | 66.11% |
| SalUn | 86.26% | 86.91% | 96.35% | 42.21% | 62.45% | 70.93% | 87.28% | 65.72% | 81.75% |
| SAeUron | 95.80% | 78.82% | 95.47% | 63.00% | 16.98% | 75.84% | 71.90% | 56.93% | 79.19% |
| SAEmnesia | 96.60% | 94.65% | 91.39% | 81.00% | 59.58% | 11.18% | 54.60% | 51.59% | 77.55% |
| SAEParate (本文) | 99.60% | 95.20% | 95.91% | 94.11% | 93.49% | 91.66% | 98.82% | 94.52% | 96.06% |
消融实验¶
在子集训练设置下对物体遗忘任务的核心模块进行系统消融，同时测量聚类质量指标（聚类率 Cluster Rate 与质心边距 Centroid Margin）：
| 配置 | 物体 UA (↑) | 物体 IRA (↑) | 物体 CRA (↑) | 物体平均 (↑) | 聚类率 (↑) | 质心边距 (↓) | 说明 |
|---|---|---|---|---|---|---|---|
| Vanilla SAE (SAeUron) | 58.14% | 96.74% | 97.45% | 84.11% | 0.93 | 0.73 | 基线重构 SAE，隐空间高度纠缠 |
| 仅加入 LSC 对比损失 | 83.82% | 95.23% | 91.55% | 90.20% | 0.98 | 0.71 | 显式分离提升遗忘，但受限于线性容量 |
| 仅加入 GeLU 非线性编码器 | 47.55% | 96.69% | 97.53% | 80.59% | 0.94 | 0.64 | 无对比引导，单纯增加容量无法自发分离 |
| SAEParate (完整模型) | 95.59% | 96.48% | 93.46% | 95.18% | 1.00 | 0.06 | 两者协同使质心边距骤降，达到最优权衡 |
关键发现¶
- 联合遗忘性能实现断层式领先：在最具挑战性的联合“风格-物体”遗忘任务中，SAEParate 将平均综合得分从既往最优 baseline（SalUn 65.72%，SAeUron 56.93%）大幅推升至 94.52%，其中风格一致性（SC）从 SAeUron 的 16.98% 飙升至 93.49%，彻底根除了因属性重叠引发的附带破坏。
- 特征重叠被彻底清零：成对类别特征重叠分析表明，传统 SAE 在不同类别间存在密集的神经元交叉复用，而 SAEParate 将选定消融特征集合之间的成对重叠计数降低为零，特征组间的最大余弦相似度也显著收缩。
- 重构微损换取极致解耦：重构-分离权衡测试显示，未解释方差分数（FVU）从 0.172 微幅上升至 0.217，但质心边距从 0.73 骤降至 0.06，CLIP 分数反而稳中有升（29.32 vs 29.15），证明在遗忘任务中一味追求极致像素级重构并无必要，结构解耦才是核心。
- 多请求连续遗忘极具韧性：在连续 5 个物体的顺序遗忘测试中，随着遗忘目标增加，SAeUron 的保留率发生雪崩式崩溃（从 97.91% 跌至 4.96%），而 SAEParate 稳健维持在 87.67%，展现出极佳的工程实用性。
亮点与洞察¶
- 扩散时间步天然构建多视角视图：无需引入随机裁剪、颜色抖动等图像级强增强（容易破坏扩散特征流），巧妙利用去噪轨迹中相邻时间步的特征连续性作为天然正样本对，既顺应生成动力学，又保证了对比学习的稳定性。
- 跨注意力引导的软掩码池化：不依赖额外的外部目标检测或分割模型，直接在 U-Net 内部 hook 出文本-图像交叉注意力图，以纯内部信号完成了概念相关像素的高信噪比加权。
- 非线性注入打破浅层容量瓶颈：敏锐指出单层线性编码器在面对严苛的对比几何推斥时存在容量缺陷，通过仅在编码器引入单层 GeLU 隐层、保持解码器纯线性，既放宽了表征流形折叠自由度，又保留了解码侧透明可解释的神经元操控属性。
局限与展望¶
- 依赖概念标签与分组监督：当前对比损失需要概念分组或成对二元标注（Forget vs Retain），在完全开放域或无监督未知概念消除场景下无法直接套用。
- 继承 SAE 固有机制局限：模型仍受制于稀疏自编码器普遍存在的特征分裂（Feature Splitting）和特征吸收（Feature Absorption）缺陷，部分复杂组合概念可能无法完全被单组离散神经元捕获。
- 仅在图像模态验证：目前工作集中在 2D 图像扩散模型（U-Net 及初步探索的 MM-DiT），尚未推广至视频生成扩散模型与多模态联合流。
相关工作与启发¶
- vs SAeUron: SAeUron 首次将 SAE 引入扩散遗忘，但其沿用无监督稀疏重构训练，导致概念在隐空间严重纠缠，抑制神经元时极易伤及非目标概念；SAEParate 引入对比学习与空间注意力聚合，主动雕琢概念解纠缠簇，消除特征重叠。
- vs SAEmnesia: SAEmnesia 采用严格的单概念对应单神经元（one-to-one）硬分配设计，在多属性复合概念场景下表达受限；SAEParate 保持分布式特征表达，通过隐空间几何分离支持更灵活的组合概念解耦。
- vs 权重微调方法 (ESD / SalUn): 参数修改方法不可逆且多次遗忘容易造成累积灾难性遗忘；SAEParate 保持基座完全冻结，仅需约 200MB 的外挂模块即可即插即用，推理耗时增加仅 2%，且在连续遗忘场景下展现出极高鲁棒性。
评分¶
- 新颖性: ⭐⭐⭐⭐⭐ 首次揭示并解决了 SAE 扩散遗忘中隐空间特征纠缠的核心瓶颈，构建了从时序视角到注意力池化、空间对比的自洽解耦体系。
- 实验充分度: ⭐⭐⭐⭐⭐ 涵盖单属性、联合组合、敏感内容（I2P 裸露分类消歧）、顺序遗忘、对抗还原攻击、Matryoshka 扩展及 MM-DiT 跨架构验证，证据极具说服力。
- 写作质量: ⭐⭐⭐⭐⭐ 逻辑线索清晰严密，从问题发现、理论建模、几何度量到可视化全方位自闭环。
- 价值: ⭐⭐⭐⭐⭐ 为大模型安全对齐与模块化可解释可信生成提供了极具启发性的技术范式，在兼顾零权改变与外科手术式遗忘上树立了新标杆。