Divide and Conquer: Improving Brain Age Predictions through Contrastive Pre-Training

1453 

Abstract Submission 

Maximilian Konowski¹, Ramona Leenings¹, Jan Ernsting², Nils Winter³, Lukas Fisch⁴, Daniel Emden⁴, Carlotta Barkhau⁴, Susanne Meinert⁵, Elisabeth Leehr⁶, Xiaoyi Jiang⁷, Udo Dannlowski⁴, Tim Hahn⁴

¹University of Münster, Münster, North Rhine Westphalia, ²University of Münster, Münster, NRW, ³University of Münster, Münster, North-Rhine Westphalia, ⁴Institute for Translational Psychiatry, Münster, North Rhine Westphalia, ⁵Institute for Translational Psychiatry, Münster, Germany, ⁶Institute for Translational Psychiatry, University of Münster, Münster, North Rhine-Westphalia, ⁷University of Münster, Münster, Northrine-Westphaila

Maximilian Konowski  
University of Münster
Münster, North Rhine Westphalia

Ramona Leenings  
University of Münster
Münster, North Rhine Westphalia

Jan Ernsting  
University of Münster
Münster, NRW

Nils Winter  
University of Münster
Münster, North-Rhine Westphalia

Lukas Fisch  
Institute for Translational Psychiatry
Münster, North Rhine Westphalia

Daniel Emden  
Institute for Translational Psychiatry
Münster, North Rhine Westphalia

Carlotta Barkhau  
Institute for Translational Psychiatry
Münster, North Rhine Westphalia

Susanne Meinert  
Institute for Translational Psychiatry
Münster, Germany

Elisabeth Leehr  
Institute for Translational Psychiatry, University of Münster
Münster, North Rhine-Westphalia

Xiaoyi Jiang  
University of Münster
Münster, Northrine-Westphaila

Udo Dannlowski  
Institute for Translational Psychiatry
Münster, North Rhine Westphalia

Tim Hahn  
Institute for Translational Psychiatry
Münster, North Rhine Westphalia

Brain Age models learn to derive the chronological age of healthy subjects from brain structure represented in T1-weighted MR images or derivatives thereof. The difference between predicted and chronological age, called brain age gap (BAG), has been used as a marker for atypical neurodegeneration and has shown significant salience in the context of different psychiatric conditions and lifestyle factors (Cole 2019, Bittner 2021, Wrigglesworth 2021, Hahn 2022, Blake 2023). Yet, prevalent imbalances in age distributions-e.g. an over representation of younger individuals-can introduce challenges to the training. A brain age model might benefit from disproportionately generating age predictions for larger age groups, as this statistically optimizes the loss function during training. Consequently, the BAG distributions among different age groups could vary significantly: The inherent meaning of a BAG of e.g. +2 years might have a training-induced shift in meaning across the age continuum. Artificially balancing the dataset by undersampling over-represented age groups, however, risks the loss of crucial information and variance that could be valuable for the learning process. To mitigate the negative effects of a skewed age distribution, we test a pre-training strategy. Specifically, we utilize samples from over-represented age groups to impart a structural understanding of MRI images to a deep learning model before fine-tuning it on a balanced dataset for chronological age prediction.

We loosely lean on Chaitanya et al. (2021) and implement a self-supervised pre-training with two focus points: First, we use a context restoration task to foster an understanding of the MRI-inherent data structure. Second, we use a contrastive loss to encourage an understanding of fine granular inter-individual differences. As our backbone, we utilize a 3D ResNet10 architecture implemented in PyTorch. We train and evaluate three brain age models, for which we obtained a total of n=8911 cat12-preprocessed T1-MRI scans from healthy control subjects (Gaser, https://neuro-jena.github.io/cat/; MACS cohort: Vogelbacher et al., 2018; publicly available studies curated by Fisch et al. 2023). For evaluation, we randomly select 5 samples from each age bin to form a balanced test set Xtest (n=260). The first brain age model was trained on the unbalanced dataset Xtrain(n=8651) to serve as benchmark performance. Second, we pre-train a model with samples from over-represented age groups Xpretrain (n=7080) , before we fine tune it to age prediction with an age-balanced data subset Xfine-tune (n=1571). Third, we train a randomly initialized model on the downsampled fine-tuning set only. Finally, we use the model's predictions to calculate the brain age gap (BAG) for the test set and compare the BAG distributions per age group.

Compared to the benchmark model, the fine-tuned model exhibited accelerated convergence and yielded superior results on the balanced test set (MAE of 9.9 to 4.68 years), despite utilizing only a fraction of the available data, see figure 1. Although the benchmark model exhibited promising performance on a validation set during training (MAE of 3.68), its performance notably deteriorated when assessed on the balanced test set. Notably, age predictions from the fine-tuned model exhibited a more uniform distribution, contrasting the benchmark model's implicit bias towards the skewed training data distribution. Consequently, the range of the BAGs of the benchmark model differs substantially over different age groups, potentially confounding later statistical analysis. Remarkably, the pre-trained model showed only slight improvements against a model trained on the fine-tuning dataset, only.

Our findings emphasize the nuanced relationship between data volume and bias, as well as the relevance of a carefully curated dataset. In the future, we aim to explore pre-training strategies which are specifically tailored to the brain age learning objective.

Aging ²

Classification and Predictive Modeling ¹

Methods Development

Multivariate Approaches

Anatomical MRI

Aging

Machine Learning

STRUCTURAL MRI