Proto-Caps: interpretable medical image classification using prototype learning and privileged information

Dataset 1: lung nodule classification

The Lung Image Database Consortium and Image Database Resource Initiative (LIDC-IDRI) dataset (Armato et al., 2015) is a comprehensive collection of computed tomography (CT) scans, annotated by expert radiologists, with the primary focus on lung nodule detection and characterization (Armato et al., 2011). The annotations contain detailed information on the medical assessment and visual appearance of pulmonary nodules. Each scan was reviewed by up to four radiologists. Identified lung nodules were segmented and annotated with a malignancy rating and visual attributes. The malignancy rating ranges from 1-highly unlikely to 5-highly suspicious. The visual appearance is described by the attributes subtlety (difficulty of detection, 1-extremely subtle, 5-obvious), internal structure (1-soft tissue, 4-air), pattern of calcification (1-popcorn, 6-absent), sphericity (1-linear, 5-round), margin (1-poorly defined, 5-sharp), lobulation (1-no lobulation, 5-marked lobulation), spiculation (1-no spiculation, 5-marked spiculation), and texture (1-non-solid, 5-solid).

Preprocessing for this dataset involves ignoring lung nodules smaller than 3 mm and those identified by less than three radiologists. In the case of 2D processing the lung nodules were cropped out of the volumes with the minimum square bounding box and the slices were resized to $32\times 32$, as in previous work (LaLonde, Torigian & Bagci, 2020). Considering each annotation and slice as a sample, the preprocessing results in a total of 27,379 samples. For 3D processing, a fixed depth was chosen with a centered bounding box, resulting in a volume size of $32\times 32\times 16$, according to existing work (Mehta et al., 2021). Considering each annotation as a sample, the preprocessing results in a total of 4,318 samples. Preprocessing was performed with the pylidc framework (Hancock & Magnan, 2016).

Each sample comes with the individual annotator’s segmentation, while the target and attribute labels take into account the annotations of all annotators, as in previous work (LaLonde, Torigian & Bagci, 2020). For the nodule’s attribute score the mean value of the radiologists’ scores is used. The target ground truth is represented with the distribution of the annotated malignancy scores.

Experiments with this dataset were performed using five-fold stratified cross-validation, with a patient-wise split and using 10% of the training data for validation.

Dataset 2: lung abnormality classification

As second dataset to validate the proposed method we chose the CheXpert dataset which contains 224,316 chest radiographs of 65,240 patients (Irvin et al., 2019). In this dataset the testing set includes labels which were obtained by expert radiologists, whereas the training labels were generated automatically from the associated radiology reports. In this work the results from the CheXbert labeler (Smit et al., 2020) were used for the training labels.

The findings include positive, negative, uncertain, and unmentioned (blank) classes. In our experiments we used binary attribute classes, considering unmentioned and negative classes as class 0 and positive and uncertain as class 1. We considered all 13 provided attribute classes which include enlarged cardiomediastinum, cardiomegaly, lung opacity, lung lesion, edema, consolidation, pneumonia, atelectasis, pneumothorax, pleural effusion, other pleural, fracture, and support devices. The binary target classes used were the absence (0) or the presence (1) of one or more of the attribute findings.

Considering only samples with a frontal projection the total number of data samples is 191,229. Preprocessing includes the resizing of the images to a size of $64\times 64$. The train and test split was adopted as defined in the dataset.

In contrast to the LIDC-IDRI dataset, the CheXpert dataset lacks segmentation masks, which is why the segmentation branch of the proposed method is disabled by excluding the respective term in the loss function.

Model

The architecture comprises an encoder and several branches stemming from the output of the encoder, similar to previous work (LaLonde, Torigian & Bagci, 2020). The encoder is a capsule network consisting of convolutional layers that produce latent capsule vectors. These vectors serve as input to multiple heads that predict the target, attribute scores and prototypes, as well as a region of interest (ROI) mask.

