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

TL;DR

This work tackles scalable Bayesian inference for high-dimensional imaging inverse problems by introducing a distributed plug-and-play MCMC framework that uses approximate data augmentation (AXDA) and a lightweight CNN prior within a distributed Gibbs sampler. The approach leverages operator locality and a Cartesian domain decomposition to enable multi-GPU inference via PnP-ULA for the image variable and PSGLA for auxiliary variables, with CNN denoisers such as DnCNN, DDFB, and DRUNet evaluated. Experimental results across inpainting and deconvolution tasks show competitive reconstruction quality and favorable strong/weak scaling, particularly for deep priors, while TV priors exhibit limited scalability. The proposed method enables uncertainty quantification on very large imaging problems (e.g., up to $N=3\times 2048^2$ or $4096^2$) and provides a practical pipeline for distributed Bayesian imaging with data-driven priors.

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.

Figures (3)

• Figure 1: Illustration of a 2D linear convolution operator $G$ defined by a $3 \times 3$ kernel $\bm{k}$, applied to an image $\bm{x}$ of size ${N} = 20 \times 20$ (hence $G(\bm{x}) \in \mathbb{R}^M$ with $M = 22 \times 22$). The input $\bm{x}$ is partitioned into $B = 4$ blocks (in dashed blue lines). Resulting blocks in $G(\bm{x})$ are delimited by dashed green lines. The block $\bm{x}_1$ assigned to worker 1 is in blue, with corresponding block $[G(\bm{x})]_1$ in green. The purple area contains pixels received by worker $1$ to compute $[G(\bm{x})]_1$. Pixels in $\bm{S}_1\bm{x}$ are in the union of the blue and purple areas.
• Figure 2: Communications and operations underlying the proposed distributed implementation of $\bm{W}_k$. Communications are represented by purple arrows, with messages in dashed purple lines. Discarded regions are shaded in green.
• Figure 3: Inpainting results -- Restoration results with \ref{['algo:dsgs_pnp_ula_psgla']} using $B = 4$ workers. First row: ground truth $\overline{\bm{x}}$ and observations $\bm{y}$. Last 4 rows, top to bottom: MMSE (left) and pixel variance of the red channel (right) using either TV, DDFB, DnCNN, or DRUNet. Dashed lines on the ground truth image show the partition considered for $\bm{x}$.

Theorems & Definitions (8)

• Definition 1: Local selection operators
• Definition 2: Localized operator
• Example 1
• Remark 1
• Example 2: DnCNN denoiser Zhang2017
• Example 3: DDFB denoiser
• Example 4: DRUNet denoiser
• Remark 2