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

2025-11-02

AI Summary

Traditional MCMC methods for high-dimensional imaging inverse problems suffer from low sampling efficiency, substantial memory and communication overhead, and challenges in uncertainty quantification. To address these issues, this paper proposes a scalable distributed Plug-and-Play MCMC (PnP-MCMC) algorithm. Our method is the first to integrate PnP-MCMC with approximate data augmentation and implements multi-GPU parallelization via the SPMD (Single Program, Multiple Data) paradigm. We employ a lightweight CNN denoiser that preserves reconstruction quality competitive with state-of-the-art PnP approaches while significantly reducing computational and inter-GPU communication costs. Experiments across multiple high-resolution imaging tasks demonstrate strong scalability, enabling efficient distributed posterior sampling and principled uncertainty quantification.

Technology Category

• Machine Learning: Calibration & Uncertainty Quantification
• Reasoning under Uncertainty: Probabilistic Programming
• Data Mining & Knowledge Management: Scalability, Parallel & Distributed Systems

Application Category

• User Modeling, Personalization and Recommendation: On-Device user modeling, personalization, and recommendation
• Graph Algorithms and Modeling for the Web: Algorithms and analysis for incomplete, noisy, or partially observed Web-related graphs
• Search and Retrieval-Augmented AI: Efficiency and scalability of Web search engines

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.

Problem

• Solving high-dimensional imaging inverse problems with uncertainty quantification
• Developing scalable distributed MCMC samplers for large-scale reconstruction
• Combining lightweight denoising networks with multi-GPU parallel computing

Innovation

• Distributed sampler uses data augmentation and PnP-ULA
• Lightweight denoising CNN efficiently exploits multiple GPUs
• Scalable approach enables uncertainty quantification for imaging