The capsule vectors represent different features that have been extracted from the input image. Each feature is mapped to a single predefined attribute using supervised learning. This combination of the encapsulated feature representation and attribute mapping enables a truthful and human-understandable description of the decision process. Furthermore, the capsule information is used to generate attribute-specific prototypes that determine the attribute scores during inference and provide visual examples.

Referring to Fig. 1, the model can be divided into individual sections: The capsule values are calculated by the backbone, which is then followed by an attribute head, target head and segmentation head. The prototype branch extends the model by learning attribute-specific visual prototypes.

The backbone extracts features from the input image data and implements a capsule attention mechanism. Firstly the input data is processed by a convolutional layer containing 256 kernels of size 9 followed by a Rectified Linear Unit (ReLu) activation. A subsequent primary capsule layer applies another convolutional layer with 256 kernels of size 9 and segregates the resulting features in separate low level feature capsules. A final dense layer implements the dynamic routing algorithm (Sabour, Frosst & Hinton, 2017) in an iterative manner, resulting in 16-dimensional capsule vectors and forming the starting point for the following prediction branches.

The target head concatenates the capsule vectors and applies a fully connected layer. The output dimension is selected based on the label representation. For the LIDC-IDRI dataset, which has a distribution of five labels, the output dimension of the linear layer is five, followed by a softmax activation. For binary class labels, as in the CheXpert dataset, the output dimension is one with a sigmoid activation function.

The segmentation head decodes the capsule features to a mask of the ROI. The idea of learning this optional branch is an attention of the object to classify. The capsule vectors are concatenated and processed by two fully connected layers with 512 and 1,024 output features and ReLU activation. A final linear layer completes the upsampling and creates an output image of the same size as the original image and is followed by a sigmoid activation.

The attribute head implements the specialization of the capsules on individual, human-understandable attributes. Each capsule vector is processed by a separate linear layer and sigmoid activation to fit a respective attribute score.

The prototype branch contains multiple sample vectors for each capsule that eventually become feature vectors of prototypical samples during the training. For each attribute value, a fixed number of prototypes are available to represent the visual variety of characteristics. The prototype vectors map real image samples and are visualized during inference.

The code of Proto-Caps is publicly available at https://github.com/XRad-Ulm/Proto-Caps. Algorithm 1 depicts the training and inference phases, described in the following.

Training phase 1—feature extraction

At the beginning of the training, the model is optimized in solely respect of the training labels. The clustering of the prototypes is disabled at this stage. This step-wise approach allows for a prioritized focus on feature extraction.

In this first training phase, the network is optimized using the following combined loss function given the training labels.

(1) $${\mathcal{L}}_{\mathrm{P}1}={\mathcal{L}}_{tar}+{\mathcal{L}}_{attr}+{\lambda }_{seg}\cdot {\mathcal{L}}_{seg}.$$

The target loss ${\mathcal{L}}_{tar}$ refers to the target label and has been selected depending on the format of the label. In case of a distribution of the target annotations, as in the LIDC-IDRI dataset, the pointwise Kullback-Leibler divergence was used to reflect the inter-observer agreement and thus uncertainty (LaLonde, Torigian & Bagci, 2020). In the case of binary classification, as in the CheXpert dataset, a binary cross-entropy loss is applied.

(2) $$\begin{array}{rl}{\mathcal{L}}_{tar}& =Y_{tar}\ast \log \frac{Y_{tar}}{{\hat{Y}}_{tar}}(\mathrm{L}\mathrm{I}\mathrm{D}\mathrm{C}-\mathrm{I}\mathrm{D}\mathrm{R}\mathrm{I}\mathrm{d}\mathrm{a}\mathrm{t}\mathrm{a}\mathrm{s}\mathrm{e}\mathrm{t})\\ {\mathcal{L}}_{tar}& =BCE(Y_{tar},{\hat{Y}}_{tar})(\mathrm{C}\mathrm{h}\mathrm{e}\mathrm{X}\mathrm{p}\mathrm{e}\mathrm{r}\mathrm{t}\mathrm{d}\mathrm{a}\mathrm{t}\mathrm{a}\mathrm{s}\mathrm{e}\mathrm{t}).\end{array}$$

