Voxel-wise Medical Images Generalization for Eliminating Distribution Shift

Abstract

Nowadays, more and more machine learning methods are applied in the medical domain. Supervised Learning methods adopted in classification, prediction, and segmentation tasks for medical images always experience decreased performance when the training and testing datasets do not follow the i.i.d(independent and identically distributed) assumption. These distribution shift situations seriously influence machine learning applications’ robustness, fairness, and trustworthiness in the medical domain. Hence, in this paper, we adopt the CycleGAN(Generative Adversarial Networks) method to cycle train the CT(Computer Tomography) data from different scanners/manufacturers, which aims to eliminate the distribution shift from diverse data terminals, on the basis of our previous work[14]. However, due to the model collapse problem and generative mechanisms of the GAN-based model, the images we generated contained serious artifacts. To remove the boundary marks and artifacts, we adopt score-based diffusion generative models to refine the images voxel-wisely. This innovative combination of two generative models enhances the quality of data providers while maintaining significant features. Meanwhile, we use five paired patients’ medical images to deal with the evaluation experiments with SSIM(structural similarity index measure) metrics and the segmentation model’s performance comparison. We conclude that CycleGAN can be utilized as an efficient data augmentation technique rather than a distribution-shift-eliminating method. While the denoising diffusion model is more suitable for dealing with the distribution shift problem aroused by the different terminal modules. In addition, another limitation of generative methods applied in medical images is the difficulty in obtaining large and diverse datasets that accurately capture the complexity of biological structure and variability. In future works, we will evaluate the original and generative datasets by experimenting with a broader range of supervised methods. We will implement the generative methods under the federated learning architecture, which can preserve their benefits and eliminate the distribution shift problem in a broader range.