NEMt: Fast Targeted Explanations for Medical Image Models via Neural Explanation Masks

Venue: NLDL_org_2025_Conference
Published: 2024-09-02 22:49:36

Abstract: A fundamental barrier to the adoption of AI systems in clinical practice is the insufficient transparency of AI decision-making. The field of Explainable Artificial Intelligence (XAI) seeks to provide human-interpretable explanations for a given AI model. The recently proposed Neural Explanation Mask (NEM) framework is the first XAI method to explain learned representations with high accuracy at real-time speed. 
NEM transforms a given differentiable model into a self-explaining system by augmenting it with a neural network-based explanation module. This module is trained in an unsupervised manner to output occlusion-based explanations for the original model. However, the current framework does not consider labels associated with the inputs. This makes it unsuitable for many important tasks in the medical domain that require explanations specific to particular output dimensions, such as pathology discovery, disease severity regression, and multi-label data classification.
In this work, we address this issue by introducing a loss function for training explanation modules incorporating labels. It steers explanations toward target labels alongside an integrated smoothing operator, which reduces artifacts in the explanation masks. We validate the resulting  Neural Explanation Masks with target labels (NEMt) framework on public databases of lung radiographs and skin images. The obtained results are superior to the state-of-the-art XAI methods in terms of explanation relevancy mass, complexity, and sparseness. Moreover, the explanation generation is several hundred times faster, allowing for real-time clinical applications. The code is publicly available at  https://anonymous.4open.science/r/NEMt-3249/.

Main paper: 
 introduction: 
 Explainable Artificial Intelligence (XAI) in the medical area is motivated by the need for trust and transparency in AI-driven medical diagnosis, prognosis, treatment planning, and report generation [1, 2] . This need is steadily increasing due to the flourishing development of deep learning (DL) models in these areas [3] . However, existing XAI methods suffer from limitations that impair their usefulness in clinical medical image analysis. Many methods focus too much on individual input features (e.g., pixels) and fail to detect complex patterns, while 047 mask-based techniques optimized for finding regions 048 in the input image are very slow [4, 5] . The re-049 cently proposed Neural Explanation Masks (NEM) 050 framework [5] has been shown to produce accurate 051 mask-based explanations when applied to the rep-052 resentations of models trained in an unsupervised 053 manner with a latency allowing for real-time infer-054 ence. Currently, it does not take image labels into 055 account. Therefore, it cannot explain a particular 056 output component, for instance, answer the ques-057 tion of which parts of the input affect the score for 058 a particular class (e.g., a pathology). This limits its 059 application in the medical domain. In this study we 060 Activation Mapping [10] ), methods based on SHAP 086 (Shapley Additive exPlaination [9] ), and LIME (Lo-087 cal Interpretable Model-agnostic Explanations [8] ). 088 However, explanations based on featured-based attri-089 bution may be misleading when features have strong 090 correlations or complicated interactions that are difficult for individual features to capture. In contrast, set-wise approaches aim at identifying subsets of informative features that have a collective impact on the model's prediction [13, 14] . This is especially helpful for comprehending intricate decision-making processes where the combination of features, rather than individual attributes, drives the prediction.However, these methods often suffer from long processing times and potentially produce poor-quality explanations [15] .Most XAI methods are designed for the supervised setting in the sense that they generate explanations for model outputs such as classifications, where the explanation is provided for an input-output pair similar to those used for supervised training of the model. However, the recent advances in unsupervised deep learning [16, 17] has increased the need for XAI methods working for the unsupervised setting leading to many new contributions in this field [12, 18, 19] . These methods seek to explain what part of an input mostly determines its (latent) representation learned by a (DL) model. Here the most prominent approach is RELAX and its variants [12] ,which provides feature-attribution explanations and has successfully been applied to medical imaging [20] . Recently, the Neural Explanation Mask (NEM) framework [5] was proposed for explaining unsupervised XAI and it was shown to compare favourably to previous unsupervised XAI methods, such as RE-LAX, in terms of the generated explanation masks while being significantly faster. The NEM framework aims to distinguish between significant and insignificant input regions. This approach is similar to methods from the field of Weakly Supervised Semantic Segmentation (WSSS), which extract segmentation labels from learned model representations, see, for example, the works by Kweon et al. [21] and Araslanov and Roth [22] . 
