PaperlayerView exact revision on arXiv

Source-linked AI summary

Picard Proximal Monte Carlo for Parallel Bayesian Imaging with Score-Based Generative Priors

Deliang Wei, Evan Bell, Wenhan Guo, Yifan Chen, Yu Sun

arXiv:2608.17666v1cs.LG

TL;DR

Bayesian imaging needs posterior samplers that capture uncertainty while overcoming the sequential structure of score-based MCMC methods. PiX-MC uses proximal Langevin dynamics and Picard-based time parallelism, achieving 2.87×–2.99× speedups with eight GPUs and N = 8 across tested Picard-based methods.

  • Problem

    Existing posterior samplers require long sequential trajectories, while available parallel methods largely target deterministic optimization rather than posterior sampling.

  • Method

    PiX-MC combines proximal Langevin dynamics with Picard time parallelism, with annealed and multi-block variants for scalable multi-GPU sampling.

  • Results

    Eight GPUs with N = 8 yielded 2.87×–2.99× speedups across four Picard-based methods relative to their sequential counterparts.

  • Takeaways & Limitations

    PiX-MC provides a principled framework for exploiting concurrent GPU computation in Bayesian imaging posterior sampling.

  • Takeaways & Limitations

    The paper’s scope is posterior sampling, whereas the cited parallelism efforts are largely limited to deterministic optimization.

Abstract

from arXiv · show

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 $512\times512\times80$ sparse-view computed tomography (CT) problem, annealed multi-block PiX-MC achieves up to a $50\times$ runtime speedup over the standard Langevin sampler using eight GPUs.

1 Introduction

Bayesian imaging inverse problems involve substantial posterior uncertainty because sparse, noisy measurements admit multiple plausible images. Score-based diffusion priors improve expressiveness, but existing posterior samplers face theoretical reliability and multi-GPU parallelism limitations.

  • Motivation: Sparse, noisy measurements can leave multiple plausible images consistent with the data, motivating characterization of the full posterior rather than a single estimate.Bayesian sampling incorporates an image prior to characterize π(x|y).
  • Expressive Priors: Score-based diffusion models offer more expressive priors than handcrafted total variation and sparsity-promoting penalties for modeling complex real-world image distributions.They have inspired posterior-sampling methods for imaging inverse problems.
  • Limitations: Guidance-based diffusion samplers are convenient and broadly demonstrated, but generally lack guarantees that generated samples follow the true posterior.This limits their interpretability and reliability in high-stakes settings.
  • MCMC-Based Sampling: Diffusion-integrated MCMC methods provide more rigorous posterior-sampling guarantees, including approaches combining Langevin dynamics with score networks.These methods nevertheless remain part of the broader sampling challenge addressed by the paper.
  • Parallelism Gap: Existing posterior samplers are not designed to exploit multi-GPU parallelism, while prior parallel solver efforts mainly target deterministic optimization rather than posterior sampling.This limits the direct applicability of those techniques to Bayesian sampling.

N Time axis

PiX-MC replaces sequential Langevin processing with batch-time posterior sampling by combining forward-backward splitting, proximal likelihood updates, and Picard iteration. Its multi-GPU time parallelism is supported by convergence analysis and validated across diverse 2D and 3D imaging problems.

  • Time-parallel sampling: Forward-backward splitting decouples the prior and likelihood, enabling modular score-based priors and proximal likelihood updates.This design targets imaging likelihoods with efficient proximal mappings.
  • Time-parallel sampling: PiX-MC uses Picard iteration to parallelize posterior sampling across time steps rather than image pixels, enabling concurrent multi-GPU processing.The sampling trajectory is formulated as a fixed-point problem, while a forward-backward formulation enforces measurement consistency through a proximal update.
  • Scalability variants: Temperature annealing accelerates sequential sampling, while multi-block partitioning preserves PiX-MC efficiency under limited GPU budgets.The two variants improve practical scalability in different computational settings.
  • Convergence analysis: Theoretical analysis bounds time-averaged relative Fisher information by an O(1/T) transient term plus explicit approximation-error terms.The guarantees accommodate imperfect learned score networks and weakly convex imaging likelihoods, with a detailed error decomposition.
  • Empirical validation: Experiments cover linear and nonlinear inverse problems from 1024 × 1024 image deblurring to 512 × 512 × 80 sparse-view CT on multi-GPU platforms.The evaluation spans both 2D and 3D imaging settings and demonstrates generalization across these problems.

