跳转至

Principled Design of Diffusion-based Optimizers for Inverse Problems

会议: NeurIPS 2026
arXiv: 2605.11506
OpenReview: https://openreview.net/forum?id=BJ9BmUd5tV
代码: 待开源(论文承诺发表后公开)
领域: 图像恢复
关键词: 扩散模型、逆问题求解、超参数重参数化、谱自回归、凸优化加速

一句话总结

本文提出面向扩散模型逆问题求解的 OptDiff 优化器框架,通过理论推导起始/终止噪声水平的容差重参数化与双端梯度归一化,消除了跨任务手动调参的痛点,并成功引入动量与预条件技术实现数倍推断加速。

研究背景与动机

基于分数的扩散模型凭借其强大的表达能力和丰富图像先验,在图像去模糊、超分辨率与欠采样 MRI 重建等各类逆问题中取得了 SOTA 级别的重构表现。传统方法往往直接利用预训练的无条件扩散骨干网络,通过贝叶斯准则引入测量似然或将采样过程重构为复合优化问题(如 RED-diff),从而避免针对特定退化任务从头微调网络权重。然而,这类方法在实际临床或工业部署中仍面临两大核心瓶颈:其一是采样迭代收敛极慢,通常需要数十至上百步神经函数评估(NFE),带来高昂的计算延迟;其二是重构质量对推断阶段超参数高度敏感,无论是反向扩散的噪声调度(起始噪声 \(\sigma_{\max}\) 与终止噪声 \(\sigma_{\min}\)),还是先验项与数据一致性项的权衡权重 \(\lambda_t\),只要稍有偏离就会导致图像严重伪影甚至发散。

在现有实践中,这些超参数往往依赖启发式规则或昂贵的网格搜索逐任务手工微调。更棘手的是,一旦退化模型发生轻微改变(如 MRI 加速倍数、欠采样掩模或噪声水平变化),旧超参数便迅速失效。尽管近期研究尝试利用信噪比(SNR)启发式设计衰减曲线,但仍需在繁琐的函数空间(如线性、平方根、对数等)中寻找最优调度,缺乏坚实的理论原则。因此,核心矛盾在于:预训练扩散先验本身具备跨任务通用性,但现有的求解器超参数却极度缺乏跨问题鲁棒性与尺度不变性。

受到大模型中最大更新参数化(\(\mu\)P)能够稳定超参数的启发,本文探索扩散逆问题求解器的第一性原理设计。核心 idea:从频域谱自回归视角推导出逆问题中噪声起止区间的容差重参数化(\(\tau_{\max}, \tau_{\min}\)),并在后验优化中采用双端归一化梯度与单标量权衡,不仅实现超参数跨任务尺度不变性,更无缝解锁动量与多项式预条件加速。

方法详解

整体框架

本文构建的 OptDiff 框架将基于分数的逆问题求解重新表述为具有物理自洽性的复合凸/非凸优化流程。面对测量模型 \(\bm{y} = \mathcal{A}(\bm{x}_0) + \bm{\eta}\),OptDiff 首先基于信号的功率谱密度与退化算子保留的低频信息,由谱容差参数 \((\tau_{\max}, \tau_{\min})\) 封闭解出最优反向扩散区间的边界 \((\sigma_{\max}, \sigma_{\min})\);随后,在后验梯度更新中,同时对数据一致性梯度与扩散先验分数梯度执行 \(L_2\) 模长归一化,消除了时间步长相关的幅度剧烈起伏,从而能够用单一全局常数 \(\hat{\lambda}\) 替代繁琐的动态调度;最后,借助少量微调样本的贪心最优求解器完成超参数解耦,并引入动量与二阶多项式预条件器大幅压缩函数评估步数。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    In["退化测量输入 y 与测量算子 A"] --> S1["基于谱自回归确定噪声调度<br/>根据容差 tau_max / tau_min 确定起止噪声"]
    S1 --> S2["双端归一化公共下降采样<br/>单标量 lambda 权衡数据与先验梯度"]
    S2 --> S3["贪心最优超参数解耦与标定<br/>少量样本 NNLS 求解最优更新权重"]
    S3 --> S4["动量与多项式预条件加速<br/>二阶多项式预条件与动量项注入"]
    S4 --> Out["高质量无伪影重构图像 x_0"]

关键设计

1. 基于谱自回归确定噪声调度:自适应匹配频域特性的启闭准则

