XProtoNet: Diagnosis in Chest Radiography with Global and Local Explanations

A number of researchers have attempted to identify diseases via chest radiography using DNNs. Wang et al. [32] and Rajpurkar et al. [23] proposed the use of a conventional convolutional neural network to localize disease through a class activation map [34]. Taghanaki et al. [31] utilized a variational online mask on a negligible region within the image and predicted disease using the unmasked region. Guan et al. [6] proposed a class-specific attention method and Ma et al. [20] used cross-attention with two conventional convolutional neural networks. Hermoza et al. [10] used a feature pyramid network [18] and an additional detection module to detect disease. Li et al. [17] proposed a framework to simultaneously perform disease identification and localization, exploiting a limited amount of additional supervision. Liu et al. [19], also utilizing additional supervision, proposed a method to align chest X-ray images and learn discriminative features by contrasting positive and negative samples. Some of these approaches localize the disease along with classification but cannot explain the predictive process of how this localized part contributes to model prediction. Herein, we aim to build a diagnostic framework to explain the predictive process rather than simply localize the disease.

There have been various post-hoc attempts to explain already-trained models [28, 27, 30, 2, 15], but some of them provide inaccurate explanations [29, 1]. Additionally, they only show the region where the network is looking within a given image [26]. To address this problem, several models have been proposed with structurally built-in interpretability [21, 16, 3, 8]. Since their prediction process itself is interpretable, they require no additional effort to obtain interpretation after training. A self-explaining neural network [21] obtains both concepts that are crucial in classification and the relevance of each concept separately through regularization, then combines them to make a prediction. Case-based interpretable models, mostly inspiring us, learn prototypes that present the properties of the corresponding class and identify the similarity of the features of a given input image to the learned prototypes [16, 3, 8]. Li et al. [16] used an encoder-decoder framework to extract features and visualize prototypes. Chen et al. [3] defined prototypes as a local feature of the image and visualized the prototypes by replacing them with the most similar patches of training data. Hase et al. [8] proposed training prototypes in a hierarchical structure. These works targeted classification tasks in general images, and there was no attempt to make an interpretable automated diagnosis framework for chest radiography. To this end, we propose an interpretable diagnosis model for chest radiography that learns disease-representative features within a dynamic area.

XProtoNet compares a given input image to learned disease-representative features to diagnose a disease. It has a set of $K$ learned prototypes $\mathcal{P}^{c}=\{\mathbf{p}_{k}^{c}\}_{k=1}^{K}$ for each disease $c$, where the prototype $\mathbf{p}_{k}^{c}$ presents a discriminative feature of disease $c$. Given an input image $\mathbf{x}$, the feature extractor extracts the feature vector $\textbf{f}_{\mathbf{p}_{k}^{c}}(\mathbf{x})$ for each prototype $\mathbf{p}_{k}^{c}$, and the prototype layer calculates a similarity score $s$ between $\textbf{f}_{\mathbf{p}_{k}^{c}}(\mathbf{x})$ and $\mathbf{p}_{k}^{c}$, which are $\rm{D}$-dimensional vectors. Similarity score $s$ is calculated using cosine similarity as

Diagnosis from chest radiography is a multi-label classification, which is a binary classification of each class. We thus derive the prediction score of target disease $c$ by considering only the prototypes of $c$, not the prototypes of the non-target diseases, in the classification layer. The prediction score is calculated from

where $w_{\mathbf{p}_{k}^{c}}$ denotes the weight of $\mathbf{p}_{k}^{c}$ and $\sigma$ represents a sigmoid activation function. Similarity score $s$ indicates how similar the feature of the input image is to each prototype, and weight $w_{\mathbf{p}_{k}^{c}}$ indicates how important each prototype is for the diagnosis. By this process, XProtoNet can diagnose the disease based on the similarity between the corresponding prototypes and the features of the input X-ray image. After the training, prototype $\mathbf{p}_{k}^{c}$ is replaced with the most similar feature vector $\textbf{f}_{\mathbf{p}_{k}^{c}}$ from the training images. This enables the prototypes to be visualized as human-interpretable training images, without an additional network for decoding the learned prototype vectors.

