跳转至

Robust Diffusion Models via Divergence-Induced Weighted Denoising

会议: NeurIPS 2026
arXiv: 2606.22521
OpenReview: https://openreview.net/forum?id=RdBtgbn9Tb
论文: https://arxiv.org/abs/2606.22521
领域: 图像生成
关键词: 扩散模型, 鲁棒训练, f-散度, 局部似然比, M-估计

一句话总结

本文证明了在高斯逆向核下扩散模型的单步似然比服从由标量失配参数化的对数正态分布,从而将标准 MSE 损失推广为由 f-散度诱导的非线性加权去噪目标,在几乎不增加计算开销的前提下显著提升了数据污染下的生成质量与训练稳定性。

研究背景与动机

扩散概率模型(DDPM 及分数量基础生成模型)已在图像、音频与跨模态生成任务中确立了前沿地位。然而,其标准训练目标本质上是对去噪残差的加权均方误差(weighted MSE),在整个训练集中对每一个样本施加无差别的二次惩罚。当训练数据不可避免地混入低质量图像、高频噪声、模糊或恶意投毒样本时,二次损失对大残差异常点极度敏感,迫使网络容量被过度分配给拟合极端离群样本,进而严重破坏正常流形上的生成质量。

这种脆弱性的核心根源在于扩散训练目标的推导假设:标准的证据下界(ELBO)在数学上严格等价于路径空间上的 KL 散度极小化。在 KL 散度下,损失对单步残差的导数恒为常数 1,没有任何抑制异常梯度贡献的机制。而在统计学 M-估计理论中,获得抗污染鲁棒性的基本条件是得分函数 \(\psi\) 具有有界性乃至再下降(redescending)特性。直接在整条生成轨迹的数据空间边缘分布上应用其他具有有界影响函数的 f-散度在计算上是极其棘手的,因为边际分布 \(f_\theta(x_0)\) 仅能通过高维逆向马尔可夫链隐式定义,无法直接计算或分解。

本文的切入角度是直接剖析高斯扩散过程在单步转移上的精细概率结构,发现无需引入复杂的判别器(critic)或对抗训练,即可在局部转移核层面建立通用的 f-散度投影。核心 idea:利用 DDPM 逆向高斯核下单步似然比严格服从单参数对数正态分布的性质,将整条扩散轨迹的散度极小化松弛为单步条件 f-散度之和,构造出只需替换一行代码的散度诱导加权去噪目标,通过影响函数衰减在残差空间实现稳健 M-估计。

方法详解

整体框架