The attribute loss ${\mathcal{L}}_{attr}$ is based on the error between the ground truth attribute score $Y_a$ and the network prediction ${\hat{Y}}_a$ for the $a$-th attribute:

(3) $$\begin{array}{rl}{\mathcal{L}}_{attr}& =b\ast \frac{1}{A}\sum _a^A\Vert Y_a-{\hat{Y}}_a\Vert (\mathrm{L}\mathrm{I}\mathrm{D}\mathrm{C}-\mathrm{I}\mathrm{D}\mathrm{R}\mathrm{I}\mathrm{d}\mathrm{a}\mathrm{t}\mathrm{a}\mathrm{s}\mathrm{e}\mathrm{t})\\ {\mathcal{L}}_{attr}& =b\ast \frac{1}{A}\sum _a^{A}BCE(Y_a,{\hat{Y}}_a)(\mathrm{C}\mathrm{h}\mathrm{e}\mathrm{X}\mathrm{p}\mathrm{e}\mathrm{r}\mathrm{t}\mathrm{d}\mathrm{a}\mathrm{t}\mathrm{a}\mathrm{s}\mathrm{e}\mathrm{t}).\end{array}$$

If the attribute scores are continuous, as in the LIDC-IDRI dataset the mean square error is applied, else if the attribute scores are binary, as in the CheXpert dataset, a binary cross-entropy loss is applied. The random binary mask $b$ can be used to control from which training sample the attribute annotations are used and from which not, allowing semi-supervised attribute learning. This version was investigated in case attribute labels are only available for part of the training dataset (see “Sparse Data Study”).

The segmentation loss ${\mathcal{L}}_{seg}$ is used when the dataset includes segmentation masks of the region of interest. By incorporating this loss term the segmentation branch of the model is activated with the mean square error between the reconstructed and the ground truth mask.

(4) $${\mathcal{L}}_{seg}=\Vert Y_{mask}-{\hat{Y}}_{mask}\Vert .$$

The hyper-parameter ${\lambda }_{seg}=0.512$ was chosen according to LaLonde, Torigian & Bagci (2020).

Training phase 2—training prototypes

After the warm-up phase, the randomly initialized prototype vectors are being unfrozen and trained to represent cluster centers for the differen attribute manifestations. Motivated by existing work about prototype learning (Chen et al., 2019), the combined loss function is being extended by two terms, defined as a cluster loss ${\mathcal{L}}_{clu}$ and a separation loss ${\mathcal{L}}_{sep}$. The cluster loss reduces the Euclidean distance between the capsule vector ${\overrightarrow{c}}^a$ and the nearest prototype ${\overrightarrow{p}}^{a,p}$ of the correct attribute score in $P_{a_s}$.

(5) $${\mathcal{L}}_{clu}=\frac{1}{A}\sum _a^A\underset{{\overrightarrow{p}}^{a,p}\in P_{a_s}}{min}{\Vert {\overrightarrow{c}}^a-{\overrightarrow{p}}^{a,p}\Vert }_2.$$

On the other hand the separation loss increases the distance between the capsule vector and prototypes, that are dedicated to other attribute scores, limited by a maximum distance:

(6) $${\mathcal{L}}_{sep}=\frac{1}{A}\sum _a^A\underset{{\overrightarrow{p}}^{a,p}\notin P_{a_s}}{\mathrm{m}\mathrm{i}\mathrm{n}}\mathrm{m}\mathrm{a}\mathrm{x}(0,\mathrm{d}\mathrm{i}\mathrm{s}{\mathrm{t}}_{max}-{\Vert {\overrightarrow{c}}^a-{\overrightarrow{p}}^{a,p}\Vert }_2).$$

The overall loss function for the second training phase is the following weighted sum, where the separation factor ${\lambda }_{sep}=0.1$ was chosen empirically:

(7) $${\mathcal{L}}_{\mathrm{P}2}={\mathcal{L}}_{tar}+{\mathcal{L}}_{attr}+{\lambda }_{seg}\cdot {\mathcal{L}}_{seg}+{\mathcal{L}}_{clu}+{\lambda }_{sep}\cdot {\mathcal{L}}_{sep}.$$

