Better Aggregation in Test-Time Augmentation

In plain words

To make a photo recognition more reliable, you can show the model several slightly changed copies of the same photo (cropped, flipped, zoomed) and average its answers. This is common, easy to add, and usually helps a bit.

The authors found that the plain average has a hidden cost. Some answers that were right on the original photo become wrong after the change, for example when a crop cuts out the one detail that identifies the object. Their fix is to let the model learn how much to trust each change instead of giving them all equal say. The learned weights are small in number, cost almost nothing to store, and can also show which changes are worth keeping.

TL;DR

Test-time augmentation (TTA) averages a classifier's outputs over transformed copies of one input. The paper shows that the average can correct many predictions but also corrupt a sizeable share of them, and that the cause depends on the model, dataset and augmentation set §3. It learns non-negative weights for the views (AugTTA, M weights) or for each view and class (ClassTTA, M×C weights) by minimising cross-entropy on a held-out split §4. Results are author-reported and mixed across settings tab. 1 tab. 2.

Problem

TTA usually improves accuracy overall, yet many individual predictions change from correct to wrong. In the paper's ImageNet analysis with a standard policy (flips, crops, scales), a little over one third of the changed labels for ResNet-18 are wrong §3.2 fig. 2. The paper's aim is to say when and why a prediction changes, and to use that to aggregate better than a simple average §1 fig. 1.

The learning task is stated as: given a black-box classifier that is not fully invariant to the augmentations, M augmentation functions, and a labelled validation set representative of the test domain, learn an aggregation function over the M predictions §4.

Method

Inputs §4:

1. A pretrained black-box classifier f that maps an image to a vector of class probabilities, and is not fully invariant to the augmentations §4.
2. M deterministic augmentation functions a_m that keep the class and change nuisance factors such as scale or colour balance §4.
3. A validation set of N labelled images that represents the test domain §4.

Aggregation function. The learned function g takes the M×C matrix of logits for one image and returns a C-vector, followed by softmax §4. Two restricted forms are used so that the cost stays low §4:

• ClassTTA (class-weighted): one weight per augmentation and class, Θ ∈ R^{M×C}, where ⊙ is element-wise eq. 1.
• AugTTA (augmentation-weighted): one weight per augmentation, Θ ∈ R^M eq. 2. The paper suggests AugTTA when labelled data is scarce §4.

g(A(x_i)) = \sum_{m=1}^{M} (\Theta \odot A(x_i))_{m,*} \qquad [\text{eq. 1}]

g(A(x_i)) = \Theta^{\top} A(x_i) \qquad [\text{eq. 2}]

Constraint and training §4 §5:

• Weights are constrained to be non-negative, so they can be read as the importance of each augmentation §4.
• Θ is learned by minimising cross-entropy between the true labels and the output of g using gradient descent §4.
• The paper states the optimiser as SGD with learning rate 0.01, momentum 0.9, weight decay 1e-4, 30 epochs, and projection by clipping to zero after each update §5.
• The variant to deploy (ClassTTA or AugTTA) is chosen per dataset on a small held-out part of the validation data §4 §5. The paper does not name the selection criterion.

# Learn aggregation weights on the 40% training split  [§4] [§5]
# Input:  outputs (N, M, C) logits for M views of N images
#         labels (N,)
# Output: theta (M, C) for ClassTTA, or (M, 1) for AugTTA

def train_weights(outputs, labels, epochs):
    # Start from uniform weights 1/M  [code: tta_agg_models.py#L21]
    theta = full((M, C), 1 / M)
    for epoch in range(epochs):
        for batch_outputs, batch_labels in minibatches(outputs, labels):
            # Weight each view per class and sum over views  [eq. 1]
            scores = sum(theta * batch_outputs, axis=1)  # (B, C)
            # Cross-entropy applies the softmax to the scores  [§4]
            loss = cross_entropy(scores, batch_labels)
            # SGD step, lr 0.01, momentum 0.9, weight decay 1e-4  [§5]
            theta = sgd_step(theta, loss)
            # Clip weights to be non-negative  [§5] [code: tta_train.py#L65]
            theta = maximum(theta, 0)
    return theta

