Do Transformers and CNNs Learn Different Concepts of Brain Age?

Abstract

"Predicted brain age" refers to a biomarker of structural brain health derived from machine learning analysis of T1-weighted brain magnetic resonance (MR) images. A range of machine learning methods have been used to predict brain age, with convolutional neural networks (CNNs) currently yielding state-of-the-art accuracies. Recent advances in deep learning have introduced transformers, which are conceptually distinct from CNNs, and appear to set new benchmarks in various domains of computer vision. Given that transformers are not yet established in brain age prediction, we present three key contributions to this field: First, we examine whether transformers outperform CNNs in predicting brain age. Second, we identify that different deep learning model architectures potentially capture different (sub-)sets of brain aging effects, reflecting divergent "concepts of brain age". Third, we analyze whether such differences manifest in practice. To investigate these questions, we adapted a Simple Vision Transformer (sViT) and a shifted window transformer (SwinT) to predict brain age, and compared both models with a ResNet50 on 46,381 T1-weighted structural MR images from the UK Biobank. We found that SwinT and ResNet performed on par, though SwinT is likely to surpass ResNet in prediction accuracy with additional training data. Furthermore, to assess whether sViT, SwinT, and ResNet capture different concepts of brain age, we systematically analyzed variations in their predictions and clinical utility for indicating deviations in neurological and psychiatric disorders. Reassuringly, we observed no substantial differences in the structure of brain age predictions across the model architectures. Our findings suggest that the choice of deep learning model architecture does not appear to have a confounding effect on brain age studies.

FIGURE 1

Overview of workflow and results: (a) We used 46.381 structural magnetic resonance imaging (sMRI) brain scans from the UK Biobank (UKBB) to train and evaluate a convolutional neural network (CNN; 3D ResNet50) and two transformers (3D simple vision transformer; sViT; 3D shifted window transformer; SwinT) for brain age prediction. Mean absolute errors (MAEs) for held‐out healthy subjects were nearly identical for ResNet (2.66 years) and SwinT (2.67 years). We define the term “concept of brain age” as the distinct brain aging effects identified by a brain age model and the way these aging effects are synthesized into scalar predictions. (b) Effect sizes between prediction errors (brain age gaps; BAGs) of patients and matched controls were similar for CNN and transformers across neurological‐ and psychiatric diseases, yielding no indication that different model architectures rely on meaningfully different concepts of brain age for their predictions.

FIGURE 2

SwinT will likely to outperform ResNet with additional training samples We trained multiple instances of each model architecture with gradually decreased training samples and found that accuracies of shifted window transformer (SwinT) and simple vision transformer (sViT) decline stronger compared to the ResNet. Extrapolating each model architecture's accuracy using power laws (Schulz, Bzdok, et al. 2024) indicates SwinT would surpass ResNet's accuracy given additional training samples. Uncertainty estimates refer to the SD across model instances.

FIGURE 3

Different brain age model architectures encode similar disease patterns. The figure shows effect sizes (Cohen's $d$) measured between BAGs of patients and matched controls. Effect sizes between model architectures were within one $\sigma$ from each other for any disease, with no indication of differences. Error bars indicate the standard error of the mean estimate derived by bootstrapping patient‐control pairs.

FIGURE 4

Association of BAG and cognitive, lifestyle and biomedical phenotypes seems not to depend on the model architecture. We fitted linear models from BAG and confounds to phenotype and report the t‐statistic for whether the BAG is a significant predictor. Error bars indicate the t‐statistic's standard error of the mean estimate, derived by bootstrapping. BAGs of different model architectures were similarly predictive for the analyzed phenotypes.

FIGURE 5

Similar brain features appear to be relevant for age predictions across different model architectures. Using Input $\times$ Gradient (IxG) Shrikumar et al. (2016), we generated feature‐relevance heatmaps for each held‐out healthy subject across ResNet, SwinT, and sViT. These heatmaps, averaged across random model architecture initializations and visualized at group‐level using a color scale (dark red = low relevance, white = high), revealed highly consistent brain regions across architectures, suggesting they capture comparable features of brain aging. Slight variations in the heatmaps likely stem from interactions between the model architectures and IxG, rather than reflecting meaningful differences in the underlying relevant features. The consistency in highlighted brain regions across ResNet, SwinT and sViT reinforces our conclusion that different model architectures are unlikely to learn different concepts of brain age. Notably, brain regions such as the cerebellum, basal ganglia, and brain stem, which were consistently identified as important, are well‐documented for their roles in aging processes (Walhovd et al. 2011), further validating their relevance as predictors of age.