Picard Proximal Monte Carlo for Parallel Bayesian Imaging with Score-Based Generative Priors
Organizations: Johns Hopkins University · University of California, Los Angeles
Abstract
Bayesian imaging inverse problems often require sampling from high-dimensional posterior distributions. While recent score-based and diffusion models provide expressive Bayesian priors, their sampling procedures remain inherently sequential and computationally expensive for large-scale imaging applications. We propose PiX-MC, a time-parallel posterior sampling framework based on proximal Langevin dynamics and Picard iteration. The proximal-likelihood formulation exploits the fact that many imaging likelihoods admit efficient, problem-specific proximal operators, while Picard refinement exposes parallelism across discretization nodes and naturally supports multi-GPU implementation. To further improve practical scalability and sampling performance, we develop multi-block and annealed variants of the proposed framework. We establish convergence guarantees under transparent assumptions, accommodating non-log-concave posteriors, imperfect learned score models, multi-block implementations, and annealing schedules. Experiments on a diverse collection of imaging inverse problems demonstrate that PiX-MC substantially reduces wall-clock time while preserving reconstruction quality. On a sparse-view computed tomography (CT) problem, annealed multi-block PiX-MC achieves up to a runtime speedup over the standard Langevin sampler using eight GPUs.
Figures & tables
| Methods | PiL-MC vs. L-MC (491s) | APiL-MC vs. AL-MC (190s) | |||||
|---|---|---|---|---|---|---|---|
| 4 GPUs | 264s / | 328s / | 550s / | 99s / | 139s / | 250s / | |
| 8 GPUs | 171s / | 175s / | 295s / | 63s / | 81s / | 151s / | |
| Methods | PiX-MC vs. X-MC (365s) | APiX-MC vs. AX-MC (157s) | |||||
| 4 GPUs | 195s / | 230s / | 440s / | 84s / | 120s / | 194s / | |
| Method | PSNR (dB) | LPIPS | SSIM |
|---|---|---|---|
| DPS | 26.43 | 0.2498 | 0.7142 |
| DAPS | 28.17 | 0.1632 | 0.7444 |
| L-MC | 26.82 | 0.2605 | 0.7273 |
| AL-MC | 27.92 | 0.2042 | 0.7737 |
| X-MC | 27.24 | 0.2686 | 0.7342 |
| AX-MC | 27.98 | 0.2138 | 0.7719 |
Appendix figures & tables1 asset
Supplementary material from the paper’s appendix.
Appendix
| Inverse problem | Method | |||||||||
|---|---|---|---|---|---|---|---|---|---|---|
| MRI | PiL-MC | – | ||||||||
| APiL-MC | – | |||||||||
| PiX-MC | ||||||||||
| APiX-MC | ||||||||||
| Rician | PiL-MC | – | ||||||||
| APiL-MC | – |