用于解决线性逆问题的精确后验分数估计

Exact Posterior Score Estimation for Solving Linear Inverse Problems

arXiv: 2606.17048v1

论文信息

标题: Exact Posterior Score Estimation for Solving Linear Inverse Problems

作者: Abbas Mammadov, Ozgur Kara, Kaan Oktay, et al.

发布日期: 2026-06-15

arXiv ID: 2606.17048v1

PDF 链接: 下载 PDF

3 分钟速览

  • 研究问题：如何利用预训练扩散/流模型的强数据先验，在线性高斯逆问题（修复、超分、去模糊等）中精确采样后验样本，而不是依赖近似的测量匹配修正。
  • 核心方法：将后验分数推导为在测量依赖的 “后验枢轴”μ⋆\mu_\star 和算符相关的各向异性噪声协方差 Σ⋆\Sigma_\star 下的去噪问题，并直接训练一个去噪器 Dθ(μ⋆,y,t)\mathcal D_\theta(\mu_\star,y,t) 来拟合 E[x0∣xt,y]\mathbb E[\mathbf x_0|\mathbf x_t,\mathbf y]，保持与标准预训练相同的输入输出结构。
  • 关键结果：EPS 在 FFHQ 与 ImageNet 的五个线性逆问题上，以仅 20 次去噪器评估（NFE） 就达到了点保真度、感知质量和分布校准指标的最好或次好水平，远低于梯度基后验采样器的 100–250 NFE（表 1、图 3）。
  • 主要局限：推导假定线性前向算子和高斯观测噪声，对非线性算子需线性化或特殊训练；应用于潜在扩散模型时，像素空间的线性算子经过解码器会变成潜在空间的非线性变换。
  • 适合读者：从事扩散模型、计算成像、医学图像重建、贝叶斯逆问题的研究者，以及需要高效且校准良好的后验采样的工程师。

论文背景和研究动机

线性逆问题（压缩感知、加速 MRI、超分辨率、去模糊等）通常需要从带噪声的测量 y=Ax0+η\mathbf y = A\mathbf x_0 + \boldsymbol\eta 中恢复未知信号 x0\mathbf x_0。由于 AA 往往是秩亏或病态的，许多信号与同一观测相容，因此目标不是点估计，而是从后验 p(x0∣y)p(\mathbf x_0|\mathbf y) 中采样。扩散模型和流模型通过学习一条从噪声到干净样本的去噪轨迹，为这一问题提供了强大的数据先验。然而，反向采样需要的是后验分数 ∇xtlog⁡p(xt∣y)\nabla_{\mathbf x_t}\log p(\mathbf x_t|\mathbf y)，而预训练给出的只是无条件分数。

现有方法分两类：免训练方法（如 DPS、DDNM、ΠGDM）通过在每个去噪步骤添加一个测量匹配修正来近似后验，但这些修正本质上是将无条件去噪器在当前状态 xt\mathbf x_t 下查询，而不是按照精确后验在 “枢轴”μ⋆\mu_\star 处查询，从而引入了偏差；训练基方法（如 Palette、条件扩散、桥方法）直接训练一个以测量 y\mathbf y 为输入的条件模型，但网络需要端到端学习算符依赖，无法利用预训练去噪器的结构。本文的贡献在于：推导出线性高斯逆问题的精确后验分数闭式，进而设计出一种既保持预训练输入/输出结构，又能将后验采样完全还原为去噪任务的方法 EPS。

核心方法和技术细节

论文的关键观察是：对于插值 xt=αtx0+βtϵx_t = \alpha_t x_0 + \beta_t \epsilon 和线性高斯观测模型 y=Ax0+η, η∼N(0,σy2I)y = Ax_0 + \eta,\ \eta\sim\mathcal N(0,\sigma_y^2 I)，后验边缘 p(xt∣y)p(x_t|y) 可通过高斯乘积（completing the square）写成一个新高斯噪声下的数据密度（定理 1）：

∇xtlog⁡p(xt∣y)=1βt2(αt DΣ⋆(t)(μ⋆(xt,y,t))−xt),\nabla_{x_t}\log p(x_t|y) = \frac{1}{\beta_t^2}\Bigl(\alpha_t\,D_{\Sigma_\star(t)}\bigl(\mu_\star(x_t,y,t)\bigr) - x_t\Bigr),

其中

