Source-linked AI summary
Better Aggregation in Test-Time Augmentation
Divya Shanmugam, Davis Blalock, Guha Balakrishnan, John Guttag
TL;DR
Test-time augmentation is widely used, but simple averaging can change correct predictions into incorrect ones even when net accuracy improves. This paper analyzes those changes and learns aggregation weights for transformed predictions, achieving consistent gains across models, datasets, and augmentation policies. The approach is lightweight, though expanded augmentation sets increase inference cost without clear accuracy justification.
Problem
Simple averaging is widely used for TTA, yet its prediction changes and cases where non-uniform aggregation is preferable require further analysis.
Method
The paper analyzes TTA-induced prediction changes and learns nonnegative weights for aggregating augmentation predictions for a given model and dataset.
Results
The learned method consistently outperforms existing TTA approaches across 4 datasets and 6 models, with Top-1 accuracy gains of up to 2.5%.
Takeaways & Limitations
Learned aggregation can serve as a lightweight replacement for simple averaging and can help smaller models match larger-model performance.
Takeaways & Limitations
Expanded augmentation sets increase inference cost, while ClassTTA is best suited to datasets with few classes and sufficient labeled data.
Abstract
from arXiv · showhide
Test-time augmentation -- the aggregation of predictions across transformed versions of a test input -- is a common practice in image classification. Traditionally, predictions are combined using a simple average. In this paper, we present 1) experimental analyses that shed light on cases in which the simple average is suboptimal and 2) a method to address these shortcomings. A key finding is that even when test-time augmentation produces a net improvement in accuracy, it can change many correct predictions into incorrect predictions. We delve into when and why test-time augmentation changes a prediction from being correct to incorrect and vice versa. Building on these insights, we present a learning-based method for aggregating test-time augmentations. Experiments across a diverse set of models, datasets, and augmentations show that our method delivers consistent improvements over existing approaches.
1. Introduction
Test-time augmentation pools predictions from transformed inputs, but simple averaging can improve net accuracy while also corrupting many correct predictions. This paper analyzes these changes and introduces a lightweight learned aggregation method with broad accuracy gains.
- Motivation: Test-time augmentation pools predictions from transformed versions of an input to improve robustness, accuracy, or uncertainty estimation.A common example averages predictions from multiple crops to reduce sensitivity to an unfavorable crop.
- Motivation: TTA is easy to deploy without changing the underlying model or requiring additional data, but its augmentation and aggregation choices are under-studied.The paper focuses on how to aggregate predictions rather than which augmentations to select.
- Motivation: Positive net accuracy gains can coexist with a sizeable number of predictions changed from correct to incorrect under standard TTA.Figure 1 distinguishes corrected, corrupted, and net-changed predictions, while prior work typically emphasizes only net improvement.
- Contributions: The paper analyzes when and why TTA changes predictions and develops a method that learns how to aggregate augmentation predictions for each model and dataset.The proposed aggregation learns weights per augmentation and is designed to work alongside other methods.
- Contributions: The learned aggregator is a lightweight replacement for simple averaging, offering up to 2.5% Top-1 accuracy improvement with little model-size, training-time, or implementation burden.The paper also reports consistent gains across architectures, datasets, and augmentation policies, and a smaller model combined with TTA can match larger-model performance.
- Contributions: Experiments on ImageNet and Flowers-102 provide practical insights into why augmentations change predictions from correct to incorrect and vice versa.These analyses motivate the design of the learned aggregation method.
2. Related Work
Prior accuracy-oriented TTA work commonly averages predictions or selects augmentations, while this paper focuses on learning the aggregation function itself. Its novelty is analyzing factors behind TTA changes and using them to learn aggregation weights.
- Scope: The paper focuses on TTA for increasing image-classification accuracy rather than its uses for uncertainty estimation, attacks, or defenses.This scope distinguishes the reviewed related work from other test-time augmentation applications.
- Prior TTA methods: Accuracy-oriented TTA research commonly presents test-time augmentation as a supplemental detail, often averaging predictions from crops, flips, or rescaled inputs.Examples include averaging softmax predictions over five patches and their reflections, with later variants adding rescaling.
- Prior TTA methods: Related methods include random augmentation sampling with averaged log class probabilities and greedy selection of useful augmentations.These approaches differ from directly optimizing the aggregation function over augmentation predictions.
- Positioning: This paper differs by studying factors that influence TTA and learning how to aggregate augmentation predictions given those factors.The focus is aggregation rather than selecting which augmentations to include.
3. Why weight augmentations differently?
Simple averaging can be suboptimal because augmentation effects depend on the label space, class-specific visual variation, and the chosen transformation set. Learning aggregation weights addresses these biases while preserving TTA’s lightweight deployment advantages.
- Motivation: TTA commonly averages predictions, but non-uniform weights can be preferable when augmentation errors vary across datasets and inputs.The analysis uses ImageNet and Flowers-102 to examine when simple averaging introduces errors.
- Setup: Standard and Expanded policies replace original predictions with averages over 30 and 128 transformed images, respectively.Standard combines flips, crops, and scales; Expanded adds binary and continuous intensity transformations.
- Overall results: TTA usually improves accuracy, but incorrect predictions remain a substantial share of changed labels; for ImageNet ResNet-18, standard TTA corrupts slightly over one third of changed labels.The magnitude of improvement varies across datasets and models.
- Class-Dependent Invariances: Crops create label-space biases involving hierarchical, multiple-class, and similar-label images, often favoring smaller or uncentered objects.These biases can either correct or corrupt predictions depending on the assigned label and the visual evidence emphasized by crops or scales.
- Class-Dependent Invariances: Flowers-102 shows class-dependent invariances: crops help roses by preserving petal texture but harm bougainvilleas when they exclude the white stamen or emphasize background.Differences in within-class consistency also help explain why TTA improves Primula but not Sword Lily.
4. Method
The method learns how to aggregate predictions from multiple augmented inputs using trainable, nonnegative weights, with either class-specific or augmentation-specific parameterizations.
- Inputs: The method takes a pretrained classifier, a set of M deterministic augmentations, and a representative labeled validation set as inputs.The classifier outputs class probabilities, augmentations preserve class-relevant information while modifying presumed class-independent variables, and the validation set supplies labels.
- Aggregation function: The aggregation function maps the matrix of predictions from M augmented images and C classes to one C-dimensional prediction.A softmax layer converts the aggregated output into class probabilities.
- Design goal: The learned aggregation sums weighted predictions over augmentations to produce the final prediction while avoiding significant model size or latency increases.The weighting scheme is designed to preserve interpretability through nonnegative parameters.
- Parameterizations: ClassTTA and AugTTA learn nonnegative weights for each augmentation-class pair or for each augmentation, respectively.The class-specific parameterization uses a trainable parameter matrix, while the reduced parameterization uses one weight per augmentation.
- Optimization: The method learns the weights by minimizing cross-entropy between the true labels and aggregated predictions using gradient descent.A small held-out portion of the validation set selects between ClassTTA and AugTTA, alongside evaluating the individual parameterizations.
5. Experimental Evaluation
Across datasets, architectures, and augmentation policies, the learned aggregation method consistently improves over standard baselines, while revealing dataset- and class-dependent weighting patterns and computational tradeoffs.
- Results: The method significantly outperforms all baselines and the original model, winning 42 of 50 individual trials.The reported significance levels are p-value=2e-7 versus all baselines and p-value=7e-10 versus the original model.
- Results: On STL-10, the method performs comparably to Mean while identifying augmentations that can be omitted to save computation.The remaining augmentations receive weights roughly equivalent to averaging, so performance does not significantly improve.
- Analysis: The method consistently selects ClassTTA for Flowers-102 and AugTTA for ImageNet, while both parameterizations learn similar weights on STL-10 and CIFAR-100.The authors attribute the ImageNet choice to its many classes and relatively few examples per class, and interpret the STL-10/CIFAR-100 result as evidence against strong class-dependent relationships for standard TTAs.
- Analysis: TTA improvements correlate with the number of examples per class, with r=.95 and p-value=.04, and TTA can let smaller models outperform larger models without TTA.ClassTTA on MobileNetV2 with ∼3.4 million parameters outperforms the original ResNet-50 with ∼23 million parameters on Flowers-102.
- Computational Cost: Expanded TTA increases inference cost, but learned weights can reduce computation by selecting only 37 of 128 augmentations on ImageNet and 20 on Flowers-102.These correspond to 28% and 16% of the expanded policy, respectively.
- Expanded TTA Policy: In expanded policies, the method favors crop, translation, and blur while assigning zero weights to contrast, cut-out, shearing, and brightness augmentations.Across ImageNet architectures, augmentations with 10% scale are excluded, and higher scale parameters are generally weighted lower.
6. Discussion
The discussion frames learned TTA aggregation as both an explanation of model sensitivity to transformations and a basis for targeted augmentation strategies.
- Discussion: The method outperforms existing TTA approaches across 4 datasets and 6 models while highlighting useful augmentations outside flips, crops, and scales.The broader analysis focuses on when TTA changes model decisions and how weighting those transformations can improve performance.
- Future Work: TTA can identify when models lack invariance to particular transforms, informing targeted train-time augmentation policies.The paper suggests using successful test-time transformations to guide class-specific transforms during training.
- Future Work: Learning augmentation weights is one possible extension; another is learning a broader set of augmentations rather than restricting the method to common transformations.The discussion presents learned augmentations as a direction for future work.