跳转至

DiffRGD: An Inference-Time Diffusion Guidance Through Riemannian Gradient Descent

会议: ECCV 2026
论文: ECCV 原文
项目: https://diffrgd.github.io/
领域: 图像生成
关键词: 扩散模型、推理时引导、黎曼梯度下降、极分解、图像逆问题

一句话总结

针对推理时扩散模型引导因破坏潜空间高斯分布导致样本质量退化的问题,DiffRGD 利用各向同性高斯的极分解几何性质构建分布诱导的球面流形,并通过黎曼梯度下降在保持先验分布的前提下实现高质量条件生成。

研究背景与动机

扩散模型在无条件和条件生成建模中展现了优异的表现,但在针对特定下游任务进行精细控制时,重新训练或微调庞大的预训练基础模型通常面临高昂的计算开销与维护成本。为了实现无需训练的即插即用控制,推理时引导技术受到了广泛关注。以 DPS 和 FreeDoM 为代表的主流方法通过可微目标函数的损失梯度直接更新反向去噪过程中的潜变量。然而,这类方法在无约束的欧几里得环境空间中执行梯度更新,由于 Jensen 不等式引起的分布偏移(Jensen gap),外加梯度往往破坏了反向去噪每一步原有的高斯边际分布,导致潜变量逐渐偏离真实数据流形,产生严重的视觉伪影与保真度退化;若刻意缩小步长以维持生成稳定性,又会导致引导力度不足,无法精确契合条件信号。

近期部分工作尝试对这一偏离流形现象施加约束,但现有约束机制仍存在显著短板。例如 MPGD 依赖局部线性流形假设,通过预训练 VAE 编码器的切空间投影来修正梯度,但该方案依赖理想编码器的假设在实践中难以成立,容易带来数值不稳定性;DSG 提出将潜变量投影到固定半径的超球面上,但这割裂了采样步进中的尺度特性,反而退化为了均匀分布;而基于交替方向乘子法的 ADMMDiff 虽然建立了约束优化框架,却需要内层近端迭代算子与繁琐的超参数调优,收敛缓慢且计算代价沉重。

面对这种生成保真度与条件引导力度的固有张量矛盾,本文的切入视角是回归扩散过程本身的统计几何本质:DDIM 反向去噪的每一步都可以形式化为一个中心可由预测均值决定的各向同性高斯分布。核心 idea:利用各向同性高斯分布在数学上可精确极分解为相互独立的径向卡方分布与球面均匀方向的几何特性,将每步采样重构为分布诱导球面流形上的约束优化问题,并通过高效收敛的黎曼梯度下降(Riemannian Gradient Descent)在切空间与收缩映射下完成潜变量微调。

方法详解

整体框架

DiffRGD 在每个反向采样时间步 \(t\) 将引导更新建模为在特定球面流形 \(\mathcal{S}_{t,r}\) 上的约束优化。输入为条件观测 \(y\) 与当前时间步噪声潜变量 \(x_t\),通过预训练去噪网络计算出预测干净样本 \(\hat{x}_0(x_t, t)\) 及去噪均值 \(\mu_t\)。随后,算法利用极分解性质独立采样径向模长与初始球面方向,构造与扩散潜变量统计特性严格一致的球面流形 \(\mathcal{S}_{t,r}\)。在此流形上,算法利用任务特定的引导目标函数计算欧氏梯度,将其正交投影到切空间以消除破坏径向尺度的分量,最后沿负黎曼梯度方向更新并通过收缩算子拉回球面,经 \(K\) 次微调后输出下一阶段的引导潜变量 \(x_{t-1}\)。

整体执行流向如下图所示:

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["输入: 观测 y 与时间步噪声 xt"] --> B["极分解几何建模<br/>采样卡方半径 r 与球面切片 St,r"]
    B --> C["正交切空间投影<br/>欧氏梯度投影剥除径向分量"]
    C --> D["流形收缩映射与快速收敛保障<br/>沿测地切线位移并归一化拉回球面"]
    D --> E["输出: 保持高斯先验的引导样本 xt-1"]

关键设计

