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

Existing posterior samplers for imaging inverse problems follow sequential trajectories and are not designed to use available multi-GPU parallelism. PiX-MC parallelizes sampling across time nodes; in large-scale deblurring, eight GPUs achieved 2.87×–2.99× speedups across four Picard-based methods.

  • Problem

    Existing posterior samplers share a sequential structure and are not designed to exploit multi-GPU resources for posterior sampling.

    Source: §1 · p. 2

  • Method

    PiX-MC uses a proximal drift and Picard iteration to process trajectory time nodes concurrently, with annealed and multi-block variants for practical scalability.

    Source: §6 · p. 27

  • Results

    2.87×–2.99× speedups over sequential counterparts were achieved across four Picard-based methods using eight GPUs and N = 8.

    Source: §5.2 · p. 22

  • Takeaways & Limitations

    The annealed variant empirically accelerates convergence to target reconstruction quality, while the multi-block variant is designed for efficient operation under limited GPU budgets.

    Source: §6 · p. 27

    Guidance-based score-model samplers generally lack guarantees that generated samples follow the true posterior, limiting interpretability and reliability in high-stakes settings.

    Source: §1 · p. 2

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 sampling characterizes uncertainty in images reconstructed from sparse, noisy measurements, but current posterior samplers are computationally costly and poorly suited to multi-GPU parallelism. Score-based priors can model complex image distributions, while guidance-based sampling lacks guarantees that generated samples follow the true posterior.

1 Introduction

  • Sparse, noisy measurements can leave substantial uncertainty about the recovered image, so Bayesian sampling characterizes the posterior rather than producing only a point estimate.This is especially useful in high-stakes biomedical and scientific imaging, where imaging results can affect conclusions, clinical decisions, or downstream actions.Source: §1 · p. 1
  • Score-based diffusion models can capture complex image distributions that handcrafted priors often fail to represent.This has motivated their use as priors for imaging inverse problems.Source: §1 · p. 1
  • Guidance-based score-model methods lack theoretical guarantees that generated samples follow the true posterior, limiting interpretability and reliability in high-stakes settings.These methods steer the generative process using measurements [37] [16] and have been demonstrated across imaging applications [17] [79] [35].Source: §1 · p. 2
  • MCMC methods with score models provide more rigorous posterior-sampling guarantees, but their sequential trajectories repeatedly evaluate the score model and forward operator, making large-scale problems costly.Examples combine Langevin dynamics with an SDM score network [40] [86] or use an SDM as an empirical proximal sampler in split Gibbs [19] [92].Source: §1 · p. 2
  • Existing posterior samplers are not designed to exploit multi-GPU resources, while explored parallel inverse-problem solvers largely target deterministic optimization rather than posterior sampling.This leaves substantial parallel computational resources underused for posterior sampling.Source: §1 · p. 2

N Time axis

PiX-MC makes posterior sampling time-parallel by expressing Langevin sampling as a fixed-point problem solved with Picard iteration, while a forward–backward split couples score-based priors with proximal likelihood updates.

N Time axis

  • PiX-MC formulates trajectory sampling as a fixed-point problem solved by Picard iteration, parallelizing time steps across GPUs rather than image pixels.Forward–backward splitting decouples the score-based prior from the likelihood and enforces measurement consistency through a proximal update suited to structured imaging operators.Source: N Time axis · p. 2, Figure 1 · p. 2
  • Annealed PiX-MC uses temperature annealing to accelerate the sequential trajectory, while multi-block PiX-MC partitions it to work efficiently with limited GPU budgets.The framework combines time-batch processing with Picard iteration across time nodes.Source: N Time axis · pp. 2–3, Figure 1 · p. 2
  • Analysis bounds time-averaged relative Fisher information by an O(1/T) transient plus approximation errors, accommodating imperfect learned scores and weakly convex likelihoods.The error decomposition characterizes contributions from forward–backward splitting, Picard iterations, and time discretization.Source: N Time axis · p. 3
  • Up to 50× runtime speedup over the standard Langevin sampler on 512 × 512 × 80 sparse-view CT, using eight GPUs, while preserving reconstruction quality.Validation spans linear and nonlinear forward models, 2D and 3D imaging, and problems from 1024 × 1024 deblurring to large-scale CT.Source: N Time axis · p. 3

2 Background and Related Works

The background traces Bayesian reconstruction and score-based priors to posterior sampling, then identifies trajectory-level parallelism as a gap PiX-MC addresses.