# Aggregate the M views of each test image  [eq. 1] [eq. 2]
# Input:  outputs (B, M, C) logits, theta from training
# Output: scores (B, C); apply softmax to get probabilities

def aggregate_classtta(outputs, theta):
    # Per-class weights, not normalised  [code: tta_agg_models.py#L30]
    return sum(theta * outputs, axis=1)

def aggregate_augtta(outputs, theta):
    # Normalise the M weights to sum to one  [code: tta_agg_models.py#L54]
    weights = theta / sum(theta)  # (M, 1)
    return einsum("bmc,m->bc", outputs, weights)

Augmentation policies §3.1:

• Standard policy: 30 views per image, the cross product of 2 flips, 5 crops (centre plus four corners) and 3 scales (1, 1.04, 1.10) §3.1.
• Expanded policy: 128 deterministic transforms, made of 8 binary PIL operations and 12 continuous operations each at 10 evenly spaced magnitudes, based on AutoAugment §3.1. The code enumerates more transforms; see Paper vs code.

Analysis that motivates the weights (not part of the method, but the paper's justification) §3.3 §3.4:

• Crops favour smaller or off-centre objects, which causes errors on images with hierarchical or multiple labels, for example "plate" vs "guacamole" or "hook" vs "cleaver" §3.3 fig. 3.
• Zoom-in scales make similar classes, such as dog breeds, look more alike §3.3 fig. 3.
• On Flowers-102, crops help "Rose" (texture is kept) and hurt "Bougainvillea" (the white stamen is cut out) §3.4 fig. 4.
• Two classes with equal accuracy respond differently to TTA, which the paper links to within-class variation §3.4 fig. 5.

Hyperparameters

Results

Author-reported. Top-1 accuracy (%), mean ± standard deviation over 5 random subsamples of the test set. Setup: ImageNet and Flowers-102 with ResNet-18, ResNet-50, MobileNetV2 and InceptionV3 §3.1; CIFAR-100 with a 7-layer CNN and STL-10 with a 5-layer CNN §5. Baselines: Original (no TTA), Mean (average of logits), Max (maximum logit), GPS (greedy policy search, N = 3) §5. The paper marks the best value per row in bold; the bold markers are omitted here.

Standard policy tab. 1

Expanded policy tab. 2

Headline claims (author text, not recomputed):

• Standard policy: significantly better than all baselines (p = 2e-7) and than the original model in all 8 Flowers-102 and ImageNet comparisons (p = 7e-10) §5.1. Better than the other baselines in 42 of 50 trials §5.1.
• Expanded policy: better than GPS (p = 5e-7) and better in 34 of 50 trials §5.2.
• The learned weights for the expanded policy keep the cost down: for ResNet-50 on ImageNet, 37 of 128 augmentations get non-zero weight §5.4.
• Intro: "an increase in Top-1 accuracy of up to 2.5%" §1.

Paper vs code

The official repository is divyashan/test-time-augmentation. It is linked from the paper's source only as a commented-out short link, which did not resolve through WebFetch. The repository was found by GitHub search. Its description names this paper, and the owner's GitHub profile name is "Divya Shanmugam", the first author (checked via the GitHub API, not in the repository). The repository README calls itself the code submission for the paper README.md. The repository has no LICENSE file at its root, so code_license is "not stated". Only the vendored ttach package carries a license (MIT) LICENSE. No code is quoted here.

Discrepancies, each with both sides:

• Epochs for AugTTA. Paper: 30 epochs for both variants §5. Code: 30 for ClassTTA aggregation.py but 20 for AugTTA aggregation.py.
• Size of the expanded policy. Paper: 128 transforms, 8 binary and 12 continuous at 10 magnitudes §3.1. Code: a list of 20 PIL operations at 10 levels each, so 200 transforms augmentation_transforms.py augmentation_transforms_hpy.py. The paper's split into 8 binary and 12 continuous operations could not be mapped onto that list.
• Temperature. Paper: no temperature term. Code: the weights are trained with temperature 1 aggregation.py aggregation.py, but the AugTTA model is loaded with a temperature fitted on the 50% pool aggregation.py aggregation.py, and the fitted value is computed by tta_utils.py. At deployment this divides the logits by a value the weights were never trained with.
• Weight normalisation. Paper: non-negativity only §4. Code: AugTTA weights are divided by their sum in the forward pass, so they always sum to one tta_agg_models.py. ClassTTA is not normalised tta_agg_models.py.
• L1 penalty. The code computes an L1 term with coefficient 0.01 but adds only the cross-entropy to the loss tta_train.py tta_train.py. The paper mentions no L1 term.
• Model selection rule. Paper: a held-out part of the validation data, with no criterion stated §4. Code: top-1 accuracy on the 10% split, keeping ClassTTA only if it is strictly better aggregation.py.
• Training script. tta_train.py imports ImprovedLR and TTARegressionFrozen tta_train.py, which tta_agg_models.py does not define. The training module cannot be imported as published. The script also moves tensors to CUDA unconditionally tta_train.py.
• Input type. Paper: the classifier outputs class probabilities §4, but the aggregator takes logits and the baselines average logits §4 §5. Code aggregates raw model outputs aggregation.py.
• Text vs tables. The expanded-policy text says the method underperforms the Mean baseline on ImageNet with MobileNetV2 §5.2. Table 2 shows 72.57 vs 71.84, so that claim matches the standard table (72.43 vs 72.69) instead tab. 2 tab. 1. The text also says ClassTTA beats the original model on all datasets §5.2, while the CIFAR-100 row shows 73.18 vs 74.37 tab. 2. The intro's "up to 2.5%" §1 is larger than the largest gap in Table 1 (92.62 − 90.28 = 2.34) tab. 1.

Consistent with the paper: the GPS baseline uses the full 50% pool, both train and validation aggregation.py and tta_utils.py; GPS uses N = 3 aggregation.py; the optimiser settings match tta_train.py §5.

Pitfalls

• Logits or probabilities. The paper calls the classifier output probabilities §4 but aggregates logits §4. The code uses logits. Follow the code, and check the output type of any classifier you wrap.
• Calibration hidden in code. A temperature is fitted on the 50% pool and used for AugTTA at test time but not in training [see Paper vs code]. Either train with the same temperature, or leave it out.
• Batch size and steps. The paper gives no batch size. The code takes steps of 1000 samples, or one full-batch step for pools under 1000 images tta_train.py. The number of optimiser steps per epoch therefore depends on the pool size, so a given epoch count is not the same training budget across datasets.
• Learned weights vs test cost. Learning weights does not reduce the inference cost by itself. All M views are still computed unless views with zero weight are skipped. For ResNet-50 on ImageNet with the expanded policy, 37 of 128 augmentations get non-zero weight, so the rest can be skipped §5.4.
• Class-level weights need data. The paper argues ClassTTA should only help when there are few classes and enough labelled data per class §5.1. It attributes AugTTA's choice on ImageNet to the large class count and few examples per class (25) §5.1.
• Statistics. The error bars come from random subsamples of a fixed test set, not from retraining with different seeds §5. The t-tests are over those subsamples, so they do not show training-seed variance.
• Test set and split. The 40/10/50 split is made from each released test set with a fixed random seed tta_utils.py. Results are not comparable with papers that use the full validation set.
• Views that move predictions to another class. The paper's analysis shows that crops and zooms can push a prediction to a different class §3.3. Check the learned weights for your own views §5.3 fig. 6.
• Appendix not available. The paper refers to an appendix for the augmentation details and top-5 tables §3.1 §5.1. The LaTeX source contains an empty supplement file, so those details are "not stated" here.

Lineage

• Compared against: Mean TTA §5, Max logit §5, and Greedy Policy Search (GPS, N = 3) §5. None has a card in this wiki.
• Builds on: the augmentation policies follow the common flips, crops and scales practice §1, and the expanded set is based on AutoAugment §3.1. Neither has a card in this wiki.
• Superseded by: not known.