1. 极分解几何建模:将先验约束投影到球面流形 针对欧氏空间直接加权更新导致潜变量偏离原有概率分布的核心痛点,本文在数学上给出了各向同性高斯分布的极分解性质(Proposition 1)。在扩散采样步 \(t\) 中,潜变量 \(x_t \sim \mathcal{N}(\mu_t, \sigma_t^2 I_n)\) 可以精确解耦为相互统计独立的径向分量与方向分量: $\(x_t = \mu_t + \sigma_t r u, \quad r \sim \chi(n), \quad u \sim \text{Unif}(\mathbb{S}^{n-1})\)$ 其中 \(r\) 服从自由度为 \(n\) 的卡方分布,方向向量 \(u\) 均匀分布在单位超球面 \(\mathbb{S}^{n-1}\) 上。传统方法如 DSG 采用经验固定的几何超球,破坏了高斯分布本身的径向概率密度结构。DiffRGD 的核心机制是在每步反向采样时显式抽取卡方随机变量 \(r\),进而将允许优化的可行域精确锁定为球心在 \(\mu_t\)、半径为 \(\sigma_t r\) 的超球面流形: $\(\mathcal{S}_{t,r} = \{ x \in \mathbb{R}^n \mid \| x - \mu_t \|_2 = \sigma_t r \}\)$ 这种设计既保持了潜变量边际概率密度的统计自洽性,又将引导的目标收敛为在保持模长严格符合先验尺度的前提下,寻找最贴合条件目标的最优方向向量 \(u\)。

2. 正交切空间投影:剥除径向分量以消除分布漂移 定义了球面流形 \(\mathcal{S}_{t,r}\) 后,若直接应用外部引导损失 \(\mathcal{L}(\hat{x}_0, y)\) 的欧几里得梯度,该梯度的径向分量会不可避免地改变潜变量的模长,引发累积性分布漂移。流形切空间 \(T_x \mathcal{S}_{t,r}\) 汇聚了所有正交于当前径向向量 \(x - \mu_t\) 的切向位移方向。DiffRGD 利用流形上的自然内积定义了正交切空间投影算子: $\(\Pi_{T_x \mathcal{S}_{t,r}}(g) = g - \frac{\langle g, x - \mu_t \rangle}{\| x - \mu_t \|_2^2} (x - \mu_t)\)$ 对于当前迭代步计算得到的欧氏梯度 \(g = \nabla_{x} \mathcal{L}(\hat{x}_0(x, t-1), y)\),投影算子完整剔除了沿径向扩张或收缩的能量分量,仅保留沿流形表面旋转的有效切向分量 \(\text{grad}_{\mathcal{S}_{t,r}} \mathcal{L}\)。这一机制从根本上阻断了优化步长导致的高斯方差失真,保证了每一步修正纯粹致力于提升语义/保真度对齐,而不损耗先验结构。

3. 流形收缩映射与快速收敛保障:确保潜变量严格重回球面 由于切空间只是局部的线性切平面,沿切向步长更新后的中间点 \(x - \eta_t \text{grad}_{\mathcal{S}_{t,r}} \mathcal{L}\) 必然会离开弯曲的球面流形表面。为了将位移后的向量精准映射回可行域,DiffRGD 引入了几何收缩算子(Retraction): $\(R_x(v) = \mu_t + \sigma_t r \frac{x + v - \mu_t}{\| x + v - \mu_t \|_2}\)$ 该操作先沿切空间移动,再径向缩放至精确半径 \(\sigma_t r\),使潜变量更新完全封闭在 \(\mathcal{S}_{t,r}\) 内部。在理论收敛性方面,论文证明了在目标函数满足李普希茨梯度的条件下,该算法以 \(\mathcal{O}(1/\sqrt{k})\) 的次线性速率收敛至一阶平稳点。相比于 ADMMDiff 需要多重对偶变量更新与昂贵的近端迭代,DiffRGD 每次流形迭代仅需两次轻量的一阶代数投影与归一化计算,在实践中仅需固定 \(K=3\) 次内层更新即可使黎曼梯度范数充分衰减,兼具极高的推理效率与理论完备性。

实验关键数据

主实验

