Existing diffusion-based methods for inverse problems sample from the posterior using score functions and accept the generated random samples as solutions. In applications that posterior mean is preferred, we have to generate multiple samples from the posterior which is time-consuming. In this work, by analyzing the probability density evolution of the conditional reverse diffusion process, we prove that the posterior mean can be achieved by tracking the mean of each reverse diffusion step. Based on that, we establish a framework termed reverse mean propagation (RMP) that targets the posterior mean directly. We show that RMP can be implemented by solving a variational inference problem, which can be further decomposed as minimizing a reverse KL divergence at each reverse step. We further develop an algorithm that optimizes the reverse KL divergence with natural gradient descent using score functions and propagates the mean at each reverse step. Experiments demonstrate the validity of the theory of our framework and show that our algorithm outperforms state-of-the-art algorithms on reconstruction performance with lower computational complexity in various inverse problems.
翻译:现有的基于扩散模型的逆问题求解方法利用分数函数从后验分布中采样,并将生成的随机样本作为解。在需要后验均值的应用中,必须从后验中生成多个样本,这一过程非常耗时。本文通过分析条件反向扩散过程的概率密度演化,证明了后验均值可以通过追踪每个反向扩散步骤的均值来实现。基于此,我们建立了一个称为反向均值传播的框架,该框架直接以后验均值为目标。我们证明RMP可以通过求解一个变分推断问题来实现,该问题可进一步分解为在每个反向步骤中最小化反向KL散度。我们进一步提出了一种算法,该算法利用分数函数通过自然梯度下降优化反向KL散度,并在每个反向步骤中传播均值。实验验证了我们框架的理论有效性,并表明我们的算法在多种逆问题中,以更低的计算复杂度在重建性能上优于当前最先进的算法。