Poisson flow consistency models for low-dose CT image denoising (2402.08159v1)
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.
- A. B. de González and S. Darby, “Risk of cancer from diagnostic X-rays: estimates for the UK and 14 other countries,” The Lancet (British edition), vol. 363, no. 9406, pp. 345–351, 2004
- D. J. Brenner and E. J. Hall, “Computed Tomography — An Increasing Source of Radiation Exposure,” The New England journal of medicine, vol. 357, no. 22, pp. 2277–2284, 2007.
- G. Wang, J. C. Ye, and B. De Man, “Deep learning for tomographic image reconstruction,” Nature Mach. Intell., vol. 2, no. 12, pp. 737–748, 2020.
- M. J. Willemink, M. Persson, A. Pourmorteza, N. J. Pelc, and D. Fleischmann, “Photon-counting CT: Technical principles and clinical prospects,” Radiology, vol. 289, no. 2, pp. 293–312, 2018.
- T. Flohr, M. Petersilka, A. Henning, S. Ulzheimer, J. Ferda, and B. Schmidt, “Photon-counting CT review,” Phys. Med. 79, P126-136, 2020.
- M. Danielsson, M. Persson, and M. Sjölin, “Photon-counting x-ray detectors for CT,” Phys. Med. Biol., vol. 66, no. 3, pp. 03TR01–03TR01, 2021.
- S. S. Hsieh, S. Leng, K. Rajendran, S. Tao, and C. H. McCollough, “Photon Counting CT: Clinical Applications and Future Developments,” IEEE Trans. Radiat. Plasma Med. Sci., vol. 5, no. 4, pp. 441–452, 2021.
- J. Wang, T. Li, H. Lu, and Z. Liang, “Penalized weighted least-squares approach to sinogram noise reduction and image reconstruction for low-dose X-ray computed tomography,” IEEE Trans. Med. Imag., vol. 25, no. 10, pp. 1272–1283, Oct. 2006.
- J. B. Thibault, K. D. Sauer, C. A. Bouman, and J. Hsieh, “A three dimensional statistical approach to improved image quality for multislice helical CT,” Med. Phys., vol. 34, no. 11, pp. 4526–4544, Nov. 2007.
- E. Y. Sidky and X. Pan, “Image reconstruction in circular cone-beam computed tomography by constrained, total-variation minimization,” Phys. Med. Biol., vol. 53, no. 17, pp. 4777–4807, Sep. 2008.
- H. Yu, S. Zhao, E. A. Hoffman, and G. Wang, “Ultra-low dose lung CT perfusion regularized by a previous scan,” Acad. Radiol., vol. 16, no. 3, pp. 363–373, Mar. 2009.
- Z. Tian, X. Jia, K. Yuan, T. Pan, and S. B. Jiang, “Low-dose CT reconstruction via edge-preserving total variation regularization,” Phys. Med. Biol., vol. 56, no. 18, p. 5949, 2011.
- J. W. Stayman, H. Dang, Y. Ding, and J. H. Siewerdsen, “PIRPLE: a penalized-likelihood framework for incorporation of prior images in CT reconstruction,” Phys. Med. Biol., vol. 58, no. 21, pp. 7563–7582, 2013.
- G. L. Zeng and W. Wang, “Does noise weighting matter in CT iterative reconstruction?” IEEE Trans. Radiat. Plasma Med. Sci., vol. 1, no. 1, pp. 68–75, Jan. 2017.
- P. J. La Riviere, “Penalized-likelihood sinogram smoothing for low-dose CT,” Med. Phys., vol. 32, no. 6, pp. 1676–1683, Jun. 2005.
- Y. Zhang, Y. Xi, Q. Yang, W. Cong, J. Zhou, and G. Wang, “Spectral CT reconstruction with image sparsity and spectral mean,” IEEE Trans. Comput. Imag., vol. 2, no. 4, pp. 510–523, Dec. 2016.
- J. M. Wolterink, T. Leiner, M. A. Viergever, and I. Isgum, “Generative Adversarial Networks for Noise Reduction in Low-Dose CT,” IEEE Trans. Med. Imag., vol. 36, no. 12, pp. 2536–2545, 2017.
- B. Kim, M. Han, H. Shim, and J. Baek, “A performance comparison of convolutional neural network‐based image denoising methods: The effect of loss functions on low‐dose CT images,” Med. Phys., vol. 46, no. 9, pp. 3906–3923, 2019.
- K. Kim, S. Soltanayev and S. Y. Chun, “Unsupervised Training of Denoisers for Low-Dose CT Reconstruction Without Full-Dose Ground Truth,” in IEEE J. Sel. Top., vol. 14(6):1112-1125, 2020,
- N. Yuan, J. Zhou, and J. Qi, “Half2Half: deep neural network based CT image denoising without independent reference data,” Phys. Med. Biol., vol. 65, no. 21, pp. 215020–215020, 2020.
- Y. Mäkinen, L. Azzari, A. Foi, 2020, “Collaborative Filtering of Correlated Noise: Exact Transform-Domain Variance for Improved Shrinkage and Patch Matching”, IEEE Trans. Image Process., vol. 29, pp. 8339-8354.
- Z. Li, S. Zhou, J. Huang, L. Yu, and M. Jin, “Investigation of Low-Dose CT Image Denoising Using Unpaired Deep Learning Methods,” IEEE Trans. Radiat. Plasma Med. Sci., vol. 5, no. 2, pp. 224–234, 2021.
- S. Wang, Y. Yang, Z. Yin, A.S Wang, “Noise2Noise for denoising photon counting CT images: generating training data from existing scans”, Proc. SPIE 12463, Medical Imaging 2023: Physics of Medical Imaging, 1246304, 2023;
- X. Liu, Y. Xie, J. Cheng, S. Diao, S. Tan, and X. Liang, “Diffusion Probabilistic Priors for Zero-Shot Low-Dose CT Image Denoising,” 2023, [Online]. Available: https://arxiv.org/abs/2305.15887.
- Niu, C., Li, M., Guo, X., and Wang, G., “Self-supervised dual-domain network for low-dose CT denoising,” in Developments in X-Ray Tomography XIV, vol. 12242. SPIE, 2022.
- J. Sohl-Dickstein, E. A. Weiss, N. Maheswaranathan, and S. Ganguli, “Deep unsupervised learning using nonequilibrium thermodynamics,” in PMLR, 2015.
- J. Ho, A. Jain, and P. Abbeel. “Denoising diffusion probabilistic models,” in Proc. NeurIPS, 2020.
- A. Q. Nichol and P. Dhariwal. “Improved denoising diffusion probabilistic models,” in Proc. ICML, volume 139, pages 8162–8171, 2021.
- Y. Song, J. Sohl-Dickstein, D. P. Kingma, A. Kumar, S. Ermon, and B. Poole, “Score-Based Generative Modeling through Stochastic Differential Equations,” in Proc. ICLR, 2021.
- J. Song, C. Meng, and S. Ermon. “Denoising diffusion implicit models,” in Proc. ICLR, 2021.
- T. Karras, M. Aittala, T. Aila, and S. Laine, “Elucidating the Design Space of Diffusion-Based Generative Models,” in Proc. NeurIPS, 2022.
- Y. Xu, Z. Liu, M. Tegmark, and T. Jaakkola, “Poisson flow generative models,” in Proc. NeurIPS, 2022.
- Y. Xu, Z. Liu, Y. Tian, S. Tong, M. Tegmark, and T. Jaakkola, “PFGM++:Unlocking the potential of physics-inspired generative models”, in Proc. ICML, 2023.
- G. Batzolis, J. Stanczuk, C.-B. Schönlieb, C. Etmann, “Conditional Image Generation with Score-Based Diffusion Models,” 2021, [Online]. Available: https://arxiv.org/abs/2111.13606.
- H. Chung, B. Sim, and J. C. Ye, “Come-Closer-Diffuse-Faster: Accelerating Conditional Diffusion Models for Inverse Problems through Stochastic Contraction,” 2021, [Online]. Available: https://arxiv.org/abs/2112.05146.
- Y. Song, L. Shen, L. Xing, and S. Ermon, “Solving Inverse Problems in Medical Imaging with Score-Based Generative Models,” 2021, [Online]. Available: https://arxiv.org/abs/2111.08005.
- C. Saharia, W. Chan, H. Chang, C. A. Lee, J. Ho, T. Salimans, D. J. Fleet, and M. Norouzi. “Palette: Image-to-image diffusion models.” In Proc. SIGGRAPH, 2022.
- C. Saharia, J. Ho, W. Chan, T. Salimans, D. J. Fleet, and M. Norouzi, “Image Super-Resolution via Iterative Refinement,” IEEE Trans. Pattern Anal. Mach. Intell., vol. 45, no. 4, pp. 4713–14, 2023
- H. Chung, E. S. Lee, and J. C. Ye, “MR Image Denoising and Super-Resolution Using Regularized Reverse Diffusion,” IEEE Trans. Med. Imag., vol. 42, no. 4, pp. 922–934, 2023.
- R. Ge, Y. He, C. Xia, Y. Chen, D. Zhang, and G. Wang, “JCCS-PFGM: A Novel Circle-Supervision based Poisson Flow Generative Model for Multiphase CECT Progressive Low-Dose Reconstruction with Joint Condition,” [Online]. Available: https://arxiv.org/abs/2306.07824.
- Y. Song, P. Dhariwal, M. Chen, and I. Sutskever, “Consistency Models,” in Proc. ICML, 2023.
- Y. Song, and P. Dhariwal, “Improved Techniques for Training Consistency Models,” 2023. [Online]. Available: https://arxiv.org/abs/2310.14189.
- AAPM, “Low dose CT grand challenge,” 2017. [Online]. Available: https://www.aapm.org/grandchallenge/lowdosect/.
- P. Dhariwal, and A. Nichol, “Diffusion models beat GANs on image synthesis neurips,” in Proc. NeurIPS, 2021.
- R. Zhang, P. Isola, A. A. Efros, E. Shechtman, and O. Wang, “The Unreasonable Effectiveness of Deep Features as a Perceptual Metric,” in in Proc. CVPR, 2018, pp. 586–595.
- Z. Wang, A. C. Bovik, H. R. Sheikh, and E. P. Simoncelli, “Image quality assessment: from error visibility to structural similarity,” IEEE Trans. Image Process., vol. 13, no. 4, pp. 600–612, 2004.
- B. Shahriari, K. Swersky, Z. Wang, R. P. Adams, and N. de Freitas, “Taking the Human Out of the Loop: A Review of Bayesian Optimization,” Proceedings of the IEEE, vol. 104, no. 1, pp. 148–175, 2016.
- M. Heusel, H. Ramsauer, T. Unterthiner, B. Nessler, and S. Hochreiter, “GANs Trained by a Two Time-Scale Update Rule Converge to a Local Nash Equilibrium,” in Proc. NeurIPS, 2017.