2.1 Bayesian Image Reconstruction

  • Bayesian image reconstruction casts measurement fitting and prior knowledge as MAP estimation, while proximal methods exploit nonsmooth regularizers and structured forward operators.The background gives Gaussian least squares as a data-fidelity example and transform-domain sparsity as a regularizer.Source: §2.1 · pp. 3–4

2.2 Score-Based Diffusion Models as Learned Image Priors

  • Score-based diffusion models learn Gaussian-smoothed image-prior scores linked to denoising, and annealing reduces noise during sampling to explore globally before refining samples.An alternative reverse-time diffusion formulation transforms Gaussian noise toward the target image distribution.Source: §2.2 · pp. 4–5

2.3 Posterior Sampling Methods for Imaging

  • Posterior methods span Langevin sampling, split-Gibbs schemes, and guided reverse diffusion; PiX-MC integrates Picard iteration into proximal Langevin sampling for parallel trajectory evaluation.Split-Gibbs methods modularize likelihood and prior updates but require costly sampling subproblems; guided reverse diffusion uses approximate time-dependent posterior scores that may lack an explicit invariant target.Source: §2.3 · pp. 5–6

2.4 Related Parallel Imaging Methods & Picard Iteration

  • Within-iteration parallelism leaves iterative sampling sequential, while Picard iteration parallelizes time nodes; prior trajectory-parallel methods do not handle imaging posteriors’ proximal likelihoods, learned priors, and annealing.PiX-MC extends trajectory parallelization to imaging posterior sampling and provides a theoretical analysis that accommodates these components.Source: §2.3 · p. 6, §2.4 · p. 7

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

PiX-MC uses proximal-likelihood Langevin dynamics and Picard refinement to parallelize trajectory computation, with multi-block and annealed variants. Its convergence analysis bounds stationarity and refinement errors, including weighted annealing and annealed-score discrepancy.

3.1 Proximal-Likelihood Drift via Forward-Backward Splitting

  • The FBS-inspired drift combines a learned prior score with a problem-specific proximal likelihood update, approximating the posterior drift with error from score estimation and splitting.Forward Euler–Maruyama yields sequential X-MC, the baseline for PiX-MC.Source: §3.1 · pp. 7–9

3.2 PiX-MC: Algorithm Formulation and Convergence Analysis

  • Picard iteration parallelizes X-MC across time nodes, reducing idealized cost scaling from N to K, and its finite-time Fisher-information guarantee includes inherited and geometrically decaying Picard errors.The bound accounts for splitting O(η^2), score mismatch O(δ^2), discretization O(γ), and finite-refinement errors that decay geometrically in K.Source: §3.2 · pp. 9–11

3.3 Multi-Block PiX-MC: Partitioning the Trajectory into Multiple Blocks

  • Multi-block PiX-MC refines locally to limit storage to N nodes and replace global contraction with qmult=τΛη, under uniformly bounded block-start second moments.It parallelizes within each block while proceeding sequentially across blocks; the local contraction condition can permit a larger stepsize than the single-block scheme.Source: §3.3 · pp. 12, 14

3.4 Multi-Block Annealed PiX-MC: Improving Sample Quality via Annealing

  • Multi-block annealed PiX-MC decreases the score-network noise level to capture coarse image structure before refining details, with the schedule continuing across the full trajectory rather than restarting at block boundaries.Global indexing makes the noise level and score weight decrease continuously over the sampling trajectory.Source: §3.4 · p. 14
  • Under Assumptions 1, 3, and 4 with sufficiently small η and γ, Theorem 3 bounds error contributions from annealed-score discrepancy and finite Picard refinement.The score-mismatch term is replaced by the time-averaged annealed-score discrepancy, while refinement error is governed by the uniform contraction factor q_ann.Source: §3.4 · p. 16
  • When α_n reaches 1 after a finite annealing stage, the score-weight contribution decreases as MN grows, and Picard refinement error decays geometrically with K when q_ann < 1.The finite-stage condition holds in the experiments.Source: §3.4 · p. 16
  • Annealing reduces the sequential effort needed for a given reconstruction quality, while Picard refinement reduces trajectory-evaluation wall-clock cost through time parallelism; their gains can compound.The passage states that APiX-MC yields additional wall-clock acceleration beyond trajectory parallelization.Source: §3.4 · p. 16

4 Numerical Validation of the Theory

In a compressed-sensing test with a closed-form posterior, PiX-MC variants empirically approach the target distribution. The analytical-score variant has the lowest final MMD2, while annealing accelerates the learned-score sampling trajectory.

