SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Saurabh Sihag†, Gonzalo Mateos‡, Corey McMillan†, Alejandro Ribeiro†
†University of Pennsylvania, Philadelphia, PA.
$^{\ddagger}$ University of Rochester, Rochester, NY.
In computational neuroscience, there has been an increased interest in developing machine learning algorithms that leverage brain imaging data to provide estimates of "brain age" for an individual. Importantly, the discordance between brain age and chronological age (referred to as "brain age gap") can capture accelerated aging due to adverse health conditions and therefore, can reflect increased vulnerability towards neurological disease or cognitive impairments. However, widespread adoption of brain age for clinical decision support has been hindered due to lack of transparency and methodological justifications in most existing brain age prediction algorithms. In this paper, we leverage coVariance neural networks (VNN) to propose an explanation-driven and anatomically interpretable framework for brain age prediction using cortical thickness features. Specifically, our brain age prediction framework extends beyond the coarse metric of brain age gap in Alzheimer's disease (AD) and we make two important observations: (i) VNNs can assign anatomical interpretability to elevated brain age gap in AD by identifying contributing brain regions, (ii) the interpretability offered by VNNs is contingent on their ability to exploit specific eigenvectors of the anatomical covariance matrix. Together, these observations facilitate an explainable and anatomically interpretable perspective to the task of brain age prediction.
Aging is characterized by progressive changes in the anatomy and function of the brain [1] that can be captured by different modalities of neuroimaging [2, 3]. Importantly, individuals can age at variable rates, a phenomenon described as "biological aging" [4]. Numerous existing studies based on a large spectrum of machine learning approaches study brain-predicted biological age, also referred to as brain age, which is derived from neuroimaging data [5-12]. Accelerated aging, i.e., when biological age is elevated as compared to chronological age (time since birth), may predict age-related vulnerabilities like risk for cognitive decline or neurological conditions like Alzheimer's disease (AD) [13, 14]. In this domain, the metric of interest is brain age gap, i.e., the difference between brain age and chronological age. We use the notation  $\Delta$ -Age to refer to the brain age gap.
Inferring  $\Delta$ -age from neuroimaging data presents a unique statistical challenge as it is essentially a qualitative metric with no ground truth and is expected to be elevated in individuals with underlying neurodegenerative condition as compared to the healthy population[12, 15]. The existing machine learning approaches for inferring  $\Delta$ -Age commonly rely on regression models trained to predict chronological age for a healthy population. Under the hypothesis that such models can detect accelerated aging, they are applied to cohorts representing adverse health conditions. From a statistical perspective, the residuals of the regression models inform the  $\Delta$ -Age estimates with the
Code: https://github.com/sihags/VNN_Brain_Age
expectation that they will degrade in a specific direction when deployed to predict chronological age for individuals with adverse health conditions. Hence, it is paramount to analyze the structure and statistics of the residuals of the model to validate whether  $\Delta$ -Age inferred using them provide biologically plausible information about the adverse health condition. A layman overview of the procedure of inferring  $\Delta$ -Age is included in Appendix D. In this paper, we focus on brain age prediction using cortical thickness features. Cortical thickness evolves with normal aging [16] and is impacted due to neurodegeneration [17, 18]. Further, the age-related and disease severity related variations also appear in anatomical covariance matrices evaluated from the cortical thickness [19].
Existing literature. The current state-of-the-art deep learning methods in the brain age prediction domain focus exclusively on the performance of the model on predicting chronological age for a healthy population as a metric for assessing the quality of their approach[20-22]. We refer to such methods as performance-driven approaches to brain age prediction. Major criticisms of such performance-driven approaches include the coarseness of  $\Delta$ -Age that results in lack of specificity of brain regions contributing to the elevated  $\Delta$ -Age; and the lack of clarity regarding the reliance on the prediction accuracy for chronological age in the design of these brain age prediction models [5, 23].
To address the criticism regarding the lack of interpretability or explainability of  $\Delta$ -Age, recent studies have utilized state-of-the-art post-hoc, model-agnostic methods, such as, SHAP, LIME [24], saliency maps [20, 25], and layer-wise relevance propagation [26] in conjunction with the performance-driven approaches. These methods commonly add anatomical interpretability to brain age estimates by assigning importance to the input features (often associated with specific anatomic regions). However, interpretability offered by such post-hoc approaches may not be conclusive if not shown to be stable to small perturbations to the input, variations in training algorithms and model multiplicity (i.e., when multiple models with similar performance may exist but offer distinct explanations) [27-29].
There exists sparse empirical evidence in the existing literature that hints at decoupling the task of brain age prediction from the performance achieved by the model in predicting chronological age for healthy population. For instance, a previous study has reported that models with a 'moderate' performance for predicting chronological age achieved a more informative brain age [21]. However, an appropriate 'moderate' fit on the chronological age that leads to the most informative brain age may not be generalizable to diverse datasets (diverse in terms of sample sizes, for example). Furthermore, a recent study of several existing brain age prediction frameworks has revealed that the accuracy achieved on the chronological age prediction task may not correlate with the clinical utility of associated  $\Delta$ -Age estimates [30]. Intuitively, the performance on chronological age prediction task is an incomplete, if not flawed, metric for assessing the quality of  $\Delta$ -Age estimate, as it cannot readily provide clarity on the correlation between the performance on predicting chronological age for healthy population and clinical utility of  $\Delta$ -Age.
Explainable perspective to brain age prediction. In this paper, we propose a principled framework for brain age prediction based on the recently studied coVariance neural networks (VNNs) [31]. VNN is a graph neural network (GNN) that operates on the sample covariance matrix as the graph and achieves learning objectives by manipulating the input data according to the eigenvectors (or principal components) of the covariance matrix. Thus, VNNs are inherently explainable models, as their inference outcomes can be linked with their ability to exploit the eigenvectors of the covariance matrix. In this context, the explainability offered by VNNs can be categorized as model-level explainability according to the taxonomy of explainability methods discussed in [32]. In general, model-level explainability can offer a more fundamental and generic understanding of the model than the aforementioned explainability methods applied in the brain age prediction application (such methods can broadly be categorized as instance-level methods [32]). A survey of explainability methods in GNNs is provided in Appendix B.
For the task of brain age prediction, the key focus of this paper is not on the accuracy in predicting chronological age, but rather (i) what properties does a VNN gain when it is exposed to the information provided by chronological age of healthy controls, and (ii) whether and how these properties could translate to a meaningful brain age estimate. While highly relevant, these aspects are often overlooked in existing studies on brain age prediction. In this context, VNNs provide novel insights beyond that possible by focusing only on model performance. Specifically, training VNNs to predict chronological age using cortical thickness features from the healthy population fine-tuned their ability to exploit the eigenvectors of the anatomical covariance matrix. Further, the statistical analyses of the outputs of the final layer of the VNN allowed us to identify the most significant contributors to elevated  $\Delta$ -Age
in AD with respect to healthy population. Mapping these contributors on the brain surface rendered an anatomically interpretable perspective to  $\Delta$ -Age estimates. Finally, the anatomical interpretability offered by VNNs to  $\Delta$ -Age prediction in AD was strongly associated with certain eigenvectors of the covariance matrix, thus, rendering an explainable perspective to brain age in terms of the ability of VNNs to exploit the eigenvectors of the covariance matrix in a specific manner. We emphasize here that in this paper, the term 'interpretability' is used in the context of anatomic interpretability of  $\Delta$ -Age and the term 'explainability' refers to explaining the VNN inference outcomes in terms of their associations with the eigenvectors of the covariance matrix.
Contributions. The contributions in this paper can be summarized as follows.
a) VNNs provide anatomically interpretable  $\Delta$ -Age:  $\Delta$ -Age in individuals with AD diagnosis was elevated as compared to healthy controls and significantly correlated with a clinical marker of dementia severity. Moreover, by analyzing the outputs at the final layer of VNN for AD and healthy population and mapping the results on the brain surface, we could identify contributing brain regions to elevated  $\Delta$ -Age in AD. Hence, VNN architecture yielded anatomical interpretability for  $\Delta$ -Age (Fig. 2).
b) Anatomical interpretability correlated with eigenvectors of the anatomical covariance matrix: Our experiments demonstrated that certain eigenvectors of the anatomical covariance matrix were highly correlated with the features that facilitated anatomical interpretability for  $\Delta$ -Age (Fig. 3). Thus, learning to predict chronological age of healthy population facilitated the ability of VNNs to exploit the eigenvectors of the anatomical covariance matrix that were relevant to  $\Delta$ -Age, yielding an explainable perspective to brain age prediction.
We focused our analysis on open access OASIS-3 dataset consisting of cortical thickness features from cognitively normal individuals and individuals in various stages of cognitive decline [33]. The findings were further validated on the baseline ADNI-1 dataset [34]. The utility of VNNs in predicting  $\Delta$ -Age has been explored previously in [35] but no insights were provided regarding their explainability or anatomical interpretability. See Appendix B for other relevant studies.
2 coVariance Neural Networks
We begin with a brief introduction to VNNs. VNNs inherit the architecture of GNNs [36] and operate on the sample covariance matrix as the graph [31]. A dataset consisting of  $n$  random, independent and identically distributed (i.i.d) samples, given by  $\mathbf{x}_i \in \mathbb{R}^{m \times 1}, \forall i \in \{1, \dots, n\}$ , can be represented in matrix form as  $\mathbf{X}_n = [\mathbf{x}_1, \dots, \mathbf{x}_n]$ . Using  $\mathbf{X}_n$ , the sample covariance matrix is estimated as
$$
\mathbf {C} \triangleq \frac {1}{n - 1} \sum_ {i = 1} ^ {n} \left(\mathbf {x} _ {i} - \bar {\mathbf {x}}\right) \left(\mathbf {x} _ {i} - \bar {\mathbf {x}}\right) ^ {\top}, \tag {1}
$$
where  $\bar{\mathbf{x}}$  is the sample mean of samples in  $\mathbf{X}_n$ . The covariance matrix  $\mathbf{C}$  can be viewed as the adjacency matrix of a graph representing the stochastic structure of the dataset  $\mathbf{X}_n$ , where the  $m$  dimensions of the data can be thought of as the nodes of an  $m$ -node, undirected graph and its edges represent the pairwise covariances between different dimensions.
2.1 Architecture
Similar to GNNs that rely on convolution operations modeled by linear-shift-and-sum operators [36, 37], the convolution operation in a VNN is modeled by a coVariance filter, given by
$$
\mathbf {H} (\mathbf {C}) \triangleq \sum_ {k = 0} ^ {K} h _ {k} \mathbf {C} ^ {k}, \tag {2}
$$
where scalar parameters  $\{h_k\}_{k=0}^K$  are referred to as filter taps that are learned from the data. The application of coVariance filter  $\mathbf{H}(\mathbf{C})$  on an input  $\mathbf{x}$  translates to combining information across different sized neighborhoods. For  $K > 1$ , the convolution operation combines information across multi-hop neighborhoods (up to  $K$ -hop) according to the weights  $h_k$  to form the output  $\mathbf{z} = \mathbf{H}(\mathbf{C})\mathbf{x}$ .
A single layer of VNN is formed by passing the output of the coVariance filter through a nonlinear activation function  $\sigma(\cdot)$  (e.g.,  $\mathrm{ReLU}$ ,  $\tanh$ ) that satisfies  $\sigma(\mathbf{u}) = [\sigma(u_1), \dots, \sigma(u_m)]$  for  $\mathbf{u} = [u_1, \dots, u_m]$ . Hence, the output of a single layer VNN with input  $\mathbf{x}$  is given by  $\mathbf{z} = \sigma(\mathbf{H}(\mathbf{C})\mathbf{x})$ . The construction of a multi-layer VNN is formalized next.
Remark 1 (Multi-layer VNN). For an  $L$ -layer VNN, denote the coVariance filter in layer  $\ell$  of the VNN by  $\mathbf{H}_{\ell}(\mathbf{C})$  and its corresponding set of filter taps by  $\mathcal{H}_{\ell}$ . Given a pointwise nonlinear activation function  $\sigma(\cdot)$ , the relationship between the input  $\mathbf{x}_{\ell-1}$  and the output  $\mathbf{x}_{\ell}$  for the  $\ell$ -th layer is
$$
\mathbf {x} _ {\ell} = \sigma \left(\mathbf {H} _ {\ell} (\mathbf {C}) \mathbf {x} _ {\ell - 1}\right) \quad f o r \quad \ell \in \{1, \dots , L \}, \tag {3}
$$
where  $\mathbf{x}_0$  is the input  $\mathbf{x}$ .
Furthermore, similar to other deep learning models, sufficient expressive power can be facilitated in the VNN architecture by incorporating multiple input multiple output (MIMO) processing at every layer. Formally, consider a VNN layer  $\ell$  that can process  $F_{\ell -1}$  number of  $m$ -dimensional inputs and outputs  $F_{\ell}$  number of  $m$ -dimensional outputs via  $F_{\ell -1} \times F_{\ell}$  number of filter banks [38]. In this scenario, the input is specified as  $\mathbf{X}_{\mathrm{in}} = [\mathbf{x}_{\mathrm{in}}[1], \ldots, \mathbf{x}_{\mathrm{in}}[F_{\mathrm{in}}]]$ , and the output is specified as  $\mathbf{X}_{\mathrm{out}} = [\mathbf{x}_{\mathrm{out}}[1], \ldots, \mathbf{x}_{\mathrm{out}}[F_{\mathrm{out}}]]$ . The relationship between the  $f$ -th output  $\mathbf{x}_{\mathrm{out}}[f]$  and the input  $\mathbf{x}_{\mathrm{in}}$  is given by  $\mathbf{x}_{\mathrm{out}}[f] = \sigma \left( \sum_{g=1}^{F_{\mathrm{in}}} \mathbf{H}_{fg}(\mathbf{C}) \mathbf{x}_{\mathrm{in}}[g] \right)$ , where  $\mathbf{H}_{fg}(\mathbf{C})$  is the coVariance filter that processes  $\mathbf{x}_{\mathrm{in}}[g]$ . Without loss of generality, we assume that  $F_{\ell} = F, \forall \ell \in \{1, \ldots, L\}$ . In this case, the set of all filter taps is given by  $\mathcal{H} = \{\mathcal{H}_{fg}^{\ell}\}, \forall f, g \in \{1, \ldots, F\}, \ell \in \{1, \ldots, L\}$ , where  $\mathcal{H}_{fg} = \{h_{fg}^{\ell}[k]\}_{k=0}^{K}$  and  $h_{fg}^{\ell}[k]$  is the  $k$ -th filter tap for filter  $\mathbf{H}_{fg}(\mathbf{C})$ . Thus, we can compactly represent a multi-layer VNN architecture capable of MIMO processing via the notation  $\Phi(\mathbf{x}; \mathbf{C}, \mathcal{H})$ , where the set of filter taps captures the full span of its architecture. We also use the notation  $\Phi(\mathbf{x}; \mathbf{C}, \mathcal{H})$  to denote the output at the final layer of the VNN. Various aspects of the VNN architecture are illustrated in Fig. 6 in Appendix E. The VNN final layer output  $\Phi(\mathbf{x}; \mathbf{C}, \mathcal{H})$  is succeeded by a readout function that maps it to the desired inference outcome.
Remark 2 (Statistical inference using VNNs). Given the eigendecomposition of  $\mathbf{C} = \mathbf{V}\pmb{\Lambda}\mathbf{V}^{\top}$ , the spectral properties of the VNN are established by studying the projection of the coVariance filter output  $\mathbf{z} = \mathbf{H}(\mathbf{C})\mathbf{x}$  on the eigenvectors  $\mathbf{V}$  (similar to that for a graph filter using graph Fourier transform [39, 40]). Theorem 1 in [31] established the equivalence between processing data samples with principal component analysis (PCA) transform and processing data samples with a specific polynomial on the covariance matrix  $\mathbf{C}$ . Hence, it can be concluded that input data is processed with VNNs, at least in part, by exploiting the eigenvectors of  $\mathbf{C}$ . Unlike simpler PCA-based inference models, VNNs offer stability [31] and transferability guarantees [35], which ensure reproducibility of the inference outcomes by VNNs with high confidence.
In the context of brain age prediction, we leverage the observations in Remark 2 to demonstrate the relationships of the inference outcomes with the eigenvectors of the anatomical covariance matrix  $\mathbf{C}$  (estimated from cortical thickness features) in Section 4.
2.2 VNN Learning
The VNN model is trained for a regression task, where the chronological age is predicted using  $m$  cortical thickness features. Since the VNN architecture has  $F$  number of  $m$ -dimensional outputs in the final layer,  $\Phi(\mathbf{x};\mathbf{C},\mathcal{H})$  is of dimensionality  $m \times F$ . The regression output is determined by a readout layer, which evaluates an unweighted mean of all outputs at the final layer of VNN. Therefore, the regression output for an individual with cortical thickness  $\mathbf{x}$  is given by
$$
\hat {y} = \frac {1}{F m} \sum_ {j = 1} ^ {m} \sum_ {f = 1} ^ {F} [ \Phi (\mathbf {x}; \mathbf {C}, \mathcal {H}) ] _ {j f}. \tag {4}
$$
Prediction using unweighted mean at the output implies that the VNN model exhibits permutation-invariance (i.e., the final output is independent of the permutation of the input features and covariance matrix). Moreover, it allows us to associate a scalar output with each brain region among the  $m$  regions at the final layer. Specifically, we have
$$
\mathbf {p} = \frac {1}{F} \sum_ {f = 1} ^ {F} [ \Phi (\mathbf {x}; \mathbf {C}, \mathcal {H}) ] _ {f}, \tag {5}
$$
where  $\mathbf{p}$  is the vector denoting the mean of filter outputs in the final layer's filter bank. Note that the mean of all elements in  $\mathbf{p}$  is the prediction  $\hat{y}$  formed in (4). In the context of cortical thickness
Figure 1: Workflow for brain age and  $\Delta$ -Age prediction using VNNs. Regions of interest (ROIs) contributing to elevated  $\Delta$ -Age in neurodegeneration were identified by mapping the results of the analyses of the outputs at the final layer of VNNs on the brain surface.
datasets, each element of  $\mathbf{p}$  can be associated with a distinct brain region. Therefore,  $\mathbf{p}$  is a vector of "regional contributions" to the output  $\hat{y}$  by the VNN. This observation will be instrumental to establishing the interpretability offered by VNNs in the context of  $\Delta$ -Age prediction in Section 3. For a regression task, the training dataset  $\{\mathbf{x}_i, y_i\}_{i=1}^n$  (where  $\mathbf{x}_i$  are the cortical thickness data for an individual  $i$  with chronological age  $y_i$ ) is leveraged to learn the filter taps in  $\mathcal{H}$  for the VNN  $\Phi(\cdot; \mathbf{C}, \mathcal{H})$  such that they minimize the training loss, i.e.,
$$
\mathcal {H} _ {\text {o p t}} = \min  _ {\mathcal {H}} \frac {1}{n} \sum_ {i = 1} ^ {n} \ell \left(\hat {y} _ {i}, y _ {i}\right), \tag {6}
$$
where  $\hat{y}_i$  is evaluated similarly to (4) and  $\ell(\cdot)$  is the mean-squared error (MSE) loss function.
3 Methods Overview for Brain Age Prediction
In this section, we provide an overview of the brain age prediction framework based on VNNs (see Fig. 1 for a summary). Our results primarily focus on the dataset described below.
OASIS-3 Dataset. This dataset was derived from publicly available freesurfer estimates of cortical thickness (derived from MRI images collected by 3 Tesla scanners and hosted on central .xnat . org), as previously reported [33], and comprised of cognitively normal individuals (HC;  $n = 611$ , age  $= 68.38 \pm 7.62$  years, 351 females) and individuals with AD dementia diagnosis and at various stages of cognitive decline ( $n = 194$ , age  $= 74.72 \pm 7.02$  years, 94 females). The cortical thickness features were curated according to the Destrieux (DKT) atlas [41](consisting of  $m = 148$  cortical regions). For clarity of exposition of the brain age prediction method, any dementia staging to subdivide the group of individuals with AD dementia diagnosis into mild cognitive impairment (MCI) or AD was not performed, and we use the label AD+ to refer to this group. The individuals in AD+ group were significantly older than those in HC group (t-test:  $p$ -value  $= 2.46 \times 10^{-23}$ ). The boxplots for the distributions of chronological age for HC and AD+ groups are included in Fig. 10 in Appendix J. For 191 individuals in the AD+ group, the clinical dementia rating (CDR) sum of boxes scores evaluated within one year (365 days) from the MRI scan were available (CDR sum of boxes  $= 3.45 \pm 1.74$ ). CDR sum of boxes scores are commonly used in clinical settings to stage dementia severity [42] and were evaluated according to [43].
Cross-validation. The findings obtained via the analyses of the OASIS-3 dataset were cross-validated on the ADNI-1 dataset (described in Appendix I).
3.1 Training the VNNs on HC group
We first train the VNN model to predict chronological age using the cortical thickness features for the HC group. This enables the VNN models to capture the information about healthy aging from the cortical thickness and associated anatomical covariance matrix. The hyperparameters for the VNN architecture and learning rate of the optimizer were chosen according to a hyperparameter search procedure [44]. The VNN model had  $L = 2$ -layers with a filter bank, such that we had  $F = 5$ , and 6 filter taps in the first layer and 10 filter taps in the second layer. The learning rate for the Adam optimizer was set to 0.059. The number of learnable parameters for this VNN was 290. The HC group was split into a 90/10 training/test split, and the covariance matrix was set to be the anatomical
covariance evaluated from the training set. A part of the training set was used as a validation set and the other used for training the VNN model. We trained 100 VNN models, each on a different permutation of the training data. The training process was similar for all VNNs and is described in Appendix F. No further training was performed for the VNN models in the subsequent analyses.
3.2 Analyses of regional residuals in AD+ and HC groups
Next, the VNN models trained to predict the chronological age for the HC group and (5) were adopted to study the effect of neurodegeneration on brain regions for  $\mathrm{AD + }$  group. Since the impact of neurodegeneration was expected to appear in the anatomical covariance matrix, we report the results when anatomical covariance  $\mathbf{C}_{\mathrm{HA}}$  from the combined cortical thickness data of HC and  $\mathrm{AD + }$  groups was deployed in the trained VNN models. Because of the stability property of VNNs [31, Theorem 3], the inference drawn from VNNs was expected to be stable to the composition of combined HC and  $\mathrm{AD + }$  groups used to estimate the anatomical covariance matrix  $\mathbf{C}_{\mathrm{HA}}$ .