Multi-Scale Energy (MuSE) plug and play framework for inverse problems

TL;DR

This paper introduces multi-scale energy models for inverse problems, improving prior estimation accuracy and convergence, and enabling uncertainty quantification in MRI image recovery.

Contribution

It proposes explicit and implicit multi-scale energy frameworks that enhance MAP estimation and sampling efficiency compared to traditional single-scale models.

Findings

Multi-scale models improve convergence and accuracy in inverse problems.

The implicit MuSE (i-MuSE) is simpler and faster, with better performance.

MuSE achieves MRI reconstruction quality comparable to end-to-end trained models.

Abstract

We introduce multi-scale energy models to learn the prior distribution of images, which can be used in inverse problems to derive the Maximum A Posteriori (MAP) estimate and to sample from the posterior distribution. Compared to the traditional single-scale energy models, the multi-scale strategy improves the estimation accuracy and convergence of the MAP algorithm, even when it is initialized far away from the solution. We propose two kinds of multi-scale strategies: a) the explicit (e-MuSE) framework, where we use a sequence of explicit energies, each corresponding to a smooth approximation of the original negative log-prior, and b) the implicit (i-MuSE), where we rely on a single energy function whose gradients at different scales closely match the corresponding e-MuSE gradients. Although both schemes improve convergence and accuracy, the e-MuSE MAP solution depends on the scheduling strategy, including the choice of intermediate scales and exit conditions. In contrast, the i-MuSE formulation is significantly simpler, resulting in faster convergence and improved performance. We compare the performance of the proposed MuSE models in the context of Magnetic Resonance (MR) image recovery. The results demonstrate that the multi-scale framework yields a MAP reconstruction comparable in quality to the End-to-End (E2E) trained models, while being relatively unaffected by the changes in the forward model. In addition, the i-MuSE scheme also allows the generation of samples from the posterior distribution, enabling us to estimate the uncertainty maps.