Vision Transformers for Breast Cancer Histology Image Classification

Giulia Lucrezia Baroni, Laura Rasotto, Kevin Roitero, Ameer Hamza Siraj, Vincenzo Della Mea

Abstract: We propose a self-attention Vision Transformer (ViT) model tailored for breast cancer histology image classification. The proposed architecture uses a stack of transformer layers, with each layer consisting of a multi-head self-attention mechanism and a position-wise feed-forward network, and it is trained with different strategies and configurations, including pretraining, resize dimension, data augmentation, patch overlap, and patch size, to investigate their impact on performance on the histology image classification task. Experimental results show that pretraining on ImageNet and using geometric and color data augmentation techniques significantly improve the model’s accuracy on the task. Additionally, a patch size of 16 \(\times \) 16 and no patch overlap were found to be optimal for this task. These findings provide valuable insights for the design of future ViT-based models for similar image classification tasks.