methodology: 
 In this section, we start by presenting the background on the NEM framework, and then introduce our proposed NEMt. Finally, we outline our new approach for artefact removal through local stochastic smoothing.The NEM framework [5] transforms a given differentiable ML model Φ, referred to as the frozen network, into a self-explaining system by augmenting it with a neural-based explanation module Ψ. The module is trained to generate occlusion-based explanations for Φ. More specifically, for a given input x, an explanation mask m = Ψ(x) is generated simultaneously with the model output y = Φ(x). The mask m serves to pinpoint the parts of x that influence y 145 most. Notably, Ψ is trained on (unlabeled) training 146 images X to directly predict m instead of performing 147 some post-training optimization over m, which was 148 the path in the existing occlusion-based explanation 149 methods [13, 14, 23] . The overall framework along-150 side the training process is illustrated in Figure 1 . 151 The choice of the neural architecture of Ψ and how 152 it is attached to the frozen model Φ lead to different 153 instantiations of the NEM framework. One specific 154 instantiation is NEM-U. By considering mask-based 155 explanation as segmenting images into areas of im-156 portance and non-importance, NEM-U takes inspi-157 ration from U-Net [24] , which is commonly used to 158 handle segmentation tasks. More specifically, the 159 NEM-U treats the frozen model Φ and the masking 160 model Ψ as the encoder and the decoder of U-Net, 161 respectively, and capitalizes on skip connections to 162 extract intermediate representations from Φ.To train a NEM model, we need a loss function 164 that defines what should be highlighted in a given 165 input. In the case of unsupervised representation 166 learning, a natural goal is to segment an image x 167 into sections that either do or do not strongly affect 168 the learned latent representation Φ(x). This can 169 be recast as finding the sparsest mask, such that 170 the model representation of the masked image is 171 still close to the model representation of the original. 172 This leads to a loss of the form173 174 L(X) = x∈X d Φ(x), Φ Ψ(x) ⊙ x 175 + λ 1 Ψ(x) 1 + λ 2 B(x, Φ, Ψ), (1) 176where the d is some distance function, B(x, Φ, Ψ) 177 is some regularization term enforcing binary masks, 178 ∥ • ∥ 1 is the l 1 norm, ⊙ is the hadamard product, and 179 λ 1 , λ 2 are hyper-parameters. We optimize over a 180 range of inputs X to allow Ψ to generalize to unseen inputs. We adopt the heuristic described in [5] and setB(x, Φ, Ψ) = -d Φ Ψ(x)⊙x , Φ (1-Ψ(x))⊙x + d Φ(x), Φ Ψ(x) ⊙ x .The NEM framework outlined in Section 3.1 was proposed for explanations in unsupervised scenarios.Here, we extend it to the supervised setting. We assume classification into C classes {0, C -1}. Let c be the class label predicted for image x. We assume that the frozen model outputs logits Φ(x) ∈ R C and predicts the label c = arg max i Φ(x) i , where Φ(x) i is the i'th component of Φ(x) corresponding to logit for class i. In the unsupervised setting, we asked which parts of x strongly affect the whole learned representation Φ(x), while in the supervised setting we are only interested in what explains Φ(x) c .In the classification setting, we only consider the effects on the scalar Φ(x) c measured by the squared l 2 norm. The loss function will be thereforeL(X) = - x∈X c=arg maxi Φ(x)i (Φ(x) c -Φ (x⊙Ψ(x)) c ) 2 +λ 1 ∥Ψ(x)∥ 1 .(EQUATIONThe loss is independent of the output dimensionality as it only accounts for the target logit. Compared to (1), the loss function is a negative sum, because the goal is to find the smallest area that determines the label. In other words, we wish to find the smallest input perturbation yielding the largest output change. We have removed the regularization term B(x), since we found a more robust method of deleting artifacts in the supervised case as described in the following section. This loss formulation allows for explaining a specific output channel t by settingc = t instead of c = arg max i Φ(x) i .stochastic weighted neighborhood averaging.Initial experiments with (2) showed that the method is more prone to generate artifacts in its explanation mask compared to what we observed in the unsupervised setting, see Figure 2 for an example. This is likely due to the fact that the term d Φ(x), Φ x ⊙ Ψ(x) in the unsupervised loss function, which enforces closeness in representation space, also has a regularizing effect. In contrast, the supervised loss function rewards larger changes, po-229 tentially also in latent intermediate representations. 230 It is not straight-forward to employ a similar regu-231 larization term in the supervised setting, and our 232 initial attempts to do so did not lead to satisfactory 233 results. Thus, we opted to regularize in a different 234 manner.Our idea to address this issue was to bias the 236 model to find explanations where the neighborhood 237 of included pixels is also included in the explanation, 238 that is, isolated pixels and other artifacts in the mask 239 should be removed. To this end, we add a stochastic 240 smoothing operator to the output of the explanation 241 module Ψ. For each pixel in the network explanation 242 mask, the filter performs a weighted average of the 243 pixel and the elements of its local neighborhood 244 via a convolution operation. To ensure that the 245 explanations are faithful, we train the network with 246 the smoothing filter on, that is, the smoothing is not 247 a post-processing step. To increase the chance for all 248 members of the neighborhood of a given important 249 pixel to be included, in each forward pass random 250 parts of the filter are "smoothly" masked out. This 251 is done by randomly picking weights for each filter 252 component so they are positive and sum to 1. A 253 visualization of the effect of the filter can be seen in 254 Figure 2 . bounding box. We also used the HAM10000 dataset, comprising 10,015 dermatoscopic divided into seven diagnostic categories and provides binary masks [27] . For evaluation, we extracted 1000 images from each dataset. For training the NEMt explanation modules, we extracted 21347 images from the RSNA dataset and 8012 images from the HAM10000 dataset, ensuring that there are no overlapping patients in the training and test sets. 
