Multi-Head Attention Multiple Instance Learning Deep Neural Classifier Enhanced with Model Uncertainty Quantification
Jakub Buler, Rafał Buler, Krystian Brzozowski, Michał Grochowski · 2025
Artificial Intelligence (AI) has demonstrated notable success in many fields, including high-stakes applications like medical diagnosis support systems. These should not only provide accurate predictions but also offer insights into the decision-making process. Additionally, AI models may encounter challenging inputs, making it essential to incorporate uncertainty quantification, especially in medical applications, thus analyzing the model's robustness. This study proposes an approach that expands a gated-attention pooling mechanism in multiple instance learning framework, within a multi-headed version, to enable multi-class attention visualization. The model is integrated with Monte Carlo dropout (MCDO), which is a convenient Bayesian Neural Network approximation, for uncertainty quantification. The designed system is capable of not only classifying mammograms but also signaling its decision confidence and providing attentional heat maps along with image regions that are sensitive or uncertain in the diagnosis process. The impact of applying MCDO during the training (in the validation phase) and inference was investigated. Two models were trained, one with applying MCDO during the validation phase and the other without. Both models were subsequently tested with and without MCDO during the inference. To evaluate the proposed methodology in a real-world context, is has been applied to the task of breast cancer classification using mammography images. This cancer remains the most prevalent cancer in women worldwide, and the analysis of mammograms is time-consuming for radiologists due to the high resolution and image complexity, which makes it a representative and clinically significant use case. The results showed that the model tested with MCDO only during inference achieved the highest$\mathbf{5}$-fold cross-validation mean accuracy$(77.32 \%)$and lower variability (standard deviation 2.73 % vs. 3.16 %). The implementation of MCDO only in attention heads did not result in significant computational time overhead. These results suggest that applying MCDO during the inference can improve model consistency and enables uncertainty quantification.