IAIA-BL: A Case-based Interpretable Deep Learning Model for Classification of Mass Lesions in Digital Mammography

Introduction

AI is revolutionizing Radiology. Conventional machine learning is currently used for computer-aided detection (i.e., to detect “lesion” vs. “no lesion”), but to make a greater clinical contribution, future approaches need to be able to assist with harder tasks, such as “should the patient get a biopsy for that lesion?” Consider mammography, which aims to detect breast cancer, a leading cause of death in the USA . In breast cancer screening, the majority of biopsies yield benign results, in the process subjecting many healthy patients to invasive testing and contribute to the societal cost of healthcare ; it is possible that machine learning might lead to improvements. In deciding how to treat a patient, radiologists must consider aspects of images that are so subtle that it is quite difficult for an untrained eye to identify even the important aspects of an image. These decisions can be challenging, even for the most experienced radiologists, as shown by relatively low inter-rater agreement scores for such difficult tasks . To ensure clinical acceptance, an AI tool will need to provide its reasoning process to its human radiologist collaborators in order to be a useful aide in these difficult and high-stakes decision-making processes .

Despite the hope of computer-aided radiology for mammography, there are serious concerns with present methods, primarily, confounding. Confounding occurs when the predictive model is using incorrect information or reasoning to make a decision, even if the decision is correct. In previous studies, researchers created models that appeared to perform well on their test sets, yet upon further inspection, based their decisions on confounding information (e.g., type of equipment) rather than medical information . If a model simply relied upon confounders such as the type of equipment used to take the mammogram or proxies for the patient’s age such as the density of the breast tissue, the model would likely fail to generalize. This problem is exacerbated by the fact that there are few publicly available mammography datasets, so many models are trained on relatively few cases.

Ideally, the reasoning process of any model would be similar to that of an actual radiologist, who looks at specific aspects of the image that are known to be important, based on the physiology of how lesions develop within breast tissue. If this reasoning process were correct, it would lead to a higher chance that: (1) The model could generalize beyond its finite training set. (2) The model’s reasoning process could be useful information for doctors, even if its prediction is sometimes incorrect. (3) It would be much easier to troubleshoot or evaluate trustworthiness of the model, since it is not a black box. (4) The model’s reasoning and reporting process could align with the structured lexicon that radiologists use to report results, such as the Breast Imaging-Reporting and Data System (BI-RADS) for breast cancer and other similar lexicons from the American College of Radiology.

Thus, unlike existing black-box systems that aim to replace a doctor , we aim to create an interpretable AI algorithm for breast lesions (abbreviated IAIA-BL) whose explicit reasoning can be understood and verified by a medical practitioner. Our novel deep learning architecture enables IAIA-BL to provide an explanation that shows the underlying decision-making process for each case. Figure 1c shows an example of how this works: the algorithm highlights parts of the image, explains that it considers these parts of images similar to prototypical cases it has seen before, and provides a score for the probability of the specific diagnosis (breast mass with mostly circumscribed margin) for this image as well as the likelihood of malignancy. IAIA-BL provides the radiologist with the means not to simply trust the AI but to check its output for plausibility, and overrule it when necessary. As far as we know, ours is the first work that applies case-based reasoning using interpretable deep learning techniques to analyzing medical images. Our approach, with its inherently interpretable reasoning process, contrasts directly with previous work that relied on posthoc explanation techniques to explain a trained black-box model, or work that relies on attention mechanisms to highlight the parts of an input image upon which the model prediction is based. The reasoning process the network explains to the humans is the reasoning process it is using to understand the image itself.

As shown partly in Figure 1c, the framework we establish aims to identify not just whether a lesion is malignant or benign, but aims to help with the full reasoning process in the decision of whether to perform a biopsy. Several prediction problems are involved in determining the mass margin and shape of a lesion, which a radiologist would need to consider as part of the required pipeline laid out by the American College of Radiology . In our framework, each of these prediction problems is handled by interpretable machine learning.

Machine learning becomes challenging with smaller datasets, such as those available for mammography. To extract more information from our dataset, we collected a small set of pixel-level (“fine”) annotations from our radiology team which permitted better generalization using a smaller number of images; that is, fine annotation on only 30 images of our limited data set (1136 annotated mammographic images from 484 patients with lesions) enabled high-quality reasoning and prediction. This novel approach can reduce the confounding in deep learning by leveraging both relatively-abundant coarsely-annotated data and a small amount of finely annotated data. Most fine-grained classification algorithms either assume the availability of fine-grained part labels for all input data and use them, or those fine-grained labels are completely ignored. Our approach provides a middle ground by using both data with and without fine-grained annotations, which takes full advantage of the information available. This approach is also practical in the sense that for real-world problems, annotated data are relatively less abundant and more expensive to obtain.

The main contributions of our paper are as follows:

We developed the first inherently interpretable ML-based system for medical imaging that goes beyond simple attention in its explanations. Our system, IAIA-BL, makes predictions for mammographic breast masses by comparing test mammograms with prototypical images of various mass margin types.

We developed a novel training scheme for our IAIA-BL which allows it to incorporate prior knowledge in the form of fine-grained expert image annotations. Using only a small number of finely annotated training data and imposing a novel fine-annotation loss on those data, IAIA-BL learns medically relevant prototypes, effectively addresses aspects of confounding issues in medical machine learning models, and sets our IAIA-BL apart from the ProtoPNet presented in and other prior works.

By changing the logic of ProtoPNet from max-pooling to top-k average pooling, we increase performance dramatically. This improvement can be used in any follow-on works that use ProtoPNet-style architecture.

We developed a framework for machine learning-based mammography interpretation in line with the goals of radiologists: in addition to predicting whether a lesion is malignant or benign, our work aims to follow the reasoning processes of radiologists in detecting specific aspects of each image, such as the characteristics of the mass margins.

Related Work

Background on Computer-Aided Detection/Diagnosis in Mammography. Computer-aided detection systems flag suspicious lesions that may prompt the radiologist to recall a patient for additional imaging. Despite widespread clinical adoption, however, an influential study from 2015 found that current systems do not improve diagnoses in clinical practice . More recent deep learning studies based on large numbers of cases have been reported to match or even exceed radiologist performance . Going beyond lesion detection, computer-aided diagnosis systems provide additional diagnostic information such as to classify the lesion as benign vs. malignant . In this study, we seek to advance the underlying technology beyond that of previous computer-aided diagnosis approaches.

Background on Interpretable ML. In spite of their promising performance, deep neural networks are difficult to understand by humans. There are two distinct approaches to address this challenge: (1) Design inherently interpretable networks, whose reasoning process is constrained to be understandable to humans. (2) Explain black box neural networks posthoc by creating approximations, saliency maps, or derivatives. Posthoc explanations can be problematic; for instance saliency maps highlight regions of the image, but can be unreliable and misleading, as they tend to highlight edges and do not show what computation is actually done with the highlighted pixels . We avoid posthoc solutions in this work. There are several types of approaches in interpretable machine learning, including case-based reasoning (which we use here), forcing the network to use logical conditions within its last layers [55, e.g.,], or disentangling the neural network’s latent space [11, e.g.,]. Case-based reasoning models in medicine retrieve existing similar cases in order to determine how to handle a new case .

