Physics-Guided Diffusion Priors for Multi-Slice Reconstruction in Scientific Imaging

TL;DR

The paper tackles the ill-posed problem of multi-slice reconstruction under undersampling by integrating partitioned diffusion priors with physics-based forward models. The authors introduce two algorithms, DART and DRIFT, to fuse learned diffusion priors with data consistency, achieving substantial memory savings and improved reconstruction quality across MRI and 4D-STEM. They demonstrate robustness to out-of-distribution data and provide ablation studies and multi-metric evaluations (SSIM, FVD, JEDi). The work offers a scalable, generalizable framework for fast, physics-consistent multi-slice imaging in scientific applications.

Abstract

Accurate multi-slice reconstruction from limited measurement data is crucial to speed up the acquisition process in medical and scientific imaging. However, it remains challenging due to the ill-posed nature of the problem and the high computational and memory demands. We propose a framework that addresses these challenges by integrating partitioned diffusion priors with physics-based constraints. By doing so, we substantially reduce memory usage per GPU while preserving high reconstruction quality, outperforming both physics-only and full multi-slice reconstruction baselines for different modalities, namely Magnetic Resonance Imaging (MRI) and four-dimensional Scanning Transmission Electron Microscopy (4D-STEM). Additionally, we show that the proposed method improves in-distribution accuracy as well as strong generalization to out-of-distribution datasets.

Figures (5)

• Figure 1: Data acquisition of a brain in MRI (k-space) and 4D-STEM (diffraction patterns) of crystalline materials, as well as proposed methods: DART (alternating update between trained diffusion prior and physics constraints) and DRIFT (diffusion priors as initialization before applying physical constraints). Physical constraints $\mathcal{G}(\emph{X})$ are adapted depending on the modality (MRI or 4D-STEM).
• Figure 2: Example of cubic crystal materials: unit cells of gallium arsenide (GaAs) with volume dimension $5.6533^3 \AA^3$ (a),(b) projection and 3D visualization, strontium titanate (SrTiO$_3$) with volume dimension $3.905^3 \AA^3$ (c), (d) projection and 3D visualization. In the multi-slice method, each slice is obtained by calculating the projection of each atomic plane along the z direction.
• Figure 3: Single slice MRI with zoomed-in region of interest from the volume reconstruction of file BraTS20 Training 338 t1ce (top) and BraTS20 Training 039 t1ce (bottom); Projection-based bangun2025reg; DART; DRIFT; CS MRI lustig2007sparse; Total Variation block2007undersampled. Phase Projection of crystalline materials CoPt$_3$ (top) and Tb$_3$InC (bottom) benchmarking with Sparse Decomposition bangun2022inverse; DART; DRIFT; Torchslice diederichs2024exact; 3PIE maiden2012ptychographic.Top right are visual quality metrics, namely SSIM.
• Figure 4: Single slice MRI from the volume reconstruction of file soybean roots; Projection-based bangun2025reg; DART; DRIFT; CS MRI lustig2007sparse; Total Variation block2007undersampled. Phase Projection of crystalline materials WSe$_2$ benchmarking with Sparse Decomposition bangun2022inverse; DART; DRIFT; Torchslice diederichs2024exact; 3PIE maiden2012ptychographic. Top right are visual quality metrics, namely SSIM.
• Figure 5: Top row: run time and SSIM of pure (vanilla) diffusion models ho2022video, DART, and DRIFT. The experiments are conducted on a distributed process with 8 GPUs for multi-slice MRI data with dimension $155 \times 240 \times 240$. Bottom row: time and memory benchmarking for DART using various numbers of GPUs.