[论文解读] Learning Diffusion Priors from Observations by Expectation Maximization

François Rozet, Gérôme Andry+2 more|arXiv (Cornell University)|May 22, 2024

Gaussian Processes and Bayesian Inference被引用 5

一句话总结

论文在使用 EM 框架将扩散模型作为贝叶斯反问题的先验，从不完整/嘈杂观测出发，介绍一种用于无条件扩散模型的矩量匹配后验采样方案。

ABSTRACT

Diffusion models recently proved to be remarkable priors for Bayesian inverse problems. However, training these models typically requires access to large amounts of clean data, which could prove difficult in some settings. In this work, we present DiEM, a novel method based on the expectation-maximization algorithm for training diffusion models from incomplete and noisy observations only. Unlike previous works, DiEM leads to proper diffusion models, which is crucial for downstream tasks. As part of our methods, we propose and motivate an improved posterior sampling scheme for unconditional diffusion models. We present empirical evidence supporting the effectiveness of our approach.

研究动机与目标

• 在仅有噪声观测可用时，为贝叶斯反问题学习先验的动机。
• 开发基于 EM 的管道，以从不完整数据中训练扩散模型。
• 确保得到的扩散先验是适当的、可用于后续贝叶斯任务的。
• 提供一种用于无条件扩散模型的新型后验采样方案，以提升采样质量。

提出的方法

• 在连续时间的 SDE 框架中对扩散模型进行表述，并采用去噪分数匹配目标。
• 使用蒙特卡洛 EM，通过迭代地从后验 q_theta(x|y) 采样并重新训练来学习先验 q_theta(x)。
• 定义正向模型 p(y|x,A)=N(y|Ax,Sigma_y) 以处理观测特定的正向过程。
• 引入矩量匹配后验采样（MMPS），其后验得分分解为 grad log p(x_t) + grad log p(y|x_t) 。
• 通过高斯近似估计似然分数，考虑 V[x|x_t] 协方差，使用 Tweedie 的协方差公式。
• 用 Tweedie 公式高效计算 V[x|x_t]，并用共轭梯度法求解得到的线性系统，而不需要构建大型矩阵。
• 用 EM 拟合的高斯分布初始化 q_0(x)，以加速收敛。
• 用去噪分数匹配训练扩散模型 q_theta(x)，其去噪网络为 d_theta(x_t,t)，如式（4）所示。

实验结果

研究问题

• RQ1 是否可以在 EM 框架下从不完整/嘈杂观测中学习扩散模型作为先验？
• RQ2 在经验贝叶斯设定中，MMPS 是否能为无条件扩散模型提供稳定且准确的后验采样？
• RQ3 基于EM的扩散先验在现实世界的被污染数据任务（如损坏的 CIFAR-10、加速 MRI）中的表现如何，与现有方法相比？
• RQ4 对后验协方差 V[x|x_t] 的准确建模对后验质量和下游推断有何影响？

主要发现

方法	已删除	FID ↓	IS ↑
0.4	0.4	18.85	7.45
AmbientDiffusion [77]	0.6	28.88	6.88
0.8	0.8	46.27	6.14
Ours w/ Tweedie	0.75	19.56	7.60
Ours w/ (I+Σ_t^{-1})^{-1}	0.75	49.45	6.99
Ours w/ Σ_t	0.75	125.58	3.98

• 基于 EM 的管道产生的扩散先验在受控实验下收敛到接近真实潜在分布的分布。
• 使用 Tweedie 协方差的矩量匹配后验采样（MMPS）比启发式协方差给出更准确的后验 q(x_t|y)，提高稳定性和覆盖度。
• 在损坏的 CIFAR-10 上，提出的方法配合 Tweedie 协方差在相当的损坏水平下可达到甚至超过 AmbientDiffusion；而启发式协方差表现较差。
• 在加速 MRI 中，扩散先验在高加速因子（最高 R=32）下能提供合理的重建，并给出详细的后验样本。
• 通过 EM 标识的高斯先验进行初始化可减少收敛迭代次数并提升效率。
• 使用共轭梯度求解器处理 A V[x|x_t] A^T 在高维下保持计算可控，通常仅需1-3次 CG 循环即获得收益。