Improving 2D Diffusion Models for 3D Medical Imaging with Inter-Slice Consistent Stochasticity

AI summary

Overview

• Research area: Medical image reconstruction (3D inverse problems in CT and MRI) using diffusion models as learned priors; listed under Computer Vision (cs.CV) on arXiv (arXiv:2602.04162v2).
• Technical level: Intermediate — the paper assumes some familiarity with diffusion forward/reverse processes and with diffusion-based inverse problem solvers, but its central idea (correlating the random noise injected across slices) is conceptually simple.
• Scope: The paper diagnoses uncoordinated stochasticity as the cause of inter-slice discontinuity when a 2D-trained diffusion prior is applied slice-by-slice to 3D medical volumes, and proposes a plug-and-play fix, Inter-Slice Consistent Stochasticity (ISCS), that correlates per-slice noise using spherical linear interpolation.

What This Paper Is About

Training diffusion models directly on 3D medical volumes is prohibitively expensive in memory, compute, and data, so a common workaround is to train on 2D slices and apply the model slice-by-slice to reconstruct 3D volumes. Because each 2D slice is denoised independently, the randomness built into diffusion sampling pulls adjacent slices along uncorrelated trajectories, producing visible discontinuities and artifacts along the slice (z) axis. The paper's goal is to fix this at the source — inside the sampling process itself — rather than by bolting on a smoothing regularizer such as Total Variation after the fact.

Key Contributions

1. Diagnosis of the root cause. The authors identify uncoordinated stochasticity in the re-noising step of slice-wise diffusion sampling as the fundamental driver of inter-slice inconsistency in 2D-prior-based 3D medical reconstruction, drawing a parallel to the temporal flickering observed when 2D image diffusion models are applied to video inverse problems.
2. The ISCS method. They propose Inter-Slice Consistent Stochasticity, which replaces independent Gaussian noise per slice with a smoothly correlated noise volume generated by Spherical Linear Interpolation (Slerp) between two anchor noise vectors sampled at the ends of the volume.
3. Plug-and-play integration with no added cost. ISCS is inserted directly into the re-noising step of existing diffusion inverse problem solvers such as DDNM and DDS, requiring no new loss terms, no hyperparameters, no retraining of the diffusion prior, and, per the authors, no additional computational cost (it is described as improving results without extra optimization steps).
4. Broad experimental validation plus a targeted ablation. The strategy is tested on sparse-view CT, limited-angle CT, and MRI isotropic super-resolution, and compared against an "identical noise" alternative (batch-consistent sampling, BCS) to show why interpolation is preferable to simply reusing one noise map for all slices.

Main Findings

• ISCS improves reconstructions across all three tasks and three viewing planes. In the reported tables, adding ISCS to DDNM and to DDS yields higher PSNR/SSIM and lower LPIPS than the unmodified solvers in the axial, coronal, and sagittal views.
• Sparse-view CT (30 views, AAPM data). DDS+ISCS reaches PSNR/SSIM/LPIPS of 36.97 / 0.937 / 0.064 (axial), 37.75 / 0.944 / 0.070 (coronal), and 38.16 / 0.942 / 0.065 (sagittal), compared with DDS at 34.76 / 0.919 / 0.069, 35.12 / 0.906 / 0.149, and 35.33 / 0.904 / 0.141, and with DDS+TV at 36.26 / 0.931 / 0.073, 37.08 / 0.938 / 0.095, and 37.50 / 0.936 / 0.088. DDNM improves similarly when combined with ISCS (axial PSNR 32.55 to 33.97).
• Limited-angle CT ([0°, 100°]). DDS+ISCS gives 31.65 / 0.911 / 0.071 (axial), 32.90 / 0.917 / 0.082 (coronal), and 32.49 / 0.920 / 0.077 (sagittal), versus DDS+TV at 31.40 / 0.898 / 0.086, 33.33 / 0.906 / 0.110, and 32.83 / 0.909 / 0.104. In this task the picture is mixed by metric and view: ISCS leads on the axial view and on most SSIM/LPIPS entries, while DDS+TV is higher on coronal and sagittal PSNR.
• MRI isotropic super-resolution (5× downsampling along z). DDS+ISCS records 40.33 / 0.968 / 0.019 (axial), 39.84 / 0.965 / 0.035 (coronal), and 39.35 / 0.948 / 0.052 (sagittal), compared with DDS+TV at 40.12 / 0.958 / 0.031, 39.36 / 0.955 / 0.038, and 38.53 / 0.932 / 0.075. Because the pretrained MRI diffusion prior was trained on coronal slices, the authors note inconsistencies mainly appear in the axial and sagittal views.
• Inter-slice difference (|Δ|) is reduced for the ISCS variants relative to their baselines. On SVCT, |Δ| falls from 0.005588 (DDS) to 0.001835 (DDS+ISCS), and from 0.009342 (DDNM) to 0.001785 (DDNM+ISCS); on LACT, from 0.011592 (DDS) to 0.001966 (DDS+ISCS) and from 0.016443 (DDNM) to 0.001899 (DDNM+ISCS). On MRI SR, DDS+ISCS records 0.002096 versus DDS at 0.001853 and DDS+TV at 0.004732 — the metric does not move in the same direction for every solver there. The metric is the absolute gap between the reconstruction's inter-slice difference and that of the ground truth, with smaller being better.
• Identical noise (BCS) is too rigid for medical volumes. In the ablation with five runs per strategy, the Slerp-based ISCS outperforms identical noise: coronal PSNR 38.16 vs 38.00, SSIM 0.941 vs 0.937, LPIPS 0.074 vs 0.081; sagittal PSNR 38.78 vs 38.24, SSIM 0.937 vs 0.933, LPIPS 0.073 vs 0.081. The authors attribute BCS's streak artifacts along the z-axis to the fact that medical volumes have a much larger axial extent (for example, more than 300 slices for CT) than the short video sequences (under 16 frames) BCS was designed for, so forcing identical noise suppresses genuine anatomical variation and produces "copying artifacts."
• Consistency is achieved early in the sampling trajectory. Performance curves over diffusion timesteps show that with ISCS the inter-slice difference drops close to the ground-truth reference early, while PSNR improves and LPIPS decreases thereafter, whereas the DDS baseline maintains larger inter-slice gaps until later steps.
• ISCS avoids the blurring trade-off of TV regularization. The qualitative comparisons report that TV-based regularization suppresses discontinuities at the cost of over-smoothing and cartoon-like textures, while ISCS reduces inter-slice inconsistency while preserving sharp edges and fine anatomy.