论文在 FFHQ 256×256 验证集(150 张样本)上对比了 1000 步 DDIM 采样下四种经典图像逆问题的重建表现(每项指标汇报 100 次 Bootstrap 抽样的均值):

任务 方法 会议 PSNR ↑ SSIM ↑ LPIPS ↓ FID ↓
Inpainting (70% 掩码) DPS ICLR 2023 30.44 0.863 0.153 42.68
MPGD ICLR 2024 27.51 0.724 0.256 68.24
DSG ICML 2024 31.03 0.866 0.144 36.30
ADMMDiff CVPR 2025 32.38 0.899 0.119 29.52
DiffRGD (本文) ECCV 2026 34.04 0.926 0.096 21.88
Super-Resolution 4× DPS ICLR 2023 26.03 0.727 0.260 80.05
MPGD ICLR 2024 24.40 0.614 0.354 101.50
DSG ICML 2024 26.71 0.737 0.256 74.67
ADMMDiff CVPR 2025 26.48 0.712 0.297 96.69
DiffRGD (本文) ECCV 2026 27.77 0.783 0.220 63.94
Gaussian Deblurring DPS ICLR 2023 25.88 0.721 0.237 69.38
MPGD ICLR 2024 24.07 0.576 0.328 95.12
DSG ICML 2024 27.45 0.751 0.259 75.85
ADMMDiff CVPR 2025 26.57 0.757 0.226 79.30
DiffRGD (本文) ECCV 2026 26.80 0.757 0.218 63.83
Motion Deblurring DPS ICLR 2023 24.47 0.685 0.271 80.75
MPGD ICLR 2024 23.15 0.569 0.357 106.99
DSG ICML 2024 26.80 0.709 0.290 87.96
ADMMDiff CVPR 2025 27.26 0.778 0.222 72.92
DiffRGD (本文) ECCV 2026 25.84 0.736 0.250 72.26

在人脸条件生成任务中(CelebA-HQ 256×256,100 步 DDIM),DiffRGD 在保持低 FID 的同时显著提高了条件对齐精度:

任务 / 指标 FreeDoM (ICCV 2023) DSG (ICML 2024) ADMMDiff (CVPR 2025) DiffRGD (本文)
分割图引导 mIoU ↑ 0.622 0.750 0.758 0.804
分割图引导 FID ↓ 156.02 117.48 101.86 96.10
线稿引导 Sketch-\(\ell_2\) ↓ 30.85 21.36 30.82 19.48
线稿引导 FID ↓ 101.90 107.00 97.52 87.82
身份引导 FaceID-\(\ell_2\) ↓ 0.557 0.340 0.346 0.303
身份引导 FID ↓ 127.05 95.27 100.81 93.80

消融与敏感度实验

在条件生成任务中,引导强度 \(\eta_t\) 的变化对生成稳定性具有决定性影响。下表记录了在 CelebA-HQ 分割图引导下,不同步长对 DSG 与 DiffRGD 的综合影响:

方法 引导强度 \(\eta_t\) 条件对齐 mIoU ↑ 生成质量 FID ↓ 现象与说明
DSG 0.05 0.53 95.69 小步长下条件控制不足
0.10 0.76 110.00 控制逐步增强,FID 出现恶化
0.20 0.90 157.82 生成严重偏离分布,伪影剧增
0.30 0.93 248.77 图像质量完全崩溃
DiffRGD (本文) 10 0.74 80.73 小步长即可保持高保真与强基础控制
30 0.82 89.03 平衡点,视觉效果自然且对齐精确
50 0.84 93.47 结构高度契合,FID 维持在极佳区间
100 0.86 108.33 极限步长下依然鲁棒,未发生流形崩溃

