Vision-Transformer-Based Transfer Learning for Mammogram Classification

Diagnostics 2023 Mammography AI 6 Explanations View Original
Original Paper (PDF)

Unable to display PDF. Download it here or view on PMC.

Plain-English Explanations
Pages 1-2
Limitations of CNN-Based Mammogram Analysis

Mammography is the primary screening tool for breast cancer, enabling detection of tumors too small to be felt. However, radiologists interpreting high volumes of mammograms are prone to error, and traditional computer-aided detection (CAD) systems have not substantially improved diagnostic accuracy, primarily because of high false-positive rates that lead to unnecessary biopsies and patient anxiety.

Convolutional neural networks (CNNs) improved upon traditional CAD systems and have been shown to diagnose approximately 85% of breast cancer cases accurately. However, CNNs have key limitations for mammogram analysis: they process images region by region rather than globally, missing relationships between distant image features; they are sensitive to lesion size; they struggle with class imbalance in datasets; and they are computationally expensive when processing high-resolution mammogram images due to multiple convolution operations.

The common workaround of patch-based analysis - cropping suspected tumor regions before feeding them to a CNN - causes loss of contextual information from the rest of the mammogram and requires precise tumor localization as a prerequisite, introducing additional error. A method that could process the entire mammogram image simultaneously while capturing both local and global image features would address these limitations.

Vision transformers (ViTs), derived from the transformer architecture originally developed for natural language processing, offer a fundamentally different approach. Rather than building up features through hierarchical convolutions, ViTs divide images into patches and use a self-attention mechanism to directly model relationships between all patches simultaneously, enabling global context awareness from the earliest layers of processing.

TL;DR: Traditional CNNs have fundamental limitations for mammogram analysis including regional focus, sensitivity to lesion size, and class imbalance issues - motivating this study's exploration of vision transformers as a more globally aware alternative.
Pages 3-7
Vision Transformer Architecture and Transfer Learning

A vision transformer (ViT) processes an image by dividing it into a grid of small non-overlapping patches (16x16 pixels in this study), flattening each patch into a vector, and embedding these vectors into a lower-dimensional space. Positional embeddings are added to preserve spatial location information since transformers inherently have no notion of order. A special learnable class token is prepended to the sequence and carries the final classification decision.

The embedded patch sequence is then processed by a transformer encoder - a stack of identical layers each containing: (1) multi-head self-attention, which allows every patch to attend to every other patch and learn which regions are most informative for classification; (2) layer normalization to stabilize training; and (3) a multi-layer perceptron for further feature transformation. The final state of the class token is passed to a classification head to produce the output.

Three state-of-the-art ViT architectures were evaluated. The original ViT applies global self-attention across all patches. The Swin transformer improves efficiency by computing self-attention within local non-overlapping windows, then shifting windows between layers to enable cross-window communication - achieving hierarchical representations with lower computational cost. The Pyramid Vision Transformer (PVT) uses spatial-reduction attention that reduces the resolution of keys and values before computing attention, further reducing complexity while maintaining positional encoding throughout all blocks.

Transfer learning was applied by initializing all models with weights pre-trained on the ImageNet dataset (containing over 1 million natural images across 1,000 categories). Only the final classification layer was replaced to accommodate the binary (benign vs. malignant) output. This approach leverages general visual features learned from large datasets - edge detection, texture analysis, shape recognition - and applies them to the specialized domain of mammography, where training data is more limited.

TL;DR: Vision transformers divide mammograms into patches and use self-attention to capture global image relationships simultaneously, and when pre-trained on ImageNet and fine-tuned on mammograms, they dramatically outperform CNNs for benign/malignant classification.
Pages 3, 4, 9
Dataset, Class Balancing, and Experimental Setup

The study used the publicly available Digital Database for Screening Mammography (DDSM), containing 13,128 mammogram images: 5,970 benign and 7,158 malignant. This imbalance (35% benign, 65% malignant) creates a risk of training bias, where models learn to favor the majority class. The dataset was split 80/20 into training and testing sets with 5-fold cross-validation.

To address class imbalance, a novel augmentation-based class balancing strategy was applied. Benign-class images (the minority) underwent five augmentation operations (color jitter, gamma correction, horizontal flip, salt-and-pepper noise, and sharpening), while malignant images underwent only one augmentation. This equalized the number of training images in each fold to 1,146 per class, preventing the model from developing systematic bias toward the more common malignant class.

All images were resized to 224 x 224 pixels as required for patch generation, and models were trained for 50 epochs with a learning rate of 0.0001 using the Adam optimizer with exponential decay and batch size of 64. GELU (Gaussian Error Linear Unit) activation was used for ViT models, while ReLU was used for CNN comparison models. The same hyperparameter settings were applied across all models to ensure fair comparison.

Five experimental comparisons were performed: (1) ViT models compared against each other, (2) transfer learning vs. training from scratch for ViTs, (3) ViT transfer learning vs. CNN transfer learning, (4) computational cost analysis using FLOPS (floating point operations per second), and (5) comparison with previously published transformer-based methods on the same dataset.

