Poisson flow consistency models for low-dose CT image denoising

Dennis Hein,Adam Wang,Ge Wang
2024-02-13
Abstract:Diffusion and Poisson flow models have demonstrated remarkable success for a wide range of generative tasks. Nevertheless, their iterative nature results in computationally expensive sampling and the number of function evaluations (NFE) required can be orders of magnitude larger than for single-step methods. Consistency models are a recent class of deep generative models which enable single-step sampling of high quality data without the need for adversarial training. In this paper, we introduce a novel image denoising technique which combines the flexibility afforded in Poisson flow generative models (PFGM)++ with the, high quality, single step sampling of consistency models. The proposed method first learns a trajectory between a noise distribution and the posterior distribution of interest by training PFGM++ in a supervised fashion. These pre-trained PFGM++ are subsequently "distilled" into Poisson flow consistency models (PFCM) via an updated version of consistency distillation. We call this approach posterior sampling Poisson flow consistency models (PS-PFCM). Our results indicate that the added flexibility of tuning the hyperparameter D, the dimensionality of the augmentation variables in PFGM++, allows us to outperform consistency models, a current state-of-the-art diffusion-style model with NFE=1 on clinical low-dose CT images. Notably, PFCM is in itself a novel family of deep generative models and we provide initial results on the CIFAR-10 dataset.
Image and Video Processing,Computer Vision and Pattern Recognition
What problem does this paper attempt to address?
### Problems Addressed by the Paper This paper aims to address the issue of noise suppression in low-dose CT (Computed Tomography) images. Specifically, the authors propose a new image denoising technique that combines the advantages of the Poisson Flow Generative Model (PFGM)++ and Consistency Models to achieve high-quality single-step sampling. The goal of this method is to significantly reduce computational costs while maintaining image quality, thereby enabling efficient noise suppression in clinical low-dose CT images. ### Background and Motivation 1. **Need for Low-Dose CT**: - Low-dose CT is widely used in diagnosis and treatment planning, but due to the potential risks of radiation dose, it is necessary to minimize radiation dose while maintaining high diagnostic quality. - Photon Counting CT (PCCT) technology, although improving the quality of low-dose imaging, results in increased image noise due to the reduced number of photons per voxel or energy bin, given its higher resolution. 2. **Limitations of Existing Methods**: - Iterative reconstruction methods, while effective in noise suppression, are computationally expensive. - Common preprocessing and postprocessing methods have limitations in computational efficiency and noise suppression effectiveness. 3. **Development of Generative Models**: - Diffusion models and Poisson flow models have achieved significant success in generative tasks and image processing tasks, but their iterative nature leads to high sampling computational costs. - Consistency models achieve high-quality data generation through single-step sampling, but their performance is somewhat inferior to multi-step sampling diffusion models. ### Proposed Method 1. **Poisson Flow Generative Model (PFGM)++**: - By training PFGM++ in a supervised manner, it learns the trajectory from noise distribution to target posterior distribution. - PFGM++ balances flexibility and robustness by adjusting the hyperparameter D (dimension of the augmented variable). 2. **Consistency Models**: - Through Consistency Distillation, the pre-trained PFGM++ is "distilled" into a single-step sampling Poisson Flow Consistency Model (PFCM). - Consistency models ensure that any two points on the same solution trajectory map to the same initial point through self-consistency constraints. 3. **Posterior Sampling Poisson Flow Consistency Model (PS-PFCM)**: - Combining the flexibility of PFGM++ and the single-step sampling advantage of consistency models, a new image denoising framework is proposed. - By adjusting the hyperparameter D, performance superior to the current state-of-the-art diffusion models can be achieved in single-step sampling. ### Experimental Results 1. **Quantitative Evaluation**: - Evaluations were conducted on a low-dose CT validation dataset using metrics such as Peak Signal-to-Noise Ratio (PSNR), Structural Similarity Index (SSIM), and LPIPS. - PS-PFCM outperformed consistency models (CD) at D=4096, despite being distilled from a slightly inferior PFGM++. 2. **Qualitative Evaluation**: - Visual inspection showed that images generated by PS-PFCM had realistic noise characteristics and excelled in detail reproduction. - Particularly in handling liver metastases, PS-PFCM was able to clearly display lesion areas with noise levels close to normal-dose CT images. ### Conclusion This paper proposes a new method combining the Poisson Flow Generative Model and Consistency Models for noise suppression in low-dose CT images. By adjusting the hyperparameter D, PS-PFCM achieves high performance in single-step sampling, significantly reducing computational costs and providing an efficient, high-quality solution for clinical applications.