2 Background and Related Works

The background frames Bayesian imaging as regularized reconstruction, introduces score-based generative priors and annealed sampling, and positions PiX-MC among posterior and trajectory-parallel methods. It emphasizes that PiX-MC combines proximal Langevin updates with Picard iteration to parallelize trajectory evaluation for large-scale imaging.

  • Bayesian imaging reconstruction: Bayesian imaging formulates reconstruction as balancing measurement fidelity with prior knowledge, with MAP estimation corresponding to regularized optimization.Gaussian likelihoods yield least-squares data fidelity, while transform-domain sparsity is a common regularizer.
  • Proximal methods: Proximal algorithms handle nonsmooth regularizers and exploit efficient imaging operators, separating likelihood and prior information through proximal and gradient steps.Fast transforms or closed-form inverses can make proximal implementations efficient, while forward-backward splitting has established convergence under suitable conditions.
  • Score-based generative priors: Score-based models learn Gaussian-smoothed image-prior scores, connecting score estimation to denoising and approximating the prior-potential gradient as noise vanishes.Langevin and reverse-diffusion samplers use these learned scores to generate image samples across natural, medical, and scientific imaging domains.
  • Score-based generative priors: Annealing gradually decreases the score-model noise level, enabling global exploration at high noise before sample refinement at lower noise.Empirical studies report substantially improved performance relative to using a small constant noise level.
  • Posterior sampling methods: Posterior samplers include Langevin, split Gibbs, and reverse-time diffusion families, while proximal variants address nonsmooth priors and asymmetric likelihood-prior updates.PiX-MC complements these methods by integrating Picard iteration into proximal Langevin sampling for parallel trajectory evaluation in large-scale imaging.
  • Parallelization approaches: Existing parallel methods accelerate computations within iterations but generally preserve iterative time dependence, whereas trajectory-parallel work targets time-node parallelism and PiX-MC applies it to posterior imaging.PiX-MC combines a proximal likelihood step, a prior gradient step, and Picard iteration to enable parallel trajectory evaluation.

3 Method: Picard Proximal Monte Carlo (PiX-MC)

PiX-MC constructs a proximal-likelihood Langevin drift and uses Picard iteration to parallelize computation across time nodes. Its analysis covers finite-horizon stationarity, multi-block refinement, imperfect scores, and annealed noise schedules.

  • Drift construction: PiX-MC combines a learned prior score with a problem-specific proximal likelihood update, enabling efficient incorporation of likelihood terms.The proximal drift approximates the ideal posterior drift while accounting for score estimation and splitting errors.
  • Picard parallelization: Picard iteration parallelizes the sequential X-MC trajectory across time nodes while sharing one Brownian-process realization across trajectory iterates.The sequential X-MC scheme serves as the baseline formulation that PiX-MC parallelizes.
  • Convergence guarantees: The Picard iteration error decays geometrically with the total number of Picard iterations, yielding rapid trajectory stabilization within only a few iterations.Theorem 1 quantifies splitting discrepancy O(η2), score mismatch O(δ2), discretization error O(γ), and finite Picard-refinement error.
  • Multi-block PiX-MC: Multi-block PiX-MC confines Picard refinements to each block, so trajectory storage and concurrent drift evaluations scale with block size rather than the full trajectory.The local contraction coefficient qmult < 1 depends on the block horizon τ rather than total sampling time T.
  • Annealed PiX-MC: Multi-block annealed PiX-MC uses a decreasing noise schedule to capture coarse image structures first and progressively refine details.Its convergence guarantee replaces the multi-block score-mismatch term with a time-averaged annealed-score discrepancy.

4 Numerical Validation of the Theory

This section validates PiX-MC against a closed-form Gaussian posterior in compressed sensing, using multi-block implementations on eight GPUs. The experiments compare learned-score, annealed, and analytical-score variants through posterior moments and squared MMD.

  • Experimental Setup: The compressed-sensing experiment uses a target posterior available in closed form, enabling direct empirical comparison with PiX-MC samples.This provides direct empirical support for the convergence results established in the theorems.
  • Experimental Setup: All PiX-MC variants use multi-block schemes and run on eight NVIDIA RTX PRO 6000 Blackwell GPUs.The single-block implementation is described as computationally demanding, motivating the multi-block setup.
  • Experimental Setup: The experiment constructs a Gaussian image prior from 70,000 grayscale, resized 32 × 32 FFHQ images and generates linear measurements with a normalized Gaussian sensing matrix and noise level σy = 0.001.Under the Gaussian prior and likelihood, the posterior is also Gaussian.
  • Distributional Evaluation: 150 independent chains per algorithm are sampled every 10 time steps, alongside 150 direct samples from the closed-form posterior.Evaluation uses final-time posterior mean and pixelwise standard deviation, plus squared MMD at each recorded time step.
  • Distribution-Level Results: All three PiX-MC instantiations accurately recover the posterior mean and uncertainty maps, while the analytical-score variant most closely matches the true posterior.The analytical score avoids score approximation error; learned-score comparisons between PiX-MC and APiX-MC illustrate annealing effects.

