Supervised Diffusion-Model-Based PET Image Reconstruction

1. Problem Statement

Positron Emission Tomography (PET) reconstruction is an ill-posed inverse problem, particularly when dealing with low-count data (low dose) to minimize radiation exposure. The measurement data follows a high-variance Poisson distribution, making reconstruction noisy and unstable.

• Limitations of Current Methods:
  • Model-based methods (e.g., MLEM, MAP-EM): Often struggle with noise and lack the ability to model complex image priors effectively.
  • Unsupervised Diffusion Models (DMs): Recent approaches use DMs trained on high-quality images as priors, conditioning on measured data via unsupervised schemes. While generalizable, they fail to explicitly learn the interaction between the DM prior and the specific noisy measurement data. This often leads to compromised quantitative accuracy and fidelity to the learned manifold.
  • Supervised Deep Learning (e.g., FBSEMnet): While accurate, these methods typically estimate only the posterior mean, failing to capture the full posterior distribution (uncertainty) and often producing blocky or blurred artifacts due to $L_2$ loss training.

The Core Challenge: How to leverage the generative power and posterior sampling capabilities of Diffusion Models while explicitly modeling the interaction between the prior and the specific Poisson-noisy PET measurements to achieve high quantitative accuracy and uncertainty estimation.

2. Methodology: PET-DEFT

The authors propose PET-DEFT (Doob's h-transform Efficient Fine Tuning), a supervised conditioning framework tailored for PET reconstruction. The method builds upon the DEFT framework but introduces specific modifications to address PET physics.

Key Components:

1. Two-Stage Training Strategy:

   • Pre-training: A score network $s_\theta$ is pre-trained on high-quality, normalized PET images to learn the data manifold (unconditional).
   • Supervised Fine-tuning: A conditional score network $s_\phi$ is trained on paired datasets of noisy sinograms ($m$) and ground truth images ($x$). This module learns to predict the residual score, effectively learning how to update the image estimate based on specific measurement data.

2. PET-Specific Modifications:

   • Measurement Normalization: PET images have a wide dynamic range. The authors introduce a scaling factor $c$ calculated from the Maximum-Likelihood Expectation-Maximization (MLEM) estimate ($x_{MLEM}^{10}$) to normalize the image space during diffusion updates and un-normalize for data fidelity. This ensures the model handles the intensity range effectively.
   • Non-Negativity Constraint: Since PET data follows a Poisson likelihood (domain $\mathbb{R}_{\ge 0}$), negative pixel values are physically impossible. The authors introduce a novel mechanism using a pReLU (parametric ReLU) on the conditional Tweedie's estimate. This forces the estimated image to be non-negative before calculating the score, ensuring the data fidelity updates remain valid.
   • Data Fidelity Integration: The network concatenates the noisy iterate ($x_t$), the unconditional Tweedie's estimate, and the gradient of the log-likelihood (data fidelity update) as input to $s_\phi$.

3. 3D Extension (2.5D Strategy):

   • To handle full 3D PET data efficiently, the authors decouple the physics operator from the neural network.
   • They train a 2.5D network that processes three consecutive transverse slices as input channels but predicts the score for the central slice only.
   • The 3D score estimate is formed by concatenating slice-wise outputs, allowing the method to scale to clinical 3D datasets without prohibitive memory costs.

4. Inference:

   • Reconstruction is performed using the DDIM (Denoising Diffusion Implicit Model) scheme with 100 steps (or as few as 5 with the non-negativity module active).
   • The method supports posterior sampling, allowing for the generation of multiple plausible reconstructions to estimate uncertainty.

3. Key Contributions

• First Supervised DM for PET: The paper presents the first supervised learning approach using Diffusion Models specifically for PET image reconstruction.
• Physics-Informed Architecture: It introduces critical PET-specific adaptations: measurement normalization to handle wide intensity ranges and a non-negativity module to satisfy Poisson likelihood constraints.
• Posterior Sampling: Unlike standard supervised DL methods that output a single mean image, PET-DEFT enables accurate posterior sampling, facilitating uncertainty quantification.
• Clinical Applicability: The method is validated on simulated phantoms and successfully extended to real-world, fully 3D $[^{18}\text{F}]\text{FDG}$ brain PET data.

4. Experimental Results

The method was evaluated on realistic brain PET phantoms and real clinical data across various dose levels.

• Quantitative Performance:
  • PET-DEFT matches or outperforms state-of-the-art supervised methods (specifically FBSEMnet) in terms of Normalized Root Mean Square Error (NRMSE) and Structural Similarity Index (SSIM) across 5 different dose levels.
  • It significantly outperforms unsupervised DM methods (PET-DDS, PET-LiSch) and traditional model-based approaches (MLEM, MAP-EM).
• Qualitative Analysis:
  • At low doses, PET-DEFT preserves fine-grained features better than FBSEMnet, which tends to produce blocky artifacts.
  • The images resemble blurred scans (due to averaging multiple samples) but retain structural integrity.
• Ablation Studies:
  • Normalization: Removing the measurement normalization step caused a significant drop in performance (NRMSE increased from 22.1% to 30.0%).
  • Non-Negativity: The non-negativity module allows the model to achieve peak accuracy with fewer diffusion steps (5 steps vs. 8 steps without the module).
  • Robustness: The model shows resilience to reduced parameter counts and smaller training datasets compared to other supervised methods.
• Uncertainty Estimation:
  • Samples generated by PET-DEFT show superior adherence to the ground truth manifold compared to unsupervised methods, as evidenced by lower LPIPS scores.
  • The method successfully generates coefficient of variation (CV) maps from real 3D data, visualizing uncertainty in specific anatomical regions.

5. Significance

This work bridges the gap between the generative power of Diffusion Models and the rigorous physical constraints required in medical imaging. By moving from unsupervised to supervised conditioning, the authors demonstrate that DMs can be trained to explicitly understand the relationship between noisy sinogram data and the underlying image distribution.

The significance lies in:

1. Improved Accuracy: Achieving state-of-the-art quantitative results while maintaining the ability to sample from the posterior distribution.
2. Uncertainty Quantification: Providing clinicians with not just a single image, but a distribution of possible images, which is crucial for diagnosing ambiguous features in low-dose scans.
3. Efficiency: The 2.5D training strategy makes the application of complex 3D diffusion models feasible on standard clinical hardware (e.g., 24GB GPUs), paving the way for real-world deployment in nuclear medicine.