CAT-DG: A Cross-Attention-Based Domain Generalization Model for Medical Image Segmentation

In medical image segmentation tasks, the performance of the trained segmentation model in the unseen domain is affected by the domain shifting problem. Therefore, improving the model’s generalization is crucial for the practical application of intelligent models for medical images. The domain generalization model for medical image segmentation is typically divided into two stages: data augmentation and segmentation training, and most existing studies focus on improving the data enhancement stage, little consideration is given to the training stage. However, it is very important to consider the model’s generalization in the segmentation training stage to alleviate the domain shifting problem. To validate the idea, we propose the CAT-DG, a cross-attention-based domain generalization model for medical image segmentation. The model pays more attention to the image content information and ignores the style information during segmentation training phase, resulting in an average improvement of 1.7%–7% compared to other methods in unseen domains. Additionally, we propose a hybrid loss function that combines Dice loss and Focal loss, and the loss calculation incorporates the distillation idea to mitigate the impact of class imbalance in medical images on model performance, which improves accuracy by 3% compared to not using the hybrid loss. The detailed experiment results not only prove the effectiveness of the CAT-DG but also demonstrate that considering the model generalizability during segmentation training phase can furtherly alleviate domain shifting problems.