IAIA-BL’s framework incorporates the architecture of the inherently interpretable neural network ProtoPNet, described in . While ProtoPNet works well with bird classification, it was not able to be directly extended to mammograms, because of the problems with confounding, which is made worse by the dearth of data and the difficulty of the overall problem. IAIA-BL overcomes these obstacles through its framework, including incorporating fine-grained labels, modified modular training, and the addition of multi-stage reasoning wherein the model first determines the mass margin feature and uses that information to predict malignancy.

Confounding and Fine Annotation. Neural networks models often use context or confounding information instead of the information that a human would use to solve the same problem in both medical and non-medical applications . The ability of these networks to use context or background information is so powerful that networks trained on images of only the background outperform networks trained on images of only the object to be classified . For high-stakes applications in medicine, model decisions must use relevant medical information rather than context or background information. To address this we introduce an attention mechanism which redirects model attention to a selected part of the input image. Ways to direct model attention include data augmentation , techniques that combat hand-selected confounders , techniques that combat learned confounders , and an approach where a human critic manually approves the attention map during training . Other techniques that show model attention but do not aim to change it are class activation maps , multi-attention CNNs , and recurrent attention CNNs . Our model uses an attention mechanism to incorporate expert annotations by adding a term to the objective function which penalizes attention outside of the regions marked as relevant by the radiologist-annotator. Mechanically, the method is most similar to that of , but differs in that our class-specific attention mechanism asks for different attention from prototypes of different classes.

Data and Methods

Models need to be collaborators in the medical decision-making process in order to be useful. In mammography, the initial clinical decision is expressed as a BI-RADS category of 1 to 5, corresponding to the recommendation of whether the patient needs a biopsy . An inscrutable model predicting malignant/benign is not useful as a decision aid, as a biopsy is recommended for every lesion with greater than 2% chance of malignancy (BI-RADS 4 and 5). To alter clinical management, an interpretable model is needed to describe its reasoning process for why the patient should or should not receive a biopsy rather than provide an inscrutable prediction of malignancy.

Our AI approach includes an explicit reasoning system that resembles that of a practicing radiologist. Existing interpretability techniques for mammography include localization as in Figure 1b, but there is no explanation of why an area is selected, what attributes of the region are used for classification, or what parts of the training set these associations are learned from. In a non-medical image analogy, though localization may provide a good interpretation for whether or not an image contains a vase (perhaps by highlighting the vase), it does not provide a good interpretation for classification of the vase pattern as Roman vs. Asian antiquity (highlighting the vase pattern provides no further insight). Many recently published AI-mammography algorithms are still entirely uninterpretable as in Figure 1a .

We train a ProtoPNet-based IAIA-BL model for classifying mass margins of breast lesions – such a model would be able to learn a set of prototypical features associated with each margin type, and predict the margin type of a previously unseen breast lesion based on the similarity of its margin with the learned features. It further uses the logits of the mass margin prediction to predict malignancy. Such a model may be integrated into a clinical support system for classifying breast lesions, because it can point to mammogram regions that resemble prototypical signs of cancerous growth (e.g., spiculated mass margin), and thereby assist doctors in making diagnoses.

2 Fine Annotations

When starting to build an interpretable model for breast lesion classification, we naïvely applied the case-based reasoning ProtoPNet architecture to medical images. Though the model appeared to be learning medically relevant features because of its high validation accuracy, the model made predictions using regions of the image that did not correspond to the medical information; in other words, the model used confounding information rather than medically relevant information. This is consistent with observations made by other groups of the dangers of confounding in medical imaging . For non-medical image classification tasks, a typical approach might be to increase the size of the training set. However, as discussed above, one major barrier to implementation of AI in the medical field is the limited availability of annotated data .

To make our limited institutional data stretch further, we designed a new training paradigm that incorporates additional expert annotation information on a subset of the existing patient cases. A radiologist (FS) annotated the area of a lesion image that indicates the mass margin for that lesion as in Figure 2(a)(b), with the most prominent and defining features marked by circles and the rest of the lesion margin highlighted by simple lines. The model incorporates the radiologist-supplied fine annotations by regularizing the activation of the prototypes over the image. It penalizes a prototype for activating anywhere on an image not of its class, or for activating outside the region of the image marked “relevant” by the radiologist. Figure 2(a)(c) shows an attention map that highlights confounding information and would be heavily penalized. Figure 2(a)(d) shows an attention map that highlights relevant information. By directing the network to the most relevant parts of the image, we set a strong prior on the network for where the useful information is centered in the image. Because these annotations are expensive to obtain, we designed the method to be able to use a small number of these finely annotated cases and a larger number of less expensive coarsely annotated cases. We include a training loss term in the objective as described in Section 3.5. When using this, our performance and explanation both improve (see Appendix E).

3 IAIA-BL Model Architecture

We use a variation on the ProtoPNet architecture from as the underlying architecture of our mass-margin classifier in IAIA-BL. Figure 3 gives an overview of our IAIA-BL model. Given a region of interest $\mathbf{x}$ in a mammogram, our IAIA-BL model first extracts useful features $f(\mathbf{x})$ for mass-margin classification, using a series of convolutional layers $f$ from a VGG-16 network pre-trained on ImageNet. Like , our IAIA-BL model has a prototype layer $g$, which follows the convolutional layers $f$. The prototype layer $g$ contains $m$ prototypes $\mathbf{P}=\{\mathbf{p}_{j}\}_{j=1}^{m}$ learned from the training set. In our experiments, each prototype is a $1\times 1$ patch with the same number (i.e., $512$) of channels as the convolutional feature maps $f(\mathbf{x})$. Since a prototype has the same number of channels but a smaller spatial dimension than the convolutional feature maps, we can interpret the prototype as representing a prototypical activation pattern of its class and we can visualize the prototype as a patch of the training image it appears in. For example, our IAIA-BL model learns prototypical representations of spicules for the spiculated masses, fuzzy borders for indistinct masses, and clearly defined borders for circumscribed masses, and stores these prototypical representations as prototypes in the prototype layer $g$ for later comparison.

Each similarity map $[s_{j,l}]_{l=(1,1)}^{(14,14)}$ between an input image and a prototype $\mathbf{p}_{j}$ is reduced to a single similarity score $s_{j}$, summarizing the degree of similarity between the input image and the learned prototype. Unlike who used max-pooling to reduce each similarity map to a single similarity score, we used top-k average pooling [25, e.g., as in] because we found that our IAIA-BL trained with the relaxed cluster and separation costs outperforms the one trained with the original cluster and separation costs on the task of margin classification. For a given similarity map $[s_{j,l}]_{l=(1,1)}^{(14,14)}$, top-$k$ average pooling of the similarity map finds the $k$ highest similarity scores from the map and computes the average of those $k$ similarity scores. We denote this operation using $s_{j}=\text{AVGPOOL}({\text{topk}}([s_{j,l}];k))$ in our paper. Note that max-pooling is a special case of top-$k$ average pooling, by using $k=1$. The top-$k$ average pooling allows the model to consider similarity between multiple parts of the input image and a mass-margin prototype, so that the similarity score after top-$k$ average pooling can be interpreted as how strong a prototypical feature is present (on average) in the $k$ most activated parts of the input image (instead of in the most activated part of the input image as in max-pooling). In the IAIA-BL, top-$5\%$ average pooling (i.e., $k=\lfloor 5\%(14\times 14)\rfloor=9$) is used to reduce each similarity map $[s_{j,l}]$ to a similarity score $s_{j}$. The similarity scores between an input image and the learned prototypes are illustrated in the “Similarity score” column in Figure 3.

