Unmasking Dementia Detection by Masking Input Gradients:
A JSM Approach to Model Interpretability and Precision

Jacobian Saliency Maps (JSM) emerged recently [1] as a highly effective tool for deciphering the decision-making mechanisms of a deep learning model. It accomplishes this goal by defining specific zones within an input image and measuring their volumetric changes, thereby providing a precise understanding of feature attribution. Feature attribution, which ascribes significance and influence to individual features, enables us to identify the particular aspects of an input that exert significant influence on the model’s output. Thus, JSM transforms data in a way that aligns with human intuition and enhances the model’s interpretability before the actual deep learning pipeline, making it a promising choice for a diagnosis model debugger.

Our approach provides interpretation by computing the gradients of input with respect to weighted elements of the input image and optimizing them toward matching the patterns of deformations highlighted by the JSM. On manipulating the input gradients, we can explore two directions: enhance their significance in relevant brain areas or reduce their significance in irrelevant brain areas. However, because the appropriate magnitudes of gradients in relevant areas are typically unknown a priori, we choose to dampen the gradients in irrelevant areas as the normal regions are a reliable reference.

We are inspired by the work done by Ross et al. [23] who introduced a method to regularize a model’s gradients with respect to input features based on a binary mask annotation matrix. However, the annotation term added to the model in [23] is a simple binary mask, which fails to capture correlations between different regions of a brain scan. This poses a significant limitation to the effectiveness because such correlations carry crucial diagnostic information. Additionally, the datasets used in their validation were small and synthetic, thus leaving considerable doubt on whether the method would perform well in real-world medical applications.

Moreover, we do not use the annotation matrix but a special weighted saliency map instead. The rationale is that we prefer a matrix to not only represent the varying degrees of sensitivity of each feature to the diagnosis accurately but also highlight brain regions that exhibit deformations; in the meantime, the normal regions have to be preserved for contrast purposes. These properties cannot be achieved by the binary annotation matrix. In addition, by using the special saliency map to characterize the deformations in the brain, it also allows us to align the constraints on gradients (a penalty we formulate later in (2), (4)) with the domain-specific knowledge related to the medical problem. We perform both registration and Jacobian using Advanced Normalization Tools (ANTs) [3].

Intuition behind JSM. Ideally, the presence of a disease can be identified by comparing an individual’s many scans over a long period of time. Since this is usually not feasible in practice, we can compare an individual’s scan to a standardized healthy brain template to assess its local changes using a deformation map. JSM is derived from non-linear image registration, which is a process designed to both minimize variations among individual subjects and align all different images with a standardized template. JSM then examines the deformations that transform the anatomical structures of individuals onto a common standard space, through which we can deduce the relative volume differences between each individual and the template. This analysis aids in pinpointing statistically significant anatomical variations across diverse populations, such as distinguishing between AD patients and healthy elderly individuals. In our investigation, we utilized the Montreal Neurological Institute’s 152 brain template (MNI152) as the reference. This process of image registration and JSM computation can be attributed to the realm of computational anatomy, specifically recognized as tensor-based morphometry [22], and has not been adequately explored in the context of dementia.

Formulation of JSM. In medical image processing, computing deformation involves first aligning a source image $M$ with a target image $F$. This is achieved by using a transformation $\phi$ that maps points in $M$ to their corresponding points in $F$. Then, the displacement between these corresponding points is represented as a Deformation Vector Field:

To transform a point $(x,y,z)$ to $\phi(x,y,z)$, it is necessary to impose a regularization constraint in order to ensure that the deformation is seamless, one-to-one, and differentiable. This is framed as an optimization problem that minimizes the following cost function:

where $L_{sim}$ is a similarity measure between two images, the transformed image $\phi(M)$ and $F$, and $L_{Reg}$ is a regularization term that enforces the desired properties on the deformation. We use Mattes Mutual Information (MI) [16] as our similarity measure, i.e.,

where $M^{\prime}$ denotes $\phi(M)$ for simplicity. $P(m^{\prime},f)$ is the joint probability distribution of the intensity of voxel $m^{\prime}$ in image $M^{\prime}$ and that of voxel $f$ in image $F$, and $Q_{1}(m^{\prime})$ and $Q_{2}(f)$ represent the marginal probability distributions of $m^{\prime}$ and $f$, respectively. For the regularizer, we use the B-spline regularization from [27]:

where $\nabla^{2}$ denotes the second-order derivative, $dV$ represent seach voxel $(x,y,z)$, and we integrate over the entire spatial domain.

After solving for the transformation $\phi$, we compute a Jacobian matrix $J$ from the deformation vector field $\vec{v}$, by calculating the first derivative of $\vec{v}$ at each voxel to encode local deformations including stretching, shearing, and rotation. That is,

Then, denoting the Jacobian determinant by $Det(J)$, we calculate it for every voxel $v(x,y,z)$ as

