Adaptive test-time augmentation via KL-regularized reinforcement learning for robust visual inference
4 affiliations
- 1Department of Information Technology, Bharati Vidyapeeth’s College of Engineering
- 2Department of Computer Science and Engineering, Bharati Vidyapeeth’s College of Engineering
- 3School of Computer Science and Engineering, Galgotias University
- 4Department of Computer Science, Faculty of Science, University of South Bohemia in České Budějovice
Abstract
Deep neural networks often suffer significant accuracy degradation when exposed to real-world image corruptions and distribution shifts. To overcome the limitations of fixed, input-agnostic test-time augmentation (TTA), an adaptive framework is proposed that learns per-sample transformations via reinforcement learning. Augmentation selection is cast as a Markov decision process and proximal policy optimization (PPO) agents are trained to choose sample-specific transforms under a composite reward combining classifier confidence gains with a self-consistency KL-divergence penalty on the model’s own softmax outputs, thereby preserving overall belief stability. On clean CIFAR-10 (1 000 samples), the adaptive ensemble raises accuracy from 88.5% (baseline) and 87.3% (static TTA) to 90.0% (+1.5 pp). On CIFAR-10-C (15 corruptions × 5 severities; 1 000 images per condition), pooled top-1 accuracy improves from 75.7% (baseline) and 74.3% (static TTA) to 76.4% (+0.7 pp), and exceeds a TENT entropy-minimization baseline (75.9%) while operating in a strictly label-free regime that updates no model weights. Per-corruption gains are consistently positive across noise, blur, weather, and compression distortions, with the adaptive policy outperforming TENT on texture and compression corruptions where input-space transforms are most effective. These findings demonstrate that learned, per-sample augmentation policies improve robustness and reliability of deep vision models under diverse image conditions, against a strong baseline classifier.
Related
- Adaptive intelligent resource optimization in wireless sensor networks using reinforcement learning and compressive sensing
- Attention-based Spatio-Temporal Graph Convolutional Networks-enhanced deep reinforcement learning for adaptive traffic signal control in urban traffic networks
- Enhancing self-supervised representation learning through lightweight learnable data augmentation
- Weighted estimation of multiple regions in CT-derived images of the pubic symphysis for age estimation using a deep learning framework with R-Mixup augmentation
- Reinforcement learning–coupled neural ode modelling reveals ferroptotic resolution failure as a shared dynamical axis linking periodontitis and coronary atherogenesis
BibTeX
@article{mittal2026adaptive,
title = {{Adaptive test-time augmentation via KL-regularized reinforcement learning for robust visual inference}},
author = {Tushar Mittal and Arun Kumar Dubey and Dharmender Saini and Sarita Yadav and Achin Jain and Arvind Panwar and Hamidreza Namazi},
journal = {Scientific Reports},
year = {2026},
month = jul,
doi = {10.1038/s41598-026-62059-4},
url = {https://doi.org/10.1038/s41598-026-62059-4},
note = {CC BY},
}
RIS
TY - JOUR
TI - Adaptive test-time augmentation via KL-regularized reinforcement learning for robust visual inference
AU - Tushar Mittal
AU - Arun Kumar Dubey
AU - Dharmender Saini
AU - Sarita Yadav
AU - Achin Jain
AU - Arvind Panwar
AU - Hamidreza Namazi
PY - 2026
DA - 2026/07/01
JO - Scientific Reports
DO - 10.1038/s41598-026-62059-4
UR - https://doi.org/10.1038/s41598-026-62059-4
AB - Deep neural networks often suffer significant accuracy degradation when exposed to real-world image corruptions and distribution shifts. To overcome the limitations of fixed, input-agnostic test-time augmentation (TTA), an adaptive framework is proposed that learns per-sample transformations via reinforcement learning. Augmentation selection is cast as a Markov decision process and proximal policy optimization (PPO) agents are trained to choose sample-specific transforms under a composite reward combining classifier confidence gains with a self-consistency KL-divergence penalty on the model’s own softmax outputs, thereby preserving overall belief stability. On clean CIFAR-10 (1 000 samples), the adaptive ensemble raises accuracy from 88.5% (baseline) and 87.3% (static TTA) to 90.0% (+1.5 pp). On CIFAR-10-C (15 corruptions × 5 severities; 1 000 images per condition), pooled top-1 accuracy improves from 75.7% (baseline) and 74.3% (static TTA) to 76.4% (+0.7 pp), and exceeds a TENT entropy-minimization baseline (75.9%) while operating in a strictly label-free regime that updates no model weights. Per-corruption gains are consistently positive across noise, blur, weather, and compression distortions, with the adaptive policy outperforming TENT on texture and compression corruptions where input-space transforms are most effective. These findings demonstrate that learned, per-sample augmentation policies improve robustness and reliability of deep vision models under diverse image conditions, against a strong baseline classifier.
SN - 2045-2322
ER -