Enhancing Modality-Agnostic Representations via Meta-learning for Brain Tumor Segmentation

Abstract

In medical vision, different imaging modalities provide complementary information. However, in practice, not all modalities may be available during inference or even training. Previous approaches, e.g., knowledge distillation or image synthesis, often assume the availability of full modalities for all subjects during training; this is unrealistic and impractical due to the variability in data collection across sites. We propose a novel approach to learn enhanced modality-agnostic representations by employing a meta-learning strategy in training, even when only limited full modality samples are available. Meta-learning enhances partial modality representations to full modality representations by meta-training on partial modality data and meta-testing on limited full modality samples. Additionally, we co-supervise this feature enrichment by introducing an auxiliary adversarial learning branch. More specifically, a missing modality detector is used as a discriminator to mimic the full modality setting. Our segmentation framework significantly outperforms state-of-the-art brain tumor segmentation techniques in missing modality scenarios.

Figure 1:

Comparison of the paradigms generally adopted by existing missing modality approaches (left) vs. ours (right) for brain tumor segmentation. $N$ and $n$ refer to the number of subjects (patients) with partial and full modalities, respectively. Previous methods either utilize full modality data ${𝒟}^f$ for all subjects or simulate partial modality data ${𝒟}^m$ from ${𝒟}^f$. On the contrary, our approach works in a limited full modality setting, i.e., $\left|{𝒟}^f\right|\le \left|{𝒟}^m\right|$.

Figure 2:

Framework overview. ${𝒟}^m$ (partial modality) and ${𝒟}^f$ (full modality) are used as inputs for encoder-decoder networks in the meta-train and meta-test phase, respectively. Partial modality representations are adapted to the full modality domain via: 1) meta-optimization of gradients in both data, and 2) adversarial learning based on predictions by a modality absence classifier.

Figure 3:

Illustration of the proposed framework. Available or full set of modalities are passed through a shared generator in the meta-train and meta-test stages respectively. The aggregation module helps to obtain a fused representation from five different levels (level $l$ and bottleneck are indicated here). Next, only the bottleneck embedding is used by the discriminator to predict which modalities are present at the input. All five fused embeddings are used by the segmentation decoder. Inner and outer loop gradient updates refer to the losses calculated in the meta-train and meta-test stages on partial modality and full modality data, respectively.

Figure 4:

Illustration of the feature aggregation module. Modality ${\mathbf{\text{F}}}_2^l$ is missing. ${\mathbf{\text{F}}}_1^l$ and ${\mathbf{\text{F}}}_3^l$ are passed through global average pooling (GAP) operation and eventually fed into an MLP to generate the shared representation ${\mathbf{\text{F}}}_{fused}^l$.

Figure 5:

Qualitative comparison. Column 1: four MRI modalities. Col 2–4: segmentation maps from three methods for different combinations of modalities. Col 5: Ground truth. Our method is able to better capture gaps/islands (rows 1,3) and boundaries (row 4) in TC segmentations.

Figure 6:

Comparison of baseline and enhanced features ACN, RFNet, and mmFormer by varying the full modality count from 100% to 40% (Fig. 7a). In order to retain sufficient samples for each combination task in meta-training, we assume that at least 50% of the subjects have partial modalities. Hence we show our results only on 50% and 40% proportions of full modality data. The fact that even with 50% full modality samples, we match the evaluation scores of SOTA at 100% setting is noteworthy. A sharp degradation can be noticed in the average WT DSC of SOTA once the number of full modality data decreases. On the other hand, our method shows only a minor drop of 0.29%. This is due to ACN being heavily dependent on the full modality for knowledge distillation. RFNet and mmFormer require full modality data as input to the network. They under-fit since their overall sample count decreases. Our method efficiently utilizes even limited samples of full modality data for feature adaptation in meta-testing. Owing to the above reasons, our approach is resilient to change in full modality proportion. Results for other tumor regions are provided in supplementary (Sec. 8).

Figure 7:

Ablation studies.