跳转至

Tempered Guided Diffusion

会议: NeurIPS 2026
arXiv: 2605.03712
OpenReview: https://openreview.net/forum?id=FdyGr3yTVZ
领域: 其他
关键词: 免训练条件扩散, 序贯蒙特卡洛, 退火后验采样, 图像反问题, 粒子剪枝

一句话总结

本文提出回火引导扩散(TGD),将免训练条件扩散采样建立在干净信号空间的退火序贯蒙特卡洛(SMC)框架上,并通过加速变体 A-TGD 在采样早期利用多粒子交互探索、中途剪枝至单一高似然轨迹,在相同计算预算下显著提升图像反问题的重建质量与采样效率。

研究背景与动机

利用预训练无条件扩散模型作为生成先验来求解逆问题(如图像修复、相位恢复等),具有无需针对新任务重新配对训练的巨大优势。然而,现有的免训练条件扩散采样器在面对高度病态或非线性的复杂反问题时表现极不稳定:从不同的随机初始噪声出发生成的重建轨迹质量差异巨大,很多轨迹会陷入与观测严重不一致或伪影频发的局部次优解。当前主流的应对策略通常采用“多轨迹独立采样后择优(Best-of-N)”的启发式重试,即独立跑完多条完整轨迹,最终依据数据一致性误差挑出最好的一条。

这种独立重试策略暴露出严峻的计算分配失衡矛盾:它为所有候选轨迹均等分配计算资源,即便某些候选轨迹在采样早期就已经明显偏离真实观测、丧失恢复希望,采样器仍会盲目将其完整推演到底。与此同时,如果在单条轨迹内盲目增加函数评估次数(NFE),也无法弥补扩散早期阶段所犯下的结构性决策错误。这表明问题的核心瓶颈不仅在于总计算量的大小,更在于计算量如何在各个候选轨迹之间进行动态、统计合理的调度与分配。

本文的切入视角是将统计物理与贝叶斯计算中的退火序贯蒙特卡洛(SMC)机制与扩散反演解耦融合:不再让各条轨迹孤立演化,而是维护一个动态交互的粒子群,使候选解在演进过程中通过似然比重赋权与重采样进行筛选。核心 idea:在干净信号空间定义显式的退火后验目标序列,将带噪扩散状态仅作为建议与传播的辅助变量,并通过早期多粒子交互探索与中途似然剪枝(A-TGD),将后续算力集中在最有希望的高似然轨迹上。

方法详解

整体框架

TGD 的核心设计在于将推断的目标分布定义在未知干净信号 \(\bm{x}_0\) 的空间,而非带噪隐变量空间。给定干净信号 \(\bm{x}_0 \in \mathbb{R}^{d_x}\)、观测 \(\bm{y} \in \mathbb{R}^{d_y}\) 与似然函数 \(p(\bm{y} \mid \bm{x}_0)\),TGD 建立了一列由外层阶段索引 \(r \in \{R, R-1, \dots, 0\}\) 驱动的退火后验目标: $$ \pi_r(\bm{x}_0) \propto p(\bm{x}_0) p(\bm{y} \mid \bm{x}0)^{\lambda_r}, \quad 0 \le \lambda_R \le \lambda \le \dots \le \lambda_0 = 1 $$ 其中 \(R\) 为外层重启阶段数,\(\lambda_r\) 逐步增强对观测的约束强度。带噪扩散状态 \(\bm{z}_r\) 仅用作辅助变量,配合任意现成的免训练扩散求解器(如 MPGD、DAPS 或 DPS)充当“单阶段重构模块”。在每个阶段,带噪粒子被映射到干净候选信号,依据增量似然比进行粒子加权与重采样,再重新加噪传播至下一阶段;针对高成本的图像反问题,加速版 A-TGD 在探索指定比例阶段后执行极大似然剪枝,仅保留最优单条轨迹演化到终点。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["初始噪声粒子群<br/>z_R ~ p_sR(z_R), w_R = 1/N"] --> B["1. 干净空间退火后验与提升 SMC 构建<br/>辅助带噪状态 z_r 经扩散求解器映射至候选 x0"]
    B --> C["2. 增量似然比重赋权与系统重采样<br/>计算权重增量 p(y|x0)^(λ_{r-1}-λ_r) 并重采样复制优质粒子"]
    C --> D["前向加噪传播<br/>向更低噪声层级 z_{r-1} ~ p_{s_{r-1}}(z_{r-1}|x0) 扩散"]
    D -->|未达到剪枝阈值 ρR| B
    D -->|达到剪枝阶段 R-K_ρ| E["3. 加速剪枝策略 (A-TGD)<br/>按最小测量误差 argmin ||A(x0)-y|| 选取单一最优粒子"]
    E --> F["单轨迹后续扩散收敛<br/>单一候选轨迹完成剩余去噪与观测修正"]
    F --> G["最终高保真重建图像 x0"]