which forms what we call a Jacobian Saliency Map $JSM$ of the source image $M$:

Volumetric changes refer to the alteration in volume at the level of individual voxels within a medical image. This JSM helps us to identify the volumetric ratio of the brain image at the voxel level before and after transformation $\phi$, which indicates the brain’s volume change.

By breaking down an input image into distinct regions and measuring how they are transformed, JSM provides precise insight into the complexities of feature attribution and thus model explainability. To this end, we take an innovative approach to explore the possibility of leveraging JSM as a powerful model debugger, for which we incorporate the JSM formulated above into the loss function $\mathcal{L}$ of a medical diagnostic model:

Given input data $\boldsymbol{x}$ and data label $\boldsymbol{y}$, the loss function $\mathcal{L}(\lambda,\boldsymbol{x},\boldsymbol{y},JSM)$ consists of the training loss (first term) and a novel, JSM-based regularization term (second term). It introduces an emphasis on the importance of anatomical changes captured by the Jacobian map values. Specifically, $JSM_{dp}$ represents the JSM values associated with the $d^{th}$ feature and $p^{th}$ spatial dimension (width, height, and depth in 3D), $log(\hat{y}_{k})$ represents the natural logarithm of the predicted probability of $x$ belonging to class $k$, and $\frac{\partial}{\partial x_{dp}}$ reflects the partial derivative with respect to the $d^{th}$ feature of $x$ in the $p^{th}$ spatial dimension. Thus, this regularization term aims to not only enhance interpretability but also rectify predictions by mitigating the influence of irrelevant cues.

With reference to Equation (7), we add a weight matrix $W$ to give more importance to areas in the JSM that have volumetric changes (expansion or compression) and discourage the input gradients from being significant in areas with no volumetric changes (marked by 1). Hence each element of $W$ is defined as follows:

Both feature_weight and debug_weight are hyperparameters that indicate the level of importance in every region. We designate a debug weight of 0.2. This choice is deliberate, allowing us to down-weight potentially misleading features while still retaining them as a reference for contrasting relevant features. On the other hand, since regions of volumetric changes are of higher importance, we assign them a feature weight of 0.8. Hence, we debug the model through weighting features according to its relevance. Note that we don’t eliminate irrelevant areas by giving a weight of zero to preserve correlations in the brain volume. More rigorous hyperparameter tuning of such weights can be incorporated in future work.

In our experiments, we utilized the recently published OASIS-3 dataset. This dataset includes multi-modality data such as MRI, PET, and CT scan images from a diverse group of 1377 participants. Among these participants, 755 were cognitively normal (CN) adults, while the remaining 622 individuals showed different levels of cognitive decline. The age range of the participants was extensive, spanning from 42 to 95 years. CT imaging was used to detect whether certain areas of the brain were shrinking, which can be an indication of AD. On the other hand, MRI provided detailed images of the body and a clear view of progressive cerebral atrophy, which is most visible through T1-weighted volumetric sequences. Consideration of PET data was deferred due to its temporal nature, making it more suitable for future spatiotemporal analyses.

During clinical assessments and diagnoses, the clinical dementia rating (CDR) scores of the participants were utilized. The scores range from 0 to 3, with 0 indicating no AD dementia and 3 indicating severe AD dementia. The very mild stage (rating 0.5) of dementia is similar to the Mild Cognitive Impairment (MCI) stage of AD. As mentioned before, we combined moderate and severe AD due to the very low number of subjects with severe AD. Ultimately, we created four classes based on CDR scores: normal, MCI, mild AD, and severe AD. Having the same subjects in multiple sets (train and test sets) can result in the model overfitting to those individuals, potentially leading to a subpar performance on new, unseen subjects. Hence, for patients who underwent multiple sessions, we preserved the initial MRI session and selected the CT session that was closest in date to that MRI.

Our preprocessing pipelines for MRI and CT scans are shown in Fig. 1. To minimize any spatially varying intensity bias that may result from factors such as magnetic field inhomogeneities and acquisition artifacts, we use the FMRIB’s Linear Image Registration Tool (FLIRT) [9] for bias field correction on MRI images. Meanwhile, we utilize a technique called contrast stretching on CT images to enhance their diagnostic value and visual perception. This process involves adjusting the pixel intensities to fully utilize the display’s dynamic range. For CT images, we followed the framework established by Kuijf et al. [11]. Also, we apply the Brain Extraction Tool (BET) [26] to eliminate non-brain portions from both MRI and CT images. Finally, we register both CT and MRI to the MNI152 brain template. Brain templates are typically created for MRI images, making it difficult to register a CT image to an MRI template. We overcame this by adhering to a method in [11] which involves identifying corresponding landmarks in the CT image and MRI template and then using these landmarks to align the images.

