[论文解读] Solving Inverse Problems by Joint Posterior Maximization with a VAE Prior

Mario González, Andrés Almansa+3 more|HAL (Le Centre pour la Communication Scientifique Directe)|Nov 14, 2019

Image and Signal Denoising Methods参考文献 35被引用 6

一句话总结

本文提出 JPMAP，一种通过联合最大化图像与潜在空间的后验分布来解决成像中病态逆问题的新方法，利用变分自编码器（VAE）作为先验。通过在联合（图像-潜在）空间中建模问题，JPMAP 实现了具有收敛性保证的交替优化，并在重建质量上优于常易陷入不良局部最优的非凸 MAP 方法。

ABSTRACT

In this paper we address the problem of solving ill-posed inverse problems in imaging where the prior is a neural generative model. Specifically we consider the decoupled case where the prior is trained once and can be reused for many different log-concave degradation models without retraining. Whereas previous MAP-based approaches to this problem lead to highly non-convex optimization algorithms, our approach computes the joint (space-latent) MAP that naturally leads to alternate optimization algorithms and to the use of a stochastic encoder to accelerate computations. The resulting technique is called JPMAP because it performs Joint Posterior Maximization using an Autoencoding Prior. We show theoretical and experimental evidence that the proposed objective function is quite close to bi-convex. Indeed it satisfies a weak bi-convexity property which is sufficient to guarantee that our optimization scheme converges to a stationary point. Experimental results also show the higher quality of the solutions obtained by our JPMAP approach with respect to other non-convex MAP approaches which more often get stuck in spurious local optima.

研究动机与目标

• 解决成像中先验为学习的神经生成模型的病态逆问题挑战。
• 克服现有基于 MAP 的方法导致高度非凸优化并易陷入虚假局部最优的局限性。
• 将先验训练与退化模型解耦，实现预训练 VAE 先验在多个逆问题中无需微调即可复用。
• 通过利用联合目标的弱双凸性，为非凸问题设计具有收敛性保证的理论基础优化方案。
• 通过在联合图像-潜在空间中采用随机编码器与交替优化，提升重建质量。

提出的方法

• 将逆问题表述为在图像 ${\bm{x}}$ 与潜在码 ${\bm{z}}$ 上的联合后验最大化，利用 VAE 定义先验 $p_X({\bm{x}}) = \text{G}(Z)$，其中 $Z \sim \mathcal{N}(0,I)$。
• 定义联合目标函数 $J_1({\bm{x}}, {\bm{z}}) = F({\bm{x}}, {\bm{y}}) + G({\bm{x}}) + \| \mathsf{E}({\bm{x}}) - {\bm{z}} \|^2$，其中 $F$ 为数据保真项，$G$ 为 VAE 重建损失。
• 采用交替优化：固定 ${\bm{z}}$ 时通过梯度下降更新 ${\bm{x}}$ 以优化 $J_1({\bm{x}}, {\bm{z}})$，随后固定 ${\bm{x}}$ 时通过 VAE 编码器更新 ${\bm{z}}$。
• 在 ${\bm{z}}$-更新步骤中引入随机编码器，提升推理速度，增强计算效率。
• 利用目标函数的弱双凸性，在较弱假设下证明其收敛至驻点。
• 通过假设 VAE 编码器的性质及数据保真项的凸性，确保部分最小化解的唯一性。

(a) Reconstruction examples

实验结果

研究问题

• RQ1 在逆成像中，对图像与潜在空间进行联合后验最大化是否能带来优于标准 MAP 方法的收敛性与重建质量？
• RQ2 使用预训练 VAE 作为先验是否可在不重新训练的情况下实现对多种退化模型的有效泛化？
• RQ3 尽管整体问题具有非凸性，该联合优化框架能否实现收敛性保证？
• RQ4 与现有插件式（plug-and-play）及端到端深度学习方法相比，该方法在重建保真度与鲁棒性方面表现如何？
• RQ5 弱双凸性对优化算法的稳定性与收敛性有何影响？

主要发现

• 所提出的 JPMAP 方法在重建质量上优于其他非凸 MAP 方法，后者常收敛至较差的局部最优解。
• 联合后验最大化框架产生弱双凸目标函数，在较弱条件下可保证收敛至驻点。
• 在潜在变量更新步骤中使用随机编码器，显著加速计算，同时保持重建质量。
• 实证验证表明目标函数满足弱双凸性条件，支持理论收敛性保证。
• 在去噪、去模糊与图像修复等图像恢复任务中，该方法表现优于现有方法，尤其在高噪声或严重退化条件下优势更明显。
• 理论分析证明，交替优化序列的所有聚点均为驻点，确保了可靠的收敛行为。

(b) Random samples