关键设计

1. 干净空间退火后验与提升 SMC 构建:解耦扩散动态与贝叶斯推断目标 以往将 SMC 引入扩散模型的方案大多直接将带噪隐空间状态作为推断目标,使得观测似然的边缘积分项 \(\int p(\bm{y} \mid \bm{x}_0) p(\bm{x}_0 \mid \bm{x}_s) \mathrm{d}\bm{x}_0\) 极度棘手且难以解析计算。TGD 摆脱了这一束缚,将退火目标 \(\pi_r(\bm{x}_0)\) 明确锚定在干净信号空间,带噪状态 \(\bm{z}_r\) 仅作为辅助变量引入提升(Lifted)联合分布 \(\gamma_r(\bm{x}_0, \bm{z}_r) = p(\bm{x}_0) p(\bm{y} \mid \bm{x}_0)^{\lambda_r} p_{s_r}(\bm{z}_r \mid \bm{x}_0)\)。由该联合分布导出的理想重构核为条件分布 \(K_r(\mathrm{d}\bm{x}_0 \mid \bm{z}_r) = p_{\lambda_r, s_r}(\bm{x}_0 \mid \bm{z}_r, \bm{y})\)。这种解耦构建带来了极大的理论与工程优势:不仅使理想条件下的粒子系统在渐近意义下收敛到精确的后验分布,而且允许将任何现有的免训练反问题算法(如 MPGD、DAPS)即插即用地作为单阶段局部重构模块,在外层享受全局粒子推断的红利。

2. 增量似然比重赋权与系统重采样:动态剔除低劣轨迹并聚集有效算力 为了在退火阶段之间平滑过渡,粒子在经历重构步骤得到候选干净样本 \(\bm{x}_0^{(i)}\) 后,依据其对观测的拟合能力进行加权。在扩展空间中推导的标准 SMC 增量权重精确呈现为干净信号上的似然比形式: $$ \tilde{w}_{r-1}^{(i)} = w_r^{(i)} p\left(\bm{y} \mid \bm{x}0^{(i)}\right)^{\lambda $$ 这一比值巧妙地抵消了复杂的扩散转移概率与先验项,仅取决于测量似然与退火指数之差。归一化后的权重直接反映了候选信号在更强约束下的存活价值。通过采用低方差的系统重采样(Systematic Resampling),算法自动舍弃那些偏离观测或陷入模糊模式的退化粒子,并复制处于高似然盆地中的潜力粒子,随后通过前向高斯核 } - \lambda_r\(p_{s_{r-1}}(\bm{z}_{r-1} \mid \bm{x}_0^{(i)})\) 加噪刷新至下一更低噪声层级。这种跨轨迹交互彻底打破了独立 Best-of-N 方案中算力浪费在“必死轨迹”上的缺陷。

