Stochastic Optimal Control Sampling for Diffusion Inverse Problems¶
会议: ECCV 2026
论文: ECCV 原文
代码: https://github.com/zjqwq01/SOCS-DIP
领域: 图像恢复
关键词: 扩散模型, 反问题, 随机最优控制, 朗之万采样, 流匹配
一句话总结¶
将扩散反问题的采样轨迹建模为由受控动力学系统引导的过程，推导单步随机最优控制显式回拉闭式解并引入有界终端惩罚能量采样，在免重训练下规避全轨迹高阶二阶导计算并实现保真度与生成先验的自适应权衡。
研究背景与动机¶
基于扩散模型的无监督图像反问题求解在近年来成为计算成像与图像恢复领域的核心范式。自然图像退化算子往往具备不可逆的多对一甚至高度非线性特征，使得从退化观测中逆向重建清晰图像天然具有严重的病态性与解的不唯一性。预训练扩散模型学习到了丰富的高阶自然图像结构先验与流形几何分布，通过在逆向去噪轨迹中注入测量一致性引导，能够在不依赖特定任务成对数据重新训练的前提下，实现高质量的后验采样。
然而，现有训练无关的引导扩散方法面临两难困境。一方面，主流基于贝叶斯近似的后验采样算法（如 DPS、DAPS）通常在经验上引入二次数据拟合项的梯度来引导去噪，隐式假定终端观测惩罚趋于无穷大；这种激进的似然注入极易在去噪早中期将中间状态强行拉离扩散先验流形，导致严重的人工伪影或过平滑细节。另一方面，现有尝试将随机最优控制引入扩散反问题的方法（如 DOC）主要依赖经典 iLQR 框架对整条扩散轨迹进行滚动优化，为了更新控制策略必须对扩散网络反传计算高阶 Hessian 矩阵，带来极高的显存占用与计算延迟，完全无法拓展至现代大分辨率模型。
这种矛盾的本质在于：无控制的扩散模型本身已具备极强的流形生成能力，反问题引导的核心目标是施加能量最小的有界修正，引导轨迹落入满足观测约束的目标集合，而非全盘推翻生成先验。因此，亟需一种既具备最优控制理论严密保障、又能完全解耦为单步显式更新的轻量采样机制。本文的核心 idea 是：将逆向去噪动力学构建为带有界终端测量惩罚的随机最优控制问题，基于确定性等价原理推导出单步控制与状态传输的显式解析闭式解，在每一步通过朗之万能量采样与显式回拉将观测一致性估计稳健映射回去噪流形，彻底消除轨迹级 Hessian 计算。
方法详解¶
整体框架¶
SOCS 将逆向去噪过程建模为一个线性随机微分方程系统，通过在扩散漂移中注入控制向量 \(u_{t,\gamma}\)，使终端状态 \(x_T\) 最小化与测量 \(y\) 之间的欧氏残差 \(\frac{\gamma}{2}\|y - \mathcal{H}(x_T)\|_2^2\)，同时最小化控制能量 \(\int_0^T \frac{1}{2}\|u_{t,\gamma}\|_2^2 dt\)。在每个离散外层去噪步 \(t_i\)，算法首先通过少步 Euler 求解器沿概率流常微分方程快速推演至终端得到粗粒度干净图像估计 \(x_T|t_i\)；随后依据算子特性选择一致性对齐策略——在线性良态任务中执行闭式仿射投影，在复杂非线性或病态任务中通过朗之万动力学执行有界能量优化；最后代入推导得出的解析传输方程，将修正后的估计直接闭式回拉至下一个去噪状态 \(x_{t_{i+1}}\)。
%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
A["输入退化观测 y 与初始先验噪声 x_t0"] --> B["阶段预演：少步 Euler ODE 求解器<br/>快速外推获得粗粒度干净估计 x_T|ti"]
B --> C{"退化算子解耦与双模态采样适配<br/>判断退化算子局部线度与适定条件"}
C -->|局部线性且良态| D["闭式仿射子空间投影<br/>直接映射至测量一致集合"]
C -->|复杂非线性或严重病态| E["有界终端惩罚调制与能量采样<br/>朗之万动力学数据一致性能量微调"]
D --> F["单步随机最优控制闭式回拉<br/>状态传输解析解将估计拉回扩散流形"]
E --> F
F -->|递推离散逆向步至终止| G["输出高保真且测量对齐的重建图像 x_N"]关键设计¶
1. 退化算子解耦与双模态采样适配：兼顾高效线性投影与复杂生成范式 针对传统采样方法难以统一处理差异巨大的退化算子这一痛点，SOCS 引入了针对算子特性的解耦分支策略，并将其泛化至潜在扩散模型与流匹配架构中。在线性且良态的退化场景下（如超分辨率），前向算子雅可比矩阵满足 \(J_\mathcal{H} \approx \mathcal{H}\)，算法令终端惩罚系数 \(\gamma \to \infty\)，直接通过伪逆投影消除迭代内层开销： $\(\hat{x}_T(x_t) = (I - \mathcal{H}^\dagger \mathcal{H}) x_T(x_t) + \mathcal{H}^\dagger(y)\)$ 该式使得重建无偏向地保持测量一致子空间的精确取值，同时将零空间交给扩散先验生成。对于潜在扩散模型（如 SD v1.5）和以 SD3-medium 为代表的流匹配模型，由于整流流在理想直线下表现为恒定速度线性 ODE，SOCS 将其状态空间对齐为线性动力学形式，从而无需修改网络即可即插即用适配高分辨率隐空间反问题求解。
2. 有界终端惩罚调制与能量采样：保留扩散先验几何的显式测量约束 常规贝叶斯采样方法通常默认 \(\gamma \to \infty\)，等价于强行施加无限大的终端约束惩罚，在严重病态或非线性任务（如相位恢复、非线性盲去模糊）中会导致采样轨迹严重畸变并产生过度平滑伪影。SOCS 明确保留了有限大小的终端惩罚系数 \(\gamma\)，将逆向控制目标转化为带有权衡的能量极小化问题。在非线性设定下，最优控制律导出的高斯-牛顿正则化正规矩阵求解转化为隐式能量极小化形式： $\(\hat{x}_T^{j+1} = \hat{x}_T^j - \eta \nabla_{\hat{x}_T^j} \left\| \left(I + \gamma(1 - e^{2\bar{f}_T}) J_\mathcal{H}(x_T^u)^\top\right) \circ \mathcal{H}(\hat{x}_T^j) - J_\mathcal{H}(x_T^u)^\top (y - e^{\bar{f}_T} \mathcal{H}(x_0)) \right\|_2^2 + \sqrt{2\eta}\epsilon_j\)$ 通过朗之万随机梯度动力学在内层展开微调，有效避免了显式矩阵求逆。有界 \(\gamma\) 的引入在数学上构筑了自适应正则化机制，既约束当前解逼近观测数据，又为扩散先验的固有结构保留了自由演化的几何空间。
3. 单步随机最优控制闭式回拉：规避全轨迹二阶 Hessian 计算的显式动力学 现有基于控制论的反问题方法需要维护整条轨迹上的动态优化，依赖二阶泰勒展开与 Hessian 矩阵滚动计算，造成严重的显存瓶颈与时间消耗。SOCS 证明了在受控线性 SDE \(dx_t = (f_t x_t + g_t u_{t,\gamma})dt + g_t dw_t\) 与二次终端代价下，由确定性等价原理可知其最优控制器与确定性控制系统一致。作者进一步在解析上推导出了最优控制律 \(u_{t,\gamma}^*\) 与状态演化 \(x_t\) 的闭式解： $\(x_t = e^{\bar{f}_t} x_0 + \gamma e^{\bar{f}_t} e^{\bar{f}_T} \bar{g}_t^2 d_{0,\gamma} J_\mathcal{H}(x_T^u)^\top \left(y - e^{\bar{f}_T}\mathcal{H}(x_0)\right)\)$ 其中 \(d_{0,\gamma} = (I + \gamma e^{2\bar{f}_T}\bar{g}_T^2 J_\mathcal{H}^\top \circ \mathcal{H})^{-1}\)。该闭式传输算子在每一步直接将经一致性校准的干净状态 \(\hat{x}_T\) 映射回逆向 SDE 去噪流形。整个过程仅需一阶雅可比向量积，彻底摆脱了整条轨迹的二次滚动优化与 Hessian 矩阵存储，使控制论方法的耗时降低至常规采样级水平。
一个完整示例¶
以带有高斯噪声的 \(4\times\) 双三次超分辨率（线性）与相位恢复（非线性）为例：
1. 初始化：输入带噪退化观测 \(y\)（噪声标准差 \(\beta_y = 0.05\)），从标准高斯噪声采样 \(x_{t_0} \sim \mathcal{N}(0, I)\)。
2. 状态前向预估：在第 \(i\) 步，扩散网络利用当前状态 \(x_{t_i}\) 执行 4 步 Euler ODE 步，向前积分得到粗粒度自然图像估计 \(x_T|t_i\)。
3. 一致性修正：
- 超分辨率任务：进入 SOCS-Linear 分支，直接利用闭式算子计算 \(\hat{x}_T|t_i = (I - \mathcal{H}^\dagger \mathcal{H}) x_T|t_i + \mathcal{H}^\dagger y\)，瞬时完成投影。
- 相位恢复任务：由于幅值观测丢失了相位且不存在有界线性伪逆，进入 SOCS-Nonlinear-γ 分支，固定超参 \(\gamma = 10^7\)，运行内层 \(N_{\text{inner}}\) 步朗之万动力学迭代更新 \(\hat{x}_T|t_i\)。
4. 单步流形回拉：将修正后的 \(\hat{x}_T|t_i\) 代入闭式状态转移方程，求得下一个逆向采样时间节点状态 \(x_{t_{i+1}}\)。
5. 循环与收敛：重复该解耦回拉过程 250 至 500 步，最终平滑收敛至与观测严格对齐的高分辨率真实图像。
实验关键数据¶
主实验¶
论文在 FFHQ 256×256 与 ImageNet 256×256 验证集（各 100 张图像）上评测了 6 项线性任务与 3 项非线性任务，评价指标涵盖 PSNR、SSIM、LPIPS 与 FID。同时在 768×768 高分辨率下采用 SD3-medium 进行了评估。
| 数据集与任务 | 评测指标 | SOCS-Linear (本文) | SOCS-Nonlinear-γ (本文) | DAPS (基线) | DPS (基线) | DDRM / DDNM (经典) |
|---|---|---|---|---|---|---|
| FFHQ: 超分辨率 4× | PSNR / SSIM / FID | 28.51 / 0.797 / 72.99 | 28.66 / 0.800 / 76.86 | 28.31 / 0.768 / 74.75 | 24.25 / 0.677 / 84.33 | 27.53 / 0.796 / 92.19 (DDRM) |
| FFHQ: 随机掩码补全 | PSNR / SSIM / FID | 29.31 / 0.834 / 77.09 | 30.63 / 0.872 / 52.73 | 29.66 / 0.800 / 54.57 | 28.30 / 0.812 / 71.74 | 29.17 / 0.850 / 64.65 (DDNM) |
| FFHQ: 方块补全 (Box) | PSNR / SSIM / FID | 24.22 / 0.794 / 53.21 | 24.45 / 0.826 / 43.79 | 24.41 / 0.753 / 51.51 | 23.49 / 0.804 / 62.18 | 24.14 / 0.837 / 48.96 (DDNM) |
| ImageNet: 超分辨率 4× | PSNR / SSIM / FID | 26.42 / 0.733 / 93.77 | 26.36 / 0.708 / 97.54 | 25.95 / 0.667 / 112.60 | 22.56 / 0.574 / 155.01 | 25.02 / 0.694 / 158.35 (DDNM) |
| ImageNet: 高斯去模糊 | PSNR / SSIM / FID | 25.38 / 0.720 / 111.49 | 26.67 / 0.740 / 110.69 | 26.37 / 0.697 / 109.61 | 22.84 / 0.587 / 115.02 | 27.81 / 0.639 / 124.55 (DDRM) |
| FFHQ: 非线性去模糊 | PSNR / SSIM / FID | 21.73 / 0.647 / 113.90 | 28.42 / 0.786 / 68.20 | 28.10 / 0.762 / 69.45 | 23.05 / 0.637 / 90.69 | 29.16 / 0.783 / 63.22 (RED-diff) |
| FFHQ: 相位恢复 | PSNR / SSIM / FID | 不适用 (无线性逆) | 29.59 / 0.808 / 44.86 | 29.88 / 0.807 / 45.33 | 16.29 / 0.442 / 137.81 | 16.83 / 0.521 / 209.75 (RED-diff) |
消融实验¶
下表整理了高分辨率（768×768）流匹配架构 SD3-medium 在超分辨率 ×12 任务下的对比消融，以及终端控制惩罚系数 \(\gamma\) 对生成质量与保真度的敏感性分析。
| 模型与超参配置 | 评测任务 | PSNR (↑) | SSIM (↑) | LPIPS (↓) | FID (↓) | 机制说明 |
|---|---|---|---|---|---|---|
| FlowDPS | SD3 ×12 (Bicubic) | 29.73 | 0.775 | 0.482 | 61.24 | 传统向量场导引，感知质量受限 |
| FlowChef | SD3 ×12 (Bicubic) | 29.56 | 0.846 | 0.346 | 152.00 | 控制过度脱离先验，FID 显著恶化 |
| LatentDAPS | SD3 ×12 (Bicubic) | 29.65 | 0.819 | 0.276 | 65.41 | 解耦退火但在隐空间缺乏最优动力学回拉 |
| SOCS-SD3 (本文) | SD3 ×12 (Bicubic) | 30.45 | 0.842 | 0.248 | 45.72 | 最优控制引导使保真度与感知质量全面领先 |
| SOCS-Nonlinear (\(\gamma \to \infty\)) | 图像反问题通用 | 略微下降 | 略微下降 | 升高 | 升高 | 惩罚过大导致强行拟合噪声，细节过平滑 |
| SOCS-Nonlinear-γ (\(\gamma = 10^7\)) | 图像反问题通用 | 最优 | 最优 | 最低 | 最低 | 自适应饱和截断，平衡流形先验与数据拟合 |
| SOCS-Nonlinear-γ (\(\gamma = 10^3\)) | 图像反问题通用 | 明显下降 | 明显下降 | 显著升高 | 明显上升 | 引导强度不足，无法约束测量一致性 |
关键发现¶
- 控制强度具有稳定宽平台区：\(\gamma\) 从 \(10^5\) 变化至 \(10^9\) 时，各项性能指标均处于宽平的高性能平台期；当 \(\gamma \le 10^3\) 时测量约束缺失，而当 \(\gamma \to \infty\) 时样本被拉离生成流形出现过度平滑，验证了有界惩罚设计的必要性。
- 算子适定性决定变体选择：在线性良态任务中，SOCS-Linear 依靠稳定的闭式仿射投影即可兼顾计算效率与高质量重建；在严重病态、大零空间或复杂非线性任务中，SOCS-Nonlinear-\(\gamma\) 通过保留雅可比项与朗之万微调，大幅超越线性投影基线。
- 计算开销低且收敛轨迹直接：得益于闭式解推导，SOCS 摆脱了整条轨迹优化，在相同偏离预算（cumulative deviation budget）下相比 DAPS 实现了更低的重构误差与更短的收敛路径。
亮点与洞察¶
- 从随机最优控制视角为后验采样经验梯度立论：首次从严密的最优控制理论证明了现有贝叶斯反问题方法中广泛使用的测量导引项 \(\nabla_{x_t} \|H(x_t) - y\|_2^2\) 本质上是终端惩罚 \(\gamma \to \infty\) 下的最优控制特例，统一了解释框架。
- 将轨迹级 iLQR 控制降维至单步闭式流形传输：基于确定性等价原理推导出解析控制律与闭式回拉公式，规避了全轨迹滚动求解与二阶 Hessian 矩阵计算，使控制论引导的复杂度降低至普通一阶采样水准。
- 无缝兼容整流流匹配等现代生成架构：将最优控制框架自 SDE 扩散系统推广至流匹配（Flow-Matching）连续 ODE 轨迹，在 SD3 等大分辨率生成模型上证明了即插即用的优异恢复能力。
局限与展望¶
- 内层朗之万步数依赖经验设定：在处理高度非线性任务时，仍需迭代运行少步内层朗之万优化，增加了函数评估次数（NFE）与时间延迟。
- 雅可比矩阵依赖局部线性假定：在高阶复杂非线性模型中，局部高斯-牛顿线性化在初始高噪声阶段可能与真实流形存在误差偏移，未来可探索自适应步长与曲率自校正机制。
- 拓展至盲反问题与非高斯噪声：当前推导严格基于加性高斯测量噪声假设与已知退化算子，针对未知退化核、泊松噪声等更复杂成像系统的控制建模有待进一步完善。
相关工作与启发¶
- vs DPS / DAPS: DPS 依赖高斯近似经验梯度引导，DAPS 引入轨迹解耦但在反向传输时未显式约束控制能量；本文从随机最优控制推导出显式传输闭式解，且在有限 \(\gamma\) 下有效抑制了过度引导导致的流形漂移。
- vs DOC / iLQR-based 控制方法: DOC 采用整条轨迹的迭代线性二次型调节器（iLQR），需计算并存储庞大的二阶 Hessian 矩阵；本文通过单步分解与解析解推导，仅使用一阶雅可比向量积即可达到相同控制目标，显存占用与计算耗时呈数量级下降。
- vs DDRM / DDNM: DDRM 与 DDNM 局限于线性可求伪逆的逆问题；本文不仅在线性任务中退化为高效的仿射投影，更通用于相位恢复、非线性去模糊等完全不具备线性逆映射的非线性反问题。
评分¶
- 新颖性: ⭐⭐⭐⭐⭐ 首次建立离散单步随机最优控制闭式解析解与扩散反问题后验采样的严密联系。
- 实验充分度: ⭐⭐⭐⭐⭐ 跨越 VE/VP-SDE、LDM 与 SD3-Flow-Matching，全面覆盖 9 种线性与非线性任务。
- 写作质量: ⭐⭐⭐⭐⭐ 数学推导严密自洽，结构清晰，理论与实验结合紧密。
- 价值: ⭐⭐⭐⭐⭐ 为免训练扩散反问题提供了一套理论完备、开销极低且即插即用的控制采样新基准。