PaperlayerView exact revision on arXiv

Source-linked AI summary

M3AE: Multimodal Representation Learning for Brain Tumor Segmentation with Missing Modalities

Hong Liu, Dong Wei, Donghuan Lu, Jinghan Sun, Liansheng Wang, Yefeng Zheng

arXiv:2303.05302v1eess.IVcs.CV

TL;DR

Brain-tumor segmentation uses complementary MRI modalities, yet one or more may be unavailable in practice. M3AE learns robust representations, substitutes a model-inverted full-modal image, and fine-tunes through self-distillation; it reports leading performance across two datasets.

Abstract

from arXiv · show

Multimodal magnetic resonance imaging (MRI) provides complementary information for sub-region analysis of brain tumors. Plenty of methods have been proposed for automatic brain tumor segmentation using four common MRI modalities and achieved remarkable performance. In practice, however, it is common to have one or more modalities missing due to image corruption, artifacts, acquisition protocols, allergy to contrast agents, or simply cost. In this work, we propose a novel two-stage framework for brain tumor segmentation with missing modalities. In the first stage, a multimodal masked autoencoder (M3AE) is proposed, where both random modalities (i.e., modality dropout) and random patches of the remaining modalities are masked for a reconstruction task, for self-supervised learning of robust multimodal representations against missing modalities. To this end, we name our framework M3AE. Meanwhile, we employ model inversion to optimize a representative full-modal image at marginal extra cost, which will be used to substitute for the missing modalities and boost performance during inference. Then in the second stage, a memory-efficient self distillation is proposed to distill knowledge between heterogenous missing-modal situations while fine-tuning the model for supervised segmentation. Our M3AE belongs to the 'catch-all' genre where a single model can be applied to all possible subsets of modalities, thus is economic for both training and deployment. Extensive experiments on BraTS 2018 and 2020 datasets demonstrate its superior performance to existing state-of-the-art methods with missing modalities, as well as the efficacy of its components. Our code is available at: https://github.com/ccarliu/m3ae.

Introduction

Four MRI modalities provide complementary information for brain-tumor sub-region analysis, but one or more are often missing in practice. M3AE is proposed as a single catch-all framework that combines masked representation learning, modality completion, and self-distillation.

Introduction

  • On BraTS 2018, M3AE achieves the best performance among four up-to-date methods while using a compact, efficient deployment model.Figure 1(c) compares model size, mean Dice scores across missing-modal situations, and GFLOPS.Source: Figure 1 · p. 1, Introduction · p. 2
  • Four MRI modalities provide complementary information: T1 and T1c highlight tumor core, while T2 and FLAIR highlight peritumoral edema.Missing modalities can arise from corruption, artifacts, acquisition protocols, contrast-agent allergy, or cost.Source: Introduction · p. 1
  • Existing dedicated approaches can require 2N −1 models for N modalities, while catch-all methods have often used complex multiple-encoder designs.The paper contrasts this with its single-model framework.Source: Introduction · pp. 1–2
  • M3AE combines multimodal masked autoencoding, model-inversion-based modality completion, and memory-efficient self-distillation in a single encoder-decoder framework.It is presented as a catch-all method for brain tumor segmentation with missing modalities.Source: Introduction · p. 2

Related Work

Prior work handles missing modalities with dedicated models, image synthesis, or catch-all models, often at computational or architectural cost. M3AE instead combines multimodal self-supervision and cross-situation distillation in a single-encoder, single-decoder design.

Multimodal Brain Tumor Segmentation with Missing Modalities

  • Dedicated methods distill from full-modal teachers to missing-modal networks, but require substantial computation and memory; M3AE uses one model for all missing-modal situations.Dedicated approaches can require 2N −1 models for N modalities.Source: Multimodal Brain Tumor Segmentation… · p. 2
  • GAN-based methods synthesize missing images, but 3D generation is difficult and may add overhead; M3AE instead optimizes a universal full-modal substitute that need not look realistic.The cited work notes that contrast images cannot be completely reproduced by generative models.Source: Multimodal Brain Tumor Segmentation… · p. 2
  • Catch-all methods commonly use multiple encoders and complex interactions, whereas M3AE uses a succinct single-encoder, single-decoder architecture.The paper describes this architecture as essentially a 3D U-Net.Source: Multimodal Brain Tumor Segmentation… · pp. 2–3