Experimental Setup

  • A compressed-sensing problem with a closed-form target posterior enables direct comparison; all variants use multi-block schemes on eight NVIDIA RTX PRO 6000 Blackwell GPUs.The multi-block implementation is used because the single-block PiX-MC implementation is computationally demanding.Source: §4 · pp. 16–17

Distribution-Level Results

  • At the final time step, analytical-score PiX-MC achieves MMD2 0.0036, versus 0.1361 for learned-score APiX-MC and 0.1480 for learned-score PiX-MC.Source: Figure 2 · p. 17, §4 · p. 18
  • Learned-score APiX-MC reduces MMD2 faster and reaches a lower final value than learned-score PiX-MC, while the analytical-score variant's advantage is consistent with removing score approximation error.The learned-score gap is consistent with the score approximation error analyzed in Section 3.Source: Figure 2 · p. 17, §4 · p. 18
  • All three PiX-MC variants recover posterior mean and uncertainty maps, and their empirical distributions move toward the target posterior over time.The trajectory behavior provides empirical support for the finite-horizon stationarity bounds in Theorems 2–3.Source: §4 · p. 18

5 Real-World Imaging Experiments

The experiments test Picard refinement on MRI and Rician denoising, practical acceleration on large-scale deblurring, and scalability on sparse-view 3D CT. They also assess reconstruction quality and the effects of GPU budgets, block sizes, proximal-likelihood drift, and annealing.

5.1 Behavior of Picard Refinement

  • After a short initial transient, Picard error decays geometrically in MRI reconstruction and Rician denoising across proximal, non-proximal, and annealed methods.Picard refinement turns initially piecewise-constant block trajectories into smooth trajectories across successive blocks.Source: §5.1 · pp. 19–20, Figure 3 · p. 19, Figure 4 · p. 20
  • APiL-MC achieves the highest PSNR and SSIM for Rician denoising, APiX-MC the lowest LPIPS there and the best PSNR and SSIM for MRI; annealing consistently improves quality over non-annealed variants.Visual comparisons show PiX-MC methods recovering sharper MRI edges with fewer residual artifacts and preserving finer structures in Rician denoising.Source: §5.1 · pp. 20–21, Table 1 · p. 20

5.2 Practical Acceleration on Large-Scale Deblurring

  • 2.87×–2.99×: with eight GPUs and N = 8, the four Picard-based methods achieve these speedups; with four GPUs and N = 32, all four are slower than their sequential counterparts.Increasing GPU count improves speedup, while oversubscribing time nodes relative to available GPUs reduces the benefit.Source: §5.2 · p. 22, Table 2 · p. 23
  • 54 seconds: APiX-MC reaches 26 dB on 1024 × 1024 deblurring, versus 157 seconds for AX-MC, 365 for X-MC, and 491 for L-MC—an overall 9× speedup over L-MC.The sequential-side combination of proximal update and annealing gives a 3.13× speedup, followed by a further 2.91× speedup from Picard parallelization.Source: §5.2 · p. 23, Figure 7 · p. 23
  • At most 0.01 dB PSNR difference: PiX-MC and APiX-MC closely match their corresponding sequential methods on deblurring while remaining competitive with DPS and DAPS.Visual comparisons under a fixed 200-second budget show sharper edges and finer structures from the Picard methods than their sequential counterparts.Source: §5.2 · pp. 23–24, Table 3 · p. 24

5.3 Scaling to 3D CT Reconstruction

  • 32.5 dB in 9 minutes: APiX-MC achieves this reference PSNR with an approximately 50× wall-clock speedup over standard L-MC on 512 × 512 × 80 sparse-view CT.The proximal-likelihood drift, annealing, and Picard time parallelization contribute complementary benefits; APiX-MC also achieves 33.64 dB in 35 minutes after convergence on the visual comparison.Source: §5.3 · pp. 25–26, Figure 9 · p. 25, Figure 10 · p. 26
  • Picard time parallelization, proximal-likelihood drift, and annealing complement one another in 3D CT, improving visual fidelity by suppressing streak artifacts and residual noise while preserving anatomical structures.These improvements are reported for both a selected axial slice and the full reconstructed volume.Source: §5.3 · pp. 26–27, Figure 10 · p. 26, Figure 11 · p. 27

6 Conclusion · A Tool Lemmas

PiX-MC parallelizes Bayesian imaging posterior sampling across time nodes using proximal likelihood operators, with annealed and multi-block variants for practical scalability. The paper establishes finite-time stationarity guarantees and reports wall-clock acceleration while maintaining reconstruction quality across imaging problems.