In IAIA-BL, we initially allocated $5$ prototypes for each of the mass-margin types represented in our dataset (circumscribed, indistinct, spiculated). The final IAIA-BL model presented has 4 prototypes for circumscribed mass margin, 3 prototypes for indistinct mass margin, and 4 prototypes for spiculated mass margin. We use $\text{class}(\mathbf{p}_{j})$ to denote the class identity of a prototype.

Our IAIA-BL uses two fully connected layers. The first fully connected layer $h_{1}$ multiplies the vector of similarity scores $[s_{1},...,s_{m}]$ by a weight matrix to produce three output scores $\hat{y}^{\text{circumscribed}}$, $\hat{y}^{\text{indistinct}}$, and $\hat{y}^{\text{spiculated}}$, one for each margin type. These are (afterwards) normalized using a softmax function to generate the probabilities that the mass margin in the input image belongs to each of the three mass-margin types. The second fully connected layer $h_{2}$ then combines the vector of (unnormalized) mass-margin scores $\hat{\mathbf{y}}^{\text{margin}}=[\hat{y}^{\text{circumscribed}},\hat{y}^{\text{indistinct}},\hat{y}^{\text{spiculated}}]$ into a final score of malignancy $\hat{y}^{\text{mal}}$, which is passed into a logistic sigmoid function to produce a probability that the input image has a malignant breast cancer.

This architecture can provide both local interpretability by explaining each prediction in terms of the similarity between a given input image and the learned prototypes, as in Figure 4, and global interpretability in terms of the clustering structure of the latent feature space (where semantically similar convolutional feature patches are clustered around prototypes representing the same semantic concepts). The set of learned prototypes is provided in Appendix D.

4 Data

Our dataset consists of 1136 digital screening mammogram images of masses in the breast from 484 patients at Duke University Health System. Each mass was coarsely annotated as a rectangular region of interest by one of four fellowship-trained breast imaging radiologists who had access to the original reports. For our training data, we cropped the region of interest as well as the surrounding area for context (rationale in Appendix F). The BI-RADS features of mass shape and mass margin were labelled by one fellowship-trained breast imaging radiologist. The ground truth for malignancy of each mass is the result of definitive histopathology diagnosis.

The 1136 masses consisted of the following mass margins: 125 spiculated, 220 indistinct, 41 microlobulated, 579 obscured, and 171 circumscribed. We excluded lesions with microlobulated margins because of the small number of lesions represented. We excluded lesions with obscured margins because this margin class is not a good indicator for classifying a lesion as benign or malignant, but instead usually indicates the need for follow-up imaging. We split each remaining margin class into 73% training, 12% validation, and 15% testing, ensuring no patient overlap between the testing set and other sets. All performances are based on the testing set alone (n=78). Given the small training set, we performed data augmentation such that each training image is randomly flipped, rotated, and undergoes random cropping with a crop size of 80% of the image’s original size. Each class is augmented to have 5000 images for the training set.

5 Model Training

The training of IAIA-BL differs from that of ProtoPNet in three major ways: (1) IAIA-BL was trained with a fine-annotation loss which penalizes prototype activations on medically irrelevant regions for the subset of data with fine annotations. (2) IAIA-BL considers the top $5\%$ of the most activated convolutional patches that are closest to each prototype, instead of only the top most activated patch as in ProtoPNet (using the single max would be equivalent to using the top 0.5%). (3) We include an additional fully connected layer to transform mass margin scores $\hat{\mathbf{y}}^{\text{margin}}$ to malignancy score $y^{\text{mal}}$ whose training is isolated from the rest of the network.

We represent the dataset of $n$ training images $\mathbf{x}_{i}$, with mass-margin labels $y^{\text{margin}}_{i}$ and malignancy labels $y_{i}^{\text{mal}}$, as $D=\{(\mathbf{x}_{i},y^{\text{margin}}_{i},y_{i}^{\text{mal}})\}_{i=1}^{n}$. A small subset $D^{\prime}\subseteq D$ in the training set comes with fine annotations. For a training instance $(\mathbf{x}_{i},y^{\text{margin}}_{i},y_{i}^{\text{mal}})\in D^{\prime}$ that comes with the radiologist’s (fine) annotations of where medically relevant information is in that training image, we define a fine-annotation mask $\textbf{m}_{i}$, such that $\textbf{m}_{i}$ takes the value $0$ at those pixels that are marked as “relevant to mass margin identification,” and takes the value $1$ at other pixels. Each fine-annotation mask $\textbf{m}_{i}$ has the same spatial dimensions (height and width) as the training image $\textbf{x}_{i}$.

The training of IAIA-BL is divided into four stages: (A1) training of the convolutional layers $f$ and the prototype layer $g$; (A2) projection of prototypes; (A3) training of the first fully connected layer $h_{1}$ for predicting mass-margin types; and (B) training of the second fully connected layer $h_{2}$ for predicting malignancy probability. Stages A1, A2, and A3 are repeated until the training loss for predicting mass-margin types converges, then we move to Stage B. By training Stage B after convergence for mass margin classification, we ensure that the mass margin classifier is not biased by the malignancy labels.

Stage A1: In the first training stage, we aim to learn meaningful convolutional features that can be clustered around prototypes that activate on medically relevant parts of a given patch. In particular, we want convolutional features that represent a particular mass-margin type to be clustered in latent space around a prototype of that particular mass-margin type, and to be far away from a prototype of other mass-margin types. As in , we jointly optimize the parameters $\mathbf{\theta}_{f}$ of the convolutional layers $f$, and the prototypes $\mathbf{p}_{1}$, …, $\mathbf{p}_{m}$ in the prototype layer $g$, while keeping the two fully connected layers $h_{1}$ and $h_{2}$ fixed. Differing from , we minimize the following training loss:

The fine-annotation loss (FineLoss) is entirely new to this paper. The purpose of the fine-annotation loss is to penalize prototype activations on medically irrelevant regions of radiologist-annotated training mammograms. The fine-annotation loss is defined by:

where $g_{\mathbf{p}_{j}}(f(\mathbf{x}_{i}))$ computes the similarity map between patches of the convolutional features $f(\mathbf{x}_{i})$ and the $j$-th prototype $\mathbf{p}_{j}$, and $\textrm{Upsample}(g_{\mathbf{p}_{j}}(f(\mathbf{x}_{i})))$ computes bilinear upsampling of the similarity map $g_{\mathbf{p}_{j}}(f(\mathbf{x}_{i}))$ to yield a prototype activation map of the same dimensions (height and width) as the fine-annotation mask.