3. 加速剪枝策略(A-TGD):平衡早期多模态探索与后期确定性精化 在实际的高分辨率图像反问题中,让完整的粒子群持续运行至最终步骤会导致总体算力成本随粒子数 \(N\) 线性倍增。作者洞察到:反问题的多模态分歧与病态分支主要发生在反向扩散的前半段(高噪声向中噪声过渡期),一旦轨迹确立了正确的几何结构与频域相位轮廓,后半段的局部细节恢复往往趋于单峰收敛。为此,Accelerated TGD(A-TGD)引入剪枝比例 \(\rho \in [0, 1]\),在完成前 \(K_\rho = \lceil \rho R \rceil\) 个退火阶段后,直接计算各个粒子的观测似然,选取与观测测量误差最小的最优粒子: $$ i^\star = \arg\min_{i \in {1, \dots, N}} |\mathcal{A}(\hat{\bm{x}}_0^{(i)}) - \bm{y}|_2^2 $$ 随后将剩余所有粒子丢弃,仅保留该优势粒子的辅助状态,以单一轨迹执行后续阶段的去噪与数据一致性修正。A-TGD 在保留早期多粒子全局探索、跨越局部极值能力的同时,将整体计算开销大幅压缩,实现了优于独立多轨迹采样的精度-耗时帕累托前沿。

损失函数 / 训练策略

TGD 与 A-TGD 属于完全免训练(Training-Free)算法,直接复用在大规模数据集(如 FFHQ、ImageNet 256×256)上预训练的无条件扩散模型,不改动也不微调模型权重。在推理配置上: - 噪声离散化:采用 EDM 连续扩散体系,外层重启阶段遵循曲率参数为 7 的 EDM 噪声网格,最大噪声 \(s_R = 100\),最小正噪声 \(s_0 = 0.1\)。 - 重构求解器:单阶段内部采用 4 步 Euler ODE 求解,图像修复选用 MPGD 模块(平滑参数 \(\gamma = 0.7\)),相位恢复选用带 100 步 Langevin 校正的 DAPS 模块(\(\gamma = 0.4\))。 - 退火与重采样策略:图像修复任务采用线性均匀退火 schedule 并开启每次重采样;高难度相位恢复任务设定 \(\lambda_r \equiv 1\)(保持满似然)以防止过度重采样削弱模式多样性。

实验关键数据

主实验

主实验在统一硬件环境(单张 NVIDIA H200 GPU, 141GB 显存)与严格对齐的实际运行时间(Matched Wall-Clock Runtime)下进行,评测基准涵盖 FFHQ 和 ImageNet 256×256 上的图像修复(Inpainting)与高难度非凸非线性相位恢复(Phase Retrieval),基于 100 张验证集图像计算均值与标准差。

任务 / 数据集 算法配置 每图耗时 (s) PSNR ↑ SSIM ↑ LPIPS ↓
图像修复 (FFHQ) A-TGD (ours, N=4, 剪枝) 6.47 25.14 ± 2.54 0.853 ± 0.025 0.122 ± 0.028
DAPS (N=1, 单轨迹) 6.44 24.93 ± 2.57 0.801 ± 0.028 0.153 ± 0.033
DAPS (N=4, 独立Best-4) 6.50 24.52 ± 2.25 0.790 ± 0.030 0.157 ± 0.034
DPS (基线) 6.34 22.12 ± 2.51 0.806 ± 0.030 0.167 ± 0.033
图像修复 (ImageNet) A-TGD (ours, N=4, 剪枝) 23.64 20.89 ± 3.67 0.781 ± 0.037 0.204 ± 0.036
DAPS (N=1, 单轨迹) 23.52 20.87 ± 3.65 0.753 ± 0.035 0.218 ± 0.044
DAPS (N=4, 独立Best-4) 23.84 20.70 ± 3.42 0.740 ± 0.033 0.227 ± 0.044
DPS (基线) 23.63 18.39 ± 2.57 0.710 ± 0.069 0.281 ± 0.069
相位恢复 (FFHQ) A-TGD (ours, N=4, 剪枝) 4.35 29.43 ± 4.07 0.811 ± 0.104 0.172 ± 0.106
DAPS (N=1, 单轨迹) 4.30 26.98 ± 8.35 0.767 ± 0.193 0.225 ± 0.185
DAPS (N=4, 独立Best-4) 4.35 27.81 ± 5.12 0.765 ± 0.128 0.207 ± 0.111
DPS (基线) 4.34 12.65 ± 2.42 0.331 ± 0.095 0.564 ± 0.063
相位恢复 (ImageNet) A-TGD (ours, N=4, 剪枝) 345 22.18 ± 8.75 0.617 ± 0.276 0.326 ± 0.218
DAPS (N=1, 单轨迹) 357 19.42 ± 9.31 0.499 ± 0.298 0.413 ± 0.222
DAPS (N=4, 独立Best-4) 352 22.41 ± 9.26 0.593 ± 0.291 0.325 ± 0.221
DPS (基线) 346 12.58 ± 2.48 0.232 ± 0.140 0.601 ± 0.064