Σ⋆(t)=(αt2βt2I+1σy2A⊤A)−1,μ⋆(xt,y,t)=Σ⋆(t)(αtβt2xt+1σy2A⊤y).\Sigma_\star(t) = \Bigl(\frac{\alpha_t^2}{\beta_t^2}I + \frac{1}{\sigma_y^2}A^\top A\Bigr)^{-1},\quad \mu_\star(x_t,y,t) = \Sigma_\star(t)\Bigl(\frac{\alpha_t}{\beta_t^2}x_t + \frac{1}{\sigma_y^2}A^\top y\Bigr).

这里 μ⋆\mu_\star 被称为后验枢轴，它是当前状态 xtx_t 和测量 yy 的高斯贝叶斯融合，Σ⋆\Sigma_\star 则是相应的后验协方差，在 AA 的零空间方向保留大方差，在观测方向上收缩。DΣ⋆(μ⋆)=E[x0∣xt,y]D_{\Sigma_\star}(\mu_\star) = \mathbb E[x_0|x_t,y] 是对该各向异性噪声的最优去噪器。换言之，后验采样就是在一个由测量决定的输入 μ⋆\mu_\star 下进行去噪，但噪声不再是各向同性的 βtI\beta_t I，而是算符依赖的 Σ⋆\Sigma_\star。

由此，论文定义 EPS 训练目标：

LEPS(θ)=Ex0,t,ϵ,η[w(t) ∥Dθ(μ⋆(xt,y,t), y, t)−x0∥2].\mathcal L_{\text{EPS}}(\theta) = \mathbb E_{x_0,t,\epsilon,\eta}\bigl[w(t)\,\|D_\theta(\mu_\star(x_t,y,t),\, y,\, t) - x_0\|^2\bigr].

这一目标在结构上与标准扩散预训练完全相同（干净目标 x0x_0，均方误差），唯一变化是网络输入从 xtx_t 换成了 μ⋆\mu_\star，并附加了观测 yy 作为条件。根据命题 3，μ⋆∣x0\mu_\star|_{x_0} 的分布恰好是 N(x0,Σ⋆(t))\mathcal N(x_0,\Sigma_\star(t))，因此训练时无需显式采样各向异性噪声，只需通过枢轴构造即可获得所需噪声。训练后，EPS 直接使用底层骨干（如 EDM）的原有采样器，每个采样步只需计算一次 μ⋆\mu_\star 的前向运算（对常见算子仅为逐元素、池化或 FFT 求解，开销远小于去噪器），无需计算似然梯度或投影。在高噪声极限（σt→∞\sigma_t\to\infty），μ⋆\mu_\star 收敛到 A†yA^\dagger y 加上 AA 零空间上的纯噪声，一次去噪器调用即近似后验均值（观察 4），这也是表 1 中 NFE=1 的 EPS 能获得高 PSNR 的原因。

论文进一步指出，现有免训练方法（如 DPS、ΠGDM 等）的近似误差本质在于：它们用无条件去噪器在 xtx_t 处的查询近似 DΣ⋆(μ⋆)D_{\Sigma_\star}(\mu_\star)，而未在正确的输入 μ⋆\mu_\star 下考虑各向异性噪声 Σ⋆\Sigma_\star。因为测量匹配分数（式 14）恰等于两个去噪器之差：后验去噪器在 μ⋆\mu_\star 的估计减去无条件去噪器在 xtx_t 的估计；免训练方法忽略了这一差异，EPS 通过训练直接学习了后项。

创新点和贡献

  1. 闭式后验分数与几何解释：首次明确给出一般高斯插值下线性高斯逆问题的精确后验分数，并用 “枢轴 + 各向异性噪声” 统一理解后验去噪的查询几何。
  2. 保留预训练结构的训练范式：EPS 训练目标与标准去噪预训练同构，因此可以从预训练检查点高效微调（Appendix D.3 显示 EPS 收敛速度远快于同等条件的 Palette），甚至随机初始化后训练。
  3. 无梯度的采样器：采样过程不需要任何测量损失函数的梯度、投影或内层优化，只要完成对 μ⋆\mu_\star 一次线性求解（开销相对于去噪器一前向可忽略，附录 D.12 测量为 <1 ms）。因此 EPS 能复用底层模型的所有采样技巧。
  4. 极低 NFE 下高性能：实验表明，EPS 在约 20 NFE 时各项指标即平坦（图 3、5、6），而免训练的 DPS、ΠGDM 等 100 NFE 仍不能达到同一水平，训练基 Palette 也需更多计算才能持平，这使 EPS 成为低采样预算下的高效选择。

