AXIAL: Attention-based eXplainability for Interpretable Alzheimer’s Localized Diagnosis using 2D CNNs on 3D MRI brain scans

Recent reviews (Wen et al., 2020; Zhou et al., 2023) highlight several reliability and reproducibility issues:

• 1.

  Wrong Data Split: Data leakage due to incorrect data splitting is recurring. When datasets are not split at the subject level, data from the same subject may appear in both training and test sets, leading to overestimated model performance. This problem is particularly pronounced in patch or slice-based approaches (Wen et al., 2020).

• 2.

  Late Split in Data Processing: Data processing steps such as data augmentation and feature selection must be conducted post-data splitting to avoid contamination of the test set. If these steps involve test data, it biases the model, rendering the evaluation unreliable (Wen et al., 2020).

• 3.

  Biased Transfer Learning: Transfer learning, while beneficial, can introduce bias, particularly when the source and target domains overlap. For instance, using a network pre-trained on an AD vs CN task for an MCI vs CN task can cause bias if the CN subjects overlap in both tasks’ training or validation sets (Wen et al., 2020).

• 4.

  Absence of an Independent Test Set: The integrity of model evaluation is compromised if the test set is used for any purpose other than final performance evaluation. Hyperparameter optimization must rely on a separate validation set. The absence of an independent test set in some studies leads to inflated performance metrics (Wen et al., 2020).

• 5.

  Differences in Diagnostic Criteria: Variability in diagnostic criteria for AD across different studies creates inconsistencies in ground truth labeling. This variability complicates the comparison of results across studies and hampers the development of universally applicable models (Zhou et al., 2023).

• 6.

  Lack of Reproducibility: A significant challenge in the field is the non-availability of frameworks and models for public use. Without access to open-source code and detailed implementation procedures, reproducibility is hampered. This lack of transparency affects the ability to validate and build upon existing work (Zhou et al., 2023).

Furthermore, Wen et al. (2020) highlights that the prevalent use of the 2D slice approach in AD analysis is often associated with a series of methodological problems. In particular, the authors highlight that among the numerous studies using this approach, only a few have successfully addressed the abovementioned challenges. This fact underlines the need for rigorous methodological standards in the field.

In response to these challenges, our work was carefully developed to avoid common weaknesses identified in previous studies. Section 5.1 details the data splitting for model selection and data augmentation used for this study. Section 4.2.2 outlines the double transfer learning strategy investigated. To ensure reproducibility, we provide an open-source repository with all necessary code and implementation procedures at https://github.com/GabrieleLozupone/AXIAL.git.

Several studies have highlighted the difficulties that prevent using DLADS on medical images due to the current challenges of XAI methods (Van der Velden et al., 2022; Singh et al., 2020; Viswan et al., 2024). XAI researchers often use self-intuition to determine a good explanation without validating with a medical professional. This missing validation is partially due to the inability to provide quantitative but only qualitative explanations. Qualitative explanations provided for small and specific subsets of subjects make a global evaluation and comparison of XAI techniques difficult. Furthermore, there is no specific correlation between the prediction and the associated brain region. Model-agnostic XAI techniques, such as GradCAM, often produce saliency maps that highlight extensive regions of the brain without sufficient specificity or consistency. Another critical limitation of GradCAM-like methods, mainly when applied to the analysis of 2D slices from MRI scans, is their focus on intra-slice features while neglecting the inter-slice relationships crucial for underlying the 3D brain regions related to AD. Aggregating these saliency maps to form a 3D visualization produces an overly generalized activation map that often highlights different and sparse areas.

As an alternative to traditional saliency-based methods, recent advancements have explored the use of attention-based models to enhance the interpretability of DL models. In this context, Jetley et al. (2018) proposed a trainable attention mechanism to highlight where and in what proportion the network paid attention to input images for classification. Attention amplifies relevant areas and suppresses irrelevant ones, improving interpretability. In Schlemper et al. (2019), grid attention captures anatomical information in medical images, and attention coefficients were used to explain which areas of the image the network focused on. Recently Wang et al. (2024), described in Section 2.3, introduced a slice-aware module to advance XAI in their methodology. Their variant of the attention mechanism enables the extraction of attention weights to determine the importance of each slice in the decision-making process, offering a more precise understanding of model focus.