When extracting feature vectors $\textbf{f}_{\mathbf{p}_{k}^{c}}$, XProtoNet considers two separate aspects of the input image: the patterns within the image and the area on which to focus to identify a certain disease. Therefore, the feature extractor of XProtoNet contains a feature module and an occurrence module for each one of the above-mentioned aspects. The feature module extracts the feature map $F(\mathbf{x})\in\mathbb{R}^{\rm{H}\times\rm{W}\times\rm{D}}$, the latent representations of the input image $\mathbf{x}$, where $\rm{H}$, $\rm{W}$, and $\rm{D}$ are the height, width, and dimension, respectively. The occurrence module predicts the occurrence map $M_{\mathbf{p}_{k}^{c}}(\mathbf{x})\in\mathbb{R}^{\rm{H}\times\rm{W}}$ for each prototype $\mathbf{p}_{k}^{c}$, which presents where the corresponding prototype is likely to appear, that is, the focus area. Both modules consist of $1\times 1$ convolutional layers. Using occurrence map $M_{\mathbf{p}_{k}^{c}}(\mathbf{x})$, feature vector $\textbf{f}_{\mathbf{p}_{k}^{c}}(\mathbf{x})$ to be compared with prototype $\mathbf{p}_{k}^{c}$ is obtained as follows:

where $u\in[0,\rm{H}\times\rm{W})$ denotes the spatial location of $M_{\mathbf{p}_{k}^{c}}(\mathbf{x})$ and $F(\mathbf{x})$ (Figure 3(b)). The values of occurrence map, which are in the range $[0,1]$, are used as the weights when pooling the feature map $F(\mathbf{x})$ so that the feature vector $\textbf{f}_{\mathbf{p}_{k}^{c}}(\mathbf{x})$ represents a feature in the highly activated area in the occurrence map.

By pooling the feature map with the occurrence map, a class-representative feature is presented as a vector of a single size, regardless of the size or shape of the area in which the feature appears. During training, the occurrence area is optimized to cover the area where disease-representative features for each disease appear, and the prototypes become disease-representative features in an adaptive area size. As mentioned in Section 3.1, prototype $\mathbf{p}_{k}^{c}$ is replaced with the most similar feature vector $\textbf{f}_{\mathbf{p}_{k}^{c}}$ after training the feature extractor, thus the prototype can be visualized as the occurrence area of the images that the prototype vectors are replaced with.

Comparison with ProtoPNet. XProtoNet differs from ProtoPNet [3] by being able to learn features within a dynamic area. In ProtoPNet, the prototypes are compared with fixed-size feature patches from an input image (Figure 3(a)). The spatial size of the prototype is $r\times r$, which is smaller than the feature map. At all spatial locations in feature map $F(\mathbf{x})$, a patch from $F(\mathbf{x})$ of the same size as prototype $\mathbf{p}_{k}^{c}$ is compared to the prototype; the maximum value of the resulting similarity map becomes the final similarity score. Since a fixed-size patch in the feature map is compared with the prototypes, the prototypes can only learn representative patterns within that patch. Thus, the size of the patch greatly affects the classification performance. The prototypes may learn an insufficient portion of the class-representative pattern if the patch is not large enough, and class-irrelevant features may be presented in the prototypes if the patch is too large. The disease-representative pattern can appear in a wide range of areas, so comparing it with a fixed-size patch may limit the performance. By contrast, the feature vector in XProtoNet represents the feature throughout the wide range of area predicted by the network, and is not limited to a fixed-size region (Figure 3(b)).

There are four losses in training XProtoNet: classification loss $\mathcal{L}_{\text{cls}}$, cluster loss $\mathcal{L}_{\text{clst}}$, separation loss $\mathcal{L}_{\text{sep}}$, and occurrence loss $\mathcal{L}_{\text{occur}}$.