实验结果分析

论文在 FFHQ‑64 和 ImageNet‑64 上评估了五种任务：随机修复（70% 缺失）、方块修复、4× 超分辨率、高斯去模糊和运动去模糊，所有实验均添加 σy=0.05\sigma_y=0.05 的观测噪声。指标涵盖 PSNR、SSIM、LPIPS、FID 以及分布校准的 CRPS 和 MMD（像素空间和 Inception 特征空间）。

表 1 与表 6 显示，在 ImageNet 和 FFHQ 上，EPS（100 NFE 和 20 NFE）在几乎所有指标上位列第一或第二，在随机修复和超分辨率这类秩亏严重的任务上优势更明显。例如，ImageNet 随机修复中 EPS‑20 获得 PSNR 24.87、FID 77.06，显著优于 ΠGDM 的 23.95 / 99.60 和 Palette 的 24.09 / 81.88（表 1）。图 3 展示的采样步数敏感性表明，EPS 很快收敛，而 DPS、DAPS、DDNM、ΠGDM 等即使加到 100 NFE 仍追不上 EPS 的渐近水平。100 NFE 的 Palette 与 EPS 的对比尤其说明：仅仅将 xtx_t 替换为 μ⋆\mu_\star 并微调就带来了持续的提升（附录 D.1 输入配置消融表明，从 Palette 到 EPS 的升级是单调收益）。此外，NFE=1 的单步后验均值估计（Tweedie 调用）得到了最高 PSNR/SSIM，但损失了感知与分布指标，这符合感知-失真权衡。

定性结果（图 2、7、8）显示，EPS 在修复和去模糊任务中能恢复清晰的纹理与边缘，同时保持与观测一致性；而免训练基线在未观测区域常出现过度平滑或伪影。在 95% 极端掩码和 16× 超分辨率的额外实验中（表 7、图 9），EPS 仍然在感知与分布指标上大幅领先，且能生成语义上多样但都符合观测的后验样本，证明了 EPS 的校准能力。

此外，论文还验证了 EPS 的多个实用特性：可以在一个模型中摊销所有五个任务而不损失性能（表 4），对掩码密度的外推也表现出一定程度的鲁棒性（表 8），在 256×256 分辨率下依旧有效（表 9）。

实践建议

EPS 非常自然地可被集成到现有的扩散/流模型工作流中，尤其是那些已经拥有预训练无条件去噪器（如 EDM）的场景。实际部署时：

  • 微调策略：保留预训练去噪器的架构，仅把输入通道从 xtx_t (3 通道) 扩展为 [μ⋆,y][\mu_\star, y]（通常 6 通道，图像问题中 yy 可经掩码或上采样对齐尺寸），额外卷积核零初始化，然后以标准 EDM 损失权重微调 10 个 epoch 左右（约 25k 次迭代）即可收敛（附录 B）。这比从头训练条件模型节省巨大算力。
  • 采样配置：直接采用 EDM 的 Euler 或 Heun 采样器，推荐 20 NFE，此时性能已接近饱和。每步需计算 μ⋆\mu_\star：对修复（逐元素求解）、超分（池化-反池化递推）、卷积模糊（FFT 对角化）都有极高效的实现；若遇到一般稠密 AA，可用共轭梯度求解对称正定系 (αt2βt2I+1σy2A⊤A)μ⋆=αtβt2xt+1σy2A⊤y\bigl(\frac{\alpha_t^2}{\beta_t^2} I + \frac{1}{\sigma_y^2}A^\top A\bigr)\mu_\star = \frac{\alpha_t}{\beta_t^2}x_t + \frac{1}{\sigma_y^2}A^\top y，每次迭代仅需 AA 和 A⊤A^\top 的矩阵-向量乘积。
  • 摊销多任务：若需支持多个逆问题（如不同掩码或模糊核），建议训练一个统一的 EPS 模型，在训练时随机采样算子，文中已证实摊销版本性能几乎不降，且大幅减少部署模型数量。
  • 潜在扩散模型：若使用 VAE 潜在空间中的扩散先验，由于解码器使像素空间的线性算子变成非线性的，EPS 的封闭形式不再直接适用；一种策略是在潜在空间内对解码器的雅可比进行局部线性近似，或训练时包含解码器的梯度，但会增加复杂度。这类场景应谨慎验证。

总之，EPS 为线性高斯逆问题提供了一种精确、高效、易集成的后验采样方案，特别适合对推断准确性和计算预算都有严格要求的应用。