This module employs a pre-trained 2D convolutional backbone to process a series of $N$ 2D slices. The convolutional backbone, which can be any architecture such as VGG or ResNet, is pre-trained on the ImageNet dataset to handle cases of small dataset sizes in AD diagnosis tasks.

The original data are 3D images with a value for each voxel, so the resulting sliced 2D images are 1-channel images. Since pre-trained weights are for 3-channel images, we summed the pre-trained convolutional filters of the backbone’s first layer. Since these filters operate linearly and sum their results across channels, we can sum all the filters over the input channel dimension in the first layer. This operation is equivalent to replicating the single-channel image three times, but it is more computationally efficient. The extracted slices were resized to $224\times 224$ and normalized.

Each convolutional backbone ends with a max-pooling operation that converts the feature map of each slice into a feature vector of dimension $f_{\text{dim}}$. This sequential input processing is achieved through sharing convolutional weights across the sequence, similar to a Recurrent Neural Network approach:

where $x_{i}$ represents the $i$-th slice in the sequence, $f_{i}$ is the resulting feature vector for that slice, and Conv denotes the convolutional operations performed by a convolutional architecture.

Models pre-trained on large datasets can help overcome challenges posed by limited data, but they may not be enough for complex tasks. As shown in Table 1, the data available for training a network to differentiate between AD and CN individuals is usually more than for distinguishing between sMCI and pMCI. Although both tasks are complex, the latter is considered more difficult. Therefore, it is more likely that transfer learning can distinguish AD from CN more effectively than recognizing the difference between sMCI and pMCI. Since doctors rely on their expertise to assess a patient’s cognitive decline, leading to an AD diagnosis, we suggest fine-tuning a model pre-trained on the AD vs. CN task to the sMCI vs. pMCI task. Following our definition of classes in Section 1, the $AD\cup CN$ set is disjoint from $(sMCI\cup pMCI)\subset MCI$. As a result, adopting this strategy under these assumptions does not imply data leakage issues.

This module enables the network to learn inter-slice dependencies and global 3D patterns through the Feature Extraction module. During backpropagation, the shared weights of the convolutional module are updated based on the entire image, thus enabling the learning of 2D patterns in relation to the importance of that slice at a global level. The Attention XAI Fusion module leverages a fully connected (FC) layer to assign importance to each section based on its feature vector. This introduces a parameter-efficient approach, with only $f_{\text{dim}}+1$ learnable parameters:

where $w_{i}$ is the computed weight for slice $i$, and FC denotes the fully connected layer operation. Normalization of these weights is performed using a softmax function, ensuring they sum to one and represent the slice-importance distribution:

where $\alpha_{i}$ is the normalized importance weight for the $i$-th slice. The module performs fusion by calculating a weighted sum of the feature vectors using their corresponding weights, generating a composite feature vector representing the entire brain image:

This module yields a feature vector $F$ that synthesizes the brain image and a set of attention weights $\{\alpha_{i}\}$ highlighting the significance of each slice.

In the final Diagnosis module, the synthesized feature vector $F$ is fed into a head network comprising a fully connected layer with two output neurons for the binary classification task (e.g., AD vs. CN):

$D$ represents the softmax output, providing a probability distribution over the diagnosis classes, and $\text{FC}_{\text{head}}$ denotes the fully connected layer specific to the diagnosis task.

The implemented XAI method yields attention weights for each of the three principal slicing planes of the brain: sagittal ($s$), coronal ($c$), and axial ($a$). We trained a dedicated Diagnosis and XAI Network for each plane to assess slice importance distribution from different views. The MRI brain image is sliced in the three planes, resulting in three sequences of 2D slices: $N_{s}$, $N_{c}$, and $N_{a}$ slices for the sagittal, coronal, and axial planes.

