Source-linked AI summary
Greedy Policy Search: A Simple Baseline for Learnable Test-Time Augmentation
Dmitry Molchanov, Alexander Lyzhov, Yuliya Molchanova, Arsenii Ashukha, Dmitry Vetrov
TL;DR
Test-time augmentation commonly uses simple transformations, while more diverse learnable augmentation has mainly been developed for training. The paper introduces greedy policy search to learn test-time policies and reports superior prediction and uncertainty results on clean and corrupted data; it does not compare TTA with model ensembles under domain shift.
Problem
Test-time augmentation is usually limited to simple transformations and typically does not benefit from more diverse policies such as those used during training.
Source: §1 · p. 2Method
Greedy policy search (GPS) builds a test-time augmentation policy based on predictive performance on a validation set.
Source: §1 · p. 2Results
Learned policies provide superior predictive performance and uncertainty estimates compared with existing test-time augmentation approaches on clean and corrupted data.
Source: §5 · p. 9Takeaways & Limitations
Transferring policies from small to large architectures is a viable strategy when computational resources are limited.
Source: §5 · p. 9The paper does not compare model ensembles with test-time augmentation under dataset shift.
Source: §4.3 · p. 7
Abstract
from arXiv · showhide
Test-time data augmentation$-$averaging the predictions of a machine learning model across multiple augmented samples of data$-$is a widely used technique that improves the predictive performance. While many advanced learnable data augmentation techniques have emerged in recent years, they are focused on the training phase. Such techniques are not necessarily optimal for test-time augmentation and can be outperformed by a policy consisting of simple crops and flips. The primary goal of this paper is to demonstrate that test-time augmentation policies can be successfully learned too. We introduce greedy policy search (GPS), a simple but high-performing method for learning a policy of test-time augmentation. We demonstrate that augmentation policies learned with GPS achieve superior predictive performance on image classification problems, provide better in-domain uncertainty estimation, and improve the robustness to domain shift.
1 INTRODUCTION
CNNs and training-time augmentation promote invariance, but test-time augmentation typically uses simple transformations and does not usually benefit from more diverse policies. The paper introduces GPS to learn test-time augmentation policies and reports improvements across image classification, uncertainty estimation, and domain-shift evaluation.
1 INTRODUCTION
- CNNs use spatially invariant operations designed to reflect label-preserving symmetries, and data augmentation further promotes invariance by expanding training data with label-preserving transformations.Earlier augmentation strategies used handcrafted transformations and required manual tuning; more recent work explored learnable, more diverse strategies.Source: §1 · p. 1
- Test-time augmentation averages predictions across multiple augmentations, but commonly uses simple transformations and usually does not benefit from more diverse policies such as those used during training.Multi-crop evaluation is a standard protocol for large-scale image classification.Source: §1 · p. 2
- Greedy policy search (GPS) learns test-time augmentation policies from predictive performance on a validation set.The paper reports that calibrated log-likelihood is crucial to policy search, while accuracy and plain log-likelihood lead to a significant drop in final performance.Source: §1 · p. 2
- GPS-learned policies outperform other augmentation baselines across image classification, in-domain uncertainty estimation, and classification under dataset shift.The reported gains span VGG-style networks and EfficientNets, models with and without advanced train-time augmentation, and ensembles; the policies also transfer across architectures.Source: §1 · p. 2
2 RELATED WORK
Prior work established test-time augmentation and explored learnable augmentation primarily for training. This section also relates test-time augmentation to ensembling and greedy selection of ensemble members.
Test-time augmentation
- Test-time augmentation averages predictions over image transformations; random crops and flips became a standard evaluation protocol, with applications including medical imaging and adversarial defense.Prior work also reported improved uncertainty estimation for medical image segmentation and applying mixup at test time to improve defense against adversarial attacks.Source: §2 · p. 2
Learnable train-time augmentation
- Learnable augmentation research primarily adapts training-time policies to a dataset, with methods ranging from costly reinforcement-learning search to more efficient policy-search approaches.AutoAugment requires repeating expensive model training for each policy-search iteration.Source: §2 · p. 2
Ensembling
- Ensembling averages predictions from a distribution or set of trained neural networks, with methods differing in time and memory costs, member diversity, and performance..Source: §2 · p. 2
Sub-ensemble selection
- The paper treats test-time augmentation as an ensemble of sub-policies and adopts greedy ensemble pruning to select members from discrete possibilities.Earlier work described heuristic and simpler greedy approaches to ensemble-member selection.Source: §2 · p. 2
3 LEARNABLE TEST-TIME AUGMENTATION
The section defines test-time augmentation (TTA) policies as sets of image-transform sub-policies whose predictions are averaged at inference, and motivates learning them directly because train-time choices may not transfer. GPS builds a policy greedily using calibrated log-likelihood.
3 LEARNABLE TEST-TIME AUGMENTATION
- A TTA policy is a set of sub-policies, each applying consecutive image transformations, and inference averages predictions across samples of different sub-policies.Source: §3 · pp. 2–3
3.1 Naive approaches to test-time augmentation
- Train-time policies need not suit test-time augmentation, and the optimal test-time augmentation magnitude can differ from the training magnitude.Source: §3.1 · p. 3
- On ImageNet, scale-crop-flip TTA outperforms RandAugment TTA for EfficientNet B2 trained with RandAugment, even at large sample counts.Source: Figure 2 · p. 3, §3.1 · p. 3
3.2 Greedy policy search
- GPS is presented as a way to improve predictive performance, uncertainty estimates, and robustness through learned TTA policies.Source: §3.2 · p. 3
- GPS draws candidate sub-policies from a prior, precomputes their predictions, then iteratively adds the candidate with the largest performance gain until the policy reaches its desired length.The authors describe candidate generation and selection as parallelizable and scalable.Source: §3.2 · p. 3
- Calibrated log-likelihood guides greedy selection because accuracy is noisy and uncalibrated log-likelihood cannot fairly assess TTA; it is better suited to policy search than either alternative.Source: §3.2 · pp. 3–4, Figure 3 · p. 4
4 EXPERIMENTS
Across image-classification experiments, GPS beats conventional and advanced test-time augmentation baselines, and calibrated log-likelihood is the best search objective in the tested comparison. The learned policies also show benefits under domain shift, transfer, ensembling, and vanilla training augmentation.
4 EXPERIMENTS
- CIFAR policies are searched on a validation split and evaluated after retraining on the full training set; ImageNet uses separate halves of its validation set for search and evaluation.GPS uses approximately 1,000 candidate sub-policies, and the evaluation procedure calibrates temperature through test-time cross-validation.Source: §4 · pp. 4–5
4.1 In-domain predictive performance
- GPS consistently beats other test-time augmentation methods in calibrated log-likelihood and accuracy on clean CIFAR-100 and ImageNet.The comparisons include conventional crops and flips, training-policy reuse, and RandAugment with magnitude grid search.Source: Figure 4 · p. 5, §4.1 · p. 5
4.3 Robustness to domain shift
- On corrupted CIFAR datasets, GPS is more robust than alternatives in most cases, and on ImageNet-C it outperforms non-learnable TTA.Policies are learned on clean data; on corrupted data, extreme augmentations can improve performance over conservative crops and flips.Source: Figure 5 · p. 6, Figure 6 · p. 6, §4.3 · p. 7
4.2 What metric to use for policy search?
- Calibrated log-likelihood optimization outperforms accuracy and log-likelihood objectives across the tested CIFAR-100 tasks and metrics.Accuracy is too noisy for reliable magnitude selection, while log-likelihood can favor conservative magnitudes; calibrated log-likelihood avoids that problem.Source: Table 1 · p. 6, §4.2 · p. 6
4.4 Policy transfer
- Transferred GPS policies outperform random crops and flips across considered CIFAR dataset-architecture transfers and outperform conventional baselines on EfficientNet-L2.Keeping the dataset fixed matters more than keeping the architecture fixed; CIFAR-to-ImageNet transfers are excluded because preprocessing differs.Source: §4.4 · pp. 7–8, Figure 8 · p. 7
4.5 Greedy policy search for ensembles
- Applying GPS directly to an ensemble improves performance more than transferring a single-model policy, though both GPS variants beat other baselines.The tested approaches use a policy found for one model or GPS applied to the whole ensemble.Source: Figure 9 · p. 8, §4.5 · p. 8
4.6 Greedy policy search for models trained with vanilla augmentation
- On CIFAR-100, GPS for vanilla-trained models significantly outperforms standard crops and flips.Compared with policies for RandAugment-trained models, these policies are less diverse, use nearly half the transformation magnitude, and choose identity more often.Source: Figure 10 · p. 8, §4.6 · p. 9
5 CONCLUSION
GPS policies improve predictive performance and uncertainty estimates over existing TTA approaches on clean and corrupted data. The paper also identifies calibrated log-likelihood as a stronger search objective and reports architecture transfer, while noting that its policies are static and object-independent.
5 CONCLUSION
- GPS policies consistently improve predictive performance and uncertainty estimates over existing TTA approaches on clean and corrupted data.The authors report significant improvements for both in-domain data and under domain shift.Source: §5 · p. 9
- Calibrated log-likelihood outperforms log-likelihood and accuracy as an objective for learning TTA strategies.The authors note possible implications for meta-learning and neural architecture search, where uncalibrated objectives are often used.Source: §5 · p. 9
- GPS policies transfer across architectures, making small-to-large policy transfer viable when computational resources are limited.Source: §5 · p. 9
- The learned policies are static and object-independent; dynamic object-dependent TTA is proposed as a future direction, potentially helpful under domain shift.Source: §5 · p. 9
A Experimental details
The experimental details specify a modified RandAugment setup for CIFAR models, the candidate-policy pools used for CIFAR and ImageNet, and image-processing choices for transformations and backgrounds.
A Experimental details
- The RandAugment modification addresses ill-defined transformation magnitudes and transformations that can become less severe as magnitude increases.For contrast, brightness, color, and sharpness, enhancement uses a random sign applied to the magnitude-scaled parameter.Source: A Experimental details · p. 11
- Transformation parameters include shear, translation, rotation, solarization, posterization, contrast, brightness, color, sharpness, and cutout.Their magnitude mappings and ranges are specified in the implementation details.Source: A Experimental details · p. 11
- Mirrored backgrounds replace RandAugment’s black backgrounds to preserve augmented-image statistics and mitigate potential issues from batch normalization on varying data.Using separate batch-normalization statistics for different transformation types is noted as an alternative.Source: A Experimental details · p. 11
- In the described RandAugment setup, each magnitude is drawn uniformly from 0 to M, and each application is followed by a random crop and flip.Source: A Experimental details · p. 11
- The CIFAR policy pool contains 1101 sub-policies sampled across three RandAugment magnitude ranges plus an identity transformation.Each GPS sub-policy application is followed by a random crop and flip.Source: A Experimental details · p. 11
- ImageNet sub-policies begin with scale-crop-flip at a learnable scaling magnitude, and the pool includes Invert and Equalize transformations.The scale-crop-flip operation is not counted toward sub-policy length.Source: A Experimental details · p. 11
B Additional experimental results
The additional results examine temperature choice, RandAugment magnitude, learned policy composition, and clean-set TTA performance. They report model-dependent magnitude preferences and show that GPS can outperform the random-crop-and-flip baseline even when magnitude grid search alone cannot.
B Additional experimental results
- The optimal temperature decreases as the number of samples drawn from a GPS policy increases, and validation-set temperature matches test-time cross-validation.Source: Figure 11 · p. 12
- RandAugment’s optimal test-time magnitude varies by model and need not match its training magnitude; increasing magnitude also widens the gap between log-likelihood and calibrated log-likelihood.For models trained with standard crops and flips, performance starts lower and degrades quickly as magnitude increases, while GPS still finds a policy that beats that baseline.Source: Figure 12 · p. 12
- GPS policies contain more diverse sub-policies when the initial model was trained with RandAugment; crop-flip-trained models use substantially more identity transformations.Figure 16 compares transformation counts and average magnitudes across CIFAR-100 model architectures.Source: Figure 14 · p. 14, Figure 16 · p. 15
- GPS outperforms other methods in calibrated log-likelihood and achieves comparable accuracy on CIFAR-10’s clean test set.Source: Figure 17 · p. 15