Influence of In-domain Pre-training on 3D Vision Transformer in Structural Neuroimaging Data

1372 

Abstract Submission 

JiHoon Kim¹, Roshan Rane^2,3,4, Kerstin Ritter^2,3

¹Freie Universität Berlin, Department of Education and Psychology, Berlin, Germany, ²Charité – Universitätsmedizin Berlin, Department of Psychiatry and Psychotherapy, Berlin, Germany, ³Bernstein Center for Computational Neuroscience, Berlin, Germany, ⁴Einstein Center for Neurosciences Berlin, Berlin, Germany

JiHoon Kim  
Freie Universität Berlin, Department of Education and Psychology
Berlin, Germany

Roshan Rane  
Charité – Universitätsmedizin Berlin, Department of Psychiatry and Psychotherapy|Bernstein Center for Computational Neuroscience|Einstein Center for Neurosciences Berlin
Berlin, Germany|Berlin, Germany|Berlin, Germany

Kerstin Ritter  
Charité – Universitätsmedizin Berlin, Department of Psychiatry and Psychotherapy|Bernstein Center for Computational Neuroscience
Berlin, Germany|Berlin, Germany

Vision Transformer (ViT) is gaining attention for predictive modeling of associations between brain structure and phenotype in neuroimaging [1]. Adapting ViT to neuroimaging is challenging due to its data-hungry nature. This nature may bring a risk of overfitting, and poor generalization with limited data [2,3]. Transfer learning is crucial to overcome these challenges [4]. Two in-domain pre-trained models: one for sex classification and the other for age estimation, show promise in neuroimaging [5,6]. However, a systematic comparison of the impact of different kinds of pre-training on different phenotypes in diverse sample sizes is lacking. Therefore, we pre-train two in-domain pre-trained models and conduct a systematic comparison of the models in predicting different phenotypes across different sample sizes.

We use T1-weighted brain images from the UK Biobank dataset (N=36,598), subjected to minimal processing: skull stripping, bias correction, and Montreal Neurological Institute (MNI) template registration. Images are downsampled to 96x96x96 voxels. We utilize randomly initialized Vision Transformer (ViT) and 'DenseNet-121' Convolution Neural Network (CNN) models, referred to as Vanilla ViT and Vanilla CNN.
We conduct two-stage experiments consisting of pre-training and fine-tuning stage (Figure 1). In the pre-training stage, we pre-train ViT models for sex classification (ViT w.SEX) and age estimation (ViT w.AGE). In the fine-tuning stage, we perform three binary classification tasks: biological sex, chronological age group (68 – 82 vs. 44 – 61), and high vs. low alcohol use. Four models: ViT w.SEX, ViT w.AGE, Vanilla ViT, and Vanilla CNN are trained per task across sample sizes from 200 to 7,000. We evaluate model performance on a holdout set (N=3,055) using balanced accuracy (BACC). Lastly, we apply Chefer's feature visualization method to obtain ViT classification relevance feature values and overlay normalized mean values of true positive samples onto the MNI templates [7].

After pre-training, ViT w.SEX achieved BACCs of 98.89% and 98.84%, and ViT w.AGE showed mean absolute errors of 2.94 and 2.88 in validation and holdout sets, respectively. Figure 2 (a) shows the model performance of three prediction tasks on the holdout set after fine-tuning. ViT w.SEX and ViT w.AGE outperformed other models in Sex and age group predictions, respectively. In alcohol consumption prediction, ViT w.SEX performed better than Vanilla ViT in small sample sizes (200 to 1,000). Also, ViT w.SEX reaches its best BACC of 61.53% in training size 7,000.
We observed different feature importance patterns for each model. Within each model, patterns are consistent across varying training sizes. Figure 2 (b) shows the feature visualization of the high alcohol use in training size 7,000. Notably, Vanilla ViT and ViT w.SEX have common features but also ViT w.SEX highlighted brain stem, right cerebellum, left corpus callosum, and left midcingulate cortex regions, setting it apart from Vanilla ViT.

We show that phenotype prediction strongly depends on the pre-training. The results of sex and age group prediction showed the importance of aligning pre-training tasks with target phenotypes. ViT w.SEX showed competitive results in alcohol consumption prediction, especially in smaller datasets. It may indicate that the prediction is confounded by sex.
We provide feature visualizations of high alcohol use. The literature consistently reported associations between alcohol intake and alterations in grey matter volume of widespread areas such as the cerebellum, frontal, and temporal lobes, as well as various subcortical structures [8, 9, 10]. Our findings echo these observations. ViT w.SEX may outperform Vanilla ViT based on the distinct highlighted features.
We demonstrate that the effectiveness of in-domain pre-training can vary based on factors such as the size of the fine-tuning dataset, and the relationship between the pre-training and target phenotypes.

Classification and Predictive Modeling ¹

Methods Development ²

Data analysis

Machine Learning

Modeling

STRUCTURAL MRI

Other - Alcohol Consumption