Modality-agnostic Domain Generalizable Medical Image Segmentation by Multi-Frequency in Multi-Scale Attention

TL;DR

This work tackles the challenge of generalizing medical image segmentation across multiple modalities by highlighting frequency variance as a key factor. It introduces MADGNet, composed of MFMSA to fuse multi-frequency and multi-scale features and E-SDM to mitigate information loss during deep supervision in multi-task learning. The approach achieves state-of-the-art performance across six modalities and fifteen datasets, including robust domain generalization to unseen clinical settings, with a reported parameter count around 31 million and efficient inference. The combination of frequency-aware channel attention, boundary-focused spatial attention, and ensemble-like decoding provides practical benefits for accurate, modality-agnostic segmentation in diverse clinical scenarios.

Abstract

Generalizability in deep neural networks plays a pivotal role in medical image segmentation. However, deep learning-based medical image analyses tend to overlook the importance of frequency variance, which is critical element for achieving a model that is both modality-agnostic and domain-generalizable. Additionally, various models fail to account for the potential information loss that can arise from multi-task learning under deep supervision, a factor that can impair the model representation ability. To address these challenges, we propose a Modality-agnostic Domain Generalizable Network (MADGNet) for medical image segmentation, which comprises two key components: a Multi-Frequency in Multi-Scale Attention (MFMSA) block and Ensemble Sub-Decoding Module (E-SDM). The MFMSA block refines the process of spatial feature extraction, particularly in capturing boundary features, by incorporating multi-frequency and multi-scale features, thereby offering informative cues for tissue outline and anatomical structures. Moreover, we propose E-SDM to mitigate information loss in multi-task learning with deep supervision, especially during substantial upsampling from low resolution. We evaluate the segmentation performance of MADGNet across six modalities and fifteen datasets. Through extensive experiments, we demonstrate that MADGNet consistently outperforms state-of-the-art models across various modalities, showcasing superior segmentation performance. This affirms MADGNet as a robust solution for medical image segmentation that excels in diverse imaging scenarios. Our MADGNet code is available in GitHub Link.

Figures (16)

• Figure 1: Scale vs Frequency distribution per modality. The scale denotes the size of lesions, measured as the ratio of foreground pixels to the total number of pixels. Frequency is calculated by the power spectrum ratio of the high-frequency and full-frequency. We observe that the frequency variance is higher than the scale, which mainly focuses on solving the various sizes in other methods.
• Figure 2: (a) The overall architecture of the proposed MADGNet mainly comprises MFMSA block and E-SDM (See Fig. \ref{['fig:parallel_vs_cascade']}). (b) MFMSA block contains $S$ scale branches ($S = 3$ in this figure) where the $s$-th branch input feature map are downsampled into $\eta^{s - 1}$ ($\eta = \frac{1}{2}$ in this figure). As our MFMSA block considers two dimensions (scale and frequency), MADGNet achieves the highest performance in various modalities and other clinical settings. Additionally, since E-SDM predicts a core task from sub-tasks, the final output is more accurate than when processed parallelly.
• Figure 3: Comparison of multi-task learning method between (a) parallel and (b) ensemble manners.
• Figure 4: Comparison of parameters (M) and inference speed (ms) vs segmentation performance (DSC) on average for all datasets.
• Figure 5: Qualitative comparison of other methods and MADGNet. (a) Input images. (b) UNet ronneberger2015u. (c) AttUNet oktay1804attention. (d) UNet++ zhou2018unet++. (e) CENet gu2019net. (f) TransUNet chen2021transunet. (g) FRCUNet azad2021deep, (h) MSRFNet srivastava2021msrf. (i) HiFormer heidari2023hiformer. (j) DCSAUNet xu2023dcsau. (k) M2SNet zhao2023m. (l) MADGNet (Ours). Green and Red lines denote the boundaries of the ground truth and prediction, respectively.