跳转至

Steering Away from Memorization: Reachability-Constrained Reinforcement Learning for Text-to-Image Diffusion

会议: NeurIPS 2026
arXiv: 2603.00140
OpenReview: https://openreview.net/forum?id=s8BKAZHA8D
代码: https://s-karnik.github.io/rads-memorization-project-page/
领域: 图像生成
关键词: 文生图扩散模型, 样本记忆化缓解, 可达性分析, 约束强化学习, 推理阶段引导

一句话总结

针对文生图扩散模型记忆并复现训练集隐私或版权图像的问题,本文提出推理阶段控制框架 RADS,将去噪轨迹建模为受控动力系统,利用后向可达管表征记忆吸引盆,并通过约束强化学习在低维文本嵌入空间施加微小扰动以主动规避记忆化。

研究背景与动机

基于扩散模型的现代文本到图像生成架构取得了令人瞩目的合成质量,但其背后暴露出的样本记忆化(memorization)风险却愈发严峻。当输入特定提示词时,模型倾向于精确复现训练集中的受版权保护或含隐私信息的原图。现有缓解方案往往通过启发式规则对输入或模型内部进行粗暴干预:部分方法在文本端替换关键词或注入随机噪声标记以破坏检索,部分方案直接修改跨注意力权重掩码或抑制关键神经元,还有方法在监测到过高分类器无条件引导(CFG)模长时直接强行截断引导。这些经验性干预不仅极易破坏最终图像的视觉真实感与细节结构,还经常导致生成内容偏离提示词语义描述(例如丢失提示词中明确指定的色彩与场景修饰)。

这种质量、语义对齐度与记忆消除之间的激烈冲突,根源在于去噪过程具有高度的路径依赖性:扩散模型在极早期的去噪阶段(前 1-2 步)就会掉入由过拟合文本引导诱发的“记忆吸引盆”(attraction basin)。一旦潜变量轨迹跌入该区域,后续无论施加何种常规修正,轨迹都会无可避免地坍缩回训练样本的几何流形中。模型权重“概念消除”(unlearning)虽然能永久抹除特定高层概念,但面对训练集数以万计缺乏明确语义归纳的偶发性实例记忆时既昂贵又容易导致灾难性遗忘。

为了在完全不改变扩散骨干权重的前提下实现无损去噪,本文引入控制理论中的安全保证工具,将反向去噪过程视为受控动力系统。核心 idea:将记忆化吸引盆形式化为后向可达管(Backward Reachable Tube, BRT),并基于拉格朗日松弛的约束强化学习在压缩的文本嵌入隐空间中学习轻量级微扰策略,在去噪极早期前瞻性地将潜变量轨迹导离记忆化危机区域,达成生成多样性、保真度与文本对齐的优越帕累托前沿。

方法详解

整体框架

RADS(Reachability-Aware Diffusion Steering)将预训练扩散模型的逐级反向去噪更新抽象为一个有限步离散受控动力系统。系统状态由当前去噪步的潜变量与其剩余时间步拼接构成,控制动作则作用在文本提示词的特征嵌入上。为了避免高维 CLIP 文本嵌入空间导致强化学习动作探索失效,RADS 先通过预训练的变分自编码器(VAE)将文本特征投影至 64 维紧凑动作隐空间。在去噪反向推进过程中,受约束的 Soft Actor-Critic(SAC)智能体依据当前潜变量和时间步输出隐空间微扰量,经解码器重构出修正后的文本嵌入送入跨注意力层,在前瞻性规避记忆化后向可达管的同时最大化图像与文本提示词的余弦相似度奖励。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["输入提示词与初始噪声<br/>c 及 x_T ~ N(0, I)"] --> B["文本动作空间低维降维<br/>Transformer-VAE 投影至 64 维隐空间"]
    B --> C["受控动力系统去噪演化<br/>状态 s_t = (x_{T-t}, T-t)"]
    C --> D["后向可达管记忆失效边界判定<br/>基于 CFG 向量模长构建目标函数"]
    D --> E["约束强化学习轨迹引导<br/>SAC 双 Critic 与拉格朗日乘子动态平衡"]
    E -->|前瞻性施加嵌入微扰 u_t| C
    E -->|完成 T=50 步安全去噪| F["安全高质量生成输出<br/>规避训练集复制并保持文本语义"]