Each network’s Attention XAI module computes a set of attention weights $\mathbf{\alpha}_{s}$, $\mathbf{\alpha}_{c}$, and $\mathbf{\alpha}_{a}$ for its corresponding plane. The computation is formalized as follows:

where $p\in\{s,c,a\}$ represents each plane, $f_{p}$ denotes the feature vectors derived from slices within that plane.

We integrate the attention weights $\mathbf{\alpha}_{s}$, $\mathbf{\alpha}_{c}$, and $\mathbf{\alpha}_{a}$ into a unified 3D attention map. The 3D attention map $A$ is constructed by applying the following operation to each voxel within the brain’s imaging volume:

For every voxel at coordinates $[i,j,k]$, its attention value is derived from the product of the corresponding sagittal, coronal, and axial attention weights. This process retains and amplifies the diagnostic importance perceived across all three anatomical planes, yielding a comprehensive 3D depiction of attention distribution.

To facilitate interpretation and comparison, we normalize the entire 3D attention map in the range [0, 1], ensuring that the highest values correspond to areas of high diagnostic relevance. We used min-max normalization:

The resulting map highlights the brain regions most significantly associated with the diagnostic output of the network, offering insights into the pathological hallmarks of AD as learned by the model through its training process.

To achieve comparability across subjects, each MRI image, denoted by $I$, was normalized to the MNI 152 standard space using the transformation function $f_{MNI152}$, thus $I_{norm}=f_{MNI152}(I)$. This step is crucial for aligning brain structures across different individuals to facilitate the identification of AD-specific biomarkers.

Generation of Binary Heatmap A binary heatmap, $H_{binary}$, was generated to isolate regions of significant structural patterns associated with AD, utilizing a threshold $\theta$ set at the 99.9th percentile. The binary heatmap is defined as:

where $\delta$ is an amplification factor set to 10 to increase the prominence of the heatmap overlay.

Identification and Analysis of Regions of Interest Employing a template atlas, we determined regions of interest (ROI) within the brain. For each ROI identified with a unique label $r$, we computed the overlap $O_{r}$ as follows:

where $T_{atlas}$ represents the template atlas data. The quantitative evaluation of each region’s overlap involved several statistical measures calculated from the MRI data. The overlap volume of the heatmap in a given region is:

which can be useful to indicate predominant regions in decision process. The mean ($\mu_{r}$) intensity of the heatmap within each identified region is given by

which represents the average activity level across the voxels of interest, allowing us to pinpoint regions with consistently high AD-related changes. The standard deviation ($\sigma_{r}$) of the heatmap intensities,

provides insight into the variability of attention within the region, indicating the heterogeneity of AD impact across different brain areas. The maximum ($A_{max,r}$) and minimum ($A_{min,r}$) values of the heatmap,

and

respectively, highlighting the extremes in activation within the regions, shedding light on the most and least affected areas within the brain’s AD-implicated regions. Lastly, the percentage of overlap ($P_{r}$) is calculated as

to provide a single value that quantifies the portion of a specific brain region involved in the AD-related patterns identified by the network. This measure is a critical indicator of the importance of each region in distinguishing AD patients from CN individuals.

We used Accuracy (ACC), Sensitivity (SEN), Specificity (SPE), and Matthews Correlation Coefficient (MCC) as performance metrics. ACC is the most common performance measure for classification, and it consists of the percentage of correctly classified samples over the total. MCC (Chicco and Jurman, 2020) is a correlation coefficient between predictions and true labels; it provides a more informative and truthful score than accuracy when evaluating binary classifications, allowing a more realistic interpretation of classifier performance, especially in the case of unbalanced datasets. SEN and SPEC serve as critical metrics for evaluating the diagnostic accuracy of our model, highlighting its ability to correctly identify positive and negative cases, respectively. In AD vs. CN task, the AD samples are in the positive class, and CN samples are in the negative one. In sMCI vs. pMCI task, the sMCI samples are in the positive class, and pMCI samples are in the negative one.

All the results obtained are averaged on the test sets from the 5-fold cross-validation.

