MRI-Based Brain Age Estimation with Supervised Contrastive Learning of Continuous Representation

TL;DR

This work tackles MRI-based brain age estimation by addressing the continuous nature of neuromorphological aging with a supervised contrastive learning framework using Rank-N-Contrast on 3D MRI. It demonstrates that RNC-highRez with Grad-RAM delivers state-of-the-art-like accuracy with far less training data, outperforming end-to-end ResNets and rivaling public SOTA models trained on much larger datasets. The approach yields meaningful brain age gaps correlated with Alzheimer's disease severity and Parkinson's disease motor scores, and Grad-RAM provides anatomically meaningful explanations that align with known aging patterns. Overall, the method offers data-efficient, explainable brain aging estimation with potential as a biomarker for neurodegenerative diseases.

Abstract

MRI-based brain age estimation models aim to assess a subject's biological brain age based on information, such as neuroanatomical features. Various factors, including neurodegenerative diseases, can accelerate brain aging and measuring this phenomena could serve as a potential biomarker for clinical applications. While deep learning (DL)-based regression has recently attracted major attention, existing approaches often fail to capture the continuous nature of neuromorphological changes, potentially resulting in sub-optimal feature representation and results. To address this, we propose to use supervised contrastive learning with the recent Rank-N-Contrast (RNC) loss to estimate brain age based on widely used T1w structural MRI for the first time and leverage Grad-RAM to visually explain regression results. Experiments show that our proposed method achieves a mean absolute error (MAE) of 4.27 years and an $R^2$ of 0.93 with a limited dataset of training samples, significantly outperforming conventional deep regression with the same ResNet backbone while performing better or comparably with the state-of-the-art methods with significantly larger training data. Furthermore, Grad-RAM revealed more nuanced features related to age regression with the RNC loss than conventional deep regression. As an exploratory study, we employed the proposed method to estimate the gap between the biological and chronological brain ages in Alzheimer's Disease and Parkinson's disease patients, and revealed the correlation between the brain age gap and disease severity, demonstrating its potential as a biomarker in neurodegenerative disorders.

Figures (7)

• Figure 1: Relative age distribution of cognitively normal subjects across datasets (n = 1618).
• Figure 2: Left: Overview of the training process for the RNC supervised learning model, with the first step consisting of training the ResNet backbone with the RNC loss (contrastive), and the second step training a FC regression layer with MAE/L1 loss. Right: Overview of the RNC loss zha2023rankncontrast with a sample batch of data (n=4 for demonstration) and their labels (in blue box). Two example positive pairs and their corresponding negative pair(s) are shown (in orange box) when the anchor is the 40 year old brain MRI sample. When the anchor forms a positive pair with a 50 year old sample, their label distance is 10, hence the corresponding negative samples are the 20 year old and the 80 year old samples, whose label distances to the anchor are greater than 10. When the 20 year old sample creates a positive pair with the anchor, the 80 year old sample is a negative sample as it has a greater label distance to the anchor.
• Figure 3: Investigation of different batch sizes on MAE for RNC models trained on 1x1x1 $mm^{3}$ resolution and 2x2x2 $mm^{3}$ resolution data.
• Figure 4: Averaged Grad-RAM heatmaps produced from the ResNet 50-highRez model based on the healthy control test set for the age groups of 20-40 yo, 40-60 yo, 60-80 yo, and 80-100 yo.
• Figure 5: Averaged Grad-RAM heatmaps produced from the RNC-highRez model based on the healthy control test set for the age groups of 20-40 yo, 40-60 yo, 60-80 yo, and 80-100 yo.