End-to-end Training for Whole Image Breast Cancer Diagnosis using An All Convolutional Design
The paper "End-to-end Training for Whole Image Breast Cancer Diagnosis using An All Convolutional Design" by Li Shen presents a significant advancement in the application of deep learning methods for mammographic breast cancer diagnosis. The approach is characterized by its end-to-end training methodology, simplified network design, and the alleviation of reliance on region of interest (ROI) annotations.
Key Contributions
The paper introduces an all convolutional design framework that leverages convolutional neural networks (CNNs) for whole-image breast cancer diagnosis, focusing primarily on mammograms. The principal contribution lies in significantly reducing the dependency on lesion annotations, which are typically available during only the initial stage of training. The subsequent stages of the training pipeline operate with image-level labels alone, addressing a crucial challenge in medical image analysis where fully annotated datasets are scarce. The approach contrasts with many extant methods that necessitate fully annotated datasets, restricting their applicability across different datasets.
Methods
The proposed method operates by first training a patch-level classifier and subsequently converting it to a whole-image classifier. This conversion involves integrating additional convolutional layers to allow the model to process whole images, generating a "heatmap" that indicates the likelihood of each patch being benign or malignant, thereby facilitating whole-image predictions.
Patch to Image Conversion:
- The conversion process capitalizes on the recursive nature of CNNs. By adding new layers to a pre-trained patch-classification network, the patch classifier effectively becomes part of a broader network that processes the entire image.
- The design eschews fully connected (FC) layers in favor of maintaining all convolutional layers, thereby preserving spatial information crucial for accurate diagnosis and potentially enhancing the model's generalization capability.
Training Strategy:
- The training is divided into two major parts. Initially, a patch classification network is trained using a staged process, starting by fine-tuning the last layer and progressively unfreezing and training more layers while decreasing the learning rate.
- The patch classifier is then converted to a whole-image classifier, with training similarly proceeding in stages to focus on learning new layers initially without disrupting the feature representations learned by the patch network.
Results
The performance of the proposed method is evaluated on two mammographic datasets: the Digital Database for Screening Mammography (DDSM) and INbreast. Strong numerical results underscore the efficacy of the approach:
- DDSM Results:
- The best single-model achieves a per-image AUC score of 0.88.
- Averaging three models enhances the AUC score to 0.91.
- INbreast Transfer Learning:
- Transfer learning from DDSM to INbreast yields a best single-model AUC score of 0.96, signifying the effective transferability of the model without requiring new lesion annotations.
Notably, the all-convolutional design outperforms architectures incorporating heatmaps and fully connected layers, which align with previous winning strategies in the DM challenge.
Transfer Learning:
- The paper explores the efficiency of transfer learning by fine-tuning models on INbreast with varying amounts of training data. Remarkably, satisfactory performance is achieved with as few as 20 patients, demonstrating that the models can adapt quickly with limited new data.
- It is observed that VGG16-based models, particularly when extended with residual blocks, display superior performance compared to Resnet50-based models during transfer learning.
Implications and Future Directions
The implications of this research are multifaceted:
- Practical Impact: The ability to train effective models with reduced annotation requirements holds substantial promise for real-world deployment since curating fully annotated medical image datasets is labor-intensive and costly.
- Improved Generalization: The success in transferring learned models from one dataset to another, with minimal additional data, indicates improved generalizability across different imaging sources, enhancing the potential utility of these models in diverse screening environments.
- Model Design: Results suggest a reconsideration of model architecture for medical imaging tasks, favoring designs that preserve spatial information through convolutional layers over traditional fully connected architectures.
Future research can robustly explore optimizing the architectural elements to further improve transfer learning efficacy, tailoring data augmentation techniques to enhance model robustness, and expanding to additional imaging modalities to validate the approach's adaptability in broader diagnostic contexts.