Anatomical Foundation Models for Brain MRIs

2024-08-07

arXiv.org

AI Summary

This work addresses the poor generalization of structural MRI-based neuroimaging diagnosis under limited labeled data. We propose AnatCL, an anatomy-guided foundation model. Methodologically, we introduce cortical thickness and other structural anatomical priors into a weakly supervised contrastive learning objective for the first time, enabling joint optimization of anatomical consistency and discriminative representation learning. The model is self-supervised pre-trained on large-scale unlabeled brain MRI data and supports cross-disease and cross-task transfer. Evaluated on diagnostic classification for 12 neuropsychiatric disorders—including Alzheimer’s disease, autism spectrum disorder, and schizophrenia—as well as prediction of 10 clinical rating scales, AnatCL achieves state-of-the-art performance, significantly outperforming existing methods. Notably, it demonstrates markedly improved robustness and generalization in data-scarce settings. The pre-trained model is publicly released.

Technology Category

• Machine Learning: Deep Neural Architectures and Foundation Models
• Computer Vision: Other Foundations of Computer Vision
• Natural Language Processing: Safety and Robustness

Application Category

• Web Mining and Content Analysis: Large pretrained models with web data
• Economics, Online Markets and Human Computation: Data quality aspects of human-annotated datasets
• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking

Abstract

Deep Learning (DL) in neuroimaging has become increasingly relevant for detecting neurological conditions and neurodegenerative disorders. One of the most predominant biomarkers in neuroimaging is represented by brain age, which has been shown to be a good indicator for different conditions, such as Alzheimer's Disease. Using brain age for weakly supervised pre-training of DL models in transfer learning settings has also recently shown promising results, especially when dealing with data scarcity of different conditions. On the other hand, anatomical information of brain MRIs (e.g. cortical thickness) can provide important information for learning good representations that can be transferred to many downstream tasks. In this work, we propose AnatCL, an anatomical foundation model for brain MRIs that i.) leverages anatomical information in a weakly contrastive learning approach, and ii.) achieves state-of-the-art performances across many different downstream tasks. To validate our approach we consider 12 different downstream tasks for the diagnosis of different conditions such as Alzheimer's Disease, autism spectrum disorder, and schizophrenia. Furthermore, we also target the prediction of 10 different clinical assessment scores using structural MRI data. Our findings show that incorporating anatomical information during pre-training leads to more robust and generalizable representations. Pre-trained models can be found at: https://github.com/EIDOSLAB/AnatCL.

Problem

• Develop anatomical foundation model for brain MRIs
• Improve diagnosis of neurological disorders via transfer learning
• Enhance MRI representation robustness using anatomical information

Innovation

• Leverages anatomical info via weakly contrastive learning
• Achieves state-of-the-art in multiple downstream tasks
• Uses brain age for weakly supervised pre-training