A Distributed Plug-and-Play MCMC Algorithm for High-Dimensional Inverse Problems

TL;DR

This paper introduces a scalable, distributed MCMC algorithm using neural network denoisers for high-dimensional inverse imaging problems, enabling efficient uncertainty quantification on large datasets.

Contribution

It proposes a novel distributed PnP-ULA-based MCMC method that leverages multiple GPUs and lightweight neural denoisers for scalable high-dimensional inverse problem solving.

Findings

Achieves high reconstruction quality comparable to existing PnP methods.

Demonstrates effective scalability on large imaging datasets.

Provides uncertainty quantification in high-dimensional inverse problems.

Abstract

Markov Chain Monte Carlo (MCMC) algorithms are standard approaches to solve imaging inverse problems and quantify estimation uncertainties, a key requirement in absence of ground-truth data. To improve estimation quality, Plug-and-Play MCMC algorithms, such as PnP-ULA, have been recently developed to accommodate priors encoded by a denoising neural network. Designing scalable samplers for high-dimensional imaging inverse problems remains a challenge: drawing and storing high-dimensional samples can be prohibitive, especially for high-resolution images. To address this issue, this work proposes a distributed sampler based on approximate data augmentation and PnP-ULA to solve very large problems. The proposed sampler uses lightweight denoising convolutional neural network, to efficiently exploit multiple GPUs on a Single Program Multiple Data architecture. Reconstruction performance and scalability are evaluated on several imaging problems. Communication and computation overheads due to the denoiser are carefully discussed. The proposed distributed approach noticeably combines three very precious qualities: it is scalable, enables uncertainty quantification, for a reconstruction performance comparable to other PnP methods.