6 Conclusion

  • Experiments across linear, nonlinear, and large-scale 2D and 3D imaging problems show practical scalability and substantial wall-clock acceleration while maintaining reconstruction quality.The conclusion does not provide numerical speedup values.Source: §6 · p. 27
  • PiX-MC uses a proximal drift based on efficient imaging-likelihood operators and processes trajectory time nodes concurrently across GPUs rather than sequentially.The framework is motivated by forward-backward splitting.Source: §6 · p. 27
  • Annealed and multi-block variants target practical scalability: annealing accelerates convergence to target reconstruction quality, while blocks support operation with limited GPU budgets.The multi-block variant partitions the trajectory into blocks.Source: §6 · p. 27
  • Finite-time stationarity guarantees are established for all PiX-MC variants under mild assumptions.Source: §6 · p. 27

A Tool Lemmas

  • Under Assumption 2(ii), the proximal drift Tη,σ is Λη-Lipschitz; when weights are non-increasing, every weighted proximal drift T_{α_nη,σ_n} is uniformly Λη,ann-Lipschitz.For α_n ≤ α_0, the weighted score is uniformly α_0L_S-Lipschitz.Source: A Tool Lemmas · p. 28
  • The Lipschitz argument uses the proximal map’s (1−ηαL)^−1 bound for αL-weakly convex L and the βL-Lipschitz gradient of L.The drift combines ∇L at the proximal point with the score; the score is LS-Lipschitz.Source: A Tool Lemmas · p. 28
  • Under Assumptions 1 and 2(i), for sufficiently small η, the squared-error bound includes 3η²δ² for score error and 3η²∥∇L(x)+∇V(x)∥² for the combined gradient.It also includes 3η²βL²∥z−x∥².Source: A Tool Lemmas · pp. 28–29

B Proof of single-block PiX-MC

The single-block PiX-MC proof establishes convergence under Assumptions 1–2 and sufficiently small step-size and splitting parameters. It also shows that the Picard refinement error decays exponentially when the contraction factor qsing is below 1.

Theorem 4

  • The detailed single-block convergence theorem assumes Assumptions 1–2 and sufficiently small η and γ satisfying the stated parameter conditions.The theorem is stated with explicit constants.Source: B Proof of single-block PiX-MC · p. 30
  • A4 tends to 0 as η and γ tend to 0, so A4 < 0.75 for sufficiently small η and γ.The proof also states that sufficient conditions ensure this bound.Source: B Proof of single-block PiX-MC · p. 31

Lemma 7

C Sufficient conditions for Assumption 3

Lemma 8 gives sufficient conditions for Assumption 3: under its hypotheses, sufficiently small step sizes and a contraction condition yield a bound on multi-block PiX-MC starting points that is uniform in the number of blocks and total time horizon.

C Sufficient conditions for Assumption 3

  • The sufficient-condition lemma assumes Assumptions 1–2, E∥x0∥2 < ∞, and constants κ > αL + αV and b ≥ 0 satisfying an additional condition.The additional condition is not specified in the supplied passage.Source: C Sufficient conditions for Assumption 3 · p. 35
  • Under Lemma 8’s hypotheses, sufficiently small η and γ, qmult = τΛη < 1, and K ≥ K0 ensure a bound on multi-block PiX-MC block starting points, establishing Assumption 3.The lemma guarantees η0, γ0 > 0 and a finite K0 such that the conclusion holds for 0 < η ≤ η0 and 0 < γ ≤ γ0.Source: C Sufficient conditions for Assumption 3 · pp. 35, 37
  • The resulting finite constant Cx is independent of the number of blocks M and the total time horizon T.The passage states T = Mτ.Source: C Sufficient conditions for Assumption 3 · pp. 35, 37

D Proof of multi-block PiX-MC

The section proves a convergence result for multi-block PiX-MC under Assumptions 1–3 and sufficiently small parameters. It also establishes geometric decay of the error during parallel refinement.

D Proof of multi-block PiX-MC

  • Theorem 5 establishes a convergence result for the law of the continuous interpolation generated by multi-block PiX-MC.The result assumes Assumptions 1–3 and sufficiently small η and γ.Source: D Proof of multi-block PiX-MC · p. 38
  • Theorem 5 requires η and γ to satisfy its stated parameter conditions.The proof notes that the conditions involving ρ and A4 match those in Theorem 4.Source: D Proof of multi-block PiX-MC · p. 38
  • The parallel-refinement error decays geometrically fast under Assumptions 1–3 and sufficiently small η and γ.This is the detailed version of Lemma 3 used in the proof of Theorem 5.Source: D Proof of multi-block PiX-MC · p. 38

