跳转至

Consistency Regularised Gradient Flows for Inverse Problems

会议: NeurIPS 2026
arXiv: 2605.07907
代码: https://github.com/LeviBorodenko/motionblur
领域: 优化/理论
关键词: 逆问题求解, 潜在一致性模型, 瓦瑟斯坦梯度流, 提示词优化, 经验贝叶斯

一句话总结

本文提出了欧几里得-瓦瑟斯坦联合梯度流框架（CWGF），利用潜在一致性模型在 VAE 隐空间中同时执行后验采样与文本提示词最大边际似然优化，无需对解码器反向传播即可在 16 次神经函数评估内实现高质量图像逆问题求解。

研究背景与动机

基于视觉-语言潜在扩散模型（Latent Diffusion Models, LDMs）的即插即用（Plug-and-Play, PnP）逆问题求解器在图像反演任务中展现出强大的生成先验能力。然而，现有基于 LDMs 的逆问题求解方案面临两个根本性的计算瓶颈：首先，传统扩散模型反演过程需要上百次甚至上千次神经函数评估（NFEs），采样耗时漫长；其次，由于物理测量似然定义在像素空间，而先验定义在变分自编码器（VAE）的隐空间，计算数据一致性项通常需要穿透高维非线性解码器进行反向传播（Backpropagation），这带来了极高的显存开销，并经常伴随解码梯度伪影。

虽然近期有方法引入少步生成的一致性模型（Consistency Models, CMs）或通过启发式提示词微调（如 P2L、TReg）来缓解问题，但提示词调整往往脱离了严格的贝叶斯推断框架。即便是引入了经验贝叶斯最大边际似然估计（MMLE）的先进方法 LATINO-PRO，也采用了随机近似近端梯度算法（SAPG），在后验采样与提示词优化之间交替迭代。这种双层嵌套结构收敛缓慢，单样本往往仍需超过 50 次 NFE，且未能从统一几何视角解决隐空间测度演化与参数优化的耦合问题。

本文的切入角度是：将图像反演中隐空间概率测度的演化与语义文本提示词的调整，统一置于乘积流形上的连续时间几何演化过程中。核心 idea：在欧几里得与瓦瑟斯坦-2 乘积流形上构建一致性正则化梯度流（CWGF），将后验采样转化为测度输运、提示词优化转化为连续流漂移，并巧妙借助 VAE 编码器与一致性流动映射避开对解码器的反向传播，实现 16 NFE 的极速高精度反演。

方法详解

整体框架