现有逆求解器通常粗暴沿用无条件生成的经验噪声端点(如 \(\sigma_{\max}=20, \sigma_{\min}=0.002\)),忽略了逆问题中低频测量信息和物理噪声底的存在,导致大量迭代步数浪费在无意义的高噪阶段或在噪声底之下过拟合。针对这一痛点,本文将反向扩散过程解释为频域谱自回归:自然图像能量集中在低频,逆问题前向观测 \(\bm{y}\) 通常已经保留了充足的低频分量。假设低频完全已知,高频不确定性由条件协方差 \(\bm{\Sigma}_{H|L}\) 描述。本文定义逆扩散第一步带来的协方差残差缩减比率为 \(\tau_{\max}\),由矩阵 Loewner 不等式推导出起始噪声下界: $$ \sigma_{\max}^2 \geq \frac{1 - \tau_{\max}}{\tau_{\max}} \nu_{\max}(\bm{\Sigma}{H|L}) = \frac{1 - \tau S(\omega) $$ 而在终止端,当测量系统存在高斯白噪声底 }}{\tau_{\max}} \max_{\omega \in \Omega_H\(\sigma_{\bm{s}}^2\) 时,继续扩散到过低噪声只会拟合测量噪声。通过控制后验残差与理论最优 LMMSE 后验的差距在容差 \(\tau_{\min}\) 内,推导出终止噪声上限: $$ \sigma_{\min}^2 \leq \frac{\tau_{\min} \sigma_{\bm{s}}^2 (\nu_{\max}(\bm{\Sigma}{\bm{s}}) + \sigma}}^2)}{\nu_{\max}(\bm{\Sigma{\bm{s}}) - \tau^2 $$ 利用容差对 } \sigma_{\bm{s}}^2} \approx \tau_{\min} \sigma_{\bm{s}\((\tau_{\max}, \tau_{\min})\) 替代具体物理数值,使超参数天然具备了随不同任务退化强度自适应缩放的理论不变性。

2. 双端归一化公共下降采样:单标量消除繁琐调度

在 RED-diff 类变分优化目标 \(\min_{\bm{\mu}} f(\bm{\mu}) + \mathbb{E}_t [\lambda_t g(\bm{\mu}; t)]\) 中,数据一致性梯度 \(\nabla f^k\) 与扩散先验梯度 \(\nabla g^k\) 的模长比值 \(r^k = \|\nabla f^k\|_2 / \|\nabla g^k\|_2\) 随着反向扩散步数 \(t\) 发生数个数量级的剧烈波动,使得静态 \(\lambda\) 极易导致发散或模糊。以往方法依赖复杂的经验 SNR 衰减曲线。本文提出双端归一化单标量重参数化更新: $$ \Delta \bm{\mu}^k = -\hat{\alpha}^k \left( \frac{\nabla f^k}{|\nabla f^k|_2} + \hat{\lambda} \frac{\nabla g^k}{|\nabla g^k|_2} \right) $$ 本文证明该更新方向与步长对于 \(f\) 和 \(g\) 的任何单调增可微重标度均保持严格不变性。更关键的是,从多目标公共下降条件出发,两梯度关于更新方向同时具备非负内积的充要条件退化为简洁的区间约束:当夹角余弦 \(\cos \vartheta^k < 0\) 时,只需全局标量满足 \(\hat{\lambda} \in (-\cos \vartheta^k, -1 / \cos \vartheta^k)\)。这从数学上严格保证了整个反向扩散过程中仅需一个固定的全局标量 \(\hat{\lambda}\),彻底摒弃了繁琐的时间依赖权重函数。

3. 贪心最优超参数解耦与标定:微调集上的非负最小二乘求解

为了在优化设计空间中彻底解耦步长、平衡权重与噪声调度搜索,OptDiff 设计了一种基于少量(5-10 张)参考样本的贪心最优标定方案。在每一步迭代中,定义梯度方向矩阵 \(\bm{U}^k = [\nabla f^k, \nabla g^k]\) 与待求超参数向量 \(\bm{w}^k = [\alpha^k, \alpha^k \lambda^k]^\top\)。通过求解非负最小二乘(NNLS)问题: $$ \hat{\bm{w}}^k = \arg\min_{\bm{w} \in \mathbb{R}^2_+} |\bm{x}^* - (\bm{\mu}^k - \bm{U}^k \bm{w})|_2^2 $$ 直接计算出当前步能让重建图像最接近真值 \(\bm{x}^*\) 的理论最优步长和权衡系数。该上界轨迹不仅能够为网格搜索提供狭窄的高效搜索邻域,更使得噪声调度的容差调优可以脱离特定求解器的步长影响独立进行。

4. 动量与多项式预条件加速:经典凸优化工具无缝接入

得益于双端归一化带来的数值稳定性,OptDiff 能够直接融合经典一阶加速与二阶近似算法,而不会引入调参噩梦。针对线性测量算子 \(\mathcal{A}\),结合 Polyak 动量与 2 阶多项式预条件器 \(\text{pol}(\bm{A}^\top \bm{A})\) 对数据一致性梯度进行频域均衡: $$ \bm{v}^{k+1} \leftarrow \alpha^k \left( \text{pol}(\bm{A}^\top \bm{A}) \nabla f^k + \lambda^k \nabla g^k \right) + \beta^k \bm{v}^k, \quad \bm{\mu}^{k+1} \leftarrow \bm{\mu}^k - \bm{v}^{k+1} $$ 在贪心最优标定中,基矩阵自然扩展为 \(\bm{U}^k = [\text{pol}(\bm{A}^\top \bm{A})\nabla f^k, \nabla g^k, \bm{v}^k]\),同样能通过 3 维非负最小二乘一并自适应求解。这使得求解器在低函数评估步数(如 NFE=20)下依旧保持极高重建保真度。

损失函数 / 训练策略

OptDiff 本身面向推断阶段优化,使用预训练冻结的无条件扩散模型,完全无需微调网络权重。其推断过程遵循改进的 RED-diff 复合能量极小化目标: $$ \min_{\bm{\mu}} |\bm{y} - \mathcal{A}(\bm{\mu})|2^2 + \mathbb{E}|_2^2 \right] $$ 在推断执行时,仅需在验证集中的 5-10 张切片/图像上使用 NNLS 快速确定全局不变超参数 }} \left[ \lambda_t |\bm{\epsilon}_{\bm{\theta}}(\bm{x}_t; t) - \bm{\epsilon\((\tau_{\max}, \tau_{\min}, \hat{\lambda}, \hat{\alpha})\),随后在全量测试集上冻结参数直接串行推断,单张图像重构仅需 20 至 100 步前向分数评估。

实验关键数据

主实验

在快速磁共振成像(fastMRI 脑部 T2 加权,加速倍数 \(R=4, R=8\))以及自然图像去模糊(ImageNet 与 FFHQ,高斯模糊核 \(k=6, \sigma=3\))的 100 步函数评估(NFE=100)下,OptDiff 与主流基线全面对比,定量表现如下:

数据集 / 任务 指标 OptDiff (本文) 之前最佳基线 (RED-diff / DDNM) 提升
fastMRI (MRI \(R=4\)) PSNR (dB) ↑ 38.45 ± 1.58 36.79 ± 1.43 (RED-diff) +1.66 dB
fastMRI (MRI \(R=4\)) SSIM ↑ 0.948 ± 0.023 0.930 ± 0.030 (DDS) +0.018
fastMRI (MRI \(R=8\)) PSNR (dB) ↑ 34.93 ± 1.38 31.92 ± 1.44 (RED-diff) +3.01 dB
fastMRI (MRI \(R=8\)) SSIM ↑ 0.921 ± 0.029 0.840 ± 0.034 (RED-diff) +0.081
ImageNet (去模糊 \(k=6\)) PSNR (dB) ↑ 33.07 ± 3.76 32.60 ± 3.89 (DDNM) +0.47 dB
ImageNet (去模糊 \(k=6\)) SSIM ↑ 0.910 ± 0.040 0.906 ± 0.045 (DDNM) +0.004
FFHQ (去模糊 \(k=6\)) PSNR (dB) ↑ 36.99 ± 1.82 36.95 ± 1.78 (DDNM) +0.04 dB
FFHQ (去模糊 \(k=6\)) SSIM ↑ 0.948 ± 0.013 0.949 ± 0.013 (DDNM) -0.001

在更严苛的高加速比 fastMRI \(R=8\) 任务下,OptDiff 相对原版 RED-diff 取得超过 3 dB 的巨幅 PSNR 跨越,同时在自然图像逆任务上稳居第一梯队。

消融实验

为验证各核心组件的有效性,在 fastMRI \(R=8\)(100 步迭代)上进行了权衡策略与优化器加速插件的消融实验:

配置 / 权重方法 调度配置 PSNR (dB) SSIM 说明
贪心最优上界 (Oracle) 标准调度 33.13 ± 1.22 0.904 ± 0.020 理论上限基准
贪心最优上界 (Oracle) 提出调度 (\(\tau\)) 34.78 ± 1.24 0.941 ± 0.021 验证提出噪声调度有效性 (+1.65 dB)
常数 \(\lambda\) (无归一化) 提出调度 (\(\tau\)) 33.56 ± 1.49 0.910 ± 0.030 未归一化梯度导致严重性能衰减
线性 SNR 调度 提出调度 (\(\tau\)) 34.39 ± 1.24 0.928 ± 0.021 常见手工衰减设计
对数 SNR 调度 提出调度 (\(\tau\)) 34.36 ± 1.25 0.928 ± 0.021 常见手工衰减设计
平方根 SNR 调度 提出调度 (\(\tau\)) 34.13 ± 1.31 0.921 ± 0.025 常见手工衰减设计
本文双端归一化 \(\hat{\lambda}\) 提出调度 (\(\tau\)) 34.51 ± 1.37 0.927 ± 0.025 超越所有经验 SNR 调度,逼近理论上界

此外,针对优化器加速插件(在 NFE=20 步下的加速评估): - Vanilla OptDiff: 29.9 ± 1.7 dB / SSIM 0.88 - + 动量 (Momentum): 31.7 ± 2.1 dB / SSIM 0.87 (+1.8 dB 显著提升早期收敛) - + 动量 + 预条件 (Precond): 32.3 ± 2.1 dB / SSIM 0.89 (相对基础版提升 2.4 dB,在仅 20 步下已匹敌 baseline 100 步性能)

关键发现

  • 噪声调度是跨求解器即插即用的通用改进:将本文提出的容差调度迁移至 DPS、DDRM、Pi-GDM、RED-diff 等现有求解器中,几乎全部带来 1~2 dB 的 PSNR 增益(如 DPS 在去模糊上由 27.51 dB 跃升至 38.17 dB),证明逆问题中尽早裁剪低效高噪区间具有普遍收益。
  • 双端归一化打破调参壁垒:经验 SNR 权重方案在不同时间窗口需要调节多种函数形态,而本文单标量 \(\hat{\lambda}\) 的归一化解法使曲线在各反向区间高度对齐,且严格满足多目标公共下降条件。
  • 预条件技术在小步数推断时收益尤为巨大:在 NFE=20 的超快速推断场景下,多项式预条件器结合动量使得 PSNR 从 29.9 dB 直接拉升至 32.3 dB,极大缓解了扩散逆求解器速度缓慢的顽疾。

亮点与洞察

  • 从频域协方差视角严格界定逆问题噪声起止点:打破了直接借用无条件扩散噪声调度的惯性思维,通过严谨的 Loewner 不等式推导出由信号 PSD 与噪声底决定的容差准则,使超参数天然具备物理尺度自适应性。
  • 双端模长归一化与公共下降理论结合:揭示了数据一致性与先验梯度模长剧烈动态失衡是超参数难以复用的根源,双端归一化配合夹角界定,使多目标优化中单标量常数即可保证单调下降。
  • 无痛接入凸优化高级工具链:将最优化领域中成熟的动量项与二阶预条件器以极低开销嵌入扩散后验采样,使扩散模型能够以 20 步极低开销输出高保真重构。

局限与展望

  • 依赖测量算子保留充足低频信息:起始噪声削减建立在算子保留可观低频的前提下(如超分辨、去模糊、平滑降采样 MRI)。对于随机大面积遮挡修复(Inpainting)等频域低频被截断的任务,起始噪声下界退化为全局极大值,加速收益受限。
  • 优化单解倾向导致极端逆问题感知质量不足:基于 RED-diff 的确定性变分最优化框架生成单一解,在 16× 超分辨率或大面积修复等严重病态多解场景下,倾向于产生像素均方误差最优的平滑解,感知指标(LPIPS)落后于随机采样求解器(如 DAPS)。
  • 优化器参数设计尚未完全泛化到非优化类求解器:尽管噪声调度能无缝迁移,但双端归一化梯度更新专为优化型求解器定制,未来需探索向一般条件反向 SDE 采样器的统一拓展。

相关工作与启发

  • vs RED-diff: RED-diff 首次将扩散逆采样形式化为变分凸优化框架,但其超参数高度依赖时间步长相关的经验 SNR 调度且步数冗长;本文通过双端归一化将其简化为单标量,并借助动量/预条件实现数倍推断加速。
  • vs DPS / Pi-GDM: DPS 与 Pi-GDM 属于基于反向 SDE/ODE 的流形似然投影与梯度引导,噪声调度依赖全局经验设置;本文推导的基于谱容差的噪声调度可作为通用插件反哺此类流形求解器并显著提质。
  • vs DDRM / DDNM: DDRM 和 DDNM 依赖退化算子 \(\mathcal{A}\) 的精确奇异值分解(SVD),在多线圈 MRI 等复杂非对角算子中计算不可行;OptDiff 仅需评估前向算子及其伴随算子,通用性更强。

评分

  • 新颖性: ⭐⭐⭐⭐⭐ 首次从谱自回归与不变性重参数化视角建立扩散逆问题求解器的第一性原理理论。
  • 实验充分度: ⭐⭐⭐⭐⭐ 涵盖 fastMRI、ImageNet、FFHQ 三大基准与四类退化任务,提供理论上界对照与跨求解器泛化验证。
  • 写作质量: ⭐⭐⭐⭐⭐ 理论推导清晰完备,从 Loewner 序到多目标梯度下降条件论证严密。
  • 价值: ⭐⭐⭐⭐⭐ 破解了扩散逆求解器“调参繁琐、收敛极慢”的工业落地痛点,为优化算法与生成先验的结合指明方向。