As the first step, we analyzed the effectiveness of different convolutional backbones in the AD vs. CN task. To this aim, two sets of experiments were performed. First, we evaluated the performance of different backbones by considering them in the Feature Extraction Module of our network. Second, to allow an evaluation independent of our method, the best networks for each family were evaluated with a 2D Majority Voting approach, which is also considered one of the baseline methods. In the Majority Voting approach, we train 2D convolutional networks by individually labelling the slices. Consequently, the weights are updated from the slice-level error during training. Then, the predictions of the slices of a single 3D image are aggregated by choosing the most recurring one in the sequence as the final label.

The first pool of experiments was conducted using a mini-batch of 16 and freezing the first 50% of the backbone’s layers. The convolutional architectures chosen are: (i) VGG16 and VGG19 (Simonyan and Zisserman, 2014), (ii) ResNet34, ResNet50 and ResNet101 (He et al., 2016), (iii) EfficientNetV2 Small and EfficientNetV2 Medium (Tan and Le, 2021) and (iv) DenseNet121 (Huang et al., 2017). The results of this pool are shown in Table 2. In this first analysis, VGG16 had the highest accuracy (ACC: 0.829) and sensitivity (SEN: 0.816), suggesting it is the better-suited backbone for AD patient identification. VGG19, while slightly less accurate, showed superior specificity (SPE: 0.875), indicating it is particularly adept at identifying CN cases. The ResNet series demonstrated that additional depth (as seen in ResNet101) does not equate to improved performance, with ResNet50 showing better accuracy and sensitivity. This could suggest diminishing returns with increased complexity for this task. EfficientNet models trailed behind their counterparts with identical accuracies (ACC: 0.764). DenseNet121, while providing good accuracy (ACC: 0.798), also did not meet the higher results obtained by other architectures. The MCC aligns with these findings, confirming VGG16 as the top-performing model in this case.

Based on the results obtained in Table 2, we selected the best backbone from each architecture family as candidates for the second pool of experiments. Since this pool comprises 2D rather than 3D images, we chose 32 as the mini-batch size. To achieve comparable results, the same percentage of backbone layers was frozen. As shown in Table 3, the performance trends of the architectures remain consistent with previous observations. VGG16 continued to lead in accuracy (ACC: 0.804) and MCC (MCC: 0.605). This result aligns with other works that design their approaches relying on a VGG as base network (Altay et al., 2021; Hu et al., 2023a). Remarkably, our method outperforms the Majority Voting one across several metrics. The superiority can be seen from the average increase across all tested backbones of ACC by 1.6% and MCC by 3.08%. Notably, a good increase was observed with VGG16, with an ACC improvement of 2.5% and an MCC improvement of 5.0%.

Although the axial plane is the most commonly used in clinical applications, there could be better choices than this one in this classification scenario. We performed a slicing plane analysis after finding that VGG16 is the best backbone in the diagnostic task. For this set of experiments, we used 100 slices for each view, a batch size of 8, a learning rate 1e-4, and froze the first half of the backbone in training. Table 4 shows the performances averaged over the five cross-validation test sets varying the slicing plane. The results show that the axial plane achieved the best performance. This finding aligns with the clinical choice.

Since VGG16 was confirmed as the best backbone among those under investigation and the axial plane provided the best results, we selected this configuration for our Feature Extraction module. The second step consisted of optimizing other parameters: the number of slices extracted $N$, the batch size, and the percentage of backbone to freeze.

The results in Table 5 indicated that varying the number of slices and the percentage of first backbone layers frozen in the network impacted performances significantly. The configuration with 70 slices, batch size of 8, and freezing 50% of the layers resulted in an ACC of 0.821 and MCC of 0.636. When the batch size was reduced to 4 under the same conditions, a slight improvement was observed with an ACC of 0.825 and MCC of 0.646. After increasing the number of slices to 80 and employing a batch size of 8, experiments with 0%, 25%, 50%, and 75% freezing were conducted. The highest accuracy was achieved at 50% freezing, yielding an ACC of 0.856 and MCC of 0.712, highlighting the best overall performance across all configurations. With a further increase to 90 slices, the performance did not improve, maintaining an ACC of 0.818 with a 50% freezing and batch size of 8. Finally, the performance slightly fluctuated with 100 and 120 slices but generally did not surpass the optimal results obtained with 80 slices end 50% freezing.