关键设计

1. 文本动作空间低维降维:构建紧凑平滑的连续控制流形
直接在扩散模型的原始文本嵌入空间执行连续强化学习控制极具挑战性。以标准 CLIP 文本编码器为例,其输出特征维度高达 \(77 \times 768 \approx 5.9 \times 10^4\),如此庞大的动作空间会导致策略梯度的方差爆炸且探索极度低效。更严重的是,在未经约束的连续空间中随意偏移嵌入很容易破坏自然语言原有的语法几何结构,导致解码出的语义发生崩溃。为此,RADS 预先训练了一个紧凑的连续隐式控制流形 \(\mathcal{Z}_{\text{act}} \subseteq \mathbb{R}^{64}\)。该组件由 4 层 Transformer 构成的编码器 \(\texttt{Enc}\) 与 4 层 Transformer 构成的解码器 \(\texttt{Dec}\) 组成。训练采用多目标联合损失: $$ \mathcal{L}{\text{VAE}} = \mathcal{L}}} + 0.1 \cdot \mathcal{L{\text{MSE}} + 2 \times 10^{-3} \cdot \mathcal{L} $$ 其中余弦嵌入损失 }\(\mathcal{L}_{\text{cos}}\) 确保文本语义朝向不被破坏,均方误差 \(\mathcal{L}_{\text{MSE}}\) 维持向量的绝对尺度与结构,KL 散度则迫使潜在动作空间呈现平滑的高斯先验分布。在实际去噪中,策略输出微扰 \(\mathbf{u}_t \in \mathbb{R}^{64}\) 并以残差形式作用于隐空间,最终向扩散模型输入的文本条件为 \(\mathbf{e}_c' = \texttt{Dec}(\texttt{Enc}(\mathbf{e}_c) + \mathbf{u}_t)\)。该设计将控制复杂度削减了近三个数量级,且保证了所有引导方向均落在合法的语义流形之内。

2. 动力系统可达性建模:以后向可达管精确刻画记忆吸引盆
扩散模型的生成过程具有极强的不可逆倾向:前 1-2 步文本引导施加的方向几乎直接决定了后续生成物体的拓扑轮廓与身份特征;一旦潜变量跌入训练集样本的吸引区域,后续无论关闭引导还是修改提示词均无法逆转复制行为。RADS 首次从控制理论的可达性分析(Reachability Analysis)切入,将生成失败集合 \(\mathcal{F}\) 定义为解码后高度重合训练样本的潜变量集合,并将不可避免跌入 \(\mathcal{F}\) 的所有中间态前兆集合定义为后向可达管(Backward Reachable Tube, BRT): $$ \mathcal{B} = {\mathbf{s}0 \in \mathcal{S} \mid \forall \mathbf{u} \in \mathcal{U}, \exists \sigma \in [0, T] \text{ s.t. } \mathbf{s}\sigma \in \mathcal{F}} $$ 为了在去噪过程中无需实际解码图像即可高效判定系统是否滑向 BRT,RADS 深入挖掘了分类器无条件引导(CFG)向量的数学特性。记忆化样本通常呈现异常膨胀的条件与无条件预测残差,表明模型在该提示词处过度拟合。因此,RADS 构造了基于引导模长偏差的瞬时安全目标函数: $$ \ell(\mathbf{s}t) = -\tanh\left(\eta \cdot \left(|\epsilon\theta(\mathbf{x}{T-t}, \mathbf{e}c') - \epsilon\theta(\mathbf{x}, \emptyset)|_2 - \beta\right)\right) $$ 通过严密统计分析,经验阈值设为 \(\beta=9.0\)、缩放因子 \(\eta=0.1\)。当 CFG 模长超过阈值时,\(\ell(\mathbf{s}_t) \le 0\),表明轨迹进入危险区域。其对应的累积安全价值函数遵循极小值 Bellman 演化:\(Q^{\text{safe}}(\mathbf{s}_t, \mathbf{u}_t) = \min\{\ell(\mathbf{s}_{t+1}), \max_{\mathbf{u}_{t+1}} Q^{\text{safe}}(\mathbf{s}_{t+1}, \mathbf{u}_{t+1})\}\),反映了在最优控制对抗下轨迹能否始终置身于记忆集合之外。

3. 约束强化学习轨迹引导:拉格朗日松弛驱动的 SAC 动态平衡
在确立了任务目标(保持提示词图文对齐)与安全边界(规避 BRT 记忆吸引盆)后,去噪过程被形式化为一个约束马尔可夫决策过程(CMDP)。任务奖励仅在去噪终止步 \(t=T\) 给出,定义为生成图与原提示词的 CLIP 余弦相似度 \(r(\mathbf{s}_T) = \textsc{clip}_{\text{im}}(\mathbf{y}) \cdot \textsc{clip}_{\text{text}}(c)\),而在中间步保持稀疏零奖励以杜绝手工塑形对去噪自然演化的干扰。为了求解该 CMDP,RADS 搭建了基于双 Q 任务评价网络 \(Q^{\text{task}}_\omega\)、安全评价网络 \(Q^{\text{safe}}_\psi\) 与随机动作策略网络 \(\pi_\phi\) 的三元架构,并利用增广拉格朗日乘子 \(\lambda\) 将硬性安全约束转化为可导目标: $$ J_\pi(\phi) = \mathbb{E}_{\mathbf{s}t \sim \mathcal{D}, \mathbf{u}t \sim \pi\phi} \left[ Q^{\text{task}}\omega(\mathbf{s}_t, \mathbf{u}t) + \alpha \mathcal{H}(\pi\phi(\cdot|\mathbf{s}t)) + \lambda \cdot Q^{\text{safe}}\psi(\mathbf{s}_t, \mathbf{u}_t) \right] $$ 安全网络参数 \(\psi\) 借助 Polyak 软更新的目标网络以极小化安全递推目标值 \(\hat{Q}_t^{\text{safe}} = (1-\gamma)\ell(\mathbf{s}_t) + \gamma \min(\ell(\mathbf{s}_t), Q^{\text{safe}}_{\psi'}(\mathbf{s}_{t+1}, \mathbf{u}_{t+1}'))\) 进行更新;拉格朗日乘子则依据对偶梯度 \(\nabla_\lambda J = \mathbb{E}[Q^{\text{safe}}_\psi(\mathbf{s}, \mathbf{u})] - \delta\)(阈值 \(\delta=0\))实时自适应调整。一旦安全 Critic 预警轨迹有滑入 BRT 的风险,\(\lambda\) 迅速拉升,强行修正策略在最初几个时间步输出微扰向量,以极微小的感知代价将扩散轨迹“拨出”危险盆地;当轨迹位于安全区时,\(\lambda\) 回落,策略全力服务于图文语义契合。

损失函数 / 训练策略

离线训练阶段,策略与 Critic 参数在包含 430 个已知记忆化提示词的离线回放池上通过交互采样进行迭代优化。优化器均采用 Adam,策略学习率设为 \(10^{-4}\),双任务 Critic、安全 Critic 以及熵温度系数 \(\alpha\) 均设为 \(3 \times 10^{-5}\)。拉格朗日乘子 \(\lambda\) 初始设为 0,对偶更新步长为 \(3 \times 10^{-5}\)。去噪步长固定为 \(T=50\) 步,批次采样 32 条去噪轨迹,训练进行 90 个 epoch(单张 A100 GPU 仅耗时约 15 小时)。每 5 个 epoch 评估一次验证目标,选取终端奖励与安全目标总和 \(r(\mathbf{s}_T) + \ell(\mathbf{s}_T)\) 最高的检查点作为推理阶段固定的即插即用导向器。

实验关键数据

主实验

实验主要在被广泛验证的 Webster (2023) 记忆化提示词数据集上基于 Stable Diffusion v1.4 和 DDIM 采样(步数 \(T=50\), 引导比例 \(g=7.5\))展开全面测评。采用自监督拷贝检测(SSCD)衡量样本复现率(\(\text{SSCD}_{\text{target}}\) 越低代表与训练集原图差异越大;\(\text{SSCD}_{\text{seeds}}\) 越低代表在不同随机噪声下的生成多样性越高;\(\text{SSCD}_{\text{prompts}}\) 越低代表跨提示词无模式坍缩)。同时报告 COCO 验证集上的 FID(越低越好)与 CLIP 语义相似度得分(越高越好)。

方法 \(\text{SSCD}_{\text{target}}^\dagger\) (↓) \(\text{SSCD}_{\text{seeds}}\) (↓) \(\text{SSCD}_{\text{prompts}}\) (↓) FID (↓) CLIP (↑)
无干预基线 (No mitigation) 0.6364 ± 0.2626 0.4910 ± 0.2983 0.0956 ± 0.1031 42.14 ± 5.91 0.3129 ± 0.0279
Wen et al. (ICLR 2024) 0.4187 ± 0.2796 0.2132 ± 0.1798 0.0613 ± 0.0686 31.78 ± 4.87 0.3056 ± 0.0329
Ren et al. (ECCV 2024) 0.3781 ± 0.2420 0.2332 ± 0.1876 0.0695 ± 0.0744 32.76 ± 4.78 0.2940 ± 0.0500
Hintersdorf et al. (NeurIPS 2024) 0.4149 ± 0.2657 0.2534 ± 0.2042 0.0573 ± 0.0584 32.50 ± 4.64 0.3119 ± 0.0296
Jain et al. (CVPR 2025) 0.1816 ± 0.1144 0.0620 ± 0.0608 0.2724 ± 0.1352 63.98 ± 16.17 0.2266 ± 0.0532
RADS (本文方法) 0.2303 ± 0.1110 0.1553 ± 0.1099 0.0409 ± 0.0227 31.57 ± 5.82 0.2917 ± 0.0366

消融与泛化实验

为了验证可达性安全约束的绝对必要性,以及在完全未见过的超大规模提示词集下的零样本外推性能,下表呈现了关键消融与 MemBench(3000 个未见记忆化提示词)的测试结果。

配置 / 评估场景 \(\text{SSCD}_{\text{target}}\) (↓) \(\text{SSCD}_{\text{seeds}}\) (↓) \(\text{SSCD}_{\text{prompts}}\) (↓) FID (↓) CLIP (↑) 说明
RADS (完整方法) 0.2303 ± 0.1110 0.1553 ± 0.1099 0.0409 ± 0.0227 31.57 ± 5.82 0.2917 ± 0.0366 帕累托最优平衡
ablation: 去除可达性约束 (\(\lambda=0\)) 0.4998 ± 0.2336 0.3449 ± 0.2453 0.0786 ± 0.0766 36.35 ± 4.17 0.3010 ± 0.0330 失去安全引导,复制严重反弹
MemBench 零样本: 无干预 0.5030 ± 0.2836 0.4751 ± 0.3080 0.0367 ± 0.0164 22.02 ± 0.89 0.3005 ± 0.0392 基线存在大量记忆复制
MemBench 零样本: Wen et al. 0.2091 ± 0.1494 0.1219 ± 0.0710 0.0412 ± 0.0178 18.78 ± 1.95 0.2998 ± 0.0388 依赖启发式替换,效果中等
MemBench 零样本: Jain et al. 0.1775 ± 0.1226 0.0726 ± 0.0741 0.1782 ± 0.0933 38.56 ± 7.22 0.2520 ± 0.0470 模式坍缩且画质劣变严重
MemBench 零样本: RADS (Ours) 0.1449 ± 0.0439 0.1006 ± 0.0598 0.0516 ± 0.0238 26.75 ± 6.55 0.2524 ± 0.0402 泛化最优复制阻断率且保真

关键发现

  • 安全约束是消除记忆的核心支柱:在消融实验中,若剥离可达性安全约束(\(\lambda=0\) 的无约束 SAC 变体),目标相似度 \(\text{SSCD}_{\text{target}}\) 暴增至 0.4998,几乎退化至无干预基线(0.6364)。这证明单纯依靠 CLIP 语义相似度奖励不仅无法抑制记忆化,反而可能激励模型去检索过拟合的高分训练集样例。
  • 破解伪低复制的“模式坍缩陷阱”:基线 Jain et al. 虽在 \(\text{SSCD}_{\text{seeds}}\) 取得了较低数值,但其 \(\text{SSCD}_{\text{prompts}}\) 异常飙升至 0.2724,且 FID 恶化至 63.98。深入分析表明,该方法频繁关闭引导导致生成过程完全被初始高斯噪声的主成分主导,使不同提示词生成出高度雷同的无意义低质画面。而 RADS 在 \(\text{SSCD}_{\text{prompts}}\) 保持 0.0409 极低水平,做到了真正多姿多彩的去记忆化。
  • 推理耗时极具工程实用价值:在单卡 A100 测试中,无干预基线耗时 2.29 秒/图,RADS 仅为 2.92 秒/图,与启发式文本干预方法耗时持平,远远优于 Hintersdorf et al. 因神经元动态定位导致的 104.09 秒/图严重延迟。

亮点与洞察

  • 跨学科理论嫁接的开创性构想:将控制理论中用于自主系统防撞与避障的安全可达性分析(Reachability Analysis / BRT)引入生成扩散模型的潜变量流形,在概念层面上将“对抗样本记忆化”严谨重塑为“动态系统避障控制问题”。
  • 抓取了记忆化发生的时间关键期:深刻洞察到“前 2 步去噪引导即可锁定记忆化”的本质,放弃笨重且易失败的后半程像素级补救,选择在前几个时间步使用文本隐空间微扰进行“四两拨千斤”的航向校正。
  • 全模型权重无侵入式即插即用:不需要昂贵且会导致灾难性遗忘的 LoRA 微调或 Concept Unlearning 重新训练,只需外挂轻量级策略网络与 VAE,即可无缝兼容 DDIM、DDPM 及 RealisticVision 等多种下游检查点与采样器。

局限与展望

  • 训练数据规模受限导致的特定语义漂移:当前强化学习策略仅在 430 个已知记忆化提示词上训练,当遇到包含罕见多义词的长尾提示词(例如“Mothers influence on her young hippo”)时,微扰可能因动作网络对局部触发词过敏,而误将小河马生成为人类母子,产生轻微语义偏移。未来需在大规模通用无害文本上引入正向正则化。
  • 存在轻量级离线训练前置步骤:相较于完全零样本的启发式硬规则(如即时截断),RADS 仍需约 15 小时的前置强化学习训练流程以学习策略与安全评价器权重。
  • 扩展至多模态安全性泛化:未来可将该框架直接推广到 NSFW 敏感内容阻断、版权 IP 风格防护及生物特征防伪等更广泛的安全可控生成任务中。

相关工作与启发

  • vs. 启发式推理干预(Wen et al., ICLR 2024; Jain et al., CVPR 2025):现有方法依靠粗暴的 token 丢弃或 CFG 截断,不仅破坏全局构图而且导致严重的模式坍缩;RADS 通过连续低维流形上的最优控制策略,在保持图像平滑度与文本对齐的同时精准脱离吸引盆。
  • vs. 神经元/注意力掩码(Ren et al., ECCV 2024; Hintersdorf et al., NeurIPS 2024):依赖神经元定位的方法在不同随机种子下表现出严重的随机性失效(stochastic failure),且计算延迟增加近 50 倍;RADS 建立在可达性 worst-case 保证之上,各随机噪声下稳定性极强且延迟极低。
  • vs. 概念消除微调(Gandikota et al., ICCV 2023; Kumari et al., ICCV 2023):消除微调需要提前定义抽象语义概念并重新训练权重,面对训练集随机散落的海量具象记忆样本力不从心且破坏基础生成能力;RADS 完全冻结骨干网络,实现纯推理端安全防护。

评分

  • 新颖性: ⭐⭐⭐⭐⭐ 首次将控制理论后向可达管(BRT)与约束强化学习引入文生图扩散模型记忆化抑制,数学建模严谨且立意极佳。
  • 实验充分度: ⭐⭐⭐⭐⭐ 涵盖 3 类扩散模型配置、多种采样器、5 组强基线对比、消融实验以及 3000 篇规模的 MemBench 零样本泛化验证。
  • 写作质量: ⭐⭐⭐⭐⭐ 逻辑严密,图文呼应紧凑,从动力系统转换到 CMDP 设计的数学推导与实验论证非常清晰。
  • 价值: ⭐⭐⭐⭐⭐ 为解决 AIGC 版权复现争议、企业级隐私合规与安全生成提供了无需重训且即插即用的前沿通用范式。