Skip to content
Open access

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

Jul 2026 · Scientific Reports · 0 citations
Medicine

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.

Read PDF

Similar papers

Conference Open access Jul 2026

Symbolic Knowledge Transfer for Sample-Efficient Deep Reinforcement Learning

This work proposes a neuro-symbolic DRL approach that incorporates background symbolic knowledge to improve both sample efficiency and generalization to more challenging, unseen tasks and demonstrates consistent performance improvements over a state-of-the-art reward machine baseline.

Celeste Veronese, Alessandro Farinelli, Daniele Meli · 0 citations
Preprint Aug 2026

Falcon Perception-HD: High Density Perception via Reinforcement Learning

This paper explores post-training reinforcement learning (RL), specifically GRPO, to directly align autoregressive perception models with their evaluation metrics, and designs an RL framework that addresses perception-specific challenges: reward design for set-structured outputs and multi-head sampling control.

Sofian Chaybouti, Yasser Dahou, Ngoc Dung Huynh et al. · 0 citations
Preprint Jul 2026

Learning the Supports for Categorical Critic in Reinforcement Learning

Value functions are an essential component in actor-critic based deep reinforcement learning (RL). Conventionally, these functions are trained as a regression task by minimising the mean squared error (MSE) relative to bootstrapped target values. Meanwhile, in distributional RL, a distribution of returns is modelled based on the distributional Bellman operator. This work investigates the Gaussian Histogram Loss (HL-Gauss), a recent approach that reframes value estimation as classification by encoding each scalar Bellman target as a Gaussian-smoothed categorical target. Despite its potential, applying histogram-based losses to RL presents inherent challenges, most notably the requirement to pre-define a fixed support interval, which is often complicated by the non-stationary and stochastic nature of target values typically found in RL tasks. In this work, we propose an approach that dynamically learns the lower and upper bounds of the support instead of assigning them beforehand. We derive an objective that jointly learns these bounds whilst learning the categorical representation of the scalar values, and we show that this objective forms an upper bound on the mean-squared Bellman error. Our theoretical analysis further shows that this bound is tighter than that of non-learned supports of HL-Gauss. Empirically, the proposed objective enables stable adaptation of the support interval and matches HL-Gauss-based actor-critic algorithms on most continuous-control tasks whilst improving on a subset, without requiring a pre-specified support interval.

Jen-Yen Chang, Takayuki Osa, Tatsuya Harada · 0 citations
Preprint Jul 2026

Mask-based Predictive Representations for Reinforcement Learning

This work proposes a self-supervised task based on mask prediction as an auxiliary task for reinforcement learning that outperforms state-of-the-art sample-efficient reinforcement learning methods on multiple continuous and discrete control benchmarks.

Kaisheng Zhao · 0 citations
Jun 2026

FedLAS: Feature-Modulated Bidirectional Label Smoothing for Neural Network Calibration

Deep Neural Network (DNN) classifiers suffer from poor calibration when their softmax outputs (predictive confidence) deviate from the empirical likelihoods. This manifests itself as either overconfident incorrect predictions or under-confident correct predictions. Label smoothing (LS) enhances model calibration by introducing entropy regularization during training through redistributing probability mass from the ground-truth label to the remaining classes. LS, including Margin-based LS (MbLS), have restrictive assumptions: they rely on predefined, uniform smoothing rules and only tackle overconfidence. In reality, samples exhibit diverse characteristics, such as difficulty/ambiguity, that interact with the evolving nature of the model being trained. In training, samples may have various degrees of under- or overconfidence. To overcome this, a mechanism that identifies the specific confidence state of each sample and determines the appropriate degree of smoothing in each training step is needed, tailoring the adjustment to the individual sample. We propose FedLAS: Feature-Modulated Bidirectional Label Smoothing, a plug-and-play algorithm for label smoothing-based losses. In FedLAS, we introduce a Feature Norm-based Confidence Indicator (NCI) to control smoothing and a Bidirectional Calibration Gating (BCG) module to detect both over and under-confidence. Our algorithm can be integrated with LS and MbLS based losses when applied to standard DNNs, enhancing performance. Extensive experiments on standard and fine-grained high-resolution vision benchmarks show that FedLAS consistently improves calibration compared to modern baselines, reducing Expected Calibration Error (ECE) and Adaptive ECE while maintaining Top-1 accuracy. Code: github.com/nadarasarbahavan/FEDLAS

Thiru Thillai Nadarasar Bahavan, Sachith Seneviratne, Saman K. Halgamuge · 0 citations
Preprint Jul 2026

Optimizing Visual Generative Models via Distribution-wise Rewards

A novel framework that finetunes generative models using distribution-wise rewards, ensuring better alignment with real-world data distributions is presented, and a subset-replace strategy that efficiently provides reward signals by updating only a small subset of a generated reference set is introduced.

Ruihang Li, Mengde Xu, Shuyang Gu et al. · 0 citations