# Test-Time Augmentation Meets Variational Bayes

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

## Key facts

- Cited by 3 papers in the map (data through 14 Oct 2026).
- Builds on 40 of its 84 references that are in the map.
- First posted on arXiv 19 Sep 2024.

## Builds on (most cited first, 30 of 40)

- [Deep Residual Learning for Image Recognition](https://mlsanity.com/paper/1512.03385/) (2015)
- [PyTorch: An Imperative Style, High-Performance Deep Learning Library](https://mlsanity.com/paper/1912.01703/) (2019)
- [Scaling Laws for Neural Language Models](https://mlsanity.com/paper/2001.08361/) (2020)
- [mixup: Beyond Empirical Risk Minimization](https://mlsanity.com/paper/1710.09412/) (2017)
- [CutMix: Regularization Strategy to Train Strong Classifiers with Localizable Features](https://mlsanity.com/paper/1905.04899/) (2019)
- [The Many Faces of Robustness: A Critical Analysis of Out-of-Distribution Generalization](https://mlsanity.com/paper/2006.16241/) (2020)
- [Making Deep Neural Networks Robust to Label Noise: a Loss Correction Approach](https://mlsanity.com/paper/1609.03683/) (2016)
- [Simple Copy-Paste is a Strong Data Augmentation Method for Instance Segmentation](https://mlsanity.com/paper/2012.07177/) (2020)
- [Generalizing to Unseen Domains via Adversarial Data Augmentation](https://mlsanity.com/paper/1805.12018/) (2018)
- [Age Progression/Regression by Conditional Adversarial Autoencoder](https://mlsanity.com/paper/1702.08423/) (2017)
- [MEMO: Test Time Robustness via Adaptation and Augmentation](https://mlsanity.com/paper/2110.09506/) (2021)
- [The Effectiveness of Data Augmentation in Image Classification using Deep Learning](https://mlsanity.com/paper/1712.04621/) (2017)
- [Aleatoric uncertainty estimation with test-time augmentation for medical image segmentation with convolutional neural networks](https://mlsanity.com/paper/1807.07356/) (2018)
- [Puzzle Mix: Exploiting Saliency and Local Statistics for Optimal Mixup](https://mlsanity.com/paper/2009.06962/) (2020)
- [Learning with Noisy Labels Revisited: A Study Using Real-World Human Annotations](https://mlsanity.com/paper/2110.12088/) (2021)
- [Adversarial Domain Adaptation with Domain Mixup](https://mlsanity.com/paper/1912.01805/) (2019)
- [Fixing Data Augmentation to Improve Adversarial Robustness](https://mlsanity.com/paper/2103.01946/) (2021)
- [Improving Out-of-Distribution Robustness via Selective Augmentation](https://mlsanity.com/paper/2201.00299/) (2022)
- [Automatic Differentiation Variational Inference](https://mlsanity.com/paper/1603.00788/) (2016)
- [Remix: Rebalanced Mixup](https://mlsanity.com/paper/2007.03943/) (2020)
- [Data augmentation using learned transformations for one-shot medical image segmentation](https://mlsanity.com/paper/1902.09383/) (2019)
- [Better Aggregation in Test-Time Augmentation](https://mlsanity.com/paper/2011.11156/) (2020)
- [Faster AutoAugment: Learning Augmentation Strategies using Backpropagation](https://mlsanity.com/paper/1911.06987/) (2019)
- [Augmentation Strategies for Learning with Noisy Labels](https://mlsanity.com/paper/2103.02130/) (2021)
- [On Mixup Regularization](https://mlsanity.com/paper/2006.06049/) (2020)
- [TransMix: Attend to Mix for Vision Transformers](https://mlsanity.com/paper/2111.09833/) (2021)
- [Understanding data augmentation for classification: when to warp?](https://mlsanity.com/paper/1609.08764/) (2016)
- [Greedy Policy Search: A Simple Baseline for Learnable Test-Time Augmentation](https://mlsanity.com/paper/2002.09103/) (2020)
- [DADA: Differentiable Automatic Data Augmentation](https://mlsanity.com/paper/2003.03780/) (2020)
- [Learning Loss for Test-Time Augmentation](https://mlsanity.com/paper/2010.11422/) (2020)

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

- [A Survey on Efficient Vision-Language Models](https://mlsanity.com/paper/2504.09724/) (2025)
- [Invariant Transformation and Resampling based Epistemic-Uncertainty Reduction](https://mlsanity.com/paper/2602.23315/) (2026)
- [Higher-Order Asymptotics of Test-Time Adaptation for Batch Normalization Statistics](https://mlsanity.com/paper/2505.16257/) (2025)

## Abstract (arXiv)

Data augmentation is known to contribute significantly to the robustness of machine learning models. In most instances, data augmentation is utilized during the training phase. Test-Time Augmentation (TTA) is a technique that instead leverages these data augmentations during the testing phase to achieve robust predictions. More precisely, TTA averages the predictions of multiple data augmentations of an instance to produce a final prediction. Although the effectiveness of TTA has been empirically reported, it can be expected that the predictive performance achieved will depend on the set of data augmentation methods used during testing. In particular, the data augmentation methods applied should make different contributions to performance. That is, it is anticipated that there may be differing degrees of contribution in the set of data augmentation methods used for TTA, and these could have a negative impact on prediction performance. In this study, we consider a weighted version of the TTA based on the contribution of each data augmentation. Some variants of TTA can be regarded as considering the problem of determining the appropriate weighting. We demonstrate that the determination of the coefficients of this weighted TTA can be formalized in a variational Bayesian framework. We also show that optimizing the weights to maximize the marginal log-likelihood suppresses candidates of unwanted data augmentations at the test phase.