# BayTTA: Uncertainty-aware medical image classification with optimized test-time augmentation using Bayesian model averaging

arXiv: https://arxiv.org/abs/2406.17640 · HTML: https://mlsanity.com/paper/2406.17640/

## Key facts

- Cited by 5 papers in the map (data through 14 Oct 2026).
- Builds on 28 of its 63 references that are in the map.
- First posted on arXiv 25 Jun 2024.

## Builds on (most cited first, 28 of 28)

- [Deep Residual Learning for Image Recognition](https://mlsanity.com/paper/1512.03385/) (2015)
- [Dropout as a Bayesian Approximation: Representing Model Uncertainty in Deep Learning](https://mlsanity.com/paper/1506.02142/) (2015)
- [Going Deeper with Convolutions](https://mlsanity.com/paper/1409.4842/) (2014)
- [Certified Adversarial Robustness via Randomized Smoothing](https://mlsanity.com/paper/1902.02918/) (2019)
- [Aleatoric and Epistemic Uncertainty in Machine Learning: An Introduction to Concepts and Methods](https://mlsanity.com/paper/1910.09457/) (2019)
- [A Review of Uncertainty Quantification in Deep Learning: Techniques, Applications and Challenges](https://mlsanity.com/paper/2011.06225/) (2020)
- [Aleatoric uncertainty estimation with test-time augmentation for medical image segmentation with convolutional neural networks](https://mlsanity.com/paper/1807.07356/) (2018)
- [PixelDefend: Leveraging Generative Models to Understand and Defend against Adversarial Examples](https://mlsanity.com/paper/1710.10766/) (2017)
- [Uncertainty Quantification and Deep Ensembles](https://mlsanity.com/paper/2007.08792/) (2020)
- [Better Aggregation in Test-Time Augmentation](https://mlsanity.com/paper/2011.11156/) (2020)
- [Repulsive Deep Ensembles are Bayesian](https://mlsanity.com/paper/2106.11642/) (2021)
- [Bayesian Deep Ensembles via the Neural Tangent Kernel](https://mlsanity.com/paper/2007.05864/) (2020)
- [Deep Ensembles Work, But Are They Necessary?](https://mlsanity.com/paper/2202.06985/) (2022)
- [Trustworthy clinical AI solutions: a unified review of uncertainty quantification in deep learning models for medical image analysis](https://mlsanity.com/paper/2210.03736/) (2022)
- [Greedy Policy Search: A Simple Baseline for Learnable Test-Time Augmentation](https://mlsanity.com/paper/2002.09103/) (2020)
- [Learning Loss for Test-Time Augmentation](https://mlsanity.com/paper/2010.11422/) (2020)
- [Uncertainty Quantification in Deep MRI Reconstruction](https://mlsanity.com/paper/1901.11228/) (2019)
- [Tractable Function-Space Variational Inference in Bayesian Neural Networks](https://mlsanity.com/paper/2312.17199/) (2023)
- [Automatic Brain Tumor Segmentation using Convolutional Neural Networks with Test-Time Augmentation](https://mlsanity.com/paper/1810.07884/) (2018)
- [Doubly Stochastic Variational Inference for Neural Processes with Hierarchical Latent Variables](https://mlsanity.com/paper/2008.09469/) (2020)
- [Dangers of Bayesian Model Averaging under Covariate Shift](https://mlsanity.com/paper/2106.11905/) (2021)
- [Tensor Networks for Medical Image Classification](https://mlsanity.com/paper/2004.10076/) (2020)
- [Soft-Attention Improves Skin Cancer Classification Performance](https://mlsanity.com/paper/2105.03358/) (2021)
- [Classification Confidence Estimation with Test-Time Data-Augmentation](https://mlsanity.com/paper/2006.16705/) (2020)
- [Uncertainty Quantification by Ensemble Learning for Computational Optical Form Measurements](https://mlsanity.com/paper/2103.01259/) (2021)
- [TAAL: Test-time Augmentation for Active Learning in Medical Image Segmentation](https://mlsanity.com/paper/2301.06624/) (2023)
- [UncertaintyFuseNet: Robust Uncertainty-aware Hierarchical Feature Fusion Model with Ensemble Monte Carlo Dropout for COVID-19 Detection](https://mlsanity.com/paper/2105.08590/) (2021)
- [Bayesian Model-Averaged Meta-Analysis in Medicine](https://mlsanity.com/paper/2110.01076/) (2021)

## Built on by (most cited first, 5 of 5)

- [Few-Shot Multimodal Medical Imaging: A Theoretical Framework](https://mlsanity.com/paper/2511.01140/) (2025)
- [Boundary-Aware Test-Time Adaptation for Zero-Shot Medical Image Segmentation](https://mlsanity.com/paper/2512.04520/) (2025)
- [A Comprehensive Inference-Time Augmentation Framework in Physiological Signals: Application to PPG-Based AF Detection](https://mlsanity.com/paper/2606.10410/) (2026)
- [Confidence is Not Reliability: Rethinking MC Dropout in Brain Tumour Segmentation](https://mlsanity.com/paper/2606.19300/) (2026)
- [Improving Calibration of Black-Box Radiology AI Using Test-Time Augmentation](https://mlsanity.com/paper/2609.29931/) (2026)

## Abstract (arXiv)

Test-time augmentation (TTA) is a well-known technique employed during the testing phase of computer vision tasks. It involves aggregating multiple augmented versions of input data. Combining predictions using a simple average formulation is a common and straightforward approach after performing TTA. This paper introduces a novel framework for optimizing TTA, called BayTTA (Bayesian-based TTA), which is based on Bayesian Model Averaging (BMA). First, we generate a prediction list associated with different variations of the input data created through TTA. Then, we use BMA to combine predictions weighted by the respective posterior probabilities. Such an approach allows one to take into account model uncertainty, and thus to enhance the predictive performance of the related machine learning or deep learning model. We evaluate the performance of BayTTA on various public data, including three medical image datasets comprising skin cancer, breast cancer, and chest X-ray images and two well-known gene editing datasets, CRISPOR and GUIDE-seq. Our experimental results indicate that BayTTA can be effectively integrated into state-of-the-art deep learning models used in medical image analysis as well as into some popular pre-trained CNN models such as VGG-16, MobileNetV2, DenseNet201, ResNet152V2, and InceptionRes-NetV2, leading to the enhancement in their accuracy and robustness performance. The source code of the proposed BayTTA method is freely available at: underline https://github.com/Z-Sherkat/BayTTA.