For a visual representation of the optimized prototype vectors with real sample images, training samples are pushed onto the prototype vectors. Two approaches of this push operation were examined. The first approach includes to find the sample with the smallest distance for each prototype vector and save it with the original image and the ground truth information. The second approach, which is similar to existing work about prototype learning (Chen et al., 2019) additionally involves replacing the prototype vector with the vector generated by this real sample.

The repetition rate of this operation is set by the hyper-parameter $pushstep$.

Inference phase—predict from prototypes

During inference, the prediction of the attribute and target scores is based on the sample prototypes. Firstly, the prototype vectors that are closest to the capsule vectors are determined. The attribute prediction is set to the attribute label of the corresponding closest prototype. The trained attribute layers are ignored in the inference phase. For the target prediction, the prototype vectors are concatenated and processed by the prediction layer.

Faithfulness

The following experiments investigate the extent to which the explanations reflect the model’s decision.

In the first experiment we analyzed the correlation between attributes and target correctness using a post-hoc logistic regression study: Based on the correctness of the attribute prediction, the correctness of the target prediction was calculated. The result on the LIDC-IDRI dataset shows a strong relationship between both with a prediction accuracy of 94.1% /1.1.

In the second experiment we analyzed the joint feature importance for the combined task in order to gain insights about the information flow within the network. The explainability of Proto-Caps relies on the capsule vectors as a foundation for decision-making. These are used to derive the attribute scores and also to classify the target. To ensure the trustworthiness of the explanations, we used the LIDC-IDRI dataset to analyze whether the information was shared across the multidimensionality of the capsule vectors for this combined task.

The joint information is derived from the weights $w$ of the linear layer of each attribute and the target prediction. The layer input values are used to standardize the weights $w^{\mathrm{\prime }}$. To achieve this, the capsule vectors $z$ for the test dataset are first standardized:

(8) $$z^{\mathrm{\prime }}=\frac{z-\tilde{z}}{{\sigma }_z}.$$

The standardized weights are calculated using the linear function:

(9) $$y=z\ast w=z^{\mathrm{\prime }}\ast w^{\mathrm{\prime }}.$$

The correlation functions are calculated from each attribute layer $w_{attr}^{\mathrm{\prime }i}$ and the respective weight values of the target layer $w_{tar}^{\mathrm{\prime }i}$:

(10) $${\mathrm{c}\mathrm{o}\mathrm{r}\mathrm{r}}_{\mathrm{P}}^i=\mathrm{P}\mathrm{e}\mathrm{a}\mathrm{r}\mathrm{s}\mathrm{o}\mathrm{n}(|w_{tar}^{\mathrm{\prime }i}|,|w_{attr}^{\mathrm{\prime }i}|)$$

(11) $${\mathrm{c}\mathrm{o}\mathrm{r}\mathrm{r}}_{\mathrm{S}}^i=\mathrm{S}\mathrm{p}\mathrm{e}\mathrm{a}\mathrm{r}\mathrm{m}\mathrm{a}\mathrm{n}(|w_{tar}^{\mathrm{\prime }i}|,|w_{attr}^{\mathrm{\prime }i}|).$$

The mean correlation over the eight attribute capsules and five target classes for the LIDC-IDRI dataset is with the Pearson calculation 0.62 and with Spearman 0.54. These values show a moderate to strong correlation.

The experiments show that the model’s decision making process is largely driven by the attribute explanation.

User-centered evaluation

In addition to the technical investigation of the explanations, we conducted a user-centered evaluation of the attribute explanations provided by Proto-Caps (Gallée et al., 2024). Six radiologists diagnosed lung nodules from the LIDC-IDRI dataset in test scenarios while having access to the model output. For each case, three variants of explanation were presented: (A) the malignancy prediction alone, (B) the prediction with attribute scores or (C) the prediction with both attribute scores and attribute prototypes.

