Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Revisiting MAE pre-training for 3D medical image segmentation

Tassilo Wald, Constantin Ulrich, Stanislav Lukyanenko, Andrei Goncharov, Alberto Paderno, Maximilian Miller, Leander Maerkisch, Paul F. Jäger, Klaus Maier-Hein

TL;DR

The study addresses the limited success of SSL in 3D medical image segmentation by leveraging a large-scale MAE pre-training regime on 39k brain MRI volumes with a Residual Encoder U‑Net within the nnU‑Net framework. A robust development framework across five development and eight test datasets enables reliable, performance-driven design choices, including sparsification, a mask-token strategy, and careful fine-tuning schedules. The approach yields consistent improvements over strong baselines and prior SSL methods, achieving approximately a 3 Dice-point advantage on 11 diverse downstream tasks, and demonstrating strong generalization across modalities and centers. This work emphasizes evaluation-driven development and demonstrates the practical potential of MAE-based pre-training for CNNs in 3D medical segmentation, with implications for data efficiency and cross-domain transfer.

Abstract

Self-Supervised Learning (SSL) presents an exciting opportunity to unlock the potential of vast, untapped clinical datasets, for various downstream applications that suffer from the scarcity of labeled data. While SSL has revolutionized fields like natural language processing and computer vision, its adoption in 3D medical image computing has been limited by three key pitfalls: Small pre-training dataset sizes, architectures inadequate for 3D medical image analysis, and insufficient evaluation practices. In this paper, we address these issues by i) leveraging a large-scale dataset of 39k 3D brain MRI volumes and ii) using a Residual Encoder U-Net architecture within the state-of-the-art nnU-Net framework. iii) A robust development framework, incorporating 5 development and 8 testing brain MRI segmentation datasets, allowed performance-driven design decisions to optimize the simple concept of Masked Auto Encoders (MAEs) for 3D CNNs. The resulting model not only surpasses previous SSL methods but also outperforms the strong nnU-Net baseline by an average of approximately 3 Dice points setting a new state-of-the-art. Our code and models are made available here.

Revisiting MAE pre-training for 3D medical image segmentation

TL;DR

The study addresses the limited success of SSL in 3D medical image segmentation by leveraging a large-scale MAE pre-training regime on 39k brain MRI volumes with a Residual Encoder U‑Net within the nnU‑Net framework. A robust development framework across five development and eight test datasets enables reliable, performance-driven design choices, including sparsification, a mask-token strategy, and careful fine-tuning schedules. The approach yields consistent improvements over strong baselines and prior SSL methods, achieving approximately a 3 Dice-point advantage on 11 diverse downstream tasks, and demonstrating strong generalization across modalities and centers. This work emphasizes evaluation-driven development and demonstrates the practical potential of MAE-based pre-training for CNNs in 3D medical segmentation, with implications for data efficiency and cross-domain transfer.

Abstract

Self-Supervised Learning (SSL) presents an exciting opportunity to unlock the potential of vast, untapped clinical datasets, for various downstream applications that suffer from the scarcity of labeled data. While SSL has revolutionized fields like natural language processing and computer vision, its adoption in 3D medical image computing has been limited by three key pitfalls: Small pre-training dataset sizes, architectures inadequate for 3D medical image analysis, and insufficient evaluation practices. In this paper, we address these issues by i) leveraging a large-scale dataset of 39k 3D brain MRI volumes and ii) using a Residual Encoder U-Net architecture within the state-of-the-art nnU-Net framework. iii) A robust development framework, incorporating 5 development and 8 testing brain MRI segmentation datasets, allowed performance-driven design decisions to optimize the simple concept of Masked Auto Encoders (MAEs) for 3D CNNs. The resulting model not only surpasses previous SSL methods but also outperforms the strong nnU-Net baseline by an average of approximately 3 Dice points setting a new state-of-the-art. Our code and models are made available here.

Paper Structure

This paper contains 39 sections, 3 figures, 12 tables.

Table of Contents

  1. Introduction
  2. Development Framework
  3. Pre-training Dataset
  4. Development Datasets
  5. Test Datasets
  6. Related Work
  7. Masked Autoencoders
  8. Masked Autoencoder In 3D Medical
  9. Other Pre-Training For 3D Medical
  10. Revisiting 3D MAEs
  11. Default parameters
  12. Sparsification
  13. Masking strategy
  14. Fine-tuning strategy
  15. Results and Discussion
  16. ...and 24 more sections

Figures (3)

  • Figure 1: Well-configured MAE pre-training for CNNs is state-of-the-art. When comparing our MAE pre-trained model on all eleven test Datasets, Spark3D achieves almost +3 DSC points compared to the strong nnU-Net baseline and outperforms current SSL methods.
  • Figure 2: S3D ranks best across all methods. In addition to absolute mean performance, we report the ranking stability of the methods through bootstrapping for all test datasets as well as the aggregated rank across all datasets.
  • Figure 3: Distribution of our pre-training dataset. The dataset stems from 44 centers and includes 8400 Patients with a 60 to 40 female-to-male ratio. Most patients were imaged with a 1.5 Tesla Philips Achieva or Ingenia scanner. The most prevalent modalities are T1 and T2-weighted images with some additional FLAIR images present. While other modalities were in the dataset, these were not used as prevalence was deemed too low.