Classification. To address the imbalance in the dataset, a weighted balance loss is used for $\mathcal{L}_{\text{cls}}$ as in [20]:

where $p^{c}_{i}=p(y^{c}|\mathbf{x}_{i})$, the prediction score of the $i$-th sample $\mathbf{x}_{i}$, and $\gamma$ is a parameter for balance. $|N_{\text{neg}}^{c}|$ and $|N_{\text{pos}}^{c}|$ denote the number of negative (0) and positive (1) labels on disease $c$, respectively. Further, $y^{c}_{i}\in\{0,1\}$ denotes the target label of $\mathbf{x}_{i}$ on disease $c$.

Regularization for Interpretability. To allow $\mathbf{p}_{k}^{c}$ to present the characteristics of disease $c$, the similarity between $\mathbf{x}$ and $\mathbf{p}_{k}^{c}$ should be large for a positive sample and small for a negative sample. Similar to [3], we define cluster loss $\mathcal{L}_{\text{clst}}$ to maximize the similarity for positive samples and separation loss $\mathcal{L}_{\text{sep}}$ to minimize the similarity for negative samples:

As in Eq. 4, $\mathcal{L}_{\text{clst}}^{c}$ and $\mathcal{L}_{\text{sep}}^{c}$ are weighted with the number of negative and positive samples when they are summed over all diseases and samples.

Regularization for Occurrence Map. To obtain prediction results with good interpretability, it is important to predict an appropriate occurrence map. Thus, we add two regularization terms to the training of the occurrence module. As in general object localization [33], since an affine transformation of an image does not change the relative location of a sign of the disease, it should not affect the occurrence map, either. We thus define the transformation loss $\mathcal{L}^{c}_{\text{trans}}$ for disease $c$ as

where $A(\cdot)$ denotes an affine transformation. We also add $L_{1}$ loss on the occurrence map to achieve locality of the occurrence area. It makes the occurrence area as small as possible to avoid covering more regions than necessary. The occurrence loss $\mathcal{L}_{\text{occur}}^{c}$ is thus expressed as

Overall Cost Function. All components of the loss are summed over all diseases, so the total loss is expressed as

where $\lambda_{\text{clst}}$, $\lambda_{\text{sep}}$, and $\lambda_{\text{occur}}$ are hyperparameters for balancing the losses.

Dataset. The public NIH chest X-ray dataset [32] consists of 112,120 frontal-view X-ray images with 14 disease labels from 30,805 unique patients. Experiments are conducted with two kinds of data splitting. In most of the experiments, we use an official split that sets aside 20% of the total images for the test set. We use 70% for training and 10% for validation from the remaining images. In comparison with recent methods using additional supervision (Table 3) and analysis with a prior condition to have specific prototypes (Section 4.5), we conduct a five-fold cross validation, similar to that in [17, 19]. In the official test set, there are 880 images with 984 labeled bounding boxes, provided for only eight types of diseases. We separate the total data into box-annotated and box-unannotated sets and conduct a cross-validation, where each fold has 70% of each set for training, 10% for validating, and 20% for testing. Note that we do not use the bounding box annotation during training, except for analysis with the prior condition. Patient overlap does not occur between the splits. We resize images to $512\times 512$ and normalize them with ImageNet [5] mean and standard deviation. We use data augmentation, by which images are rotated up to $10^{\circ}$ and scaled up or down by 20% of the image size, similar to that in [10].

Evaluation. We evaluate the diagnostic performance of XProtoNet using the area under the receiver operating characteristic curve (AUC) scores.

