Explainable Artificial Intelligence Architecture for Melanoma Diagnosis Using Indicator Localization and Self-Supervised Learning

Figure 2, Visualize the architecture of our proposed model, which aims to generate a heatmap for the final convolution layer that closely resembles the regions expected by human experts when the network is trained for diagnostic classification. Our architecture consists of three key subnetworks:

• (1) A pretrained Resnet50 network for attributes classification; we denote the subnetwork (1) as $f_{Res}$ where $f_{Res}(\bm{x}_{i})=E_{1}(\bm{x}_{i})+fc_{1}(\bm{x}_{i})$ , $\bm{x}_{i}\in D^{L}$ and $E_{1}(\cdot)$ represents the encoder, which consists of four consecutive blocks, each block consists of a downsampling and two ResNet residual blocks, $fc_{1}(\cdot)$ represents a fully connected layer which produces probabilistic outputs.

• (2) An architecture based on U-Net for biomarker indicator locations. We replaced the U-Net encoder with the encoder $E_{1}(\cdot)$ of the $f_{Res}(\cdot)$ ; we denote the subnetwork (2) as $f_{Seg}(\cdot)$ where $f_{Seg}(\bm{h}_{i})=E_{1}(\bm{h}_{i})+D_{1}(\bm{h}_{i})$, where $\bm{h}_{i}$ denotes the heatmaps that are generated using Grad-Cam, $D_{1}(\cdot)$ denotes the decoder which outputs biomarker indicator location.

• (3) We add a ResNet50-based projection head for self-supervised learning; we denote the subnetwork (3) as $f_{CLR}(\cdot)$, where $f_{CLR}(\bm{x}_{i}^{\prime})=E_{1}(\bm{x}_{i}^{\prime})+P_{1}(\bm{x}_{i}^{\prime})$ and $\bm{x}_{i}^{\prime}\in D^{UL}$, $P_{1}(\cdot)$ is a projection head consisting of two fully connected layers that outputs a feature vector.

In summary, our architecture learns to classify the input images. Then, we use Grad-Cam to generate attention heatmaps. These heatmaps do not necessarily show interpretable indicators, but are useful for classification which means that they should have correlations with features expert clinicians use. Our idea is to feed these heatmaps to the segmentation subnetwork and generate the indicator biomarker localization maps using the Grad-Cam heatmaps. As a result, the architecture learns to generate explanations along with diagnosis labels. We can see that our full architecture can only be trained using Dataset A.

Our Bio-Unet baseline consists of networks $f_{Res}(\cdot)$, $f_{Seg}(\cdot)$ . At first, We train the network $f_{Res}(\cdot)$ to do the skin lesion classification task using simple supervised learning. After this stage, the classification network can predict diagnosis labels with high accuracy. We then apply Grad-CAM on the network to generate attention heatmaps. Grad-CAM is generally used by computing the gradient of the classification score with respect to the last convolutional feature map to identify the classification score of the selected target in the most influential part of the input image. However, the final convolutional feature map primarily provides a high-level region of the input image which corresponds to the area of interest. Our experiments demonstrate that this mostly will just generate the full lesion mole similar to Figure 1 which is usually much larger than the region delineated by experts for a specific indicator and usually cannot reflect the location and area of an indicator. However, the earlier convolutional feature maps contain low-level information, e.g., the boundary of a region of interest. Hence, combining attention maps at all convolutional layers looks like an option that may generate a good estimation for the location of an indicator, but we empirically observed that simply averaging all heatmaps will result in poor output. Therefore, we would like to benefit from binary masks provided by experts and reconstruct these heatmaps using the segmentation network $f_{Seg}(\cdot)$. We learn to use appropriate weights for each heatmap so that all bottleneck blocks of $E_{1}(\cdot)$ contribute to reconstruction of the final heatmap. Because the encoder $E_{1}(\cdot)$ has been used for training, in order not to affect its parameters, we created an encoder $E_{3}(\cdot)$ with the same architecture as $E_{1}(\cdot)$. Specifically, the Resnet encoder $E_{3}(\cdot)$ consists of 12 bottleneck blocks, each loaded with optimal checkpoint parameters.Then we input a training image and a bottleneck block in Grad-CAM, get a heatmap, and then replace the bottleneck block in turn, getting a total of 12 localization masks. We denote $f_{Grad\_CAM}(b_{i},j)$ where $b_{i}$ is the bottleneck block and j is the index of the selected label.