消融实验

在 FFHQ 上对粒子数 \(N\)、剪枝比例 \(\rho\) 以及退火与重采样组合进行系统性消融。下表汇总了核心设计因子的敏感性验证:

消融维度 具体配置 耗时 (s/img) 图像修复 PSNR / LPIPS 相位恢复 PSNR / LPIPS 核心说明
粒子数量 \(N\) \(N=1\) (无粒子群) 1.29 / 1.30 24.61 / 0.126 22.26 / 0.303 单粒子缺乏多样性,相位恢复崩溃
\(N=2\) 1.92 / 1.94 24.73 / 0.125 24.76 / 0.248 性能显著攀升
\(N=4\) (本文默认) 3.26 / 3.27 25.01 / 0.124 26.38 / 0.206 效果与速度的最佳平衡点
\(N=16\) 11.14 / 11.27 24.92 / 0.124 27.24 / 0.189 边际收益递减,耗时成倍增加
剪枝比例 \(\rho\) \(\rho=0.10\) (极早剪枝) 1.76 / 1.80 24.88 / 0.123 23.25 / 0.292 图像修复对早剪枝不敏感,相位恢复不够充分
\(\rho=0.50\) (本文默认) 3.26 / 3.28 25.01 / 0.124 26.38 / 0.206 兼顾两类任务的最优剪枝点
\(\rho=1.00\) (全流程不剪枝) 5.12 / 5.13 24.93 / 0.124 27.19 / 0.188 质量略微提升但耗时大幅增加
退火与重采样 无退火 (\(\lambda_r \equiv 1\)) 5.05 / 5.05 24.36 / 0.130 27.19 / 0.188 相位恢复在此配置下表现最佳
有退火无重采样 5.06 / 5.06 24.42 / 0.127 26.98 / 0.191 缺少交互,指标提升有限
退火 + 系统重采样 5.05 / 5.07 24.93 / 0.124 24.80 / 0.254 图像修复最佳;相位恢复中重采样会削弱多样性

关键发现

  • 任务病态性决定粒子探索价值:图像修复任务观测约束强、解空间单峰性较高,从 \(N=1\) 增至 \(N=4\) 或采用较早剪枝(\(\rho=0.1\))便能快速饱和;相反,相位恢复存在严重的符号与相位不适定性,多粒子在早期的充分探索是逃离劣质解盆地的决定性因素(\(N=4\) 相比 \(N=1\) 使 PSNR 狂揽超 4.1 dB 增益)。
  • 早期剪枝消除计算冗余:A-TGD 在相同显存与墙上时钟预算下完胜四轨迹独立的 DAPS (N=4)。在 FFHQ 相位恢复中,达到相同感知质量阈值(LPIPS=0.20)时,A-TGD 相比 DAPS (N=4) 带来了近 2 倍的速度加速。
  • 重采样的双刃剑效应:在线退火结合重采样对多模态解空间的均匀覆盖有益,但在观测非凸极强的相位恢复任务中,激进的似然重采样过早将粒子收拢到某些看似残差小但语义错误的局部陷阱,因而在该任务上维持无退火且不重采样(仅在 \(\rho=0.5\) 处按最终误差剪枝)最具鲁棒性。