Experimental Details. We use ImageNet [5] pretrained conventional neural networks as a backbone (\eg, ResNet-50 [9] and DenseNet-121 [11]). The feature extractor consists of convolutional layers from the backbone network, feature module, and occurrence module. The feature and occurrence modules each consist of two $1\times 1$ convolutional layers with ReLU activation between them. The occurrence module has an additional sigmoid activation function to rescale the occurrence value to $[0,1]$. The weights of the classification layer are initially set to 1 so that high similarity scores with the prototypes would result in a high score for the disease. $K$ and $\rm{D}$ are set to 3 and 128, respectively. The batch size is set to 32. We set $\lambda_{\text{clst}}$, $\lambda_{\text{sep}}$, and $\lambda_{\text{occur}}$ to 0.5. Balance parameter $\gamma$ for $\mathcal{L}_{\text{clst}}$ is set to 2. We use random resizing with ratios 0.75 and 0.875 as affine transformations for $\mathcal{L}_{\text{trans}}$ in Eq. 6.

We follow the training scheme of ProtoPNet [3]: 1) training the model, except for the convolutional layers from the pretrained network and the classification layer, for five epochs; 2) training the feature extractor and the prototype layer until the mean AUC score of the validation set does not improve for three consecutive epochs; 3) replacing the prototypes with the nearest feature vector from the training data; and 4) training the classification layer. The training steps, except for the first step, are repeated until convergence. To retain only supporting prototypes for each disease, prototypes with negative weights are pruned. More details are explained in the supplementary material.

Visualization. The occurrence maps are upsampled to the input image size and normalized with the maximum value for visualization. The prototypes are marked with contours, which depict regions in which the occurrence values are greater than a factor of 0.3 of the maximum value in the occurrence map.

Table 1 shows the comparison of the diagnostic performance of XProtoNet with various baselines that use different methods of comparison with the prototypes. ResNet-50 [9] is used as the backbone. Baseline $\text{Patch}_{r\times r}$ refers to the method that follows ProtoPNet [3] with prototypes of spatial size $r\times r$, as in Figure 3(a); baseline GAP refers to the method where the feature vector $\textbf{f}_{\mathbf{p}_{k}^{c}}(\mathbf{x})$ is obtained by global average pooling (GAP) of the feature map $F(\mathbf{x})$ without an occurrence map. The different performances of the baselines $\text{Patch}_{r\times r}$ show that the performance varies greatly depending on the size of the patch. In addition, the performance of baseline GAP is similar and at times lower than that of baseline $\text{Patch}_{r\times r}$. By contrast, because XProtoNet predicts the adaptive area to compare, it achieves higher performance in all classes than the baselines: the mean AUC score of $0.820$ is $3.5\%$ higher than the highest baseline $\text{Patch}_{r\times r}$ mean AUC score, which is $0.792$. Especially, the improvement in hernia is significant ($>25\%$). This confirms that our proposed method of learning disease-representative features within a dynamic area is effective for diagnosis of medical images. Moreover, $\mathcal{L}_{\text{trans}}$ is also helpful in improving the performance.

Figure 4 shows the comparison of the predictions between XProtoNet and the baseline $\text{Patch}_{3\times 3}$ which shows the best diagnostic performance among the baselines $\text{Patch}_{r\times r}$. The cardiomegaly prototype of the baseline presents only a portion of the heart, resulting in a high similarity score (0.775) with the negative sample (Figure 4(a)). By contrast, the prototype of XProtoNet presents almost the whole area of the heart; this is more interpretable than the baseline, and the similarity score between the two occurrence areas is low (-0.369). Note that the similarity score takes a value in the range $[-1,1]$. Given the positive sample of nodule (Figure 4(b)), XProtoNet successfully detects the small nodule with a high similarity score (0.936) to the prototype, while the baseline fails. In addition, the occurrence area corresponding to the nodule prototype of XProtoNet is consistent with the ground-truth bounding box. This confirms that our proposed method shows more interpretable visualizations of the prototypes and more accurate predictions than the baseline.

