This repository contains the code for the paper "Probabilistic smooth attention for deep Multiple Instance Learning in medical imaging". The original code will be made available soon. In the meantime, check the torchmil implementation for ProbSmoothABMIL and TransformerProbSmoothABMIL.
The Multiple Instance Learning (MIL) paradigm is attracting plenty of attention in medical imaging classification, where labeled data is scarce. MIL methods cast medical images as bags of instances (e.g. patches in whole slide images, or slices in CT scans), and only bag labels are required for training. Deep MIL approaches have obtained promising results by aggregating instance-level representations via an attention mechanism to compute the bag-level prediction. These methods typically capture both local interactions among adjacent instances and global, long-range dependencies through various mechanisms. However, they treat attention values deterministically, potentially overlooking uncertainty in the contribution of individual instances. In this work we propose a novel probabilistic framework that estimates a probability distribution over the attention values, and accounts for both global and local interactions. In a comprehensive evaluation involving eleven state-of-the-art baselines and three medical datasets, we show that our approach achieves top predictive performance in different metrics. Moreover, the probabilistic treatment of the attention provides uncertainty maps that are interpretable in terms of illness localization.
If you find our work useful, please consider citing our paper:
@article{castro2025probabilistic,
title={Probabilistic smooth attention for deep multiple instance learning in medical imaging},
author={Castro-Mac{\'\i}as, Francisco M and Morales-{\'A}lvarez, Pablo and Wu, Yunan and Molina, Rafael and Katsaggelos, Aggelos K},
journal={Pattern Recognition},
pages={112097},
year={2025},
publisher={Elsevier}
}