在贝叶斯反演设定下，已知正向退化算子 \(A\) 及观测 \(y = A(x) + n\)（其中噪声 \(n \sim \mathcal{N}(0, \sigma_y^2 I)\)），图像 \(x\) 由预训练 VAE 解码器 \(p_{\phi}(x|z) = \mathcal{N}(x; \mathcal{D}_\phi(z), \sigma_{\text{dec}}^2 I)\) 从隐变量 \(z\) 生成。文本提示词 \(c \in \mathcal{C}\) 调节先验分布 \(p_c(z)\)。CWGF 将最优化目标定义在乘积空间 \(\mathcal{C} \times \mathcal{P}_2(\mathcal{Z})\) 上： $$ \min_{(c, \mu) \in \mathcal{C} \times \mathcal{P}2(\mathcal{Z})} \mathcal{F}[c, \mu] := D_\mu [\log p(y|z)], $$ 其中正则项 }}(\mu \parallel p_c) - \mathbb{E\(\mathcal{R}[c, \mu] = D_{\text{KL}}(\mu \parallel p_c)\) 促使隐空间测度逼近一致性模型诱导的先验，似然项 \(\mathcal{L}[\mu] = -\mathbb{E}_\mu[\log p(y|z)]\) 促使测度匹配观测数据。

CWGF 沿着由欧几里得度量（作用于 \(c\)）与瓦瑟斯坦-2 度量（作用于 \(\mu\)）诱导的联合梯度流进行演化，并通过 Lie-Trotter 算子分裂法将连续时间演化解耦为三个交替更新的子步骤：基于得分的提示词演化流、基于一致性流动映射的先验输运流、以及基于编码器代理的似然漂移流。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["退化观测 y 与初始文本嵌入 c0"] --> B["基于得分的提示词流<br/>Fisher 散度梯度与时间自适应权重更新 c"]
    B --> C["先验诱导输运流<br/>LCM 终点松弛驱动与粒子核化排斥"]
    C --> D["基于编码器的似然流<br/>像素共轭高斯后验与 VAE 编码器代理更新"]
    D -->|循环迭代 K 步| E["最终重构隐变量 zK 与优化提示词 cK"]
    E --> F["VAE 解码输出清晰图像 D(zK)"]

关键设计

1. 基于得分的提示词流：避免逆向扩散模拟的边际似然优化

在欧几里得参数空间中，直接对 KL 散度 \(\mathcal{R}[c, \mu] = D_{\text{KL}}(\mu \parallel p_c)\) 关于提示词嵌入 \(c\) 求导极其昂贵，因为计算先验密度 \(p_c(z)\) 的梯度通常需要逆向模拟整条 SDE 或 ODE 轨迹。CWGF 借助连续时间方差保持（VP）正向扩散路径，利用得分匹配恒等式将 KL 散度等价转化为沿扩散路径的 Fisher 散度积分： $$ \nabla_c D_{\text{KL}}(\mu \parallel p_c) = \frac{1}{2} \nabla_c \int_0^T \beta(t) \mathbb{E}{z_0 \sim \mu, z_t \sim q_t(\cdot|z_0)} \left[ \left| \nabla \log q_t(z_t | z_0) - s_\theta(z_t, t; c) \right|_2^2 \right] dt, $$ 在离散实现中，CWGF 复用潜在一致性模型（LCM）在时间步 \(t_k\) 处的重参数化得分输出 \(s_\theta(z_t, t; c) = -\sigma_t^{-1} \hat{\epsilon}_\theta(z_t, t; c)\) 作为代理得分，并结合信噪比倒数权重 \(w(t_k)\) 进行蒙特卡洛估计。提示词向量据此执行平滑的欧几里得梯度下降，且每一步投影到以初始提示词为中心的球空间内，使得生成模型的主题语义能自动、渐进地校正并与观测图像高度匹配。

2. 先验诱导输运流：利用流动映射恒等式消除先验得分奇异性

对于隐变量概率测度 \(\mu\) 针对先验正则项 \(\mathcal{R}[c, \mu]\) 的瓦瑟斯坦-2 梯度流，其漂移速度场由 \(-\nabla_z \frac{\delta \mathcal{R}}{\delta \mu} = \nabla_z \log p_c(z) - \nabla_z \log \mu(z)\) 决定。然而，在经验粒子测度下，经验得分 \(\nabla \log \mu\) 与高维未加噪先验得分 \(\nabla \log p_c\) 均存在数值奇异性。本文证明了严格的流动映射积分恒等式（Proposition 1）：对于概率流动 ODE 诱导的逆映射 \(G_t^r(z_t)\)，加权积分 \(\frac{\sqrt{\alpha_t} G_t^r(z_t) - z_t}{1 - \sqrt{\alpha_t}} = \int_0^t \gamma_t(s) \nabla_{z_s} \log r_s(z_s) ds\) 精确等于沿扩散特征轨迹的正向得分凸平均。

据此，CWGF 将难以处理的得分差替换为流动映射差异。对于先验部分，利用预训练 LCM 的预测函数 \(g_\theta(z_t, t, c)\) 逼近 \(G_t^{p_c}(z_t)\)；对于粒子测度部分，利用核加权去噪均值 \(\sum_m \pi_{nm}(t) z_0^{(m)}\) 逼近粒子后验均值。推导出的离散粒子更新为： $$ \bar{z}^{(n)}k = z^{(n)}k + \eta_R \left( g\theta(z^{(n)}}, t_k, c_k) - z^{(n)k \right) + \eta_R \left( z^{(n)}k - \sum^N \pi_k \right), $$ 式中前两项构成了粒子向一致性模型生成流形的几何松弛牵引，最后一项则是粒子间远离邻域质心的局部核化排斥力。在多粒子采样时，该排斥项能够有效拓宽粒子覆盖范围，防止后验模式坍缩。}(t_k) z^{(m)

3. 基于编码器的似然流：利用变分共轭消除解码器反向传播

传统 PnP 扩散求解器处理似然项 \(\mathcal{L}[\mu] = -\mathbb{E}_\mu[\log p(y|z)]\) 时，漂移向量 \(g(z; y) = \nabla_z \log p(y|z) = \mathbb{E}_{p_c(x|z,y)}[\nabla_z \log p_\phi(x|z)]\) 必须计算解码器 \(\mathcal{D}_\phi(z)\) 的雅可比矩阵，产生巨大显存开销与梯度畸变。CWGF 基于 VAE 训练目标的变分最优性推导证明，对于优化良好的 VAE，解码器似然梯度可以被编码器高斯分布的得分所近似：\(g(z; y) \approx \lambda \int [\nabla_z \log q_\phi(z|x)] p_c(x|z, y) dx - \lambda \nabla_z \log p_0(z)\)。

在正向线性退化模型与高斯解码假设下，像空间条件后验 \(p_c(x|z, y)\) 存在闭式高斯解 \(\mathcal{N}(x; m(z), \Sigma_y)\)，其均值点为： $$ m(z) = \Sigma_y \left( \sigma_{\text{dec}}^{-2} \mathcal{D}\phi(z) + \sigma_y^{-2} A^\top y \right), \quad \Sigma_y = \left( \sigma. $$ 在最大后验（MAP）近似下，似然项诱导的隐变量速度场简化为前向编码计算： $$ \dot{z}}}^{-2} I + \sigma_y^{-2} A^\top A \right)^{-1\tau = \lambda \Sigma\phi^{-1} \left( \mathcal{E}\phi(m(z\tau)) - z_\tau \right) + \lambda z_\tau. $$ 整个数据保真度更新完全由解码前向传播 \(\mathcal{D}_\phi\)、在像空间计算数据滤波解析解 \(m(z)\)、以及编码器前向提取 \(\mathcal{E}_\phi(m(z))\) 串联完成，全程无需对解码器或编码器网络进行反向自动微分，兼具理论严密性与轻量高效性。

实验关键数据

主实验

在 FFHQ-512 与 ImageNet-512 验证集上，评估高斯去模糊（Gaussian Deblur, \(61 \times 61, \sigma=3.0\)）、运动去模糊（Motion Deblur, \(61 \times 61, \text{intensity}=0.51\)）以及 \(8\times\) 超分辨率（SR \(\times 8\)，双三次插值），测量噪声为 \(\sigma_y = 0.01\)。主实验以单粒子 \(N=1\) 运行 16 次迭代。

数据集 / 任务 指标 CWGF (本文) LATINO-PRO (前SOTA) P2L TReg 提升幅度 (vs LATINO-PRO)
FFHQ-512 (Gaussian Deblur) NFE \(\downarrow\) / FID \(\downarrow\) / PSNR \(\uparrow\) / LPIPS \(\downarrow\) 16 / 21.95 / 29.42 / 0.309 65 / 32.48 / 28.49 / 0.325 400 / 45.12 / 26.47 / 0.433 200 / 48.73 / 24.27 / 0.418 -75.4% NFE, -10.53 FID, +0.93 dB
FFHQ-512 (Motion Deblur) NFE \(\downarrow\) / FID \(\downarrow\) / PSNR \(\uparrow\) / LPIPS \(\downarrow\) 16 / 23.18 / 27.98 / 0.338 65 / 39.72 / 25.16 / 0.410 400 / 55.73 / 23.92 / 0.503 200 / 44.97 / 22.82 / 0.432 -75.4% NFE, -16.54 FID, +2.82 dB
FFHQ-512 (SR \(\times 8\)) NFE \(\downarrow\) / FID \(\downarrow\) / PSNR \(\uparrow\) / LPIPS \(\downarrow\) 16 / 37.37 / 27.07 / 0.366 65 / 75.97 / 26.00 / 0.445 400 / 52.14 / 23.76 / 0.485 200 / 62.40 / 22.68 / 0.460 -75.4% NFE, -38.60 FID, +1.07 dB
ImageNet-512 (Gaussian Deblur) NFE \(\downarrow\) / FID \(\downarrow\) / PSNR \(\uparrow\) / LPIPS \(\downarrow\) 16 / 28.19 / 25.40 / 0.394 65 / 35.18 / 25.28 / 0.357 400 / 59.77 / 23.05 / 0.487 200 / 56.54 / 22.38 / 0.441 -75.4% NFE, -6.99 FID, +0.12 dB
ImageNet-512 (Motion Deblur) NFE \(\downarrow\) / FID \(\downarrow\) / PSNR \(\uparrow\) / LPIPS \(\downarrow\) 16 / 41.81 / 23.86 / 0.409 65 / 67.44 / 22.51 / 0.478 400 / 159.32 / 14.29 / 0.683 200 / 78.75 / 20.99 / 0.501 -75.4% NFE, -25.63 FID, +1.35 dB
ImageNet-512 (SR \(\times 8\)) NFE \(\downarrow\) / FID \(\downarrow\) / PSNR \(\uparrow\) / LPIPS \(\downarrow\) 16 / 48.66 / 23.40 / 0.458 65 / 68.89 / 23.06 / 0.498 400 / 55.04 / 22.56 / 0.468 200 / 75.91 / 20.85 / 0.504 -75.4% NFE, -20.23 FID, +0.34 dB

消融实验

为了检验对抗性/错配提示词下的校正能力，在 FFHQ-512 数据集上故意初始化为错误提示词 "A photo of a cat"（对抗提示词测试），对比提示词优化机制的作用：

退化任务 配置 NFE FID \(\downarrow\) PSNR \(\uparrow\) LPIPS \(\downarrow\) 说明
Gaussian Deblur CWGF (无提示词优化) 16 32.91 28.64 0.354 初始提示词固定为猫，人脸重建产生语义伪影
Gaussian Deblur LATINO-PRO 65 44.67 28.63 0.331 迭代更新缓慢
Gaussian Deblur CWGF (完整联合流) 16 27.58 29.59 0.322 自动漂移对齐，FID 改善 5.33，PSNR 提升 0.95 dB
Motion Deblur CWGF (无提示词优化) 16 32.20 27.86 0.357 严重缺乏高频语义细节
Motion Deblur CWGF (完整联合流) 16 25.99 28.16 0.339 FID 改善 6.21，PSNR 提升 0.30 dB
SR \(\times 8\) CWGF (无提示词优化) 16 92.43 26.24 0.456 超分因语义严重错配导致结构崩塌
SR \(\times 8\) CWGF (完整联合流) 16 50.07 27.39 0.415 FID 大幅降低 42.36，重建出真实人脸结构

关键发现

  • 计算效率显著跨越：CWGF 仅需 16 次 NFE 即可稳定收敛，相比基于扩散模型的传统 PnP 方案（200-400 NFE）减少了 90% 以上计算量；相比少步基线 LATINO-PRO（65 NFE），计算耗时减少 75.4%，同时在各项指标上全面超越。
  • 提示词自动纠错韧性：在对抗性错配提示词实验中，CWGF 展现出极强的语义自愈能力。例如在 \(8\times\) 超分任务中，完整 CWGF 使 FID 从 92.43 骤降至 50.07，证明欧几里得提示词梯度流能迅速脱离局部错误模式，自发向匹配真实观测的文本嵌入空间靠拢。
  • 隐空间维度与多粒子交互：在 SD1.5 的高维隐空间中，\(N=1\) 与 \(N=8\) 的多粒子表现差异较小（高维下粒子转移核权重矩阵接近单位阵）；但在 MNIST 低维隐空间（\(d=10\)）的修复实验中，多粒子（\(N=1024\)）联合演化展现出显著的核化质心排斥效应，能够发掘出独立单粒子采样遗漏的后验边缘模式（如数字 5 的解），证明了瓦瑟斯坦相互作用粒子流在刻画复杂后验多峰分布上的理论潜力。

亮点与洞察

  • 乘积流形上的统一步调演化：打破了过去将"后验采样"与"提示词优化"视作内外两层嵌套算法的固有思维，将其统一表述为欧几里得参数空间与瓦瑟斯坦测度空间上的连续时间联合梯度流，理论框架严谨优雅。
  • 完全免于解码器反向传播：巧妙结合 VAE 最优编码器性质与正向线性观测的共轭高斯解，仅通过解码前向、数据空间闭式滤波与编码前向，便获得了精确的隐空间似然漂移，彻底根除了反向穿透解码器带来的显存溢出与梯度高频噪声。
  • 普适性技术迁移价值：基于流动映射近似正向得分的凸平均恒等式（Proposition 1），为任意基于一致性模型或流匹配（Flow Matching）的少步生成先验处理无导数逆问题提供了通用数学范式。

局限与展望

  • 文本嵌入空间的欧几里得假设：作者在推导中假设提示词嵌入空间 \(\mathcal{C}\) 具备标准欧几里得几何结构，但近期研究表明 CLIP 等视觉-语言文本特征空间往往呈现非欧几何或双椭球流形特性，未来探索黎曼流形下的提示词梯度流有望进一步加速语义收敛。
  • 强依赖线性前向观测假设：似然流中像空间高斯后验的闭式解依赖于正向算子 \(A\) 的线性结构；虽然论文在附录中针对非线性 JPEG 压缩提供了投影梯度代理，但针对极端非线性退化（如散射介质、非线性盲去噪）仍需探索更通用的变分近似机制。

相关工作与启发

  • vs LATINO / LATINO-PRO: LATINO 依赖固定先验提示词且在像空间使用近端校正；LATINO-PRO 引入经验贝叶斯，但采用外层 SAPG 提示词更新与内层 Langevin 链交替的随机逼近结构，需 65 NFE 且收敛缓慢。CWGF 将其合成为单层连续梯度流，在 16 NFE 内实现双流同步加速演化。
  • vs PSLD / P2L: PSLD 与 P2L 依赖迭代扩散模型（200-400 NFE），且必须对整个非线性 VAE 解码器进行反向自动微分，显存占用极大且耗时严重。CWGF 借助 LCM 与编码器代理，完全消除了反向传播，推理速度高出数十倍。

评分

  • 新颖性: ⭐⭐⭐⭐⭐ 首次将欧几里得-瓦瑟斯坦联合梯度流引入潜在一致性模型逆问题求解，理论推导严谨精妙。
  • 实验充分度: ⭐⭐⭐⭐⭐ 涵盖人脸与通用图像数据集的三类标准退化任务、对抗错配提示词消融、非线性问题测试及低维后验分布粒子分析。
  • 写作质量: ⭐⭐⭐⭐⭐ 理论脉络清晰，数学证明完备，算法伪代码与参数设置透明严谨。
  • 价值: ⭐⭐⭐⭐⭐ 为超少步（16步）生成模型高保真求解图像逆问题树立了新的效率与质量标杆。

相关论文