亮点与洞察

  • 将推断锚定于干净信号空间:传统扩散逆求解器常陷于带噪边缘似然积分不可解的泥潭,TGD 创新地在干净图像空间构造退火目标,带噪状态退化为纯粹的提议与扰动工具,从而以极其优雅的方式兼容现有任意单步或多步扩散求解算子。
  • 动态算力分配颠覆暴力重试:不仅在理论上证明了完整 TGD 渐近收敛于真实后验分布,工程上提炼的 A-TGD 模式更是为生成式反问题提供了“早期粒子广播探索、中途收敛剪枝深耕”的标准算力优化范式。
  • 具备极强的算法即插即用性:TGD 的外层框架与内层重构解耦,无论是使用基于评分预测梯度的 DPS/MPGD 还是基于 Langevin 校正的 DAPS,均能作为内部模块直接装配入该粒子交互外壳。

局限与展望

  • 理论一致性与工程加速的取舍割裂:定理 3.1 证明的后验弱收敛性依赖于“单阶段精确重构”与“粒子全流程不剪枝”两个严格假设;而在实际高维图像任务中,内层使用粗粒度 ODE 近似且 A-TGD 实施了硬性剪枝,使得实际落地版本丧失了严格的后验一致性分布保证。
  • 超参配置依赖任务特性:实验表明,退火方案与重采样开关对反问题的非凸程度高度敏感(图像修复与相位恢复最佳配置完全相反),目前尚缺乏根据前向算子性质自适应调节退火调度与剪枝时机 \(\rho\) 的自动化准则。
  • 显存并行开销:在早期维持多粒子交互需要并发执行多次网络前向计算,若将粒子规模扩大至十数个以上,高分辨率图像生成对 GPU 显存的峰值需求会显著攀升。

相关工作与启发

  • vs DAPS (Decoupled Annealing Posterior Sampling): DAPS 提出了带噪声退火的解耦逆向求解流程,并依赖独立 Best-of-N 挑选最终图像;TGD 证明了 DAPS 本质上是 TGD 在退火指数恒为 1 且省略重采样时的特例,TGD 通过粒子间交互重采样与早期剪枝,显著打破了 DAPS 全程运行多条独立轨迹的算力浪费。
  • vs DPS (Diffusion Posterior Sampling) / MPGD: DPS 等方法主要关注单条扩散轨迹内部如何通过一阶梯度项近似条件得分;TGD 将其降维视为可替换的局部重构模块,并在外层引入 SMC 机制,用全局退火后验引导粒子跳出单轨迹易陷的局部极小。
  • vs TDS / SMCDiff: TDS 将 SMC 粒子权重定义在带噪扩散潜在空间;而 TGD 的粒子权重与退火势函数直接建立在未知干净信号上,避免了带噪边缘似然的复杂估算。

评分

  • 新颖性: ⭐⭐⭐⭐⭐ 首次从干净信号空间建立起扩散反问题的严格退火 SMC 框架,视角独到且理论自洽。
  • 实验充分度: ⭐⭐⭐⭐⭐ 涵盖 2D 可控基准的分布收敛验证,以及 FFHQ/ImageNet 严格对齐耗时的主实验与细致消融。
  • 写作质量: ⭐⭐⭐⭐⭐ 数学推导严谨清晰,附录完备给出了与多种前序算法的严格约简关系。
  • 价值: ⭐⭐⭐⭐⭐ 为免训练扩散逆问题求解提供了极具启发性的粒子交互与算力分配新范式。