TLDR: Poisson Flow Consistency Training (PFCT) is a novel method that allows Poisson Flow Consistency Models (PFCMs) to be trained independently, eliminating the need for a pre-trained model. By leveraging the perturbation kernel and introducing a sinusoidal discretization schedule and Beta noise distribution, PFCT enhances model stability and sample quality. Applied to low-dose CT image denoising, PFCT demonstrated improved image quality metrics (LPIPS, SSIM, PSNR) and achieved competitive results against other generative models, establishing it as a flexible and effective approach for medical imaging and opening new avenues for future research.
Deep Generative Modeling has made significant strides in creating realistic data across various fields, including audio and image generation. Among these, Diffusion Models and their physically inspired counterparts, like the Poisson Flow Generative Model++ (PFGM++), have set new standards. These models are powerful but can be computationally intensive, especially during the sampling process.
To address the need for faster sampling while retaining the benefits of Diffusion Models, Consistency Models (CMs) were developed. These models can generate high-quality images in a single step. The Poisson Flow Consistency Model (PFCM) extends this concept, applying single-step generation to PFGM++ and showing promise in areas like computed tomography (CT) image denoising.
However, a significant limitation of PFCMs has been their training process. Traditionally, PFCMs could only be trained through a method called distillation, which requires a pre-trained PFGM++ model. This dependency meant that if a pre-trained PFGM++ was unavailable or too costly to train, PFCMs could not be utilized. Until now, there hasn’t been a practical way to train a PFCM independently.
This research introduces a groundbreaking methodology called Poisson Flow Consistency Training (PFCT), which enables the training of PFCMs in isolation. This innovation removes the reliance on pre-trained models, offering much greater flexibility for researchers and practitioners. The core of PFCT involves leveraging the perturbation kernel to sample adjacent noisy data without needing a pre-trained PFGM++. It also incorporates advanced optimization strategies, including a sinusoidal discretization schedule and a Beta noise distribution, to enhance model stability and improve the quality of generated samples.
The methodology for PFCT involves a few key adaptations. To generate adjacent perturbed data, the model samples a single uniform angle while varying the radius component of the perturbation kernel. Drawing inspiration from Improved Consistency Training (iCT), PFCT removes the exponential moving average from training and utilizes the Pseudo-Huber Loss as a distance metric. A sinusoidal discretization schedule, which offers a smoother progression of timesteps compared to exponential schedules, is introduced to facilitate adaptability. Furthermore, a Beta noise distribution is used to select noise levels, aligning with the principle that high noise scheduling is crucial for effective training.
The effectiveness of PFCT was rigorously evaluated using the Low-Dose CT Image Denoising and Projection Dataset. Low-dose CT scans are vital for reducing patient exposure to radiation, but they often result in noisy images that can hinder diagnosis. Denoising these images without compromising diagnostic quality is a critical task in medical imaging. The model was implemented using a U-Net style architecture, enhanced with a Weighted Attention Gate for improved conditional image generation.
The results demonstrated that PFCT significantly improved the quality of low-dose CT images. When compared to the original low-dose CT scans, PFCT achieved a lower LPIPS (Learned Perceptual Image Patch Similarity), indicating better perceptual quality, and higher SSIM (Structural Similarity Index) and PSNR (Peak Signal Noise Ratio) scores, reflecting increased pixel-wise similarity and structural integrity. While PFCT did not outperform all existing generative AI methods, such as certain diffusion-based models or the distillation-trained PFCM, it showed competitive performance with other single-step generative models like ICCM and even surpassed the original Consistency Model in SSIM and PSNR. This establishes PFCT as a viable and promising method for medical image denoising.
Also Read:
- Unlocking Image Diversity with Latent Graphs and GFlowNets
- Unpacking Diffusion Models: A Journey from Core Concepts to Cutting-Edge Techniques
The introduction of PFCT marks a significant step forward, allowing the PFCM to function as a standalone Consistency Model. This newfound ability to train PFCMs in isolation opens up possibilities for their application in various data modalities beyond CT image denoising, such as unconditional image generation, which has yet to be fully explored. Further research will focus on optimizing PFCT’s parameters, such as the hyperparameter D, and conducting ablation studies to understand the impact of individual optimizations. This framework provides greater flexibility for how PFCMs can be developed and applied in the broader field of generative modeling. You can read the full research paper here.