Since the fine-annotation mask $\mathbf{m}_{i}$ and the prototype activation map (denoted $\textrm{Upsample}(g_{\mathbf{p}_{j}}(f(\mathbf{x}_{i})))$) have the same dimensions, we can compute a Hadamard (component-wise) product between them. For a given training instance $i\in D^{\prime}$ that comes with a fine-annotation mask $\mathbf{m}_{i}$ and has a mass-margin type $y^{\text{margin}}_{i}$, and for a mass-margin prototype $\mathbf{p}_{j}$ with $\text{class}(\mathbf{p}_{j})=y^{\text{margin}}_{i}$, the Hadamard product between $\mathbf{m}_{i}$ and the prototype activation map (denoted $\textrm{Upsample}(g_{\mathbf{p}_{j}}(f(\mathbf{x}_{i})))$) gives a map that shows the prototype activations in the medically irrelevant regions of the training image (because the fine-annotation mask $\mathbf{m}_{i}$ takes the value $1$ at medically irrelevant pixels, and $0$ at medically relevant pixels). Hence, the first sum in the parentheses of the fine-annotation loss tends to reduce the amount of prototype activations in medically irrelevant regions when those prototypes are of the same class as the training image $i\in D^{\prime}$. This, in turn, reinforces the training algorithm to learn prototypes that encode medically relevant mass-margin features for the prototypes’ designated classes. On the other hand, for a given training instance $i\in D^{\prime}$, the second sum in the parentheses of the fine-annotation loss penalizes any amount of prototype activation for a mass-margin prototype $\mathbf{p}_{j}$ with $\text{class}(\mathbf{p}_{j})\neq y^{\text{margin}}_{i}$. This promotes the learning of prototypes that stay away from any features that could appear in classes that are not the prototypes’ designated classes, so that the prototypes of a particular class represent distinguishing features of that class.

To incorporate the training data with fine annotations into model training, we optimize the convolutional layers $f$ and the prototype layer $g$ by minimizing the training objective in Equation (1), using stochastic gradient descent with $75$ training examples with lesion-scale annotation and $10$ training examples with fine annotations. The fine-annotation loss on a lesion-scale annotation penalizes activation outside of the area marked as the lesion, whereas the fine-annotation loss on a finely annotated image penalizes activation outside of the region “relevant to the mass margin class” as marked by the radiologist.

The prototype layer was initialized randomly using the uniform distribution over a unit hypercube (because the convolutional features from the last convolutional layer all lie between $0$ and $1$).

Stage A2: As in , we project the prototypes $\textbf{p}_{j}$ onto the nearest convolutional feature patch from the training set $D$, of the same class as $\textbf{p}_{j}$. See for a detailed description of how a prototype is visualized.

Stage A3: After the previous two training stages, the (medically relevant) convolutional features have been clustered in latent space around mass-margin prototypes that are identical to some (medically relevant) convolutional features from training images, which can be visualized in the original image space. In this stage, we fine-tune the first fully connected layer $h_{1}$ to further increase the accuracy in predicting mass-margin types. In particular, we fix the parameters $\mathbf{\theta}_{f}$ of the convolutional layers $f$ and the prototypes $\mathbf{p}_{1}$, …, $\mathbf{p}_{m}$, and minimize the following training objective with respect to the parameters $\mathbf{\theta}_{h_{1}}$ of the first fully connected layer $h_{1}$:

The first time we enter stage A3, we initialize connections in fully connected layer $h_{1}$ to a value of 1 for prototypes that are positive for that mass margin, -1 otherwise.

Stage B: In this stage, we train the second fully connected layer $h_{2}$ for predicting malignancy probability, using a logistic regression model whose input is the (unnormalized) mass-margin scores produced by the first fully connect layer $h_{1}$, and whose output is the probability of malignancy. To prevent the malignancy information from biasing the mass margin classification, we train the model in a modular style and it is not trained completely end-to-end in any stage, i.e., there is no return to Stage A from Stage B.

Experiments and Results

We use the AUROC (area under receiver operator characteristic curve) for each of the three mass margin classes as the performance metric for both mass-margin prediction and malignancy prediction. An image-weighted average of these AUROCs to measures overall performance. 95% confidence intervals were derived using Delong’s method .

Cohen $\kappa$ shows the agreement between our model’s predictions and the physician-annotator’s labels for the mass margin prediction task. We use Cohen $\kappa$ to compare our model’s agreement to the agreement of physicians with each other from previous studies . 95% confidence intervals were derived using non-parametric bootstrap resampling with 5000 samples each equal to the size of the test set.

2 Interpretability Metric

We designed the interpretability metric activation precision to quantify what proportion of the information used to classify the mass margin comes from the relevant region as marked by the radiologist-annotator. Using the notations defined in Section 3.5, the activation precision for a single prototype $\mathbf{p}_{j}$ on a single image $\mathbf{x}_{i}$ that has mass-margin type $y^{\text{margin}}_{i}$ and comes with a fine-annotation mask $\mathbf{m}_{i}$, is defined as:

where $T_{\tau}$ is a threshold function that returns the top $(1-\tau)\times 100\%$ of the input values as $1$ and the bottom $\tau\times 100\%$ as $0$. Activation precision is only defined where the prototype has the same class identity as the image. The fraction in Equation (6) gives a proportion of highly activated pixels that are medically relevant. We elaborate on this in Appendix G. To evaluate activation precision for GradCAM and GradCAM++ , we calculate as in Equation (6) but replace the prototype activation map $\textrm{Upsample}\left(g_{\mathbf{p}_{j}}(f(\mathbf{x}_{i}))\right)$ with the normalized gradient map for the correct class.

Note that we do not compute the proportion of medically relevant pixels that are highly activated, i.e., the denominator in Equation (6) is not the number of medically relevant pixels (given by $\sum(1-\mathbf{m}_{i})$). This is because we do not require each prototype to detect the entire mass margin (that was annotated by a doctor), but rather, we expect each prototype to detect a differentiating feature that may only be present at parts of a mass margin. Since a prototype may only focus on parts of a margin, intersection over union or measuring the proportion of medically relevant pixels on which the prototype activates highly would not be appropriate metrics.

We can extend the above definition to a dataset $D^{\prime}$, as:

and further to a set of prototypes $\mathbf{P}=\{\mathbf{p}_{j}\}_{j=1}^{m}$:

Activation precision is a measure of interpretability, in the sense that the higher the activation precision, the better a prototype (or a set of prototypes) is at detecting medically relevant features for mass-margin classification. In our experiments, we used $\tau=0.95$ because IAIA-BL uses the top 5% of activated patches in its predictions. 95% confidence intervals were derived using non-parametric bootstrap resampling with 5000 samples each equal to the size of the test set.

Activation precision can be measured both at lesion-scale (i.e., is the activation within the lesion area and not the added context window?) and at fine-scale (i.e., is the activation on the specific part of the margin marked relevant by the radiologist?).

3 Mass Margin Prediction

IAIA-BL. We used ProtoPNet with VGG-16 pre-trained on ImageNet as the base architecture trained for 50 epochs because model training converges between 40 and 50 epochs. The final model is trained on the combination of the training set and validation set, and tested on a test set never before seen in training. See Appendix E for hyperparameters. Our model can be fully trained on one P100 GPU in 50 hours.

Baselines 2a and 2b: VGG-16 with GradCAM and GradCAM++ . We trained a VGG-16 model with two added fully connected layers to account for the larger number of parameters in our model. Pre-trained on ImageNet, it was trained for 250 epochs and the epoch with the highest test accuracy is selected for comparison. There is no native way to incorporate our fine annotation into VGG-16. VGG-16 provides no inherent interpretability or localization. Using the posthoc GradCAM and GradCAM++ techniques we show localization information and calculate activation precision.