The third step involved investigating the utility of transfer learning for predicting AD progression with the double transfer learning strategy proposed in Section 4.2.2. This approach leverages a backbone pre-trained on the ImageNet dataset, further fine-tuned through a subsequent task distinguishing AD from CN individuals. As mentioned in Section 4.2.2, it is crucial to note that the patient cohorts classified as AD or CN differ from those labeled as sMCI or pMCI, given that the latter are diagnosed as MCI. Consequently, this strategy does not cause data leakage problems.

First, we fine-tuned our model with VGG16 pre-trained on ImageNet in the Feature Extraction module on the entire AD vs. CN dataset. We froze the first 50% of the layers as this configuration provided the best results in Table 5. The final model pre-trained on the AD vs. CN task was selected using 10% of the dataset as a validation set. This step provides a model in which the first convolutional part is highly specialized in extracting fine-grained low-level feature representations and the rest in analyzing AD-related high-level patterns. Second, we ulteriorly fine-tuned the AD pre-trained model on the sMCI vs. pMCI classification task at a 36-month forecast interval. The learning rate for this pool of experiments was reduced to $1\times 10^{-5}$ to fine-tune the pre-trained model weights, ensuring that the disease knowledge is preserved during the training process. Table 6 demonstrates the efficacy of this approach in this task. With baseline ImageNet transfer learning and 50% of the network layers frozen, the model achieved an ACC of 0.4822 and an MCC of 0.065. The MCC value of 0.065 indicates that the model could not find useful correlations to distinguish the two cases. Incorporating AD vs. CN knowledge with a 50% freezing resulted in a marked enhancement across metrics: an ACC of 0.703 and an MCC of 0.396. Advancing to a 75% freezing of the network layers under the double transfer learning paradigm, we observed a peak performance of an ACC of 0.725, SPE of 0.763, SEN of 0.678, and an MCC of 0.443.

We performed a comparison with current 2D slice attention-based state-of-the-art models. These methods include Attention Transformer (Altay et al., 2021) and AwareNet (Wang et al., 2024). Attention Transformer uses a multi-head self-attention to perform a cross-attention operation resulting in a slice feature maps fusion. AwareNet represents the diagnosis network of a joint learning framework that uses a slice-aware module and a slice-shift module. The slice-aware module leverages the attention mechanism to interpret significant slices and regions. While we have utilized the official implementation of AwareNet, as it is available and provided by the authors, we have implemented the architecture for the Attention Transformer ourselves based on the paper details, as the authors did not provide the official implementation. For both methods, the hyperparameters were optimized to maximize performance. As in our case, Attention Transformer obtained the best results by freezing 50% of the base network pre-trained on ImageNet and a batch size of 8. The AwareNet network was trained from scratch with a batch size of 2. The network is initialized by the Kaiming method (He et al., 2015) and trained using the Adam optimization algorithm with $\beta_{1}=0.48$ and $\beta_{2}=0.999$. In both cases, the learning rate was set to $1\times 10^{-5}$ to prevent overfitting.

In the sMCI vs. pMCI task, we adopted the double transfer learning strategy proposed in this work for the Attention Transformer method. The first 75% of the Attention Transformer base network layers were frozen. Since the AwareNet is trained from scratch on the AD vs. CN task, we froze the first 50% of the backbone.