Methodology in Plain English

• Setup. Start from a diffusion prior trained only on 2D image slices. To reconstruct a 3D volume, run the model on each slice independently inside an existing inverse-problem solver, which alternates three steps: predict a clean image from the current noisy one, correct that prediction so it agrees with the measured data, and then re-noise it back to the next timestep.
• The problem with re-noising. The re-noising step injects fresh random Gaussian noise. When the measurements are highly undersampled, they impose weak constraints, so this random noise is free to steer each slice wherever it likes — and independent noise per slice means neighboring slices drift apart.
• The fix. Instead of drawing independent noise for every slice, the authors draw only two anchor noise maps, for the first and last slice, and then fill in every intermediate slice by interpolating between them along the surface of a high-dimensional sphere (Slerp). The interpolation parameter is the slice's normalized position, (i - 1)/(S - 1).
• Why a sphere and not a straight line. In high dimensions, Gaussian samples concentrate tightly in a thin shell around a hypersphere of radius the square root of d (the Gaussian Annulus Theorem). Straight-line interpolation would cut through the interior of that shell and leave the typical set, so Slerp is used to travel along the geodesic and preserve vector norms and distributional statistics. The result is noise that is strongly correlated between nearby slices and decorrelates with distance, while each individual slice's noise map still follows a standard Gaussian.
• Why not just reuse one noise map. That is the BCS approach from video restoration. It is too strong a constraint for volumes with hundreds of anatomically varying slices, causing features to be improperly replicated across distinct slices.
• No retraining, no tuning. The correlated noise volume simply substitutes for the independent noise in the existing sampler update, so any 2D diffusion-based 3D reconstruction pipeline can adopt it.
• Experimental scope. Three tasks are evaluated: sparse-view CT with 30 views sampled over [0°, 360°), limited-angle CT with 100 views over [0°, 100°], and MRI isotropic super-resolution from 5× z-axis downsampling. CT uses the AAPM 2016 low-dose CT grand challenge dataset (5936 slices from 10 patients; 5410 slices from 9 patients for training and one patient, L506, held out; evaluation volume 256×256×300, resized to 256×256, HU clipped to [-1000, 1600] and rescaled to [0, 1]). MRI uses the public IXI dataset (T1-weighted; evaluation volume 256×256×150 at 1.2×0.9375×0.9375 mm³ spacing, resampled to isotropic 1 mm³ and padded to 256×256×256). The diffusion prior is a Variance Exploding model in the ncsnpp architecture, trained for 70 epochs with batch size 16, learning rate 2×10⁻⁴, Adam (β₁ = 0.9, ε = 10⁻⁸) without weight decay, on a single NVIDIA A100 GPU; the MRI prior is the pretrained coronal-slice checkpoint from Lee et al. (2023). All methods use 30 NFEs for CT and 20 NFEs for MRI with the same pretrained prior.
• Baselines and metrics. DDNM and DDS represent state-of-the-art diffusion solvers, augmented as DDNM+ISCS and DDS+ISCS; FDK and AD

Authors’ abstract

3D medical imaging is in high demand and essential for clinical diagnosis and scientific research. Currently, diffusion models (DMs) have become an effective tool for medical imaging reconstruction thanks to their ability to learn rich, high-quality data priors. However, learning the 3D data distribution with DMs in medical imaging is challenging, not only due to the difficulties in data collection but also because of the significant computational burden during model training. A common compromise is to train the DMs on 2D data priors and reconstruct stacked 2D slices to address 3D medical inverse problems. However, the intrinsic randomness of diffusion sampling causes severe inter-slice discontinuities of reconstructed 3D volumes. Existing methods often enforce continuity regularizations along the z-axis, which introduces sensitive hyper-parameters and may lead to over-smoothing results. In this work, we revisit the origin of stochasticity in diffusion sampling and introduce Inter-Slice Consistent Stochasticity (ISCS), a simple yet effective strategy that encourages interslice consistency during diffusion sampling. Our key idea is to control the consistency of stochastic noise components during diffusion sampling, thereby aligning their sampling trajectories without adding any new loss terms or optimization steps. Importantly, the proposed ISCS is plug-and-play and can be dropped into any 2D trained diffusion based 3D reconstruction pipeline without additional computational cost. Experiments on several medical imaging problems show that our method can effectively improve the performance of medical 3D imaging problems based on 2D diffusion models. Our findings suggest that controlling inter-slice stochasticity is a principled and practically attractive route toward high-fidelity 3D medical imaging with 2D diffusion priors. The code is available at: https://github.com/duchenhe/ISCS