Prediction Results: Treating the radiologist annotations as the ground truth, pruned IAIA-BL achieves AUROCs as reported in Table 2 and an accuracy for the overall mass margin classification task of 83% (n=78, 95% CI: 0.75%, 0.92%). (Without pruning, there is a 0.004 increase in AUROC for mass margin prediction.) Figure 5 shows ROC curves for all prediction tasks and all methods.

Though there are many papers on computer vision with applications to mammography, few papers attempt to classify masses by margin type. We found only one study, who report an accuracy of mass margin prediction at 70.6% and include more margin types than we do, but their provided results are not separated into different margin classes so we cannot directly compare. Further reducing comparability, that study used digitized mammography images from the DDSM database. Their model is not publicly available.

Performance of IAIA-BL is better than that of ProtoPNet (Baseline 1), which does not have the stabilization of the gradient provided by the average pooling improvement of IAIA-BL. VGG-16 (Baselines 2a and 2b) performed comparably to IAIA-BL for AUROC. Remember that the baseline models are permitted to use confounding information that IAIA-BL is not encouraged to use, and we will see that when we consider the interpretability results. As we know, it is easy to perform well on training data despite using logic that a radiologist would claim is incorrect .

Another measure of performance we calculated is the Cohen $\kappa$ agreement between IAIA-BL and our human mass margin annotator on the test set. We found “substantial” agreement with a $\kappa$ value of 0.74 (n=78, 95% CI: 0.60, 0.86) , further broken down into circumscribed at 0.76 (95% CI: 0.58, 0.90), indistinct at 0.69 (95% CI: 0.51, 0.84), and spiculated at 0.78 (95% CI: 0.61, 0.93). For this task of characterizing the mass margin in mammography, our performances were higher than the interobserver agreement of radiologists with each other, e.g., 0.61-0.65 in , 0.58 in , and 0.48 in .

Interpretability Results: To measure interpretability, we used the interpretability metric activation precision from Section 4.2, shown in the lower two rows of Table 2. For the unpruned IAIA-BL model (not shown in the table, because it is almost identical to IAIA-BL), the lesion-scale activation precision of the learned prototypes is 0.93 (95% CI: 0.91, 0.96) and the fine-scale activation precision of the learned prototypes is 0.41 (95% CI: 0.39, 0.43).

Both ProtoPNet (Baseline 1) and VGG-16 (Baselines 2a and 2b) show lower activation precision than IAIA-BL. Both use information from image regions entirely outside the region that contains the lesion. The baseline models are not restricted from using confounding information, and thus do so freely. These models should not be used in practice for this reason. A visual comparison of activation maps (defined in Appendix C) is shown in Figure 2(b).

To summarize, IAIA-BL’s predictive performance was as good or better than the analogous black-box model. Its performance in mimicking our annotator was better than the typical interannotator agreement between radiologists. Its interpretability, measured by how well its attention agreed with a radiologist annotator’s hand-drawn attention maps, exceeded that of existing methods and does not resort to post-hoc analysis.

4 Malignancy Prediction

Even though IAIA-BL is constrained to using only the results of the mass margin outputs to predict malignancy (rather than extra information that may be contained within the raw pixels of the image), IAIA-BL predicts mass malignancy with AUROC of 0.84 (n=75, 95% CI: 0.74, 0.94). These results are interpretable in that they only use the mass margin scores to make their predictions. In the equation below, ${y}_{i}^{\text{circumscribed}},{y}_{i}^{\text{indistinct}},{y}_{i}^{\text{spiculated}}$ correspond to raw, unnormalized mass margin scores for circumscribed, indistinct and spiculated margin respectively. The conversion between mass margin scores and malignancy score ${y}_{i}^{\text{mal}}$ is the following concise linear model:

As expected, a high spiculated score results in a high probability of malignancy, while high circumscribed or indistinct margin scores indicate a benign lesion. Each mass margin score is explained as in Figure 4.

We remark that the prediction of whether a mass has 5% or 95% probability of being malignant would not alter the clinical management, since all lesions with >2% probability of malignancy would be recommended to undergo breast biopsy.

There are a variety of malignancy performance values reported in the literature, though not necessarily from the same population as ours, which means the results are not directly comparable. Some studies have reported better performance in predicting malignancy from BI-RADS features . If our dataset were larger, and if we had non-imaging features such as patient age, it could potentially boost performance.

Baseline 3: Radiologist estimate. During data collection, we asked radiologists to estimate the probability that the lesion will be malignant. There are several caveats for this estimate: radiologists do not perform this task in standard practice, instead they only provide a categorical recommendation for biopsy; the annotations were completed as part of a separate study that used consumer-grade monitors without the necessary specifications or calibrations of medical-grade displays. Nonetheless, these estimates represent the radiologist’s “best guess” when given even more information than the model is provided. The radiologists predicted mass malignancy on the test set with AUROC of 0.91 (n=75, 95% CI: 0.85, 0.97). These radiologists are from Duke Hospital, and thus represent an extremely high quality of care for patients. Using this as a reference standard, IAIA-BL is approximately 7% in AUC away from the physicians.

Baseline 4: Unrestricted end-to-end VGG-16. The uninterpretable VGG-16 baseline given the same image data, but not restricted to predicting on only mass margin results, achieves an AUROC of 0.87 (n=75, 95% CI: 0.82, 0.93). Again, it is possible that VGG-16 uses confounding information; e.g., the age of the patient could be inferred from the density of the normal breast tissue and could be a useful predictor of malignancy.

We cannot compare with papers focused on detection because our technique works on diagnosis of an already detected lesion .

Discussion

The high performance of uninterpretable models that appear to be leveraging mainly confounding information is a point of concern when incorporating models into clinical practice. Though a radiologist may not choose to view an explanation for every prediction, interpretable models still provide value over uninterpretable models. Because we know that AI systems fail , we designed a system that can alert a radiologist to faulty reasoning at the time the prediction is made instead of only after the consequences of misprediction have been realized. The global interpretability (namely, the set of prototypes) allows the trained model to be fine-tuned by domain experts through pruning of prototypes that do not correspond to medically relevant features. The explanations provided can also be used for debugging a model and for retrospective analysis of model failures.

Our technique could be expanded with little change to include other BI-RADS features (e.g., mass shape). The technique might be able to be expanded to microcalcification clusters, the other main type of breast lesions, but there are more categories of calcification morphologies and the different types of cluster distributions can translate into lesions with extreme differences in scale which might pose interesting technical challenges. The underlying logic of the technique could be extended to digital breast tomosynthesis by representing a prototype as either a 2-dimensional part of a reconstructed slice image, or as a 3-dimensional portion of a tomosynthesis volume.

Future work with this model might include reader studies in which we measure any improvements in accuracy and radiologists report their trust in our system. Given the increased benefit of other AI assistance to less-experienced readers , it might be valuable to compare the benefit of this system to both sub-specialists and community radiologists who might be called on to do this work only occasionally. This work might help to extend the quality of care that patients receive at Duke (with highly-trained Duke radiologists) to patients that do not have access to this level of care.

The fine annotation techniques we developed to reduce the use of confounding information can be extended to other computer vision applications. The fine annotation technique could also be used on datasets with known confounders to see how effectively it reduces (or reveals) use of the confounders in its classification decisions.

Conclusion