The evaluation incorporated objective measures—such as the diagnostic accuracy of the radiologists across the different model support variants—as well as subjective assessments of trust and perceived helpfulness.

Figure 3 illustrates that the explanations provided by Proto-Caps are persuasive: when the model prediction is correct, they lead to higher diagnostic accuracy, whereas in cases of incorrect predictions, the explanations tend to increase the likelihood of an incorrect diagnosis. Radiologists reported that the explanations aligned with their diagnostic criteria and served as useful guidance during diagnosis.

To mitigate these negative effects, we suggest using Proto-Caps selectively—primarily during development and testing—and, in practice, as a warning system that activates and intervenes only when the model suspects a false negative diagnosis by the radiologist. For a thorough discussion of the experimental design, survey methodology, and analysis of user feedback, readers are referred to Gallée et al. (2024).

Evaluation metric

For the LIDC-IDRI dataset Proto-Caps is evaluated using the Within-1-Accuracy metric. A prediction is considered correct if it falls within a tolerance of 1 score of the ground truth label. As in previous work on the LIDC-IDRI dataset (LaLonde, Torigian & Bagci, 2020), this metric is used for the ordinal labels of malignancy (ascending from 1-highly unlikely to 5-highly likely).

For the CheXpert dataset, where possible, area under the curve (AUC) values are calculated as an evaluation measure to compare the results of Proto-Caps with other studies. However, the calculation by including prediction probabilities in the Proto-Caps evaluation is only possible for the target classification, not for the attribute classification. This is due to the prototype-based inference for the attributes. The predicted attribute classification is determined by the ground truth scores of the closest prototypes.

Comparison with SOTA

We compare our proposed method with literature values of existing methods. For the LIDC-IDRI dataset these include methods that predict either no attributes (non-explainable) or some attributes (explainable). Table 1 displays the accuracy of the 2D models, while Table 2 presents the accuracy of the 3D models.

The results show that Proto-Caps outperforms existing explainable approaches in predicting both nodule malignancy and visual attributes. X-Caps (LaLonde, Torigian & Bagci, 2020) in Table 1 is also a capsule-based architecture similar to Proto-Caps. It first learns attributes and then predicts malignancy based on them. However, it does not learn attribute prototypes.

Compared to methods that do not implement attribute prediction for improved explainability (3D-CNN+MTL, TumorNet, CapsNet), our method performs better in the 2D case. In the 3D case, as shown in Table 2, it is important to note that the data preprocessing and evaluation metrics differ among the compared methods. Our proposed method achieves a higher malignancy prediction accuracy in the 3D setting (2D: 93.1 vs. 3D: 94.3) as well as a better mean prediction accuracy for the eight attributes (2D: Ø = 92.2 vs. 3D: Ø = 95.3). These results demonstrate the methodological feasibility of Proto-Caps on 3D data.

Proto-Caps also excels in the second validation dataset, CheXpert, in predicting fractures, pneumothorax, pneumonia and lung lesions, but underperforms in other attributes, see Table 3.

Number of prototypes

In the following we present the results of the experiments where we used a different number or prototypes that are initialized per attribute class. Besides the original setting with 16 prototypes per attribute class, we tested four, 32 and 64 prototypes per attribute class. The results are summarized in Fig. 4 and in the following:

Accuracy: In terms of the target accuracy 93.4/1.2 (4) 93.1/0.9 (16), 92.9/1.4 (32), 93.4/1.24 (64), and of the mean attribute accuracy 92.9/1.7 (4) 92.2/2.4 (16), 92.3/2.5 (32), 91.9/3.6 (64), there is no advantage in using more prototypes.

Training time: When changing the number of prototypes, the survey of the training time needs to be considered. No change in duration is detected during the weight adjustment phase (18 s). However, one push operation of the prototypes takes significantly longer the more prototypes are learned: 17 s (four), 37 s (16), 51 s (32), and 69 s (64).

Number of prototypes used in inference: This experiment analyzed the total number of attribute prototypes actually used in the test dataset: 107/160 (4) 311/640 (16), 418/1,280 (32), and 569/2,460 (64).