关键发现

  • 步数缩减下的鲁棒性提升:当 DDIM 步数从 1000 步压缩至 100 步时,DPS 与 ADMMDiff 因近端迭代不足或累积误差暴增出现明显性能退化;而 DiffRGD 在 100 步下仍取得优于基线的指标(例如 ImageNet 256×256 修复任务 PSNR 达 31.16 dB,大幅超越第二名 ADMMDiff 的 27.97 dB)。
  • 引导强度耐受范围极宽:如消融表所示,DSG 在 \(\eta_t\) 增大到 0.2 时 FID 便暴增至 157.82,发生严重的离流形崩溃;而 DiffRGD 得益于严格的球面流形约束与收缩拉回机制,在步长跨越数量级(\(\eta_t=10\) 到 \(100\))的变化中,FID 仅温和上升至 108.33,展现了卓越的优化稳定性。
  • 高倍率超分辨率的保真度权衡:基于 Stable Diffusion 进行高倍率人脸超分辨率实验时(8× 至 12×),12× 重建能恢复更多高频纹理细节,取得更高的结构相似度 SSIM(0.709 vs 0.700),但微观纹理与真值的局部差异导致像素级 PSNR 略有下降(26.39 dB vs 26.41 dB)。

亮点与洞察

  • 高斯统计特性与黎曼流形的理论契合:巧妙地借助各向同性高斯变量的极分解定理,将繁杂的高斯分布约束精确转化为固定半径超球面的黎曼流形,不仅化繁为简,而且理论证明闭环完备。
  • 正交投影与收缩映射的计算极简化:避开了 ADMM 等方法中繁复的双层迭代求解,用简单的解析向量内积投影和除以模长的径向重缩放,以 \(K=3\) 次极小开销实现了稳健收敛。
  • 广泛的即插即用迁移性:设计不依赖任何特定模型架构或潜在先验假设,可以直接无缝外挂到任何基于 DDIM 框架的连续去噪模型(包括像素级 DDPM 和 Latent Diffusion / Stable Diffusion)。

局限与展望

  • 各向同性高斯假设的约束边界:方法推导强依赖 DDIM 反向去噪中每一步边际分布为各向同性高斯的性质。若在采用非各向同性协方差矩阵或新型 Flow Matching(连续流匹配)采样器时,流形不再是简单超球,极分解形式需要重新进行几何推导。
  • 运动去模糊场景的次优表现:在高度非线性的强运动去模糊实验中,DiffRGD 的 PSNR(25.84 dB)略逊于基于交替迭代的 ADMMDiff(27.26 dB),反映出纯一阶黎曼梯度对严重非局部退化算子的全局逆向能力仍有提升空间。
  • 未来改进方向:可进一步探索自适应步长黎曼动量优化,或将该流形约束理论推广至扩散变换器(DiT)与流匹配(Flow Matching)架构。

相关工作与启发

  • vs DPS (ICLR 2023) / FreeDoM (ICCV 2023):DPS 和 FreeDoM 直接在环境欧几里得空间利用自动微分更新潜变量,毫无概率分布保护,步长敏感且极易破坏图像保真度;DiffRGD 引入球面流形正交投影和收缩映射,消除了梯度对高斯先验的破坏。
  • vs MPGD (ICLR 2024):MPGD 试图在数据流形上进行局部线性投影,但强行依赖不完美的预训练 VAE 编码器,极易引起数值崩溃;DiffRGD 扎根于反向去噪过程本身的统计流形,完全无需外置辅助网络。
  • vs DSG (ICML 2024):DSG 采用固定的经验球面投影,导致采样分布退化为均匀分布;DiffRGD 基于极分解准确引入卡方分布采样半径,严格保留了高斯的尺度分布特性。
  • vs ADMMDiff (CVPR 2025):ADMMDiff 构建了算子分裂的优化框架,但计算沉重且超参繁多;DiffRGD 仅用 \(K=3\) 步黎曼迭代即可快速收敛,时间成本与易用性大幅领先。

评分

  • 新颖性: ⭐⭐⭐⭐⭐ [首次利用各向同性高斯极分解性质将扩散引导规范为球面黎曼流形优化,视角优美且数学自洽]
  • 实验充分度: ⭐⭐⭐⭐⭐ [覆盖 4 种逆问题任务与 3 种条件生成任务,涵盖 FFHQ、ImageNet 及 Stable Diffusion 高分辨率实验]
  • 写作质量: ⭐⭐⭐⭐⭐ [理论推导严密、问题针对性极强,配图与算法流程清晰易懂]
  • 价值: ⭐⭐⭐⭐⭐ [作为通用即插即用的推理期引导框架,为解决扩散生成中的偏离流形伪影问题提供了优雅通用的范式]