MRI Volume-Based Robust Brain Age Estimation Using Weight-Shared Spatial Attention in 3D CNNs

TL;DR

The paper tackles brain age estimation from structural MRI by proposing a voxel-based 3D CNN augmented with a weight-shared spatial attention module that targets invariant brain regions across diverse datasets. The method uses seven 3D convolutional layers with a shared spatial attention mechanism after each layer, followed by five dense layers, achieving a MAE of 1.622 on ADNI and 2.265 on OASIS3, while demonstrating robustness to cross-dataset variation. Grad-CAM analysis suggests the model focuses on the corpus callosum region, and ablation shows that sharing attention across layers improves generalization compared to per-layer attention. This approach offers a data-efficient path to robust brain-age biomarkers across multi-site data and paves the way for cross-dataset clinical applicability, with code forthcoming on GitHub.

Abstract

Important applications of advancements in machine learning, are in the area of healthcare, more so for neurological disorder detection. A crucial step towards understanding the neurological status, is to estimate the brain age using structural MRI volumes, in order to measure its deviation from chronological age. Factors that contribute to brain age are best captured using a data-driven approach, such as deep learning. However, it places a huge demand on the availability of diverse datasets. In this work, we propose a robust brain age estimation paradigm that utilizes a 3D CNN model, by-passing the need for model-retraining across datasets. The proposed model consists of seven 3D CNN layers, with a shared spatial attention layer incorporated at each CNN layer followed by five dense layers. The novelty of the proposed method lies in the idea of spatial attention module, with shared weights across the CNN layers. This weight sharing ensures directed attention to specific brain regions, for localizing age-related features within the data, lending robustness. The proposed model, trained on ADNI dataset comprising 516 T1 weighted MRI volumes of healthy subjects, resulted in Mean Absolute Error (MAE) of 1.662 years, which is an improvement of 1.688 years over the state-of-the-art (SOTA) model, based on disjoint test samples from the same repository. To illustrate generalizability, the same pipeline was utilized on volumes from a publicly available source called OASIS3. From OASIS3, MRI volumes 890 healthy subjects were utilized resulting in MAE of 2.265 years. Due to diversity in acquisitions across multiple sites, races and genetic factors, traditional CNN models are not guaranteed to prioritize brain regions crucial for age estimation. In contrast, the proposed weight-shared spatial attention module, directs attention on specific regions, required for the estimation.

Figures (4)

• Figure 1: Spatial Attention Module
• Figure 2: Proposed framework for weight shared spatial attention in 3D CNNs
• Figure 3: Illustration of Weight Shared Spatial Attention across the convolution layers
• Figure 4: Grad-CAM Interpretation across Convolution Layers in Sagittal, Coronal, and Axial Planes.