Our work shows that we are able to create interpretable mass margin prediction models with equal or higher performance to their uninterpretable counterparts. Using only a small dataset, we were able to provide an interpretable network that performs comparably with radiologists on mass margin classification and malignancy prediction. The gradient stabilization improvement to the ProtoPNet training can be added into any future use of its codebase.

Author Contributions

Idea and model development: A.J.B., F.S., C.T., C.C., J.L., C.R.. Code and code review: C.T., A.J.B., C.C.. Data collection: Y.R., A.J.B., F.S., J.L.. Data preprocessing: Y.R., C.T., A.J.B..

Acknowledgements

This study was supported in part by MIT Lincoln Laboratory, Duke TRIPODS and the Duke Incubation Fund.

We would like to acknowledge breast radiologists Michael Taylor-Cho MD, Lars Grimm MD, Connie Kim MD, and Sora Yoon MD, who annotated the dataset used in this paper. This study was supported in part by NIH/NCI U01-CA214183 and U2C-CA233254.

References

Appendix A Results Table

Appendix B Sample explanations

Figures 6, 7, 8 and 9 show explanations of mass margin classification automatically generated by IAIA-BL.

Appendix C Comparing Visual Explanations

The IAIA-BL and ProtoPNet class activation visualizations shown in Figure 2(b) are produced by taking the weighted average of the prototype activations maps for every prototype in the correct class. The weight for each prototype is the similarity score between the prototype and the original image. $\textrm{CAV}^{a,b}$, the value of the $a$-th row and $b$-th column of class activation visualization CAV, is defined as:

where $\textrm{Upsample}\left(g_{\mathbf{p}_{j}}(f(\mathbf{x}_{i}))\right)$ is the prototype activation map for prototype $p_{j}$ on image $\mathbf{x}_{i}$. CAV will have the same dimension as $\textrm{Upsample}\left(g_{\mathbf{p}_{j}}(f(\mathbf{x}_{i}))\right)$. CAV is normalized using mix-max normalization so that its values fall between 0 and 1.

Compared to the class activation visualizations produced by baselines with similar predictive performance, the prototype activation maps produced by IAIA-BL are more likely to highlight the lesion and more likely to highlight the relevant part of the mass margin. This is shown quantitatively by the activation precision metric results from Section 4.3.

Appendix D Learned prototypes

Figures 10, 11 and 12 show unpruned sets of prototypes learned by IAIA-BL. Duplicated prototypes 4, 6, 7 and 9 are pruned with a negligible loss in performance (0.001 decrease in AUROC).

Figure 13 shows an unpruned set of prototypes learned by a variant of IAIA-BL that is constrained to select prototypes only from images with fine annotation, dataset $D^{\prime}$. For this variant there is less prototype variety, but higher unpruned AUROC at 0.965 (compared to unpruned IAIA-BL at 0.955).

Appendix E Hyperparameter tuning

For the IAIA-BL model presented in this paper, we use the hyperparameters: fine annotation coefficient 0.001 and top 5% average pooling. These are selected based on hyperparameter tuning experiments trained on the training dataset and tested on the validation set, shown in Figure 14.

Appendix F Context Window

While performing preliminary experiments with the publicly available dataset CBIS-DDSM, there is an option to download only the regions of interest (ROI). In preliminary work, we compared using only the ROI provided, and using the ROI provided but including an additional 100 pixels on each side of the ROI to provide context. Figure 15 shows that for the spiculated vs. other mass margins task, the algorithm was able to achieve a AUROC of 0.80 on images that include context, but only 0.64 for the ROI images that do not. The reason for this may be that the provided boxes are too tight to convey the information needed for mass margin classification. This preliminary result informed our choice to include context around the ROIs provided by our radiologist annotations.

Appendix G Figure for Activation Precision

A visual definition for activation precision is in Figure 16.

Appendix H Expanded Discussion of Fine Annotation

Despite the promise of an interpretable mass margin classifier promised by a ProtoPNet, we found that a direct application of the ProtoPNet model and the training procedure in failed to find medically relevant prototypes of various mass margin types. This was primarily because there was no extra supervision for prototype learning during the training of a ProtoPNet, which means that the ProtoPNet was free to choose any region of any mammogram in the training set as a “prototypical case” of a specific mass margin type, even if the region it had chosen did not contain the mass margin or the breast lesion. For example, a ProtoPNet we trained using the techniques in yielded spiculated mass margin prototypes that contain healthy tissue as in Figure 2(a)(b), and the ProtoPNet was using the similarity with healthy tissue to decide if a lesion has a spiculated mass margin. This is an example of how a machine learning model could use confounding information to make predictions. The use of confounding information is fairly common in complex black-box models such as neural networks, and is especially dangerous in black-box models used for high-stakes decisions because those models do not explain their reasoning processes, and hence it is not easy to know whether they rely on confounding information or not .

In contrast, IAIA-BL addresses the confounding issues in a ProtoPNet trained for mass margin classification, by integrating additional supervision during prototype learning. As discussed, we annotated a small number of the training mammograms, by highlighting the mass margin in each of those training mammograms, and turning the highlighting into a mask where $1$ indicates that the pixel is in the highlighted mass margin and $0$ otherwise. We introduced a fine-annotation loss during training, which encourages IAIA-BL to choose prototypes that are similar to the highlighted mass margins from the annotated mammograms of the training set. The fine-annotation loss we used resembles the gender attention loss used in . However, the gender attention loss is not class/gender specific, in the sense that there are no separate terms for gender attention loss for male and that for female. In contrast, our fine-annotation loss takes into account the class identity of each annotated training mammogram, in the sense that it not only encourages the prototypes belonging to the class of an annotated training mammogram to be similar to only the highlighted mass margin, but also requires the prototypes belonging to other classes to be dissimilar to any part of the mammogram. With our fine-annotation loss imposed on a small subset of annotated training mammograms, IAIA-BL for mass margin classification is not only able to achieve an area under receiver operating characteristic curve (AUROC) of $0.951$, but also able to learn medically relevant prototypes that capture the essential features of various mass margin types.

Let us walk through the example of learning a spiculated prototype, which detects the presence of spicules in the lesion margins. Suppose that we are interested in creating a network that predicts the presence of spiculated margin. To explain why an unknown lesion is identified as spiculated, we create explanations of the form “this looks like that” where the spicules of the unknown lesion look like the spicules of the prototypical lesion. When we train the network, we train for both which lesions become prototype lesions as well as which part of the prototypical lesion is the relevant part. An ideal spiculated lesion prototype would be a lesion with a spiculated margin where the margin itself is the most relevant part of the prototype. To train for this behaviour, we introduce a loss term during training that penalizes any spiculated prototype that has the most relevant part of the prototype being on an area other than the spicules. This constraint has the effect of directing model attention away from areas of the image that do not contain spicules, and encourages the training algorithm to select prototypes from image regions with spicules. This constraint on model attention penalizes the model for using confounding information, that is, information other than the spicules.

Appendix I Expanded IAIA-BL Model Architecture

