Enhanced Masked Image Modeling to Avoid Model Collapse on Multi-modal MRI Datasets

2024-07-15

AI Summary

Masked image modeling (MIM) in self-supervised pretraining of multimodal MRI is prone to model collapse—specifically, complete collapse (trivial constant outputs) and dimensional collapse (degenerate feature representations)—hindering representation learning. Method: This work first systematically identifies and distinguishes these two collapse mechanisms, proposing an enhanced MIM framework: (1) a Hybrid Masking Policy (HMP) to increase reconstruction difficulty and suppress complete collapse; and (2) a Pyramid Barlow Twins (PBT) module that enforces multi-scale feature consistency to improve representation uniformity and mitigate dimensional collapse. The framework jointly integrates MIM with contrastive learning regularization. Contribution/Results: Evaluated on three multimodal MRI datasets, the approach significantly improves pretraining stability and achieves state-of-the-art performance gains on downstream segmentation and classification tasks.

Technology Category

• Machine Learning: Multimodal Learning
• Computer Vision: Multi-modal Vision
• Intelligent Robots: Multimodal Perception & Sensor Fusion

Application Category

• Search and Retrieval-Augmented AI: Retrieval-Augmented Generation (RAG) and multi-modal RAG
• Web Mining and Content Analysis: Large pretrained models with web data
• User Modeling, Personalization and Recommendation: Attacks and countermeasures in recommendation systems

Abstract

Multi-modal magnetic resonance imaging (MRI) provides information of lesions for computer-aided diagnosis from different views. Deep learning algorithms are suitable for identifying specific anatomical structures, segmenting lesions, and classifying diseases. Manual labels are limited due to the high expense, which hinders further improvement of accuracy. Self-supervised learning, particularly masked image modeling (MIM), has shown promise in utilizing unlabeled data. However, we spot model collapse when applying MIM to multi-modal MRI datasets. The performance of downstream tasks does not see any improvement following the collapsed model. To solve model collapse, we analyze and address it in two types: complete collapse and dimensional collapse. We find complete collapse occurs because the collapsed loss value in multi-modal MRI datasets falls below the normally converged loss value. Based on this, the hybrid mask pattern (HMP) masking strategy is introduced to elevate the collapsed loss above the normally converged loss value and avoid complete collapse. Additionally, we reveal that dimensional collapse stems from insufficient feature uniformity in MIM. We mitigate dimensional collapse by introducing the pyramid barlow twins (PBT) module as an explicit regularization method. Overall, we construct the enhanced MIM (E-MIM) with HMP and PBT module to avoid model collapse multi-modal MRI. Experiments are conducted on three multi-modal MRI datasets to validate the effectiveness of our approach in preventing both types of model collapse. By preventing model collapse, the training of the model becomes more stable, resulting in a decent improvement in performance for segmentation and classification tasks. The code is available at https://github.com/LinxuanHan/E-MIM.

Problem

• Multi-modal MRI
• Self-supervised Learning
• Masked Image Modeling

Innovation

• E-MIM
• HMP
• PBT