By performing downsampling on the input heatmaps, the encoder will focus more on important features. During training, we add only the heatmap of the currently predicted feature by the encoder as input to $f_{Seg}(\cdot)$, making the classification model $f_{Res}(\cdot)$ more accurate for the current attribute’s region of interest. Therefore, explanations can improve the diagnostic accuracy.

The proposed Bio-Unet baseline architecture can learn boundaries efficiently, but this is only effective when the number of annotated images is large for biomarker indicators, e.g., pigment network. or the area corresponding to the indicators on a lesion is contiguous and large. To enable classification models to generate accurate heatmaps when Dataset A is small or when the area for an indicator is very small or is scatterred on the image, we benefit from self-supervised learning on Dataset B to improve the baseline of our proposed network architecture that is obtained by training on Dataset B. Specifically, we use SimCLR [4] which uses contrastive learning for improved visual representation. As shown in Figure 2, there are two independent data augmenters $T_{1}(\cdot)$ and $T_{2}(\cdot)$ which are randomly selected from rotation, scaling, cropping, brightness, contrast, saturation, and flipping transforms to generate an augmented version samples of Dataset B so that we can compute the contrastive learning loss. Each training image $\bm{x}^{\prime}_{i}\in D^{UL}$is passed through two data augmenters to produce two augmented images. The two augmented images will then pass through our shared weight encoder $E_{1}(\cdot)$ and projection head $P_{1}(\cdot)$, resulting in two 128-length features. In a minibatch of $N$ input images, $2N$ augmented images will be produced, for each pair of augmented images we treat as positive pairs, the other $2(N-1)$ are negative examples.We adopt contrastive loss as semi-supervised loss in Eq. (2)

where $sim(u,v)=\frac{u^{T}v}{\|u\|\|v\|}$ and k $\not=$ i and $\tau$ denotes a temperature parameter. Upon training the encoders on Dataset B using self-supervised learning, we can benefit from transferring obtained knowledge across Dataset A and Dataset B. Due to the space limit, the full training procedure for our architecture is described in Algorithm 1.

After training the ResNet50 subnetwork and the Bio-UNet architecture for classification, we obtained $76\%$ and $82\%$ classification accuracy rates, respectively. We observe that both architecture have high melanoma diagnosis rates and using the additional subnetworks for segmentation and localization has led to a significant diagnosis performance boost. Table 2 presents results for localizing the five indicators when measured using the Continuous DICE metric. As it can be seen, utilizing Bio-UNet architecture along with our proposed training scheme enhances localization results by 0.95%, 2.33%, 5.11%, -8.24%, 9.01% over the standard architecture for the five indicators listed in Table  1, respectively. Each feature is critical, and Bio-UNet enables the encoder to better capture three features that are difficult for ResNet to detect. Every feature is vital for melanoma diagnosis, so when each feature can be localized uniformly, the accuracy of cancer diagnosis increases significantly. Table 2 presents results for localizing the five indicators when measured using the Continuous DICE metric.

We find that when the model’s decisions are close to what is expected by humans, the accuracy would increase. We conclude from our results that in order to make AI explainable, we may need to incorporate intermediate-level human-interpretable annotations in the deep learning end-to-end training pipelines and design DNN architectures that learn to perform a downstream task using human-interpretable intermediate indicators.

To provide an intuition behind the quality of results presented in Table 2, we have visualized samples of heatmaps that are generated using different methods in Figure 3 for visual inspection. In this figure, we have selected an example image such that it is annotated with a corresponding indicator biomarker and presented the localization map that the AI pipelines generate. We can observe that the three CAM-based techniques generate features maps that are far larger than the ground truth mask and pretty much include the majority of the input mole. This means that the generated maps are not interpretable because they point to the whole mole which even a beginner knows that it should be the primary area of attention. In contrast, a close visualize comparison between columns two and six, demonstrates that our method generate binary feature maps that are quite similar to the ground truth, focusing on a subarea of the mole that in reality pertains to the clinical indicator. We also conclude that although the DICE metric is the predominant metric for segmentation, it is a sensitive metric when the semantic classes in the images are imbalanced.