Figure 5 shows some examples of the global and local explanations of XProtoNet. The global explanation of XProtoNet in the diagnosis of mass can be interpreted as follows: the prototypes of mass present an abnormal spot as a major property of mass for XProtoNet; this agrees with the actual sign of lung mass [7]. In terms of the local explanation of the X-ray image (top left in Figure 5), XProtoNet predicts that the prototypes of mass are likely to appear in the large left areas of the image, which are consistent with the ground-truth bounding box. XProtoNet outputs high similarity scores between these parts and the corresponding prototypes (0.996 and 0.993), resulting in a high prediction score (0.957) for the mass. For the diagnosis on the bottom left of Figure 5, XProtoNet identifies a small region on the right within the image as the occurrence area, which is different from the first example but consistent with the actual sign. This shows that XProtoNet can dynamically predict the appropriate occurrence area.

To see whether the learned prototypes align with actual signs of diseases, we find the image that is the most similar to the prototype among the images annotated with bounding boxes. Note that those annotations are not used during training. Figure 6 shows that the occurrence area in the image is consistent with the locus of the actual sign of each disease (green boxes). This shows that the prototypes have been well-trained to present proper disease-representative features.

We compare the diagnostic performance of XProtoNet with recent automated diagnosis methods [32, 6, 20, 10]. Table 2 shows that XProtoNet achieves state-of-the-art performance on both ResNet-50 [9] and DenseNet-121 [11] backbones while ensuring interpretability. In comparison with recent methods implemented on ResNet-50, XProtoNet achieves the best performance for 10 out of 14 diseases. Note that Ma et al. [20] use two DenseNet-121 and Hermoza et al. [10] use a feature pyramid network [18] and DenseNet-121 as the backbone: these provide better representation than a single DenseNet-121. Compared with Guan et al. [6], who use a single DenseNet-121, the mean AUC score is improved from $0.816$ to $0.822$.

We also compare the diagnostic performance of XProtoNet to that of two recent automated diagnosis methods [17, 19] using bounding box supervision, which use ResNet-50 [9] as the backbone. Table 3 shows the performances based on a five-fold cross-validation. Despite having no additional supervision, XProtoNet achieves the best performance for most diseases.

As XProtoNet provides predictions based on prototypes that are exposed explicitly, we can instruct it to diagnose using specific signs of diseases by forcing the prototypes to present those signs. We conduct analysis with the prior condition that the prototypes of XProtoNet should present the features within the bounding box annotations.

XProtoNet is trained with both box-annotated and box-unannotated data. We set both $\lambda_{\text{clst}}$ and $\lambda_{\text{sep}}$ to 1.5 for the box-annotated data and both to 0.5 for the box-unannotated data. To utilize the bounding box annotations during training, we extract the feature vectors from the feature maps within the bounding boxes as $\textbf{f}_{\mathbf{p}_{k}^{c}}^{\text{bbox}}(\mathbf{x})=\sum_{u\in\text{bbox}}{M_{\mathbf{p}_{k}^{c},u}(\mathbf{x})F_{u}(\mathbf{x})}$, where bbox denotes the spatial location inside the bounding box. We also change $L_{1}$ loss on the occurrence map for the box-annotated data to $\sum_{\mathbf{p}_{k}^{c}\in\mathcal{P}^{c}}{\sum_{u\not\in\text{bbox}}{M_{\mathbf{p}_{k}^{c},u}(\mathbf{x})}}$ to suppress the area outside the bounding box from being activated in the occurrence map. To enable the prototypes to present the features within the bounding boxes, the prototype vectors are replaced with their most similar feature vectors $\textbf{f}_{\mathbf{p}_{k}^{c}}^{\text{bbox}}$ from the box-annotated data, instead of the feature vectors $\textbf{f}_{\mathbf{p}_{k}^{c}}$ from the box-unannotated data.

Figure 7 shows the learned prototypes of XProtoNet trained with and without the prior condition. Owing to the constraint, the prototypes of XProtoNet trained with the prior condition present disease-representative features within the bounding box annotations. Although this can be a strong constraint for the model, there is no significant difference in the diagnostic performance: the mean AUC scores over 14 diseases of XProtoNet trained with and without the prior condition are 0.850 and 0.849, respectively. Therefore, using the prior condition, we enable XProtoNet diagnoses based on the specific features, thus rendering the system more trustworthy.