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

Alina Jade Barnett; Fides Regina Schwartz; Chaofan Tao; Chaofan Chen; Yinhao Ren; Joseph Y. Lo; Cynthia Rudin

Source-linked AI summary

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

Alina Jade Barnett, Fides Regina Schwartz, Chaofan Tao, Chaofan Chen, Yinhao Ren, Joseph Y. Lo, Cynthia Rudin

arXiv:2103.12308v1Mar 23, 2021cs.LGcs.AIcs.CV

TL;DR

Mammography involves high-stakes biopsy decisions, small datasets, confounding information, and difficult lesion assessment. IAIA-BL addresses these challenges with interpretable case-based reasoning that combines coarse labels with limited pixel-wise annotations. The resulting models provide clinically relevant explanations while maintaining mass-margin performance comparable to uninterpretable counterparts.

Problem
Mammography requires difficult biopsy-related judgments from small datasets where models may rely on confounding information instead of medically relevant features.
Method
IAIA-BL uses interpretable case-based prototypes, radiologist-relevant mass features, and a training scheme combining whole-image labels with fine pixel-wise annotations.
Results
IAIA-BL provides interpretable mass-margin models with equal or higher performance than uninterpretable counterparts and predicts malignancy with AUROC 0.84 (n=75, 95% CI: 0.74, 0.94).
Takeaways & Limitations
The models are decision aids that let radiologists inspect and verify reasoning rather than simply trust an AI prediction.
Takeaways & Limitations
Malignancy results are not directly comparable across studies, and larger datasets or non-imaging features could potentially improve performance.

Abstract

from arXiv · show

Interpretability in machine learning models is important in high-stakes decisions, such as whether to order a biopsy based on a mammographic exam. Mammography poses important challenges that are not present in other computer vision tasks: datasets are small, confounding information is present, and it can be difficult even for a radiologist to decide between watchful waiting and biopsy based on a mammogram alone. In this work, we present a framework for interpretable machine learning-based mammography. In addition to predicting whether a lesion is malignant or benign, our work aims to follow the reasoning processes of radiologists in detecting clinically relevant semantic features of each image, such as the characteristics of the mass margins. The framework includes a novel interpretable neural network algorithm that uses case-based reasoning for mammography. Our algorithm can incorporate a combination of data with whole image labelling and data with pixel-wise annotations, leading to better accuracy and interpretability even with a small number of images. Our interpretable models are able to highlight the classification-relevant parts of the image, whereas other methods highlight healthy tissue and confounding information. Our models are decision aids, rather than decision makers, aimed at better overall human-machine collaboration. We do not observe a loss in mass margin classification accuracy over a black box neural network trained on the same data.

1 Introduction

IAIA-BL targets mammography’s difficult biopsy-related decisions by combining interpretable lesion classification with radiologist-relevant mass-feature reasoning. It uses fine annotations, case-based prototypes, and top-k pooling to address confounding and limited data while preserving performance.

Mammography requires decisions beyond lesion detection, including whether a patient should undergo biopsy despite subtle findings and frequent benign biopsy results.
Confounding can make apparently accurate models rely on equipment or other incorrect information rather than medically relevant image features.
IAIA-BL explains predictions by localizing image regions, associating them with medical features, and using explained evidence for classification.
Fine annotations from 30 images supplement 1,136 coarsely annotated mammographic images from 484 patients, enabling reasoning and prediction with limited data.
The model compares mammograms with prototypical mass-margin images and incorporates expert annotations through a novel training scheme.
Top-k average pooling replaces max-pooling in ProtoPNet-style logic and is reported to increase performance dramatically.

2 Related Work

Related work distinguishes inherently interpretable networks from posthoc explanations and situates IAIA-BL as a case-based approach adapted to mammography’s confounding and data limitations.

Interpretable networks constrain their reasoning to remain human-understandable, whereas posthoc methods explain black boxes using approximations, saliency maps, or derivatives.
Table 1 compares selected AI mammography techniques with IAIA-BL, while prior mass-margin results were limited and not directly comparable.
Saliency-based posthoc explanations can highlight edges without revealing what computation uses the highlighted pixels, motivating IAIA-BL’s inherent interpretability.
ProtoPNet provides the architectural basis, but its direct extension to mammograms was hindered by confounding, scarce data, and task difficulty.
Medical models must use relevant medical information rather than context or background information because confounding can produce misleadingly strong performance.

3 Data and Methods

IAIA-BL combines case-based prototype reasoning, radiologist fine annotations, and staged prediction to classify mass margins and malignancy while exposing medically relevant evidence. Its architecture compares image patches with learned margin prototypes and produces interpretable similarity maps and predictions.