implementation details: 
 To explore the feasibility of NEM-U for medical image analysis, we leverage a pre-trained variant of the MedCLIP model [6] using a RESNET50 [28] architecture. We applied the MedCLIP model to the RSNA pneumonia dataset and explained the generated latent representations. We compared our method to RELAX [12] and the related U-RELAX method, where RELAX explanations were thresholded by uncertainty estimates.To judge the performance in the supervised setting, we explained two classifiers based on the RESNET50 Architecture. The models are trained on the RSNA dataset and the HAM10000 dataset, respectively, using the same training splits as used for training the NEMt models. Given we wish to find areas of illness, we explain the channel belonging to the true target label using the targeted variant of (2). The RESNET50 architecture is extracted from the Pytorch Image Models library [29] . We compare our methodology against Integrated Gradients [30] , Gradient SHAP [9] , GradCAM++ [25] , RISE++ [31] , and Smooth Pixel Mask [23] with a area constraint of 5%.We implemented all NEMt and NEM-U explanation architectures using U-Net decoders with 9 million trainable parameters. We leverage the stochastic smoothing filter described in Section 3.2 with a neighborhood size of (21, 21) when running NEMt in the supervised case. For the loss function for explaining representations, we chose the distance function d to be negative cosine similarity (see [5] ).In (1) we set λ 1 = 1 8 and λ 2 = 1 4 as is done in prior work [5] . In (2), we set λ 1 = 1 2 based on visual inspection of the masks generated on the validation split of the training data. All models were trained with the Prodigy optimizer [32] . We provide an ablation study of various implementation choices in appendix A.3.For evaluation, we computed the complexity [33] and sparseness [34] metrics. These metrics favor explanations that highlight the minimum number of important pixels. Furthermore, we calculated relevance rank and relevance mass [35] . Relevance mass is determined by calculating the ratio of pos-322 itive attributions located inside the ground truth 323 bounding box to the total positive attributions. A 324 high relevance mass score shows that the explana-325 tion significantly focuses on areas highlighted by 326 human annotations while minimizing the focus on 327 irrelevant areas. Relevancy rank is measured as the 328 IoU of the k highest rated pixels according to attri-329 bution and the ground truth bounding box, where k 330 is equal to the size of the bounding box. Thus, the 331 relevancy rank measures the overall ranking ability 332 of a given explanation method. The metrics im-333 plementation was taken from the Quantus library 334 [36] . We also considered the faithfulness [37] 415 VOC [39] , or COCO [40] .  
 
 Appendix:  


Decision: The paper proposes a method called Neural Explanation Masks with target labels (NEMt), an extension of the existing NEM framework tailored for supervised settings to improve interpretability for neural nets. The method addresses a weakness of the existing NEM and optimizes for target labels, incorporates a stochastic smoothing filter to remove artifacts to provide more reliable explanations for classification. The paper demonstrates improvements over existing baselines across multiple metrics. The authors tested the approach on two datasets and provided a code repository for reproducibility. All reviewers regarded the paper as an interesting contribution but also raised some comments. The reviewers noted that the paper would benefit from a broader evaluation on additional datasets, comparison with more related methods and more details on certain implementation aspects. Most of the comments from the reviewers were addressed during the rebuttal and the authors added some more experiments and ablation studies and gave some additional details for clarification which they plan to add to the manuscript prior to publication. This was well received by the reviewers. Given all strengths and weaknesses, all reviewers rated the paper very positively and suggested accepting it for NLDL.

