PSI3D: Plug-and-Play 3D Stochastic Inference with Slice-wise Latent Diffusion Prior

TL;DR

PSI3D introduces a plug-and-play 3D stochastic inference method using slice-wise latent diffusion priors, enabling high-quality reconstruction of large-scale volumetric data like OCT with improved consistency and efficiency.

Contribution

The paper presents a novel MCMC-based algorithm for 3D inference that operates on large volumes by sampling 2D slices with a latent diffusion model and incorporating stochastic TV regularization.

Findings

Significantly improves OCT super-resolution reconstructions.

Outperforms traditional and learning-based baselines.

Provides robust and credible 3D reconstructions.

Abstract

Diffusion models are highly expressive image priors for Bayesian inverse problems. However, most diffusion models cannot operate on large-scale, high-dimensional data due to high training and inference costs. In this work, we introduce a Plug-and-play algorithm for 3D stochastic inference with latent diffusion prior (PSI3D) to address massive ($1024\times 1024\times 128$) volumes. Specifically, we formulate a Markov chain Monte Carlo approach to reconstruct each two-dimensional (2D) slice by sampling from a 2D latent diffusion model. To enhance inter-slice consistency, we also incorporate total variation (TV) regularization stochastically along the concatenation axis. We evaluate our performance on optical coherence tomography (OCT) super-resolution. Our method significantly improves reconstruction quality for large-scale scientific imaging compared to traditional and learning-based baselines, while providing robust and credible reconstructions.