TL;DR: The study used the DDSM mammography dataset with augmentation-based class balancing and compared six vision transformer models against six CNN models, all using transfer learning from ImageNet pretraining.
Pages 9-12
Perfect Performance and Computational Efficiency

All six vision-transformer-based transfer learning models (ViT-base, ViT-large, Swin-small, Swin-base, PVT-medium, PVT-large) achieved accuracy, AUC, F1 score, precision, recall, MCC, and kappa scores of 1.0 with zero variance on the DDSM test set - representing perfect classification performance across all metrics. This result substantially exceeds all previously published transformer-based methods on the same dataset.

In direct contrast, vision transformer models trained from scratch (without ImageNet pretraining) achieved AUCs ranging from only 0.72 to 0.77, confirming that transfer learning is essential for achieving strong performance on limited medical imaging datasets. Training from scratch led to overfitting given the relatively small mammography dataset - demonstrating the critical role of pretraining on large general datasets.

CNN-based transfer learning models achieved strong but inferior performance. The best CNN models (ResNet50, ResNet101) reached accuracy and AUC of 0.95, while EfficientNet and InceptionNet variants achieved 0.93-0.95 AUC. The vision transformer models surpassed all CNN benchmarks by a clear margin, supporting the hypothesis that the global self-attention mechanism is more effective than hierarchical convolutions for mammogram analysis.

Among all ViT models, PVT-medium offered the best computational efficiency: with only 44 million parameters and 7G FLOPS, it achieved the same perfect performance in 2,900 seconds of training time - more than 2.5 times faster than ViT-large (309 million parameters, 59G FLOPS, 7,400 seconds). PVT-medium's efficiency makes it the most practical choice for clinical deployment where both accuracy and processing speed matter.

TL;DR: All six vision transformer transfer learning models achieved perfect 1.0 AUC on the DDSM mammography dataset, dramatically outperforming CNNs (best AUC 0.95) and transformer models trained from scratch (best AUC 0.77).
Page 13
Why Vision Transformers Excel at Mammogram Classification

The superior performance of vision transformers over CNNs for mammogram classification is primarily attributed to their global attention mechanism. From the very first layer, ViTs can model relationships between any two patches of a mammogram regardless of their spatial distance. This allows the model to leverage subtle contextual cues across the entire image - for example, comparing tissue patterns in the tumor region to patterns elsewhere in the breast - that CNNs can only capture after building up many convolutional layers.

The critical importance of transfer learning over training from scratch (1.0 vs. 0.72-0.77 AUC) confirms a fundamental principle in medical AI: the features learned from large natural image datasets are broadly useful even for specialized medical tasks. Low-level features like edge and texture detection are relevant to identifying mass boundaries and structural abnormalities in mammograms, providing a vastly better starting point than random initialization.

The augmentation-based class balancing strategy also played an important role. Without addressing the malignant-to-benign imbalance in the DDSM dataset (65:35 ratio), models would tend to favor predicting the majority class. By equalizing class representations through augmentation, the model was exposed equally to both categories during training, preventing systematic diagnostic bias toward over-predicting malignancy.

The comparison with published literature shows the magnitude of improvement: prior transformer-based approaches on the DDSM dataset achieved AUCs ranging from 0.65 to 0.803. The current approach's AUC of 1.0 represents a substantial advance, likely attributable to the combination of larger ViT models, systematic class balancing, and optimized training procedures not employed in earlier work.

TL;DR: Vision transformers outperform CNNs because their self-attention mechanism captures global mammogram context from early processing layers, and transfer learning from ImageNet provides essential feature knowledge that compensates for the limited size of medical imaging datasets.
Page 14
Conclusions and Limitations

This study demonstrates that vision-transformer-based transfer learning achieves perfect classification performance (AUC 1.0) for distinguishing benign from malignant mammographic masses on the DDSM dataset. The approach outperforms CNN-based transfer learning, ViT models trained from scratch, and all previously published transformer methods on the same dataset, establishing a new state of the art for this classification task.

The PVT-medium model provides the optimal balance of accuracy and computational efficiency (7G FLOPS, 2,900 second training time), making it the recommended architecture for clinical deployment scenarios where speed and resource constraints matter alongside performance. This finding provides practical guidance for implementation choices beyond simply comparing accuracy metrics.

An important limitation is that all results were obtained from a single dataset (DDSM) from a single source. Perfect performance on one dataset may reflect characteristics specific to that dataset rather than genuine generalizability across all mammographic presentations. Validation on additional datasets from different imaging equipment, patient demographics, and clinical settings is needed to confirm the approach's real-world utility.

Future work should evaluate the approach on diverse international mammography datasets, investigate the optimal patch size and ViT architecture depth for medical images, and explore the use of attention maps for explainability - generating visualizations that highlight which image regions the model focused on for its classification decision. Such explainability tools would be essential for clinical acceptance, allowing radiologists to verify that the model's attention aligns with clinically meaningful anatomical features.

TL;DR: Vision transformer transfer learning achieves perfect mammogram classification on the DDSM dataset, with PVT-medium offering the best efficiency for clinical deployment, though multi-dataset validation is essential before real-world implementation.
Citation: Open Access, 2023. Available at: PMC9857963.