Self-Supervised Multimodal Representation Learning for Medical Image Analysis

  • M3AE masks both random modalities and patches of the remaining modalities, targeting global inter-modal correlations and local intra-modal anatomical integrity.The paper contrasts this with modality reconstruction that models global correlations but overlooks local structural integrity.Source: Self-Supervised Multimodal… · p. 3

Knowledge Distillation

  • Self-distillation transfers shared semantics across heterogeneous missing-modal situations within one network, unlike co-training’s fixed full-to-specific-missing pairing and dual-network memory cost.The paper reports that this strategy uses fewer training resources than prior co-training methods.Source: Knowledge Distillation · p. 3

Method

M3AE combines self-supervised multimodal reconstruction and model inversion in pretraining, then fine-tunes a single network with self-distillation for segmentation across missing-modality patterns.

Training and Inference Procedures

  • Training first jointly pretrains M3AE and the substitute image, then fine-tunes a segmentation head with supervised loss and consistency loss.For each case, the fine-tuning stage instantiates two inputs by randomly dropping zero to three modalities.Source: Method · p. 3, Training and Inference Procedures · pp. 4–5

Self-Supervised Multimodal Representation Learning via M3AE

  • M3AE masks random modalities and patches so reconstruction learning captures anatomical structure and inter-modal correlations.Whole-modality recovery uses global inter-modal correlation; patch recovery uses intramodal structure and local inter-modal correlation.Source: Self-Supervised Multimodal… · p. 4

Model Inversion based Modality Completion

  • Model inversion optimizes a full-modal substitute image by replacing masked content with location-matched content and minimizing reconstruction error.The substitute is updated alongside network parameters, adding no extra module and only marginal optimization cost.Source: Model Inversion based Modality… · p. 4

Fine-Tune with Heterogeneous Missing-Modal Self Distillation for Tumor Segmentation

  • Self-distillation aligns bottleneck features from two randomly sampled modality-dropout situations, transferring knowledge between heterogeneous inputs within one network.The situations can include full modalities; transfer from more modalities supports recovery of lost information, while reverse transfer enhances modality-specific features.Source: Fine-Tune with Heterogeneous… · p. 4

Training and Inference Procedures

  • At inference, the substitute fills any missing modalities before the trained catch-all segmentation network produces its output.The same model is intended for any subset, as well as the full set, of modalities.Source: Method · p. 3, Training and Inference Procedures · p. 5

Experiments and Results

Across missing-modality tests, M3AE achieved the highest mean DSC across modal combinations on both BraTS datasets. Full-modal evaluation and component ablations also support the reported performance of the framework and its components.

Datasets and Evaluation Metrics

  • Evaluation used BraTS 2018 and 2020 with DSC and HD95; training, validation, and testing splits contained 199:29:57 and 219:50:100 cases, respectively.Source: Datasets and Evaluation Metrics · p. 5

Comparison with State of the Art (SOTA)

Full-Modal Performance

Ablation Study

  • Removing M3AE pretraining or patch masking lowered average performance; optimized inversion beat zero-filling, while the population mean sharply hurt and self-distillation modestly improved both metrics across regions.Self-distillation saved approximately 4.7 million parameters compared with dual-network co-training.Source: Ablation Study · p. 7

Conclusion

M3AE establishes a new state of the art for brain tumor segmentation with missing modalities on two public benchmarks, using a single model for all modality combinations.

Conclusion

  • M3AE established a new state of the art for brain tumor segmentation with missing modalities on two public benchmark datasets.The experiments also found it competent in multimodal representation learning.Source: Conclusion · p. 7
  • As a catch-all model, M3AE accommodates all possible combinations of missing modalities with a single trained model.The conclusion frames this as one of the framework’s central properties.Source: Conclusion · p. 7
  • M3AE combines multimodal masked autoencoding, model inversion-based modality completion, and memory-efficient self-distillation between heterogeneous missing-modal situations.The components target robust representations, modality completion, and knowledge distillation, respectively.Source: Conclusion · p. 7
  • Ablative experiments validated the efficacy of M3AE’s three novel components.Source: Conclusion · p. 7
  • Future work will apply M3AE to different modalities, such as MRI and CT, and benchmarks beyond BraTS.Source: Conclusion · p. 7

Related papers

Loading 2303.05302v1…