Source-linked AI summary
Toward a Foundation Plug-and-Play Prior for Computed Tomography Reconstruction via a Multimodal Diffusion Model
Haley Duba-Sullivan, Patxi Fernandez-Zelaia, Obaidullah Rahman, Amirkoushyar Ziabari
TL;DR
CT reconstruction must handle sparse-view and low-dose measurements, yet learned priors commonly require retraining when scan settings change. This paper uses one diffusion model trained across imaging domains as a frozen plug-and-play prior, with acquisition physics handled by the forward operator. Across three heterogeneous CT datasets, it substantially improves analytic reconstructions without retraining or test-time adaptation.
Problem
Learned CT priors often target one scan setting, creating a need for retraining or adaptation when modality, geometry, or material changes.
Method
A diffusion ViT trained across eight imaging domains serves as a shared frozen plug-and-play denoiser, while the forward operator represents acquisition physics.
Results
PnP-Diffusion-ViT substantially improves analytic reconstructions across all three datasets and achieves the lowest HFEN on each.
Takeaways & Limitations
A diffusion prior trained across several imaging domains can support CT reconstruction across changes in modality, beam geometry, material, and degradation type without retraining.
Abstract
from arXiv · showhide
Computed tomography (CT) throughput is limited by scan time, which grows with both the number of projections acquired and the detector integration time for each. Reconstructing high-quality volumes from sparse-view or low-dose measurements therefore depends on an informative prior, typically a neural network trained for one specific scan setting and retrained whenever the modality, geometry, or material changes. We investigate whether a single diffusion model trained across several imaging domains can instead serve as a prior for many CT problems simultaneously. We evaluate the proposed method using the same frozen model on three datasets that differ in modality, beam geometry, material, and degradation type, spanning flaw analysis in additively manufactured metal parts imaged with cone-beam X-ray CT and concrete microstructure imaged with parallel-beam neutron CT. Our proposed method out-performs analytic reconstructions in all three cases, providing a step toward a reusable foundation prior for heterogeneous CT reconstruction problems.
I. INTRODUCTION
CT throughput is constrained by acquisition time, while sparse-view and low-dose measurements make reconstruction ill-posed. The paper investigates a single frozen diffusion prior across heterogeneous CT settings without retraining or fine-tuning.
- Scan time grows with both the number of acquired projections and detector integration time per projection.
- Fewer projections create ill-posed reconstructions with streaking artifacts in analytic methods such as FBP and FDK.
- Shorter integration times reduce recorded photons and the signal-to-noise ratio of measured projections.
- Learned priors often target one modality, geometry, and material, requiring retraining or test-time adaptation when scan settings change.
- The proposed foundation prior uses one diffusion ViT trained across eight imaging domains and applies identical frozen weights and hyperparameters to three distinct CT datasets.
II. RELATED WORK
Plug-and-play reconstruction separates the denoiser from the acquisition model, while diffusion models provide noise-level-conditioned priors. The paper adopts a one-step denoising interpretation to reduce the cost of diffusion-based CT reconstruction.
- PnP methods replace a prior proximal operator with a denoiser inside an iterative algorithm such as ADMM.
- Diffusion models learn to reverse progressive Gaussian corruption and predict noise across a continuum of noise levels.
- Full reverse diffusion is costly because a volume may contain about a thousand slices and the prior is evaluated for every slice at each outer iteration.
- The method uses one broadly trained frozen diffusion model as a shared denoiser while leaving acquisition physics to the forward operator.
III. METHODOLOGY
The CT problem recovers a 3D attenuation volume from noisy projections using a regularized weighted least-squares formulation. Geometry and modality enter through the forward operator and measurement weighting, while regularization supplies the prior.
- CT reconstruction estimates a 3D volume of linear attenuation coefficients x from noisy measured projections y by inverting the measurement process.
- The model-based formulation combines a weighted data-fidelity term with a regularization function g(x).
- The forward operator A is specific to acquisition geometry and imaging modality, while Λ is derived from measurement statistics.
B. PnP-ADMM
PnP-ADMM decouples data fidelity from regularization through variable splitting and alternates physics-based, prior-based, and dual updates. The prior update can be implemented by a denoiser rather than an explicit regularizer.
- Variable splitting converts the reconstruction into a constrained problem that separates data fidelity and regularization with an equality constraint.
- ADMM alternates a data-fidelity update, a prior update, and a dual update, with u as the scaled dual variable and γ as a positive step-size parameter.
- The data-fidelity subproblem is determined entirely by acquisition physics and is solved from the corresponding normal equations.
- The regularization subproblem is a Gaussian-denoising MAP estimate and can be solved with Dθ(·, γ) without explicitly specifying g.
- Because the prior update excludes measurements and the forward operator, whether one denoiser transfers across scan settings becomes a central assumption.
C. Multimodal Diffusion Model
The method constructs a plug-and-play denoiser from a multimodal diffusion ViT and applies it across heterogeneous CT problems. Training combines simulated and experimental data from multiple imaging domains, while reconstruction uses slice-wise patch processing.
- Denoiser construction: The network predicts noise at discrete diffusion levels, so the paper builds a denoiser by applying a deterministic DDIM reverse map.The construction resolves the mismatch between the model’s normalized noise schedule and the physical denoising strength required by reconstruction.
- Denoiser construction: Input intensities are normalized using s = std(x0), and the requested denoising strength is represented in model space as σ = γ/s.The scale statistic is taken from the whole initial reconstruction and reflects object contrast rather than measurement noise.
- Inference: Each PnP iteration evaluates only one diffusion time step rather than a full reverse diffusion trajectory.This makes the per-iteration cost comparable to that of a conventional PnP denoiser.
- Training data: The diffusion model is trained across eight imaging domains spanning four modalities with both experimental and simulated data.The training domains include cell microscopy, powder micro-XCT, metal AM XCT, and concrete XCT, with equal mixtures of experimental and simulated examples.
- Inference: The 2D prior is applied independently to axial slices using overlapping 128×128 patches recombined by weighted overlap-add.A squared-sine window tapers patch boundaries to reduce seams.
A. Datasets
The evaluation uses three CT datasets spanning X-ray and neutron modalities, cone- and parallel-beam geometries, and nickel, steel, and concrete materials. Their degradations include sparse angular sampling, short integration time, noise, and scatter.
- Dataset scope: The three datasets cover X-ray and neutron CT, cone- and parallel-beam geometries, and nickel, steel, and concrete materials.The collection includes both simulated and measured acquisitions.
- Simulated Ni AM part: The simulated nickel AM dataset uses cone-beam X-ray projections with sparse angular sampling plus added noise and simulated scatter.The reference full scan has neither the added noise nor simulated scatter.
- Measured steel AM part: The measured steel AM dataset compares 10-minute and 60-minute acquisitions with different detector integration times.Both are 1000-view cone-beam X-ray scans; FDK reconstructs the input and MBIR reconstructs the reference.
- Concrete microstructure: The concrete neutron dataset uses 34 of 546 projections, corresponding to roughly a 16× reduction in angular sampling.This dataset’s degradation is purely angular and uses parallel-beam geometry over 180°.
B. Baselines and Metrics
The proposed PnP-Diffusion-ViT method is compared with analytic input reconstructions and PnP-BM3D under a matched iterative reconstruction setup. Performance is measured against reference reconstructions using PSNR, SSIM, and HFEN.
- Baselines: The comparison includes the proposed PnP-Diffusion-ViT method, the analytic input reconstruction, and PnP-BM3D.PnP-BM3D uses the same PnP-ADMM algorithm, iteration count, and nominal denoising strength as the proposed method.
- Evaluation protocol: Metrics are computed against the reference reconstruction over a masked center sub-volume comprising 16 slices.The mask is obtained by Otsu thresholding and hole filling so internal voids and flaws remain included.
- Metrics: PSNR and SSIM measure agreement with the reference, whereas lower HFEN indicates better reconstruction quality.Higher PSNR and SSIM are preferred.
C. Algorithmic Parameters
The algorithm uses a shared normalized denoising strength across the three datasets despite large differences in physical step size. This shared setting remains near independently optimized values and causes only a small PSNR reduction at worst.
- Parameterization: The physical step size γ differs by two orders of magnitude across the three datasets, so the algorithm uses normalized strength σ = γ/s.The shared value is selected as σ = 0.25 from a parameter sweep on the nickel dataset and then reused unchanged.
- Parameter transfer: All reconstructions use 5 PnP iterations.The shared denoising strength is evaluated without per-dataset retuning.
- Parameter transfer: The independently optimal normalized strengths are σ⋆= 0.25 for nickel and σ⋆= 0.29 for both steel and concrete.Thus, the shared σ = 0.25 remains close to each dataset’s independently optimized value.
- Results: The shared denoising strength gives the best HFEN on every dataset and the best PSNR and SSIM on the two measured datasets.For simulated nickel, PnP-BM3D leads on PSNR and SSIM because its assumed additive noise matches that dataset.
- Parameter transfer: 0.5 dB is the maximum PSNR reduction from using the shared value instead of each dataset’s optimum.The shared setting therefore transfers with limited PSNR loss across the tested datasets.
V. EXPERIMENTAL RESULTS
Across three datasets, the frozen PnP-Diffusion-ViT prior improves analytic reconstructions under different degradation types and preserves fine-scale structure better than PnP-BM3D in most comparisons. It also requires substantially less reconstruction time than PnP-BM3D.
- PnP-Diffusion-ViT suppresses noise, ring artifacts, and sparse-view streaking while retaining small flaws and fine structure that PnP-BM3D smooths or leaves artifacted.The remaining degradation is most apparent for concrete, which differs most strongly from the training distribution in modality and material.
- 15.8 dB, 13.0 dB, and 2.1 dB PSNR increases over analytic inputs are obtained for nickel, steel, and concrete, respectively.HFEN decreases by 80%, 83%, and 45% on the same datasets using one frozen prior and denoising strength.
- PnP-Diffusion-ViT achieves the best PSNR, SSIM, and HFEN on the measured steel and concrete datasets relative to PnP-BM3D.Its PSNR advantage is 2.1 dB on steel and 0.17 dB on concrete.
- PnP-BM3D achieves 0.48 dB higher PSNR and 0.086 higher SSIM on simulated nickel, while PnP-Diffusion-ViT achieves the lowest HFEN on all three datasets.The nickel advantage is consistent with the approximately additive white Gaussian noise used in that simulation, which matches BM3D’s assumed noise model.
- 2.1–9.6 s per slice is required by PnP-Diffusion-ViT, compared with 27.8–47.7 s for the standard CPU PnP-BM3D implementation.With K = 5, the proposed method uses one diffusion time step per iteration and therefore five network evaluations per slice.
VI. CONCLUSION
The paper demonstrates that one diffusion model trained across imaging domains can serve as a frozen PnP prior for heterogeneous CT problems without retraining or test-time adaptation. Across three datasets, it improves analytic reconstructions, preserves fine-scale structure, and supports transfer of denoising strength across datasets.
- A diffusion model trained across several imaging domains serves as a PnP prior across CT problems differing in modality, beam geometry, material, and degradation type without retraining or test-time adaptation.
- A single deterministic DDIM step per PnP iteration avoids a full reverse diffusion trajectory while retaining the diffusion model as an effective denoiser.
- Across all three datasets, the proposed prior substantially improves analytic reconstructions and achieves the lowest HFEN, indicating preservation of fine-scale structure.
- A normalized denoising strength selected on nickel transfers unchanged to steel and concrete with at most a 0.5 dB PSNR loss relative to per-dataset tuning.The corresponding physical step sizes differ by two orders of magnitude.
- Future work will compare the multimodal prior with architecture-matched single-domain models to isolate the benefit of multi-domain training.