5 Real-World Imaging Experiments

Experiments across MRI, Rician denoising, large-scale deblurring, and sparse-view CT verify geometric Picard refinement, reconstruction quality, and substantial wall-clock acceleration. Proximal likelihood, annealing, and Picard parallelization complement one another, with quality preserved relative to sequential methods and diffusion-based baselines.

  • Experimental scope: Experiments evaluate Picard refinement on MRI and Rician denoising, acceleration under GPU and block configurations, and large-scale sparse-view CT reconstruction.The methods use shared score networks unless task-specific choices are specified, and annealed variants are denoted with an “A” prefix.
  • MRI and Rician denoising: After a short transient, the in-block Picard error decreases geometrically for both accelerated MRI reconstruction and Rician denoising.The observed behavior empirically supports the predicted convergence independently of whether the inverse problem is linear or nonlinear.
  • MRI and Rician denoising: APiL-MC achieves the highest PSNR and SSIM for Rician denoising, while APiX-MC attains the lowest LPIPS; APiX-MC also achieves the best MRI PSNR and SSIM.The proposed methods are competitive with diffusion-based baselines, and annealing consistently improves the reported comparisons.
  • Large-scale deblurring: 54 seconds: APiX-MC reaches average PSNR 26 dB, versus 157 seconds for AX-MC, 365 seconds for X-MC, and 491 seconds for L-MC.The comparison shows that Picard parallelization compounds with proximal updates and annealing, reducing sequential-side runtime from 491 to 157 seconds, a 3.13× speedup.
  • Reconstruction quality: PiX-MC and APiX-MC closely match their sequential counterparts, with PSNR differences of at most 0.01 dB, while producing sharper structures under fixed runtime budgets.Under 200 seconds, they recover sharp edges and fine structures with less blur and noise than DPS and DAPS, which require 334 and 741 seconds to converge.
  • Sparse-view CT: 3.14×: PiX-MC reaches its reference PSNR in 37 minutes versus 116 minutes for X-MC, while APiX-MC requires 9 minutes versus 28 minutes for AX-MC, a 3.11× speedup.Picard acceleration similarly yields 3.08× and 2.96× speedups for PiL-MC versus L-MC and APiL-MC versus AL-MC, respectively.

6 Conclusion · A Tool Lemmas

The paper introduces PiX-MC as a time-parallel posterior sampler for Bayesian imaging, using proximal drifts and concurrent processing across time nodes and GPUs. The tool lemmas establish Lipschitz properties and bound score-, gradient-, and proximal-drift-related terms under stated assumptions.

  • 6 Conclusion: PiX-MC is a principled time-parallel posterior sampling framework for Bayesian imaging.It constructs a proximal drift motivated by forward-backward splitting and exploits efficient imaging-likelihood proximal operators.
  • 6 Conclusion: PiX-MC parallelizes long sampling trajectories across time nodes, enabling concurrent multi-GPU processing.This replaces sequential advancement along the trajectory with parallel computation across discretization nodes.
  • A Tool Lemmas: Under Assumption 1 and 1−ηαL > 0, the proximal drift Tη,σ is shown to be Lipschitz under Assumption 2(ii).The proof uses the (1−ηαL)^−1-Lipschitz property of proxηL and the score model’s Lipschitz continuity.
  • A Tool Lemmas: When αn is non-increasing, every weighted proximal drift T αn η,σn is uniformly Λη,ann-Lipschitz.The argument replaces the score with αnSθ(·;σn), whose Lipschitz constant is uniformly bounded by α0LS.
  • A Tool Lemmas: Lemma 5 states that ∇L(x)+∇V(x) is (βL+βV)-Lipschitz under Assumption 1.The lemma also applies to any probability measure µ absolutely continuous with respect to π.
  • A Tool Lemmas: Lemma 6 assumes Assumptions 1 and 2(i) and requires sufficiently small η so that specified stability conditions hold.Its bounds decompose the update error into proximal-gradient, score-model, and total-gradient terms, including the score error δ.
  • A Tool Lemmas: The proximal optimal condition gives z+η∇L(z)=x−ηSθ(x;σ) and Tη(x)=∇L(z)+Sθ(x;σ).These identities provide the starting point for the proximal-drift proof.