We use a variation on the ProtoPNet architecture from as the underlying architecture of our mass-margin classifier in IAIA-BL. Figure 3 gives an overview of our IAIA-BL model. Given a region of interest $\mathbf{x}$ in a mammogram, our IAIA-BL model first extracts useful features $f(\mathbf{x})$ for mass-margin classification, using a series of convolutional layers $f$. In our experiments, the convolutional layers $f$ in our IAIA-BL model consist of all convolutional and max-pooling layers, excluding the last max-pooling layer, in a VGG-16 network , followed by two additional $1\times 1$ convolutional layers each with $512$ filters. We removed the last max-pooling layer from the base VGG-16, so that the resulting convolutional feature maps $f(\mathbf{x})$ will have a larger spatial dimension and a finer resolution. Since we resized all input images to $224\times 224$, the spatial dimension of the convolutional feature maps is $14\times 14$ (instead of $7\times 7$ if the last max-pooling in the base VGG-16 were retained). The higher-resolution convolutional feature maps are able to represent finer details, such as variations in mass margin appearance, and are better suited for mass margin classification, because the mass margin is a relatively small part of the input image. Since the last convolutional layer has $512$ filters, there are $512$ convolutional feature maps in $f(\mathbf{x})$. In addition, the last convolutional layer also uses a sigmoid activation function, so that the convolutional features all lie between $0$ and $1$.

where we have $d_{j,l}=\|\mathbf{p}_{j}-f(\mathbf{x})_{l}\|_{2}^{2}$, $l\in\{(1,1),...,(1,14),(2,1),...,(14,14)\}$ indexes the $1\times 1$ patches of the $14\times 14$ convolutional feature maps $f(\mathbf{x})$, and $f(\mathbf{x})_{l}$ is the $l$-th $1\times 1$ patch of the convolutional feature maps $f(\mathbf{x})$. Since each patch of the convolutional feature maps $f(\mathbf{x})$ has a similarity score with each prototype $\mathbf{p}_{j}$, the similarity scores between patches of $f(\mathbf{x})$ and a prototype $\mathbf{p}_{j}$ can be organized spatially into a similarity map, denoted $[s_{j,l}]_{l=(1,1)}^{(14,14)}$, which can then be upsampled to the size of the input image to produce a prototype activation map that identifies which parts of the input image are similar to the learned prototype. This is shown by the “Sim. maps” column of Figure 3.

Each similarity map $[s_{j,l}]_{l=(1,1)}^{(14,14)}$ between an input image and a prototype $\mathbf{p}_{j}$ is reduced to a single similarity score $s_{j}$, summarizing the degree of similarity between the input image and the learned prototype. Unlike who used max-pooling to reduce each similarity map to a single similarity score, we used top-k average pooling. This idea is first found in the CNN literature as a subsection of Kalchbrenner et al. where dynamic top-k average pooling is used for sentiment analysis. For a given similarity map $[s_{j,l}]_{l=(1,1)}^{(14,14)}$, top-$k$ average pooling of the similarity map finds the $k$ highest similarity scores from the map and computes the average of those $k$ similarity scores. We denote this operation using $s_{j}=\text{AVGPOOL}({\text{topk}}([s_{j,l}];k))$ in our paper. Note that max-pooling is a special case of top-$k$ average pooling, by using $k=1$. The top-$k$ average pooling allows the model to consider similarity between multiple parts of the input image and a mass-margin prototype, so that the similarity score after top-$k$ average pooling can be interpreted as how strong a prototypical feature is present (on average) in the $k$ most activated parts of the input image (instead of in the most activated part of the input image as in max-pooling). We use top-$k$ average pooling in IAIA-BL because the margin is distributed between multiple parts of the input image, so it makes semantic sense. In the IAIA-BL, top-$5\%$ average pooling (i.e., $k=\lfloor 5\%(14\times 14)\rfloor=9$) is used to reduce each similarity map $[s_{j,l}]$ to a similarity score $s_{j}$. The similarity scores between an input image and the learned prototypes are illustrated in the “similarity score” column in Figure 3.

In our IAIA-BL, we initially allocated $5$ prototypes for each of the mass-margin types represented in our dataset. We use $\text{class}(\mathbf{p}_{j})$ to denote the class identity of a prototype.

Fully connected layer $h_{1}$ multiplies the vector of similarity scores $[s_{1},...,s_{m}]$ by a weight matrix to produce three output scores $\hat{y}^{\text{circumscribed}}$, $\hat{y}^{\text{indistinct}}$, and $\hat{y}^{\text{spiculated}}$; one for each margin type. These are (afterwards) normalized using a softmax function to generate the probabilities that the mass margin in the input image belongs to each of the three mass-margin types. The second fully connected layer $h_{2}$ then combines the vector of (unnormalized) mass-margin scores $\hat{\mathbf{y}}^{\text{margin}}=[\hat{y}^{\text{circumscribed}},\hat{y}^{\text{indistinct}},\hat{y}^{\text{spiculated}}]$ into a final score of malignancy $\hat{y}^{\text{mal}}$, which can be passed into a logistic sigmoid function to produce a probability that the input image has a malignant breast cancer.

This architecture provides both local interpretability by explaining each prediction in terms of the similarity between a given input image and the learned prototypes, as in Figure 4, and global interpretability in terms of the clustering structure of the latent feature space (where semantically similar convolutional feature patches are clustered around prototypes representing the same semantic concepts). The set of learned prototypes is provided in Appendix D.

Appendix J Expanded IAIA-BL Training

The training of IAIA-BL differs from that of ProtoPNet in three major ways: (1) Our IAIA-BL was trained with a fine-annotation loss which penalizes prototype activations on medically irrelevant regions for the subset of data with fine-scale annotations. For images with only lesion-scale annotation, fine-annotation loss penalizes prototype activations on regions other than the lesion. (2) Our IAIA-BL considers the top $5\%$ of the most activated convolutional patches that are closest to each prototype, instead of only the top most activated patch (using just the top activated patch is equivalent to top 0.5% in our implementations) as in ProtoPNet. (3) We include an additional fully connected layer to go from mass margin scores $\hat{\mathbf{y}}^{\text{margin}}$ to malignancy score $y^{\text{mal}}$ whose training is isolated from the rest of the network. Some of the training regime from ProtoPNet is explained here to provide context for our changes.

We represent the dataset of $n$ training images $\mathbf{x}_{i}$, with mass-margin labels $y^{\text{margin}}_{i}$ and malignancy labels $y_{i}^{\text{mal}}$, as $D=\{(\mathbf{x}_{i},y^{\text{margin}}_{i},y_{i}^{\text{mal}})\}_{i=1}^{n}$. A small subset $D^{\prime}\subseteq D$ in the training set comes with fine-scale annotations. For a training instance $(\mathbf{x}_{i},y^{\text{margin}}_{i},y_{i}^{\text{mal}})\in D^{\prime}$ that comes with doctor’s annotations of where medically relevant information is in that training image, we define a fine-annotation mask $\textbf{m}_{i}$, such that $\textbf{m}_{i}$ takes the value $0$ at those pixels that are marked as “relevant to mass margin identification” by a radiologist, and takes the value $1$ at other pixels. Each fine-annotation mask $\textbf{m}_{i}$ has the same spatial dimensions (height and width) as the training image $\textbf{x}_{i}$.

The learnable parameters of our IAIA-BL includes: (1) the parameters of the convolutional layers $f$, collectively denoted by $\mathbf{\theta}_{f}$, (2) the prototypes $\mathbf{p}_{1}$, …, $\mathbf{p}_{m}$ in the prototype layer $g$, (3) the parameters of the first fully connected layer $h_{1}$, collectively denoted by $\mathbf{\theta}_{h_{1}}$, and (4) the parameters of the second fully connected layer $h_{2}$, collectively denoted by $\mathbf{\theta}_{h_{2}}$.

