Source-linked AI summary
End-to-end Training for Whole Image Breast Cancer Diagnosis using An All Convolutional Design
Li Shen
TL;DR
Whole-image mammography diagnosis is difficult because clinically relevant labels depend on small, variable regions and many datasets lack lesion annotations. The paper converts an annotated patch classifier into an all-convolutional whole-image model that can then be trained with image-level labels. It reports competitive performance on DDSM and INbreast, including AUC 0.91 from three-model averaging on DDSM and AUC 0.96 from averaging two augmented models on INbreast.
Problem
Mammography diagnosis must identify small, variable regions from whole images, while many databases lack complete ROI annotations.
Method
The method first trains a localized patch classifier with ROI annotations, converts it to an all-convolutional whole-image classifier, and trains the latter with image-level labels.
Results
The models achieve competitive performance across DDSM and INbreast, including per-image AUC 0.91 from three-model averaging on DDSM and 0.96 from averaging two augmented models on INbreast.
Takeaways & Limitations
The approach reduces reliance on lesion annotations and supports transfer from DDSM to INbreast using only a small amount of training data.
Takeaways & Limitations
Direct comparison with the top-performing challenge team was not allowed when the manuscript was written.
Abstract
from arXiv · showhide
We develop an end-to-end training algorithm for whole-image breast cancer diagnosis based on mammograms. It requires lesion annotations only at the first stage of training. After that, a whole image classifier can be trained using only image level labels. This greatly reduced the reliance on lesion annotations. Our approach is implemented using an all convolutional design that is simple yet provides superior performance in comparison with the previous methods. On DDSM, our best single-model achieves a per-image AUC score of 0.88 and three-model averaging increases the score to 0.91. On INbreast, our best single-model achieves a per-image AUC score of 0.96. Using DDSM as benchmark, our models compare favorably with the current state-of-the-art. We also demonstrate that a whole image model trained on DDSM can be easily transferred to INbreast without using its lesion annotations and using only a small amount of training data. Code availability: https://github.com/lishen/end2end-all-conv
1 Introduction
Mammography diagnosis is difficult because whole-image labels depend on small, highly variable regions, while down-sampling can obscure those regions. The proposed approach first learns localized patch recognition with annotations, then converts the model into an end-to-end whole-image classifier trained without ROI annotations.
- Problem: Whole-image mammography labels are determined by small, localized regions that vary substantially among patients.These properties make breast cancer diagnosis challenging from mammograms.
- Problem: Down-sampling a mammogram to 224x224 for classification can make regions of interest difficult to detect or classify.
- Problem: Many prior methods assume databases are fully annotated with regions of interest, limiting transfer to databases without those annotations.The passage notes that fully annotating all databases is unrealistic because of the time required.
- Contribution: The proposed method first trains localized patch recognition on an annotated database, then converts it into a whole-image classifier trained end-to-end without ROI annotations.The approach uses an all-convolutional design intended to improve on the winning challenge strategy.
2 Methods
The method converts a patch classifier into a whole-image model that scans overlapping regions in one pass, then connects convolutional features to image-level predictions without fully connected layers. Training uses staged unfreezing to preserve pretrained and patch-level features while learning whole-image top layers.
- Patch-to-image conversion: Changing the input from patches to whole images lets the patch classifier generate predictions for all overlapping patches in one forward pass.These predictions form a heatmap of background, benign, or malignant likelihoods.
- All-convolutional top layers: The converted model adds convolutional top layers that reduce feature-map sizes and uses global average pooling for image-level output.The design removes fully connected layers.
- All-convolutional top layers: The method discards the three-dimensional heatmap layer because its small dimension may create an information barrier for subsequent layers.
- All-convolutional top layers: Unlike the comparison strategy, the proposed fully convolutional design preserves spatial information through every layer except the last.The comparison strategy uses a heatmap and fully connected layers to connect with image-level output.
- Training strategy: Patch-network training uses ImageNet-pretrained weights and three stages that progressively unfreeze the network.The stages train the last layer first, then the top layers, and finally all layers.
- Training strategy: Whole-image training changes the input size, proportionally enlarges convolutional feature maps, adds top layers, and trains those layers first under a two-stage strategy.Training the top layers first is intended to avoid unlearning important patch-network features.
- Training strategy: Pixel-wise mean centering and on-the-fly data augmentation are applied to both patch and whole-image training.No other preprocessing is applied.
3 Results
The experiments develop whole-image classifiers from annotated patch networks and evaluate architecture choices, model averaging, and transfer learning on DDSM and INbreast. Results favor the all-convolutional approach, with ensemble AUC reaching 0.91 on DDSM and 0.96 on INbreast.
- Dataset and patch classifiers: 1249 patients and 2584 mammograms from CBIS-DDSM were split by patient, with each view treated as a standalone image.The dataset used an 85-15 stratified train-test split, with an additional 10% of training patients held out.
- Whole-image classifier development: 15 points: the all-convolutional design outperformed the heatmap, max-pooling, fully connected, and shortcut design by as much as 15 points.The comparison design was inspired by the winning method but used a heatmap followed by max pooling and two fully connected layers.
- Whole-image classifier development: 0.80: adding a heatmap to a Resnet50-based whole-image classifier scored lower than the same design without the heatmap.The authors conclude that removing the heatmap facilitates information flow during whole-image training.
- Model averaging: 0.91: averaging three augmented models raised DDSM per-image AUC beyond the individual augmented-model scores.The individual score changes were 0.86→0.88, 0.83→0.86, and 0.85→0.88; augmentation used horizontal and vertical flips.
- Transfer learning: 0.89: the hybrid model reached this validation AUC using only 20 patients or 79 training images, with scores quickly saturating as training size increased.The authors connect this rapid adjustment to reduced whole-image training-set construction burden.
- Transfer learning: 0.96: averaging two augmented models produced the INbreast score, while individual VGG16 and hybrid models reached 0.94 and 0.96.The finetuned models used DDSM-trained whole-image networks and required no INbreast ROI annotations.
Appendix 1: Patch to whole image conversion
The appendix presents conversion strategies that extend a patch-classification network into a whole-image classifier using additional top-level components.
- The primary structure adds convolutional layers on top of a patch classifier to produce whole-image predictions.The example uses two residual blocks with structure [512-512-1024] x 2.
- An alternative strategy adds heatmaps and fully connected layers above the patch network.
- Another alternative places a random forest classifier on top of the heatmap.
Appendix 2: Overfitting of VGG structures
The appendix compares VGG structures through training and validation loss curves, showing that the more complex structure suffers from overfitting.
- Training and validation loss curves are shown for two VGG structures with different complexity.
- The more complex VGG structure suffers from overfitting.