B Proof of single-block PiX-MC

The single-block PiX-MC analysis establishes convergence through an explicit-constant version of Theorem 1 under sufficiently small step-size and refinement parameters. A technical lemma shows exponential decay of Picard refinement error when the contraction factor qsing is below one.

  • Convergence theorem: The proof introduces explicit constants, including the contraction factor qsing = TΛη, and ensures A4 < 0.75 for sufficiently small η and γ.The analysis notes that A4 tends to zero as (η, γ) tends to (0, 0).
  • Convergence theorem: Theorem 4 gives a detailed convergence result for single-block PiX-MC under Assumptions 1–2 and sufficiently small η, γ.The parameters must satisfy the stated positivity conditions, including 1 − ηαL > 0 and 1 − 3η^2βL^2 > 0.
  • Picard refinement lemma: When qsing < 1, the Picard iteration error decays exponentially fast across parallel refinements.Lemma 7 establishes this refinement-error behavior under Assumptions 1–2 and sufficiently small η, γ.

C Sufficient conditions for Assumption 3

Lemma 8 gives sufficient conditions for Assumption 3: under Assumptions 1–2, a finite second moment for x0, and a suitable coercivity condition, multi-block PiX-MC satisfies the required bound for sufficiently small step sizes and sufficiently many Picard iterations. The resulting constant is independent of the number of blocks and total time horizon.

  • Sufficient conditions: Lemma 8 assumes Assumptions 1–2, E∥x0∥2 < ∞, and constants κ > αL + αV and b ≥ 0.These conditions provide the stated sufficient setup for Assumption 3.
  • Parameter restrictions: For sufficiently small η and γ satisfying qmult = τΛη < 1, a finite K0 exists such that the result holds for every K ≥ K0.Specifically, there exist η0, γ0 > 0 with 0 < η ≤ η0 and 0 < γ ≤ γ0.
  • Convergence bound: For every K ≥ K0, the multi-block PiX-MC starting points satisfy the required bound with finite constants independent of the block index.The proof obtains a contraction after sufficiently many Picard iterations.
  • Conclusion: The bound uses a finite constant Cx that is independent of the number of blocks M and total time horizon T = Mτ, establishing Assumption 3.This independence is the concluding property of Lemma 8.

D Proof of multi-block PiX-MC

The multi-block PiX-MC analysis establishes convergence under Assumptions 1–3 and sufficiently small parameters, using a geometrically decaying Picard-refinement error. The resulting estimate is uniform across blocks and independent of the total time horizon.

  • Convergence theorem: Theorem 5 gives a convergence result for multi-block PiX-MC under Assumptions 1–3 and sufficiently small η and γ.The parameter conditions include 1 −ηαL > 0, 1 −3η2β2 L > 0, and 1 −18γ2(βL + βV)2.
  • Picard error control: The detailed Picard-refinement lemma shows that the error decays geometrically fast across parallel refinements.The result applies for each block m = 0, . . . , M −1 and refinements k = 1, . . . , K −1.
  • Picard error control: The global Picard-error estimate is uniform over all blocks and independent of the total time horizon T.The proof obtains the global estimate after bounding the initial Picard error and using the uniformity over m = 0, . . . , M −1.

E Proof of multi-block APiX-MC

The multi-block APiX-MC analysis proves convergence under stated assumptions and sufficiently small step-size parameters. Its effective discrepancy separates score-model error, weighted-annealing transients, and discretization/refinement contributions, with annealing error vanishing after finite post-annealing sampling.

  • Convergence theorem: Theorem 6 establishes convergence for multi-block APiX-MC under Assumptions 1, 3, and 4, provided η and γ satisfy the displayed stability conditions.The conditions include positivity constraints involving ηαL, η^2βL^2, and γ^2(βL + βV)^2.
  • Error decomposition: The effective discrepancy decomposes into a score approximation term and a transient error from weighted annealing, alongside discretization and refinement contributions.The local bound uses the weighted score error δn + (αn −1)RS and includes terms proportional to γd and d(βL + βV).
  • Annealing error: The weighted-annealing error vanishes as the post-annealing sampling horizon increases when αn = 1 after finitely many iterations.This identifies finite-time termination of annealing as the condition eliminating the annealing transient asymptotically.
  • Picard refinement: In-block Picard refinement errors decay geometrically for every block and refinement step, yielding the same bound uniformly over all global nodes.The result is obtained from a blockwise contraction and the equality of each block’s starting point with its preceding global node.

