Source-linked AI summary
Learning Loss for Test-Time Augmentation
Ildoo Kim, Younghoon Kim, Sungwoong Kim
TL;DR
Test-time corruptions and deformations can reduce neural-network performance, while many augmentation methods focus on training. This paper predicts losses for candidate transformations and selects low-loss transformations per test input; experiments report improved robustness against deformations at low computational cost.
Problem
Neural networks can lose performance under test-time deformations, and training-time augmentation may not suffice to handle deformed test data.
Source: §1 · p. 1Method
A separate loss predictor estimates relative losses for candidate transformations, and the method applies low-predicted-loss transformations and averages the resulting predictions.
Source: §1 · p. 2, §3.4 · p. 5Results
Takeaways & Limitations
Instance-level selection can reduce the number of augmentations used for ensembling while targeting transformations for each test input.
Source: §1 · p. 2The authors could not train appropriate loss predictors for some target models, including models trained with CutMix.
Source: §5 · p. 9
Abstract
from arXiv · showhide
Data augmentation has been actively studied for robust neural networks. Most of the recent data augmentation methods focus on augmenting datasets during the training phase. At the testing phase, simple transformations are still widely used for test-time augmentation. This paper proposes a novel instance-level test-time augmentation that efficiently selects suitable transformations for a test input. Our proposed method involves an auxiliary module to predict the loss of each possible transformation given the input. Then, the transformations having lower predicted losses are applied to the input. The network obtains the results by averaging the prediction results of augmented inputs. Experimental results on several image classification benchmarks show that the proposed instance-aware test-time augmentation improves the model's robustness against various corruptions.
1 Introduction
Neural networks can be vulnerable to small image deformations and other test-time corruptions, even when trained with advanced augmentation. The paper proposes selecting test-time transformations for each input using predicted losses, aiming to improve robustness without the large ensembles or fixed policies used by prior approaches.
1 Introduction
- Test-time deformations can reduce performance even for networks trained with advanced data augmentation.The authors report both a performance drop from simple test-time deformations and remaining room for improvement with appropriate test-time transformations.Source: §1 · p. 1
- Prior test-time augmentation uses fixed transformations or policies and can require many augmented inputs, increasing inference cost.A cited learnable policy searches over the test set rather than optimizing choices for each input.Source: §1 · p. 2, Figure 1 · p. 2
- The method predicts transformation losses for each input, selects low-loss augmentations, and averages the target network’s outputs across them.The predictor is compact, and instance-level selection can reduce the number of augmentations used for ensembling.Source: §1 · p. 2
- The proposed method improves robustness against deformations in image-classification experiments compared with predefined test-time augmentation methods.The introduction reports comparisons on two benchmark datasets and describes the method as computationally efficient.Source: §1 · p. 2
2 Related Works
Related work covers augmentation during training, test-time augmentation, uncertainty estimation, and robustness to corruptions. The paper positions its approach as instance-specific loss prediction for choosing test-time transformations.
Data Augmentation
- Training-time augmentation includes sample mixing, policy search, and randomly generated augmentation mixtures, but does not eliminate test-time deformation concerns.The cited work notes that training-phase augmentation can be insufficient for deformations in test data.Source: §2 · pp. 2–3
Test-Time Augmentation
- Prior test-time strategies average predictions from crops, mixup inputs, multiple resolutions, or flipped and rotated variants.Some methods use many transformed inputs, while a learned greedy policy selects a static policy based on validation performance.Source: §2 · p. 3
Uncertainty Measure
- Loss estimation has been used to measure uncertainty, including auxiliary loss predictors, while this paper predicts relative losses across transformations for test-time augmentation.The cited loss-prediction approach addresses uncertainty in unlabeled data; the paper distinguishes its goal from active learning.Source: §2 · p. 3
Robustness in Convolutional Neural Network
- Prior robustness studies examine corruption-specific fine-tuning, specialized experts, mixup, and model capacity, with increased capacity requiring substantial computation.The cited studies report that robustness to one corruption type may not generalize to others.Source: §2 · p. 3
3 Method
The method selects a small set of transformations for each test image using a separate loss predictor. It defines a discretized transformation space, predicts relative losses with a compact network, and trains that predictor while keeping the target model fixed.
3 Method
- The method section introduces the overall test-time procedure, transformation space, loss-prediction module, and its training method.Implementation details are deferred to Appendix A.3.Source: §3 · p. 3
3.1 Test-Time Augmentation
- Conventional test-time augmentation averages target-network predictions over candidate transformations; the proposed procedure selects the k transformations with lowest predicted losses for each input.An idealized setting uses only the best transformation to improve performance with the smallest computation.Source: §3.1 · pp. 3–4
3.2 Test-Time Augmentation Space
- The candidate space contains 12 individual transformations: rotations, zooms, color and contrast changes, sharpness adjustments, and identity.The paper uses a moderately diverse discretized space and applies transformations individually rather than in combinations.Source: §3.2 · p. 4
3.3 Loss Prediction Module: Architecture
- A separate small loss predictor takes the untransformed image and outputs relative losses for each candidate transformation.The predictor uses EfficientNet-B0 with multi-level features; a separate module avoids an additional target-network inference for selection.Source: §3.3 · p. 5
3.4 Loss Prediction Module: Training Method
- The target network is frozen while the loss predictor is trained using a data split reserved for its training and validation.Ground-truth losses come from evaluating every transformed image with the target network; softmax-normalized relative losses and a ranking-loss surrogate train the predictor, with the latter reported to yield more stable training.Source: §3.4 · p. 5
4 Experiments
Experiments on CIFAR-100-C and ImageNet-C show that instance-aware test-time augmentation improves corruption robustness. On ImageNet, its generalization is similar to 5-Crops at less than half the computational cost.
4 Experiments
- The evaluation trains loss predictors with 15 corruption types and validates on four additional types, excluding training augmentations that overlap with corruptions.Source: §4 · pp. 5–6
4.1 CIFAR-100 Classification
- 45.15% to 41.92% average error on 15 corruptions for Fast AutoAugment-trained Wide-ResNet-40-2, with k=1 and 1% added inference computation.Source: §4.1 · p. 6
- With 15 corruptions added to training augmentation, AugMix+ fails to generalize improvements to held-out corruptions and degrades clean-set performance, while the proposed method generalizes better and preserves clean performance.Source: §4.1 · p. 6, Table 2 · p. 7
- When blur corruptions are included in training, the proposed method improves blur robustness while maintaining clean-set performance, unlike training-time blur augmentation alone.Source: §4.1 · p. 6, Table 3 · p. 8
4.2 ImageNet Classification
- The loss predictor uses 64-by-64 inputs on ImageNet, and its inference cost is negligible relative to target networks.Source: §4.2 · p. 6
- On ImageNet-C, the method outperforms GPS and GPS†; GPS policies include the method’s augmentation policies, and the reported gains come from instance-specific transformations.Source: §4.2 · pp. 7–8
- For AugMix-trained ResNet-50, k=2 lowers (known-)corruption mCE from 67.72% to 64.55%, with generalization similar to 5-Crops at less than half the cost.Source: §4.2 · pp. 6–7, Table 4 · p. 8
5 Discussion
The discussion reports strong potential when the lowest-loss augmented image can be selected, while noting limits in loss prediction and in the tested transformation set. It also identifies predictor-training failures for some target models and corruptions.
Results
- Assuming the lowest-loss augmented image can be selected, the method beats existing state-of-the-art models on clean data and can approach clean-data performance on ImageNet-C and CIFAR-100-C.The authors report the clean-data comparison as a clear margin and describe the corrupted-set performance as close to clean-data performance.Source: §5 · pp. 8–9
Examples
- Figure 3 presents correctly classified examples with corrupted inputs, prediction-versus-actual-loss scatter plots, and the method's augmented images.The red circle marks an augmented image matching the label; the blue triangle does not, and corruptions in examples (b) and (d) are exaggerated.Source: Figure 3 · p. 9
Augmentation possibilities
- Sharpness enhancements contributed to performance on noise corruptions, and the authors expect purpose-designed operations such as deblurring to yield further gains.They suggest that sophisticated augmentations with a clear purpose could improve results.Source: §5 · p. 9
Learning Loss Prediction Module
- Loss predictions correlate 0.38 on validation and 0.30 on test, and the method works well in ensembles when k ≤ 2.The authors characterize predicting transformation-induced loss changes for the same data as more difficult than predicting losses for unseen data.Source: §5 · p. 9
Limitations
- The predictor was ineffective for some CutMix-trained models and difficult to train for some corruptions included in training-time augmentation.The authors conjecture that small perturbations can make loss values noisy to predict.Source: §5 · p. 9
6 Conclusion
The conclusion presents instance-aware test-time augmentation as a low-cost way to improve classification accuracy and robustness against test-time deformation. It also identifies limited transformation coverage and broader validation as open work.
Conclusion
- The method uses a loss-prediction network to choose image transformations dynamically at test time and is described as computationally efficient and practical for real-world use.The conclusion says the loss prediction module improves classification accuracy and robustness against deformation at low computational cost.Source: §6 · p. 9
Future work
- The experiments used a limited number of transformations and specifications, motivating study of more general loss-prediction settings.The authors also propose examining reinforcement learning for efficient loss-predictor optimization.Source: §6 · p. 9
- The authors plan to validate the framework on object detection and segmentation and to examine how training-time and test-time augmentation can be combined.They also identify reducing costs while including more augmentations as future work.Source: §6 · p. 9
Broader Impact
The broader-impact discussion argues that identical preprocessing may be ineffective across inputs and suggests instance-aware augmentation could improve deployment stability and performance. It stresses that stability and generalization still require verification.
Broader Impact
- Applying the same preprocessing regardless of the input's state can be ineffective for performance and stability.The passage does not specify the conditions under which this occurs.Source: Broader Impact · p. 10
- The authors suggest the proposed method could increase stability and performance when a deep learning model is deployed.This is presented as an expectation rather than a verified deployment result.Source: Broader Impact · p. 10
- The method's stability and generalization remain to be verified, and performance may degrade in unexpected situations.The passage does not identify specific unexpected conditions.Source: Broader Impact · p. 10
A.1 Conventional Test-Time Augmentation
The appendix describes conventional test-time augmentation as crop-, flip-, or resolution-based aggregation. In the reported experiments, multi-resolution inputs were ineffective, while existing methods generally aggregate transformations considered safe during training.
A.1 Conventional Test-Time Augmentation
- Center-Crop resizes an image and crops its central area; for ImageNet ResNet-50, the setup resizes to 256 pixels and crops 224 pixels.CIFAR images are used at their original 32-by-32 resolution without test-time modification.Source: §A.1 · p. 13
- Horizontal-Flip ensembles the original and horizontally inverted images, while 5-Crops and 10-Crops combine center and corner crops, with 10-Crops adding flipped images.The cited ImageNet methods use 5 or 10 crops.Source: §A.1 · p. 13
- Multi-resolution test-time inputs were ineffective in these experiments, and existing methods generally aggregate results from transformations considered safe during training.Source: §A.1 · p. 13
A.2 Test-Time Augmentation Space
The augmentation space contains 12 image operations, with transformations applied to a test image illustrated through relative losses and classification outcomes. The examples show that different transformations can help a corrupted image, while the displayed blur was exaggerated for visibility.
A.2 Test-Time Augmentation Space
- The test-time augmentation space contains 12 operations, including rotation, zoom, contrast, sharpness, and color adjustments.Rotation and zoom use bicubic resampling; blank areas from transformations are filled with mirrored image content.Source: §A.2 · p. 13
- For a noisy swing image, zoom-out, blur, or rotation into the normal orientation can be advantageous at test time.The blur corruption for two displayed sharpness settings was exaggerated to make the difference visible.Source: Figure 4 · p. 13
- Table 5 compares random-TTA and oracle-TTA Top-1 accuracy on CIFAR-100 and ImageNet, with oracle-TTA indicating an upper bound for the proposed transformation space.Its caption describes this as evidence of performance-improvement potential relative to state-of-the-art models without external data.Source: Table 5 · p. 14
A.3 Implementation
The experiments compare hypothetical loss predictors, examine which transformations are selected across corruption types and severities, and evaluate clean-set augmentation. They also describe the computational cost of preparing predictor training data and measures used to reduce overfitting.
A.3 Implementation
- Generating predictor-training targets requires transforming images and running target-network inference, so distributed data loading reduces learning time but not total computational cost.The implementation also perturbs training data with repeated same-content images and Cutout to help avoid overfitting.Source: §A.3 · p. 14
A.4 Test-Time Augmentation Oracle
- Oracle-TTA improves performance on clean and corrupted datasets, while Random-TTA reduces performance, indicating the potential of instance-aware test-time augmentation.Oracle-TTA assumes accurate relative-loss prediction and selects the lowest-loss transformation; Random-TTA selects randomly.Source: §A.4 · p. 14
A.5 Analysis of Selected Test-Time Augmentation for ResNet-50 on ImageNet
- Selection varies with corruption and severity: high-severity Gaussian noise favors stronger blur, while defocus blur favors sharpening or zoom-in, with selection rates increasing as severity rises.At relatively low Gaussian-noise severity, blur or zoom-out is selected; selection also varies for clean images, where identity is chosen for 75%.Source: §A.5 · pp. 14–15, Figure 5 · p. 15
A.6 Test-time Augmentation for The Clean Set
- On clean ImageNet, one selected crop performs almost as well as the 5-crop ensemble, while two crops perform better with less computation; selection from GPS policies further improves over static GPS ensembles.The predictor selects promising crop regions and chooses among policies tailored to each test instance.Source: Figure 6 · p. 16, Figure 7 · p. 16, §A.6 · p. 16