The results of our comparative analysis are summarized in Table 7. To assess the effectiveness of our attentional module in enhancing diagnostic performance, we introduced three additional approaches: Majority Voting, Attention-Guided Majority Voting, and Majority Voting 3D. We detailed the Majority Voting approach implemented in Section 5.2.1. The Attention-Guided Majority Voting approach utilizes the network trained in the majority voting case. However, predictions are made during testing on a subset of slices identified as most informative by our attentional fusion module. Specifically, we identified a contiguous range where the attention values exceeded the 75th percentile. This range was calculated for the axial plane by averaging the attention weight distribution across the five validation sets of the fold cross-validation, resulting in the range [0, 25]. The results of this method validate the module’s ability to select feature maps with high information content. Finally, the 3D Majority Voting approach replaces the attentional fusion module with an averaging module, where feature maps from individual slices are averaged to form a single feature map for making subject-level decisions. This approach further helps to evaluate the impact of slices’ importance weighting compared to considering all slices as equally important.

Our model demonstrates effectiveness in disease diagnosis and prediction tasks, achieving (i) an accuracy (ACC) of 0.856 and a Matthews correlation coefficient (MCC) of 0.712 in the AD vs. CN comparison, and (ii) an ACC of 0.725 and MCC of 0.443 in the sMCI vs. pMCI task. While AwareNet shows promise in the diagnosis task with an ACC of 0.832 and an MCC of 0.659, it struggles in the disease prediction task, yielding an MCC of 0.039, likely due to insufficient data for training. On the other hand, Attention Transformer, leveraging double transfer learning, provides a more balanced alternative, with an MCC of 0.651 in the diagnosis task and an MCC of 0.238 in disease prediction. However, its increased parameterization from self-attention prevents it from bridging the gap with lightweight models like ours and AwareNet in the diagnosis task. Majority voting, lacking a mechanism to learn correlations between sections, underperforms compared to other methods. Incorporating subject-level error evaluation through the Majority Voting 3D approach the MCC improves from 0.605 to 0.667, even surpassing overfitting-prone approaches such as AwareNet and Attention Transformer. Furthermore, leveraging our attentional mechanism to identify informative slices while using networks pre-trained without considering slice relationships yields an MCC of 0.683. This result underscores our model’s ability to select crucial slices for decision-making.

Many studies focused on model interpretability by generating saliency maps for randomly chosen test subjects, enabling local qualitative analysis. However, this approach does not support the evaluation of interpretability consistency when the training and testing data vary. To examine this consistency in our approach, we considered the variability of attentional weight distributions across the five different folds used for cross-validation. Our evaluation leverages a clinical hypothesis specific to AD, which is believed to affect similar brain areas across different patients. For each fold, we calculated the weight distributions for each sample in the test set and derived an average distribution specific to that test set. By analyzing these average distributions across the five folds, we could assess whether our interpretability remained consistent, thereby supporting the generalizability of our findings across different subsets of data.

Fig. 5 presents our diagnosis network’s average attentional weight distributions for each fold in the axial, coronal, and sagittal views. Upon examination of the histograms for each view, we observe a remarkable consistency in the distribution shapes across all five folds, indicating that our interpretability approach is stable despite the variation in the train/test set data. Specifically, the axial distributions reveal a consistent concentration of attentional weights around the initial slices. This trend suggests the model’s recurrent focus on the brain’s inferior regions, notably the areas where degenerative changes first manifest in AD, such as the hippocampus. In the coronal view, attentional weights are notably centered, indicating that the model consistently identifies the central part of the brain as significant. This central focus might correspond to the medial temporal lobe, including the hippocampus and the surrounding regions, further substantiating the axial findings. The sagittal view is the only bimodal distribution, suggesting that the model pinpointed symmetrical areas along this plane. We hypothesize that the network was focusing on the hippocampus since it adheres to all the constraints: situated in the inferior part of the brain, centrally located, and symmetrical. The consistency and specificity of these findings across multiple data folds strengthen the argument that our network could reliably identify specific brain regions as a critical biomarker for distinguishing between AD and CN subjects.

We also analyzed the consistencies of AwareNet distributions (Wang et al., 2024) to compare the robustness and interpretability of different attention mechanisms. However, the average distributions produced by this model, as seen in Fig. 6, present sparsely distributed peaks and do not allow the identification of predominant slice ranges. Furthermore, the distributions between the folds are inconsistent. These results indicate that in this case study, AwareNet could not produce consistent attentional weights useful for contributing to the interpretability of its decisions.