Self Pre-training with Topology- and Spatiality-aware Masked Autoencoders for 3D Medical Image Segmentation

2024-06-15

arXiv.org

AI Summary

Existing MAE-based pretraining methods, designed for ViT architectures, struggle to capture the critical geometric structures and spatial relationships inherent in medical images, thereby limiting 3D segmentation performance. To address this, we propose a topology- and spatially-aware self-supervised pretraining framework: (1) a novel topological signature-based loss function that explicitly preserves anatomical structural integrity; (2) two new auxiliary tasks—3D cropping center localization and octant-point regression—to enhance spatial position understanding; and (3) joint pretraining of a ViT backbone with state-of-the-art segmentation networks in a hybrid architecture. Evaluated on five public 3D medical segmentation benchmarks, our method achieves average Dice score improvements of 2.1–4.7 percentage points over prior MAE approaches, with significantly enhanced generalization and robustness.

Technology Category

• Computer Vision: Segmentation
• Machine Learning: Unsupervised & Self-Supervised Learning
• Knowledge Representation and Reasoning: Geometric, Spatial, and Temporal Reasoning

Application Category

• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs
• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• Graph Algorithms and Modeling for the Web: Foundation models and LLMs for Web-related graphs

Abstract

Masked Autoencoders (MAEs) have been shown to be effective in pre-training Vision Transformers (ViTs) for natural and medical image analysis problems. By reconstructing missing pixel/voxel information in visible patches, a ViT encoder can aggregate contextual information for downstream tasks. But, existing MAE pre-training methods, which were specifically developed with the ViT architecture, lack the ability to capture geometric shape and spatial information, which is critical for medical image segmentation tasks. In this paper, we propose a novel extension of known MAEs for self pre-training (i.e., models pre-trained on the same target dataset) for 3D medical image segmentation. (1) We propose a new topological loss to preserve geometric shape information by computing topological signatures of both the input and reconstructed volumes, learning geometric shape information. (2) We introduce a pre-text task that predicts the positions of the centers and eight corners of 3D crops, enabling the MAE to aggregate spatial information. (3) We extend the MAE pre-training strategy to a hybrid state-of-the-art (SOTA) medical image segmentation architecture and co-pretrain it alongside the ViT. (4) We develop a fine-tuned model for downstream segmentation tasks by complementing the pre-trained ViT encoder with our pre-trained SOTA model. Extensive experiments on five public 3D segmentation datasets show the effectiveness of our new approach.

Problem

• Capturing geometric shape information in 3D medical images
• Aggregating spatial relationships for medical image segmentation
• Enhancing masked autoencoders for hybrid segmentation architectures

Innovation

• Topological loss preserves geometric shape information
• Pre-text task predicts 3D crop positions for spatial awareness
• Hybrid co-pretraining combines ViT with SOTA segmentation architecture