Medical images, especially those related to the brain, are typically 3D. This introduces a computational burden into complex deep neural networks. Our approach seeks to harness the guidance provided by the JSM, and by incorporating such insights from JSM, we aim to develop lighter convolutional neural networks (CNNs) that alleviate computational burdens without compromising model performance. In view of the multimodal nature of our study, we incorporated two data fusion techniques: 1) Late Fusion and 2) Early Fusion. As shown in Fig. 2, late fusion adopts a dual-branch structure, treating each modality independently and subsequently aggregating their predictions through an averaging mechanism. This approach is particularly advantageous for JAL, where debugging is performed separately for each modality through its own JSM. On the other hand, early fusion involves concatenating the input images as well as their corresponding JSM maps, which allows the model to glean correlations between the two modalities and concurrently debug predictions for the input holistically.

Our lightweight CNN model [20] contains two convolutional layers coupled with batch normalization, ReLU activation, dropout, and max-pooling operations, which capture spatial hierarchies and correlation patterns in the input data. The convolutional layers use a kernel size of 3x3x3, stride of 1, and padding, with a dropout of rate 0.2 for regularization. The max-pooling operation reduces spatial dimensions to 2x2x2. The 3D tensor is reshaped into a 1D tensor by a flattening layer, preparing it for the subsequent fully connected subnetwork, which performs the overall feature integration and classification. All the specifications are presented in Table 1. Finally, a softmax layer converts the predicted scores into probabilities. In the late fusion setup, the final prediction is computed by averaging the probabilities from both branches.

We trained our model on a 40GB A100 GPU with batch size 10. Our model was able to quickly converge within only 20 epochs. To address the class imbalance problem in AD, we utilized the Adaptive Synthetic (ADASYN) [7] oversampling algorithm to generate synthetic samples for minority classes during training. Another challenge is that OASIS-3 included identical subjects from multiple sessions across training and test sets may contribute to overfitting, as the model can become excessively attuned to the characteristics of those particular subjects and sessions [2]. To address this, we take meticulous care to ensure that only subjects from distinct sessions are grouped within either the training set or the test set (but not both).

Table 2 summarizes the performance comparison in terms of accuracy, sensitivity, and specificity between our approach and the state-of-the-art. Please note that while precision and recall are more general terms used in machine learning, sensitivity and specificity are often preferred in medical and diagnostic fields due to their direct interpretation in the context of disease detection and diagnosis. Table 2 shows that our testing accuracy across four classes surpasses all the baselines that employ the same dataset. Massalimova et al. [15] achieved marginally higher sensitivity and specificity, but it is crucial to note that our model handles a larger number of classes, making it more challenging to achieve higher accuracy. In addition, [15] uses ResNet18 while our model is significantly lighter. In fact, our model is lighter than nearly all the baselines. Basheer et al. [4] used features like CDR, MMSE, age, gender, etc along with MRI images, and found that age and gender had substantial positive impact on performance. In our case, we achieved superior performance solely with images as we aimed to test our model using spatial features.

Ablation Study. To provide an in-depth assessment of the impact of JAL on our model’s performance, we conducted an ablation study on models with and without JAL. This was achieved by setting the JSM term in (8) to zero. The results are presented in Fig. 3, which provides histograms of model performance distribution across multiple mini-batches in the test set. The overlap represents the intersection between the with JAL and without JAL conditions, indicating the extent to which the model’s performance remains consistent regardless of the presence or absence of JAL. The histogram demonstrates that the model performance significantly improves in terms of all the metrics (accuracy, sensitivity, and specificity) when JAL is incorporated, as evidenced by the rightward shift of the distributions in the blue bars in comparison with the yellow bars. This compellingly demonstrates the impact of incorporating JAL on the model’s decisions. For a more comprehensive evaluation, Table 3 dissects the four classes and provides more detailed results. It shows substantial improvements over all the AD stages (CN, MCI, MLD, SEV), further affirming the efficacy of JAL. Table 3 also allows us to see that performance for CN and SEV (severe) are relatively higher than MCI and MLD, which is because MCI and MLD have more subtle differences in dementia patterns, making them more intricate to discern. Nevertheless, our model with JAL exhibits evenly promising outcomes for all stages. Scores in Table 2 are the macro average of the scores in Table 3.

Interpretability. We examined the similarity between the volumetric changes characterized by JSM and the decision-making process of the neural network. By plotting gradients overlaid on their corresponding input images, we observed that they closely aligned with the patterns highlighted by the JSM. This can be seen from Fig. 4 which provides samples from the dataset showing the axial, sagittal, and coronal views of each MRI and CT images juxtaposed with the corresponding JSM views that highlight brain deformations. This desirable alignment is attributed to our JAL loss function which incorporates the JSM during the model debugging process, contributing to the model trustworthiness by promoting transparency.

Furthermore, incorporating JSM in JAL during model debugging also enables the model to learn and adapt to the highlighted deformations, and hence serves as a powerful tool for refining the model performance as well, fostering a more accurate and informed decision-making process.