E Proof of multi-block APiX-MC

The proof establishes convergence control for multi-block APiX-MC with weighted annealing. It separates score-approximation error from a transient annealing error that vanishes when annealing ends after finitely many iterations.

E Proof of multi-block APiX-MC

  • Weighted scores replace the score-error term δ_n with δ_n +(α_n −1)R_S in the local interpolation bound, which contributes to integrated discrepancy.The integrated bound also includes discretization terms involving γ, dimension d, and the Lipschitz constants.Source: E Proof of multi-block APiX-MC · pp. 41, 43, 44
  • Under Assumptions 1, 3, and 4 and sufficiently small parameters, the theorem bounds effective discrepancy by score error and a transient annealing term that vanishes when α_n becomes 1 after finitely many iterations.The theorem’s parameter conditions include positivity constraints on expressions involving η and γ.Source: E Proof of multi-block APiX-MC · p. 41
  • Under the stated assumptions, in-block Picard refinement error decays geometrically, with the same bound holding across all global nodes.The proof establishes a geometric recursion for each block and refinement, then uses uniformity across blocks to obtain the global estimate.Source: E Proof of multi-block APiX-MC · pp. 41–42

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

The experiments use shared image scaling and task-specific score-network implementations, likelihood balancing, and annealing schedules. Hyperparameters are tuned by method and task, with paired sequential and Picard methods matched on key settings.

F.1 Common Implementation Details and Hyperparameter Selections

  • Images are scaled to [0, 1], and score-network inputs and outputs are rescaled between [0, 1] and [−1, 1].Pretrained natural-image, MRI, and CT models use different output parameterizations.Source: §F.1 · p. 44
  • Natural-image experiments use a pretrained DDPM [36] model, MRI uses the pretrained PMC score network [86], and CT uses an EDM-trained MMSE denoiser [44].The passage specifies that the MRI network directly estimates the positive score; model-specific parameterizations are described separately.Source: §F.1 · p. 44
  • A task-dependent likelihood weight λ balances the differing numerical scales of likelihoods and score networks.The likelihood is described for MRI, deblurring, and CT, with A denoting a task-specific forward operator.Source: §F.1 · p. 45
  • Annealing uses geometric schedules for noise and score-network weights, with a non-increasing schedule that reaches one at σmin.The schedule parameter a is 2 for MRI and CT, and 5/2 for deblurring and Rician denoising.Source: §F.1 · p. 45
  • Deblurring and CT use adaptive Picard refinement, stopping at K = 20 or when the normalized residual εm,k/d falls below 3 × 10−4.The residual is computed empirically within each block.Source: §F.1 · p. 45
  • Hyperparameters are tuned separately for each method and task, while paired sequential and Picard methods share stepsize, score schedule, likelihood weight, and time-node count.Gradient-based and proximal variants receive method-specific stepsizes; the implicit likelihood treatment in proximal variants may improve stability and permit larger stepsizes.Source: §F.1 · p. 45

F.2 Likelihood and Proximal Implementations

This section specifies task-specific likelihood and proximal implementations for MRI, Rician denoising, image deblurring, and 3D CT, using Fourier-domain computations or iterative solvers as appropriate.

MRI reconstruction

  • For radial-mask MRI, the Fourier-plus-mask model yields a proximal inverse evaluable elementwise in k-space, avoiding a large-scale linear solve.The forward operator is the discrete Fourier transform followed by a radial subsampling mask.Source: §F.2 · pp. 45–46

Rician denoising

  • For Rician denoising, the likelihood is Lipschitz smooth [69], its gradient uses B(x)=I1(x)/I0(x), and pixelwise proximal subproblems use IRL1.IRL1 is capped at 10 iterations, though it typically solves the proximal subproblem within 5.Source: §F.2 · pp. 46–47

Image deblurring

  • For image deblurring, circular convolution is diagonalized by the discrete Fourier transform, enabling the proximal operator to use elementwise multiplication and division.The likelihood and its gradient are specified for a circular-convolution blur model [69].Source: §F.2 · p. 47

3D CT reconstruction

  • For 3D CT, the proximal linear system uses five conjugate-gradient iterations, with inter-slice Huber-TV promoting consistency between adjacent slices.The experiments set δH=10^-8 and βH=0.28 for all methods.Source: §F.2 · pp. 47–48

Related papers

Loading 2608.17666v1…