Multi-GPU distributed PnP-ULA for high-dimensional imaging inverse problems

Markov Chain Monte Carlo (MCMC) methods enable uncertainty quantification in inverse problems, making them valuable in various applications, especially when no ground-truth is available. Optimization-inspired Plug-and-Play MCMC algorithms, like PnP-ULA, have been developed to incorporate rich neural networks as priors, improving drastically the estimation quality. However, scaling MCMC samplers remains challenging, as generating and storing a sufficient number of high-dimensional samples, typically high-resolution images, can be computationally prohibitive. This work proposes a distributed implementation of PnP-ULA to target much larger problems without compromising on estimation quality. The proposed approach leverages a lightweight deep denoiser within a Single Program Multiple Data that permits an efficient exploitation of a multi-GPU architecture. Synthetic imaging tasks demonstrate the scalability and performance of this strategy to deal with very high dimensional inverse problems. This approach achieves competitive estimation quality along with uncertainty quantification compared to a typical state-of-the-art PnP optimization method, with much higher scalability.