Fine Annotations: Fine annotations regularize prototype activations toward radiologist-identified lesion-margin regions and away from medically irrelevant image regions.The training scheme uses a small set of finely annotated cases alongside more numerous coarsely annotated cases.
IAIA-BL Model Architecture: IAIA-BL extends case-based reasoning to mammography by comparing lesions with learned prototypes representing mass-margin features.Examples include spicules, fuzzy borders, and clearly defined borders associated with different margin types.
IAIA-BL Model Architecture: The prototype network compares convolutional feature patches with each prototype, converts distances into similarity maps and scores, and feeds those scores into margin and malignancy prediction layers.Similarity maps can be upsampled to the input-image size to identify regions resembling learned prototypes.
IAIA-BL Model Architecture: Top-k average pooling aggregates multiple highly activated patches, allowing similarity scores to reflect the average presence of a prototypical feature across several image regions.The method uses the top 5% of activated patches rather than only the single most activated patch used by ProtoPNet.
Model Training: IAIA-BL differs from ProtoPNet through fine-annotation loss, top-5% pooling, and a separate fully connected malignancy layer trained after mass-margin prediction.These design choices target confounding and preserve the modular relationship between margin and malignancy predictions.

4 Experiments and Results

IAIA-BL achieved competitive prediction performance while providing interpretable, medically focused explanations for mass-margin and malignancy classification. Its interpretability exceeded comparison methods and its agreement with the annotator surpassed reported radiologist interobserver agreement.

Interpretability Metric: Activation precision measures the proportion of highly activated pixels that are medically relevant, with lesion-scale and fine-scale variants.The experiments used τ=0.95, corresponding to the top 5% of activated patches.
Prediction Results: 83% accuracy was achieved for overall mass-margin classification with pruned IAIA-BL (n=78; 95% CI: 0.75%, 0.92%).Without pruning, mass-margin AUROC increased by 0.004.
Prediction Results: IAIA-BL achieved substantial agreement with the human annotator for mass-margin classification, with Cohen κ=0.74 (n=78; 95% CI: 0.60, 0.86).Class-specific κ values were 0.76 for circumscribed, 0.69 for indistinct, and 0.78 for spiculated margins.
Interpretability Results: IAIA-BL’s activation maps were more likely than baseline maps to highlight the lesion and the relevant mass-margin region.The comparison included ProtoPNet and VGG-16 with GradCAM or GradCAM++, which could use regions outside the lesion and confounding information.
Prediction Results: IAIA-BL’s mass-margin predictive performance was as good as or better than the analogous black-box model and exceeded the reported interobserver agreement of radiologists.Reported radiologist interobserver agreement values ranged from 0.48 to 0.65 in cited studies.
Malignancy Prediction: IAIA-BL predicted malignancy with AUROC 0.84 (n=75; 95% CI: 0.74, 0.94) using only mass-margin outputs.A high spiculated score increased predicted malignancy, whereas high circumscribed or indistinct scores indicated a benign lesion.

5 Discussion

The discussion emphasizes that interpretable models can expose faulty reasoning and support model debugging, while outlining extensions and needed reader studies.

Interpretable models can alert radiologists to faulty reasoning when predictions are made.Their explanations can also support debugging and retrospective analysis of model failures.
Prototype-based global interpretability allows domain experts to prune prototypes that lack medically relevant features.
The technique could expand to other BI-RADS features, microcalcification clusters, and digital breast tomosynthesis.Microcalcifications introduce more morphology categories and scale variation, while tomosynthesis could use 2D slice or 3D volume prototypes.
Future reader studies should measure accuracy improvements and radiologists’ trust in the system.Comparisons could include subspecialists and community radiologists with less frequent experience in this work.
The fine-annotation techniques could extend to other computer vision applications and datasets with known confounders.Such applications could test whether the technique reduces or reveals confounder use in classification decisions.

6 Conclusion

The work reports interpretable mass-margin models that match or exceed uninterpretable counterparts, while remaining comparable with radiologists on mass-margin and malignancy tasks.

Equal or higher performance was achieved by interpretable mass-margin prediction models compared with uninterpretable counterparts.
Using a small dataset, the interpretable network performed comparably with radiologists on mass-margin classification and malignancy prediction.
Gradient stabilization improvements to ProtoPNet training can be added to future uses of its codebase.

7 Author Contributions

The author-contribution statement assigns responsibility across idea and model development, code, data collection, and preprocessing.

Idea and model development were contributed by A.J.B., F.S., C.T., C.C., J.L., and C.R.
Code and code review were contributed by C.T., A.J.B., and C.C.
Data collection involved Y.R., A.J.B., F.S., and J.L., while data preprocessing involved Y.R., C.T., and A.J.B.

A Results Table

Table 3 compares IAIA-BL, radiologist estimates, and VGG-16 across a stated classification task.

The table’s task is the classification task.
IAIA-BL denotes the model described in Section 4.4.
Radiologist Estimate and VGG-16 are comparison entries described in Section 4.4.

B Sample explanations

IAIA-BL explains mass-margin predictions by showing which image regions resemble learned prototypes. Examples include correct classifications for circumscribed, indistinct, and spiculated lesions, plus a failure where spiculated evidence is missed.