本文构建了从全局路径空间散度上界到局部条件 f-散度加权去噪的严谨理论链条。首先,通过数据处理不等式证明联合路径空间上的完整散度 \(D_G(Q \| P_\theta)\) 构成了观测边际散度 \(D_G(g \| f_\theta)\) 的切上界;随后针对非 KL 散度中全局路径似然比端点项难以分解的问题,构建了局部单步条件散度求和目标 \(\mathcal{L}_G(\theta) = \sum_{t=1}^T \mathbb{E}_{q(x_0, x_t)} [D_G(q(\cdot \mid x_t, x_0) \| p_{\theta, t}(\cdot \mid x_t))]\)。关键在于,利用高斯逆向核协方差匹配的性质,证明单步似然比对数服从均值为 \(-m_t\)、方差为 \(2m_t\) 的正态分布,从而将每一步高维转移核之间的 f-散度坍缩为一个仅依赖于标量去噪残差失配 \(m_t\) 的一维解析或拟合变换 \(h_G(m_t)\)。在反向传播中,\(h_G'(m_t)\) 天然充当了自适应样本影响权重。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["加噪样本输入<br/>(x0, t, ϵ)"] --> B["前向扩散与残差计算<br/>xt = √α¯t x0 + √(1-α¯t) ϵ"]
    B --> C["高斯逆向核对数正态似然比<br/>log rt ~ N(-mt, 2mt)"]
    C --> D["局部条件 f-散度诱导目标<br/>LG(θ) = ∑t E[hG(mt)]"]
    D --> E["M-估计影响权重自适应加权<br/>梯度加权: hG'(mt) wt ∇θ ||ϵ - ϵθ||²"]
    E --> F["稳健去噪网络参数更新<br/>抑制大失配异常样本梯度"]

关键设计

1. 高斯逆向核对数正态似然比:将高维核散度坍缩为标量失配函数

现有扩散模型在计算一步转移的 KL 散度时,仅关注两个高斯分布均值差的加权二次型,掩盖了似然比本身的精确分布形态。本文指出,在 DDPM 的真实后验核 \(q_t(x_{t-1} \mid x_t, x_0) = \mathcal{N}(\tilde{\mu}_t, \Sigma_t)\) 与模型参数化逆向核 \(p_{\theta, t}(x_{t-1} \mid x_t) = \mathcal{N}(\mu_{\theta, t}, \Sigma_t)\) 协方差匹配的前提下,似然比随隐变量的涨落具有极为纯粹的代数结构。令单步马氏失配标量为 \(m_t := \frac{1}{2}(\tilde{\mu}_t - \mu_{\theta, t})^\top \Sigma_t^{-1} (\tilde{\mu}_t - \mu_{\theta, t})\)。在模型逆向核测度下,单步似然比 \(r_t := q_t / p_{\theta, t}\) 的对数严格服从对数正态分布: $\(\log r_t \sim \mathcal{N}(-m_t, 2m_t)\)$ 由此,\(r_t\) 是满足 \(\mathbb{E}[r_t \mid x_0, x_t] = 1\) 且方差为 \(\mathrm{Var}(r_t \mid x_0, x_t) = e^{2m_t} - 1\) 的一维对数正态随机变量。这一发现打破了以往必须借助对抗生成网络或 Fenchel 对偶判别器估计复杂概率比的壁垒,意味着任意由凸生成函数 \(G(\delta)\) 定义的条件 f-散度 \(D_G(q_t \| p_{\theta, t})\),都能在条件期望下完全表征为关于单步失配 \(m_t\) 的一维确定性变换 \(h_G(m_t) = \mathbb{E}[G(U_{m_t} - 1)]\),其中 \(\log U_m \sim \mathcal{N}(-m, 2m)\)。

2. 局部条件 f-散度诱导目标:基于路径上界的拟合代理与渐近乘积求和约化

在边际数据分布上直接优化 \(D_G(g \| f_\theta)\) 因 \(f_\theta\) 不可解析而不可行。定理 2.1 基于数据处理不等式证明了联合轨迹空间上的完整散度 \(D_G(Q \| P_\theta) \ge D_G(g \| f_\theta)\) 构成严格的切上界,当且仅当模型逆向路径测度与真实前验路径测度完全吻合时取等。然而对于非 KL 的非线性凸函数 \(G\),路径似然比包含了未对齐的端点项 \(g(x_0) q(x_T \mid x_0) / p(x_T)\),无法直接因子分解。为此,本文直接构造了局部单步条件散度之和目标: $\(\mathcal{L}_G(\theta) = \sum_{t=1}^T \mathbb{E}_{q(x_0, x_t)} [h_G(m_t(x_0, x_t))]\)$ 定理 2.2 和推论 2.1 提供了这一局部构造的渐近合理性论证:当扩散步数充分细密、单步失配均匀趋于 0(\(\max_t m_t \to 0\) 且 \(\sum_t \mathbb{E}[m_t^2] = o(\sum_t \mathbb{E}[m_t])\)),且满足类似 Lindeberg 条件的高阶矩控制与似然比增量非对角不相关时,整体乘积似然比的散度与单步散度之和的误差为更高阶无穷小。定理 2.3 和 2.4 更进一步证明了局部变换的普适性质:对任意满足正规化条件 \(G'(0)=0, G''(0)=1\) 的平滑 f-散度,\(h_G(m) = m + \mathcal{O}(m^{3/2})\) 且在 \([0, \infty)\) 上严格单调递增。这意味着在模型充分收敛或干净数据的小残差区域(\(m \to 0\)),所有 f-散度诱导的目标均自动退化回标准 DDPM 的加权 MSE,而在大失配区域才显现出非线性分流。

3. M-估计影响权重自适应加权:指数衰减实现再下降鲁棒性

对参数 \(\theta\) 求梯度可知,诱导目标的一阶更新方向为: $\(\nabla_\theta \mathcal{L}_G(\theta) = \sum_{t=1}^T h_G'(m_t) \, w_t \, \nabla_\theta \|\epsilon - \epsilon_\theta(x_t, t)\|^2\)$ 其中 \(w_t = \frac{\beta_t^2}{2\alpha_t(1-\bar{\alpha}_t)\sigma_t^2}\) 是仅与噪声调度相关的标准扩散标度。导数项 \(h_G'(m_t)\) 严格对应于稳健 M-估计中的影响权重函数。对于 KL 散度,\(h_{\mathrm{KL}}(m) = m \implies h_{\mathrm{KL}}'(m) = 1\),所有样本不论偏离多远均施加全额梯度权重,导致残差空间得分 \(\psi(r) \propto r\) 无界;而对于平方 Hellinger 距离(HD),解析解为 \(h_{\mathrm{HD}}(m) = 4(1 - e^{-m/4})\),其影响权重呈现清晰的单调指数衰减: $\(h_{\mathrm{HD}}'(m) = e^{-m/4}\)$ 对于负指数散度(NED, \(G(\delta) = e^{-\delta} - 1 + \delta\)),尽管其 \(h_{\mathrm{NED}}(m) = e\,\mathbb{E}[e^{-U_m}] - 1\) 无初等闭式解,命题 2.1 借助 Girsanov 测度变换与 Laplace 渐近展开证明其大失配衰减率为: $\(h_{\mathrm{NED}}'(m) \sim \frac{e}{4\sqrt{m}} e^{-m/4}\)$ NED 在保持与 HD 相同的主导指数衰减速度的同时,带有一个额外的 \(m^{-1/2}\) 多项式收缩项,因而对极端离群值产生更为决绝的压制。在工程实现上,由于 \(h_{\mathrm{NED}}'\) 仅是一维积分,利用 64 节点甚至 16 节点的 Gauss-Legendre 正交数值积分即可在 GPU 上向量化瞬间完成,使训练循环只需替换一步损失计算代码,额外开销低于 1%。

损失函数 / 训练策略

在训练阶段,算法保持前向加噪过程不变:采样 \(x_0 \sim \mathcal{D}, t \sim \mathrm{Uniform}\{1, \dots, T\}, \epsilon \sim \mathcal{N}(0, I)\) 生成 \(x_t = \sqrt{\bar{\alpha}_t} x_0 + \sqrt{1-\bar{\alpha}_t}\epsilon\)。计算预测噪声残差并乘以扩散权重得到标量失配 \(m_t = w_t \|\epsilon - \epsilon_\theta(x_t, t)\|^2\)。随后直接以 \(h_G(m_t)\) 作为标量损失进行反向传播(在自动微分引擎中,其梯度反传自动乘以 \(h_G'(m_t)\))。基准实验采用 35M 参数的 DDPM U-Net(基础通道 128,通道倍率 1,2,2,2,在 \(16\times 16\) 分辨率引入自注意力机制),1000 步线性加噪调度,Adam 优化器学习率 \(2\times 10^{-4}\),批大小 128,bfloat16 混合精度训练 100K 步。

实验关键数据

主实验

实验在无条件 CIFAR-10 上系统评估了不同数据污染比例 \(\varepsilon \in \{0.0, 0.1, 0.2, 0.3\}\) 下各散度诱导目标的生成质量。污染图像由 4 类均匀混合破坏生成:\(8\times 8\) 局部 Patch 噪声、高斯模糊(\(\sigma=1.5\))、椒盐噪声(\(p=0.05\))以及对比度偏移(\(\pm 0.4\))。测试采用 50 步 DDIM 采样器生成 5K 样本计算对比干净测试集的 FID,并跨 3 个随机种子(42, 123, 456)统计均值与标准差。

污染比例 \(\varepsilon\) KL(标准 DDPM MSE) HD(平方 Hellinger) NED(负指数散度) 相对 KL 改善幅度
0.00(干净基准) 35.4 ± 12.3 31.8 ± 1.7 32.4 ± 2.4 FID 相当,方差缩减 80%
0.10(轻度污染) 48.8 ± 11.0 40.5 ± 5.7 39.8 ± 3.8 FID 下降 9.0,稳定性倍增
0.20(中度污染) 67.7 ± 32.6 66.5 ± 27.9 62.2 ± 18.3 FID 下降 5.5,方差大幅收敛
0.30(重度污染) 93.0 ± 29.6 91.9 ± 33.4 77.5 ± 19.7 FID 下降 15.5,显著抗崩塌

消融实验:与传统鲁棒损失函数对比及正交积分阶数分析

在单种子(seed 42)受控环境下,将散度诱导目标与鲁棒统计中经典的 Huber 损失(截断阈值 \(\delta=1.0\))和 Clipped MSE(硬截断阈值 \(\tau=10\))进行横向对比;同时测试了 NED 数值积分节点数对 CIFAR-10 生成性能的影响。

污染比例 \(\varepsilon\) Huber (\(\delta=1.0\)) Clipped MSE (\(\tau=10\)) HD(本文) NED(本文)
0.00 27.8 27.7 32.4 31.6
0.10 53.6 54.0 43.3 42.3
0.20 87.2 67.5 45.2 49.8
0.30 94.0 90.5 56.2 55.1
NED 正交积分节点数 干净 CIFAR-10 FID 计算开销 / 延迟增加 稳定性表现
16 节点 25.7 < 0.3% 极低分辨率下表现优异
32 节点 30.0 < 0.5% 波动处于随机种子噪声内
64 节点(默认) 29.7 < 1.0% 全局动态数值极其稳健
128 节点 29.7 ~ 1.5% 与 64 节点完全饱和一致

关键发现

  • 大残差抑制带来了决定性的污染隔离能力:在 \(\varepsilon=0.30\) 的严重污染下,标准 KL/MSE 的 FID 退化至 93.0,生成的图像充斥着明显的马赛克噪声伪影;而 NED 维持在 77.5,在单 seed 实验下更达到 55.1(优于 Huber 的 94.0 和 Clipped MSE 的 90.5,提升达 35%~40%)。这证明简单的截断或线性过渡无法等同于由对数正态似然比严格推导出的平滑指数降权。
  • 干净数据下的训练方差显著压缩:在干净数据集上(\(\varepsilon=0.00\)),KL 的跨种子方差高达 12.3,而 HD 和 NED 仅为 1.7 和 2.4。有界影响函数避免了训练初期偶尔出现的极端梯度对参数更新方向的颠簸冲击,起到了内在的训练稳定器作用。
  • CIFAR-100 上表现出一致的鲁棒泛化性:在类别更分散的 CIFAR-100 上,\(\varepsilon=0.30\) 污染时 KL 的 FID 恶化为 98.0,而 NED 保持在 75.3(降低了 22.7 点),验证了理论在复杂分布下的有效性。

亮点与洞察

  • 似然比的对数正态分布表征是连接扩散过程与统计发散的桥梁:Lemma 2.1 首次在扩散模型理论中证明,协方差匹配的高斯去噪单步似然比的对数精确服从 \(\mathcal{N}(-m_t, 2m_t)\)。这使得不需要任何神经网络判别器来估计概率密度比,便可精确解析计算任意 f-散度的期望。
  • 一行代码的非侵入式工程落地方案:该方法将抽象复杂的生成流形稳健估计问题,化简为训练循环中仅需对均方误差施加 \(h_G(m)\) 非线性变换,前向反向完全兼容标准流水线,兼具极强的理论自洽性与极高的工业可用性。
  • 向连续流匹配(Flow Matching)的通用拓展潜力:由于速度场残差 \(\|v_\theta - v^*\|^2\) 在高斯流中具备对称结构,该散度诱导加权机制可天然平移至连续归一化流的训练中,构建更广义的抗噪生成模型。

局限与展望

  • 局部条件目标与边际路径散度之间缺乏严格的双向紧致界:局部条件散度之和是基于切上界与小误差渐近展开所构造的训练代理,对于有限失配(finite mismatch)情况,局部代理与真实边际散度 \(D_G(g \| f_\theta)\) 之间的理论差值界限尚未严格闭合。
  • 强有界权重在训练极早期的信号衰减权衡:指数衰减权重 \(e^{-m/4}\) 虽然能压制脏数据,但在模型刚初始化的极早期(此时即使干净样本的失配 \(m_t\) 也较大),可能会削弱有益的学习梯度,需要结合自适应温度或退火调度进一步调优。
  • 高斯逆向核假设的束缚:全文推导严格建立在分析型逆向高斯核与协方差精确匹配(matched covariance)的前提下,对于非高斯逆向过程或可学习各向异性方差的扩散模型,对数正态分布特性将被破坏。

相关工作与启发

  • vs DDPM (Ho et al., 2020) & 似然最大化 (Song et al., 2021b):Ho 等人将 ELBO 展开为单步均值匹配的 MSE 损失,Song 等人证明了其对应加权 KL 散度。本文则指出 KL 仅是 f-散度家族中影响权重为常数 1 的特例,通过引入具有指数衰减影响函数的 HD 与 NED,克服了原始扩散模型对坏样本极端敏感的缺陷。
  • vs f-GAN (Nowozin et al., 2016):f-GAN 必须引入额外的参数化判别网络通过 Fenchel 共轭对偶求解变分下界,训练极易出现鞍点不稳定性;本文利用逆向高斯结构的解析分布特性,彻底省去了判别器网络,实现了确定性的闭式/数值评估。
  • vs 传统鲁棒统计损失 (Huber, 1964; Clipped MSE):传统方法是基于启发式设计的几何截断,缺乏针对扩散逆向生成动力学的概率解释;本文的影响函数是严谨的统计发散测度在路径空间对数正态增量下的导出解,在抗重度污染上全面超越了传统启发式截断。

评分

  • 新颖性: ⭐⭐⭐⭐⭐ 首次揭示扩散逆向核单步似然比的对数正态特性,建立了完备的局部 f-散度加权去噪框架。
  • 实验充分度: ⭐⭐⭐⭐ CIFAR-10/100 上的多污染模式与多随机种子验证扎实,对比了传统鲁棒基准,但缺乏大规模高分辨率数据集(如 ImageNet)的大模型实测。
  • 写作质量: ⭐⭐⭐⭐⭐ 理论推导清晰优美,从测度论投影、中心极限定理型渐近展式到 Laplace 逼近一气呵成。
  • 价值: ⭐⭐⭐⭐⭐ 兼具优雅的统计理论深度与极其轻量的工业实用性,仅需修改一行损失计算代码即可提升生成鲁棒性。