F Additional experimental details · F.1 Common Implementation Details and Hyperparameter Selections

The experiments use task-specific score-network parameterizations, likelihood balancing, annealing schedules, and stopping criteria. Hyperparameters are tuned separately by method and imaging task, with proximal variants potentially allowing larger stepsizes through implicit likelihood treatment.

  • F.1 Common Implementation Details and Hyperparameter Selections: All images are scaled to [0, 1], while score-network inputs are rescaled to [−1, 1] and outputs mapped back to [0, 1].Pretrained natural-image, MRI, and CT models use different output parameterizations.
  • F.1 Common Implementation Details and Hyperparameter Selections: Natural-image experiments use a pretrained DDPM diffusion model, with its score network related approximately to the noise predictor.The relation is expressed using the cumulative diffusion coefficient ¯αt.
  • F.1 Common Implementation Details and Hyperparameter Selections: MRI reconstruction uses the pretrained PMC score network, whose positive-score output is converted to the negative-score function used in the formulation.The network output is specified in network coordinates as SMRIθ(2x −1; σ).
  • F.1 Common Implementation Details and Hyperparameter Selections: CT experiments train an MMSE denoiser Dθ using the standard EDM framework and relate it to the score network Sθ.The denoiser-score relation defines the score representation used for CT.
  • F.1 Common Implementation Details and Hyperparameter Selections: A task-dependent λ > 0 balances likelihood and score-network scales, with A denoting the task-specific forward operator.This balancing parameter is introduced for MRI, deblurring, and CT.
  • F.1 Common Implementation Details and Hyperparameter Selections: The annealing schedule uses global node index n = mN + i and geometric schedules for both noise levels and score-network weights.The schedule is non-increasing and becomes one once σn reaches σmin; a = 2 for MRI and CT, while a = 5/2 for deblurring and Rician denoising.
  • F.1 Common Implementation Details and Hyperparameter Selections: Deblurring and CT use adaptive Picard refinement, terminating when k reaches K = 20 or the normalized residual εm,k/d falls below 3 × 10−4.The method computes an empirical in-block residual before applying this stopping rule.
  • F.1 Common Implementation Details and Hyperparameter Selections: Hyperparameters are tuned separately for each method and imaging task, with method-specific stepsizes and potentially larger proximal stepsizes from implicit likelihood treatment.Paired sequential and Picard-based methods use the same La...

F.2 Likelihood and Proximal Implementations

The paper implements task-specific likelihood gradients and proximal operators for MRI reconstruction, Rician denoising, image deblurring, and 3D CT reconstruction, exploiting Fourier-domain structure where available and iterative solvers otherwise.

  • MRI reconstruction: MRI uses a Fourier transform followed by radial subsampling, enabling its proximal operator to be evaluated elementwise in k-space without solving a large-scale linear system.The forward operator is A = MF, with F unitary and MᵀM diagonal in the Fourier domain.
  • Rician denoising: Rician denoising evaluates the gradient through B(x) = I1(x)/I0(x) and solves each pixel-wise proximal subproblem with IRL1.The maximum number of IRL1 inner iterations is set to 10, although the subproblem typically converges within 5 iterations.
  • Image deblurring: Image deblurring models circular convolution, whose proximal operator is computed efficiently in the Fourier domain through elementwise multiplication and division.The convolution operator is diagonalized by the discrete Fourier transform.
  • 3D CT reconstruction: 3D CT reconstruction solves the proximal linear system with five conjugate-gradient iterations and adds an inter-slice Huber-TV regularizer.The HTV regularizer promotes consistency between adjacent slices in 512×512×80 CT volumes.
  • 3D CT reconstruction: The CT experiments set the Huber-TV smoothing parameter to δH = 10^-8 and its balancing parameter to βH = 0.28 for all methods.HTV is treated as an inter-slice prior characterizing correlations between adjacent slices.
Loading 2608.17666v1…