The training of IAIA-BL is divided into four stages: (A1) training of the convolutional layers $f$ and the prototype layer $g$; (A2) projection of prototypes; (A3) training of the first fully connected layer $h_{1}$ for predicting mass-margin types; and (B) training of the second fully connected layer $h_{2}$ for predicting malignancy probability. Stages A1, A2, and A3 are repeated until the training loss for predicting mass-margin types converges, then we move to stage 4. By training stage B after converging mass margin classification, we ensure that the mass margin classifier is not biased by the malignancy labels.

Stage A1: In the first training stage, we aim to learn meaningful convolutional features that can be clustered around prototypes that activate on medically relevant part of a given mammogram. In particular, we want convolutional features that represent a particular mass-margin type to be clustered around a prototype of that particular mass-margin type, and to be far away from a prototype of other mass-margin types. As in , we jointly optimize the parameters $\mathbf{\theta}_{f}$ of the convolutional layers $f$, and the prototypes $\mathbf{p}_{1}$, …, $\mathbf{p}_{m}$ in the prototype layer $g$, while keeping the two fully connected layers $h_{1}$ and $h_{2}$ fixed. Our training loss differs from that in . In particular, we minimize the following training loss:

The cross-entropy loss in Equation (12) penalizes misclassification of mass-margin types on the training data (using the subnetwork that excludes the second fully connected layer $h_{2}$). It also ensures that the learned convolutional features and the learned prototypes are relevant for predicting mass-margin types. The cluster cost is defined by:

The minimization of the separation cost encourages the average of the $k$ smallest squared distances, between the convolutional patches of a training image and a prototype not of the same class as the training image, to be large. This encourages convolutional feature clusters of different mass-margin types to separate.

The fine-annotation loss (FineLoss) is new to this paper, and has not been previously explored by . The purpose of the fine-annotation loss is to penalize prototype activations on medically irrelevant regions of doctor-annotated training mammograms. The fine-annotation loss is defined by:

where $g_{\mathbf{p}_{j}}(f(\mathbf{x}_{i}))$ computes the similarity map between patches of the convolutional features $f(\mathbf{x}_{i})$ and the $j$-th prototype $\mathbf{p}_{j}$, and $\textrm{Upsample}(g_{\mathbf{p}_{j}}(f(\mathbf{x}_{i})))$ computes bilinear upsampling of the similarity map $g_{\mathbf{p}_{j}}(f(\mathbf{x}_{i}))$ to yield a prototype activation map of the same dimensions (height and width) as the fine-annotation mask.

Since the fine-annotation mask $\mathbf{m}_{i}$ and the prototype activation map (denoted $\textrm{Upsample}(g_{\mathbf{p}_{j}}(f(\mathbf{x}_{i})))$) have the same dimensions, we can compute a Hadamard (component-wise) product between them. For a given training instance $i\in D^{\prime}$ that comes with a fine-annotation mask $\mathbf{m}_{i}$ and has a mass-margin type $y^{\text{margin}}_{i}$, and for a mass-margin prototype $\mathbf{p}_{j}$ with $\text{class}(\mathbf{p}_{j})=y^{\text{margin}}_{i}$, the Hadamard product between $\mathbf{m}_{i}$ and the prototype activation map (denoted $\textrm{Upsample}(g_{\mathbf{p}_{j}}(f(\mathbf{x}_{i})))$) gives a map that shows the prototype activations in the medically irrelevant regions of the training image (because the fine-annotation mask $\mathbf{m}_{i}$ takes the value $1$ at medically irrelevant pixels, and $0$ at medically relevant pixels). Hence, the first sum in the parentheses of the fine-annotation loss encourages the amount of prototype activations in medically irrelevant regions to be small, for those prototypes that are of the same class as the training image $i\in D^{\prime}$. This, in turn, encourages the training algorithm to learn prototypes that encode medically relevant mass-margin features for the prototypes’ designated classes. On the other hand, for a given training instance $i\in D^{\prime}$, the second sum in the parentheses of the fine-annotation loss penalizes any amount of prototype activation for a mass-margin prototype $\mathbf{p}_{j}$ with $\text{class}(\mathbf{p}_{j})\neq y^{\text{margin}}_{i}$. This encourages the training algorithm to learn prototypes that stay away from any features that could appear in classes that are not prototypes’ designated classes, so that the prototypes of a particular class represent distinguishing features of that class.

To incorporate the training data with fine annotations into model training, we optimize the convolutional layers $f$ and the prototype layer $g$ by minimizing the training objective in Equation (12), using stochastic gradient descent with $75$ training examples with lesion-scale annotation and $10$ training examples with fine annotation. The fine-annotation loss on a lesion-scale annotation penalizes activation outside of area marked as the lesion, whereas the fine-annotation loss on a finely annotated image penalizes activation outside of the region “relevant to mass margin classification” as marked by the radiologist.

We initialized IAIA-BL using the parameters of a VGG-16 pre-trained on ImageNet for the base convolutional layers. The additional convolutional layers were initialized randomly using Kaiming normal initialization . The prototype layer was initialized randomly using the uniform distribution over a unit hypercube (because the convolutional features from the last convolutional layer all lie between $0$ and $1$). For the first ten training epochs of IAIA-BL, we warmed up those randomly initialized layers, by only optimizing the training objective in Equation (12) with respect to the parameters in the two additional convolutional layers and the prototypes in the prototype layer.

Stage A2: As in , we project the prototypes $\textbf{p}_{j}$ onto the nearest convolutional feature patch from the training set $D$, of the same class as $\textbf{p}_{j}$:

The projection of each prototype onto a convolutional patch from a training image allows us to visualize a prototype $\mathbf{p}_{j}$, by cropping out the region of that training image corresponding to the highest activations from the prototype activation map obtained from comparing the convolutional features of that training image with the prototype $\mathbf{p}_{j}$ (see for a detailed description of how a prototype can be visualized).

Stage A3: After the previous two training stages, the (medically relevant) convolutional features have been clustered around mass-margin prototypes that are identical to some (medically relevant) convolutional features from training images, which can be visualized in the original image space. In this stage, we fine-tune the first fully connected layer $h_{1}$ to further increase the accuracy in predicting mass-margin types. In particular, we fix the parameters $\mathbf{\theta}_{f}$ of the convolutional layers $f$ and the prototypes $\mathbf{p}_{1}$, …, $\mathbf{p}_{m}$, and minimize the following training objective with respect to the parameters $\mathbf{\theta}_{h_{1}}$ of the first fully connected layer $h_{1}$:

The first time we enter stage A3, we initialize connections in fully connected layer $h_{1}$ to 1 for prototypes that are positive for that mass margin, -1 otherwise.

Stage B: In this stage, we train the second fully connected layer $h_{2}$ for predicting malignancy probability, using a logistic regression model whose input is the (unnormalized) mass-margin scores produced by the first fully connect layer $h_{1}$, and whose output is the probability of malignancy. Training an additional layer $h_{2}$ to predict malignancy allows us to compare our interpretable mass-margin classifier with previous non-interpretable classifiers that only output a malignancy label/probability. To prevent the malignancy information from biasing the mass margin classification, we train the model in a modular style and it is not trained completely end-to-end in any stage.