Adaptive test-time augmentation via KL-regularized reinforcement learning for robust visual inference

Area of Science:

• Computer Vision
• Machine Learning
• Deep Learning

Background:

• Deep neural networks (DNNs) exhibit reduced accuracy with real-world image corruptions and distribution shifts.
• Fixed, input-agnostic test-time augmentation (TTA) has limitations in addressing these accuracy degradations.

Purpose of the Study:

• To propose an adaptive framework that learns per-sample transformations for enhanced DNN robustness.
• To overcome the limitations of static TTA methods.

Main Methods:

• An adaptive framework using reinforcement learning (RL) was developed.
• Augmentation selection was modeled as a Markov decision process with proximal policy optimization (PPO) agents.
• A composite reward function combined classifier confidence gains and a self-consistency KL-divergence penalty.

Main Results:

• On clean CIFAR-10, accuracy improved from 88.5% (baseline) to 90.0%.
• On CIFAR-10-C, pooled top-1 accuracy increased from 75.7% (baseline) to 76.4%, outperforming static TTA and a TENT baseline.
• Consistent accuracy gains were observed across various corruption types (noise, blur, weather, compression).

Conclusions:

• Learned, per-sample augmentation policies significantly improve the robustness and reliability of deep vision models.
• The adaptive framework operates effectively in a label-free regime, without updating model weights.
• The approach demonstrates superior performance on corruptions where input-space transforms are most beneficial.