Having a diverse set of prototypical samples used on the testing dataset increases the quality of the local explainability meaning that the prototypes are adapted better on the different manifestations within one attribute class. A small set of prototypes provides a global explanation of the entire network, while a large number of prototypes offers a more local explanation for each sample.

High target and attributes accuracies were achieved when using four, 16, and 32 prototypes per attribute class. As the average number of prototypes (16 per attribute class, i.e., 80 per attribute) reflects the variance of the appearance well, we recommend the use of 16 prototypes per attribute class.

Capsule vector dimension

This section of experiments explores the effect of the dimension of the attribute capsules. The motivation for reducing the dimension is to reduce the number of learnable parameters, which raises the prospect of a more generalized representation of the attributes. In addition to the original capsule setting of 16 dimensional vectors, dimensions 2, 4, 8 and 32 were tested.

Accuracy: In terms of the target accuracy 93.5/1.3 (2), 93.0/1.6 (4), 93.3/0.9 (8), 93.1/0.9 (16), 92.6/0.9 (32), and of the mean attribute accuracy 93.1/2.5 (2), 93.1/2.0 (4), 92.3/2.2 (8), 92.2/2.4 (16), 91.0/3.4 (32), the prediction accuracies are comparable.

Number of prototypes used in inference: Of the 640 available prototypes during inference 181 (2), 265 (4), 330 (8), 311 (16), and 265 (32) were used.

Based on high prediction accuracy and the number of prototypes used, see Fig. 5, we recommend a capsule dimension of 8 or 16, as similar results were achieved with both.

Number of capsules

In previous work (LaLonde, Torigian & Bagci, 2020), the number of capsules was set to the number of attributes. This model architecture supports the idea that the target prediction is based on the defined attribute capsules. Since one could allow the prediction to be based on more than just the predefined attributes, we investigated how the model changes when the model has more capsules. The model architecture is changed by adding additional capsules that are used for the target prediction branch but are not followed by attribute prediction layers. For this experiment, we tested the variants with 4, 8, and 16 extra capsules.

Accuracy: The target’s accuracy was 78.9/10.9 (+4), 90.4/1.1 (+8), and 89.0/3.8 (+16) compared to the original accuracy of 93.1/0.9 (+0) when no additional capsules were used besides the supervised attribute capsules. The mean attribute accuracy was 91.2/2.9 (+4), 92.2/2.4 (+8), and 92.0/2.3 (+16) compared to 92.2/2.4 with the original setting.

The results indicate a decrease in target accuracy and a reduced level of robustness with the inclusion of additional capsules.

Push operation

In the following experiment on the LIDC-IDRI dataset, we compare two versions of the push operation. This operation occurs during training phase 2, when the prototype branch is being optimized. For each prototype vector, the push operation stores a real sample image to which its capsule vector is closest.

In the first variant, which was used in the experiments above, the push operation does not involve any additional steps.

In the second variant, which is based on the ProtoPNet prototype learning strategy by Chen et al. (2019), the capsule vectors of the most similar image samples replace the prototype vectors in the push operation. This additional step changes the learning progress of the prototype vectors, as they are not only adapted by the loss function, but are also replaced during the push operations.

The push variants were tested with the original capsule dimension of 16 and with 8.

It was observed that the training accuracy after the first push operation is higher with the replacing-push method. By replacing the prototype vectors with real samples, they are adapted faster initially, in contrast to the slow adaptation caused by only the change from backpropagation.

With the replacing-push method, the final accuracy for target prediction is 93.3/1.2 (16), and for mean attribute prediction it is 92.5/2.4 (16) for the capsule dimension 16. For the capsule dimension of 8 the target prediction accuracy is 90.5/2.1 (8) and the mean attribute prediction accuracy is 92.9/1.7 (8).

When comparing accuracies, there are no differences between the proposed push operation and the replacement-push variant, except for a lower target prediction accuracy when using a capsule vector dimension of 8. We recommend using the proposed push variant, where the prototype vectors are fitted through backpropagation with the prototype loss function only.