IAIA-BL automatically generates mass-margin classification explanations using prototype activations.
Correctly classified circumscribed lesions can activate multiple prototypes associated with different image regions.
For an indistinct lesion, the indistinct margin activates an indistinct prototype while a circumscribed margin portion activates circumscribed prototypes.
The model correctly identifies the illustrated spiculated lesion as spiculated.
A misclassified spiculated lesion is labeled circumscribed because the explanation highlights only the circumscribed margin portion and misses the spiculated portion.
Class activation visualizations are weighted averages of correct-class prototype activation maps, with weights given by prototype-image similarity scores.

D Learned prototypes

The learned prototype sets represent mass-margin classes and can be pruned with negligible performance loss. Restricting prototype selection to finely annotated images increases AUROC but reduces prototype variety.

Figures 10–12 show unpruned prototypes learned for circumscribed, indistinct, and spiculated lesions.
Duplicated prototypes 4, 6, 7, and 9 are pruned with a 0.001 decrease in AUROC.
The fine-annotation-only variant has less prototype variety but reaches unpruned AUROC 0.965 versus 0.955 for unpruned IAIA-BL.
The fine-annotation-only prototype set prunes duplicated prototypes and prototype 6 based on class-connection criteria.
The reported IAIA-BL configuration uses fine annotation coefficient 0.001 and top 5% average pooling, selected through training-validation hyperparameter tuning.

F Context Window

Context around the lesion region substantially affects the spiculated-margin experiment. The discussion also motivates fine annotation as supervision for avoiding prototype reliance on healthy tissue and other confounders.

Including 100 pixels of context around each ROI yields AUROC 0.80, compared with 0.64 using ROI images alone for spiculated versus other margins.
The authors attribute the context difference possibly to provided boxes being too tight to convey information needed for mass-margin classification.
Activation precision measures whether highly activated regions overlap medically relevant fine annotations.
Direct ProtoPNet application failed to learn medically relevant mass-margin prototypes and produced spiculated prototypes containing healthy tissue.
Healthy-tissue similarity could drive ProtoPNet’s spiculated-margin decision, illustrating a confounding failure in high-stakes classification.
IAIA-BL uses class-specific fine-annotation loss to encourage prototypes toward highlighted mass margins and away from irrelevant regions.

I Expanded IAIA-BL Model Architecture

IAIA-BL combines convolutional features, prototype comparisons, top-k pooling, and fully connected prediction layers. This design supports local explanations through prototype similarity and global interpretability through latent-space prototype structure.

The model extracts convolutional features from an input ROI before comparing feature patches with learned prototypes.
Each prototype is a 1 × 1, 512-channel activation pattern representing a prototypical mass-margin feature such as spicules, fuzzy borders, or clearly defined borders.
Prototype comparisons compute squared ℓ2 distances between each prototype and all 1 × 1 patches, then transform distances into similarity scores.
Similarity maps are organized spatially and upsampled to produce prototype activation maps identifying image regions similar to learned prototypes.
Top-k average pooling averages the k highest similarity scores, allowing similarity across multiple image parts; max-pooling is the k = 1 special case.
The first fully connected layer predicts three mass-margin scores, while the second combines mass-margin scores for downstream prediction.
The architecture provides local interpretability through prototype similarity and global interpretability through clustering of semantically similar feature patches.

J Expanded IAIA-BL Training

IAIA-BL uses staged case-based training to learn mass-margin prototypes from mammograms, then uses those margin scores to predict malignancy. Fine annotations constrain prototype activations to medically relevant regions, while projected prototypes remain visualizable in the original images.

Four-stage training: IAIA-BL training proceeds through convolutional and prototype learning, prototype projection, mass-margin classification, and malignancy prediction stages.Stages A1–A3 are repeated until the training loss for mass-margin classification converges before stage B trains malignancy prediction.
Four-stage training: The model learns convolutional features and prototypes that cluster same-type mass-margin features while separating prototypes associated with different mass-margin types.Cross-entropy supports mass-margin prediction, while cluster and separation costs encourage within-class proximity and between-class separation.
Fine-annotation loss: FineLoss penalizes prototype activations in medically irrelevant regions of doctor-annotated mammograms, extending prototype-based training beyond lesion-scale annotations.The mask marks radiologist-identified relevant pixels as 0 and other pixels as 1; the loss uses the upsampled prototype activation map and the mask.
Fine-annotation loss: IAIA-BL combines lesion-scale and fine annotations during optimization, using 75 lesion-scale and 10 fine-annotation training examples.For lesion-scale data, the fine-annotation loss penalizes activation outside the marked lesion area; for fine annotations, it targets medically irrelevant regions.
Prototype projection and prediction: After learning, each prototype is projected onto the nearest same-class training feature patch, enabling visualization by cropping the region with the highest prototype activation.Stage A3 then fine-tunes the first fully connected layer for mass-margin prediction, while a later fully connected layer maps margin scores to malignancy scores.