Reviews: review: title: title: An interesting XAI module for artificial neural networks
paper_summary: paper_summary: The paper proposes a variant to a recently introduced Neural Explanation Mask (NEM) framework, called NEMt, where the additional "t" in the name corresponds to the notion of "target labels".  The idea is that the original NEM framework works in an unsupervised setting, whereas the proposed framework works for artificial neural networks that operate in a supervised manner. The authors experimentally validate the performance of their approach using two datasets from the medical domain and moreover evaluate the performance of the proposed method using several metrics (relevance rank, relevance mass, complexity, sparseness) that have been introduced and provide justification for the "goodness" of the provided explanations.  Furthermore, the method is very lightweight as exhibited by the experiments, thus providing explanations in time comparable to one baseline method (Grad Shape) or, faster compared to the other baselines explored (Smooth Pixel Mask, GradCAM, Integrated Gradients).
main_review: main_review: Reading the reviews of others has increased my confidence on the rating of the paper.  I believe that this is a paper that has sufficient novelty and should be accepted, especially given the changes and small additions that the reviewers have requested and the authors have promised to include in the final version of the paper.
strength_weakness: strength_weakness: - No source code provided.  I would urge the authors to provide a link with the source code for their experiments upon acceptance of their work.  This can only benefit the community and the authors.

- Other than the above, I cannot think of any real weaknesses for this work from a scientific point of view.  Even though the novelty is not very high since it relies on the NEM framework, nevertheless, the idea is natural and very much important since supervised neural networks are used in so many contexts.

Finally, below are some comments for the presentation.

- Figure 1 seems to be out of order and I do not believe it is referenced from the main text.

- In Section 3.1 where you describe the NEM framework I was a little bit confused in retrospect.  In line 137 you say that the output of model $\Phi$ is $y$, but what is really $y$?  If we are talking about unsupervised learning, then this $y$ cannot be a label - and $y$ is (almost?) universally accepted as a "label"; so, most likely you want to output the identify function (e.g., as in an auto encoder) or some sort of embedding of the original input?  Please note that I was not familiar with the NEM framework before reading this paper and perhaps this is why it is not so clear to me what you are trying to say and how that would fit in an "unsupervised" setting where NEM operates.  Along these lines, in line 170, you mention that $d$ is some distance function.  Given that $d$ takes as input such $y$'s, it would be nice if you can provide an example of the kind of functions that you have in mind.

The above weaknesses are the reason that I am giving a rating of 4 for this work and not a 5.
title: title: Review of NEMt: Fast Targeted Explanations for Medical Image Models via Neural Explanation Masks
paper_summary: paper_summary: The work introduces NEMt, a modification of the NEM method to specifically explain image classification predictions, in contrast to the original model which explains latent representations.
The work also extends the NEM method with a new final layer to remove artifacts dubbed "stochastic weighted neighborhood averaging" which was needed as standard NEM produced a lot of artifacts in classification explanation.
NEMt works by essentially training a segmentation network that explains a given classification network by outputting an attribution score for each pixel.

During inference, the image is predicted by the classification network and the segmentation network is given the image and some representations from the classification network and predicts the attribution map.
During training additional steps are applied where the pixels indicated by the attribution map are occluded and the new image is fed through the classification network.
The segmentation network is then trained to find the attributions that maximize the difference in prediction while a regularization loss forces it to segment a smaller area.

The method is compared to several others using many relevant metrics in the new classification setting as well as the original latent space setting.
The results show that NEMt is still competitive in the original latent space setting, though no comparison to the original NEM is made which leaves it uncertain whether the new artifact removal is useful in this case.
The results also show that NEMt is better than the other methods in the classification setting for most metrics on two different datasets.
main_review: main_review: The original work contributes an evaluation of a novel method (NEM) on a new problem (image classification). Additionally, the method introduced novel solutions (artifact removal) needed for the method to work in the new domain and might be generally useful for other similar methods. There were some missing details like comparisons to similar ideas in other domains, ablation evaluation of their new solutions, a relevant existing method not being compared to, and a tiny bit of implementation details missing. These were all clarified in the rebuttal and promised to be included in the revised version. For these reasons, the work is a clear fit for the conference.
strength_weakness: strength_weakness: The NEM and NEMt methods do share a lot of similarities with methods developed by the Weakly Supervises Semantic Segmentation (WSSS) field which has also developed methods to extract masks using pretrained classification models, however, these are not mentioned at all in work.
The artifact removal is described as useful and necessary but its impact is not shown through evaluation, something that could have been achieved by evaluating how the method performs without it or with the original NEM regularizing loss instead (especially since one of the test cases is specifically in the domain NEM was developed for).
NEMt is compared against RELAX in the latent space setting. Since RELAX can be described to be a latent space adaption of RISE, it is strange to not also compare it against RISE in the classification setting.
While relevance mass is a useful metric it is biased towards metrics that produce positive attributions for relevant areas and negative otherwise, this is not strictly speaking needed for an attribution. This can be alleviated with normalization, but whether this is done is not described.