Test-time augmentation

Overview

Test-time augmentation runs a trained classifier on several transformed copies of one input (flips, crops, scales, colour changes) and combines the outputs into one prediction. The usual combination is the mean of the outputs. The invariant the method relies on is that each transform keeps the label valid, so the copies should agree on the class 2011.11156 §4.

The paper's starting point is that the average helps overall but still changes many individual predictions from correct to wrong. For ImageNet with a standard set of flips, crops and scales, a little over one third of the changed labels for ResNet-18 were wrong 2011.11156 §3.2. The authors trace this to specific transforms: crops can cut out the detail that identifies a class, and zoom-in can make similar classes look more alike 2011.11156 §3.3. Their response is to learn how much each view contributes to the final score 2011.11156 §4. The learned weights need a labelled held-out split from the target domain 2011.11156 §4.

The same averaging appears outside classification. A second-stage filter for CT detections crops each candidate into many random views and averages the network's probabilities 1505.03046 §II-C 1505.03046 eq. 1. There the views are random in scale, position and rotation rather than a fixed list of flips and crops.

Variants

Mean of outputs (baseline)

Average the logits of all views with equal weight. This needs no training and is one of the paper's baselines 2011.11156 §4 2011.11156 §5. Card: Better Aggregation in Test-Time Augmentation.

Max logit (baseline)

Take the maximum over views for each class. In every row of both results tables it is below the mean 2011.11156 tab. 1 2011.11156 tab. 2.

Greedy policy search, GPS (baseline)

A baseline the paper calls greedy policy search, run with N = 3 policies 2011.11156 §5. The paper's description of its procedure is not reproduced here; see the card for the paper's details.

AugTTA: one weight per view

A learned, non-negative weight for each of the M augmentations, shared across classes. The paper suggests it when labelled data is scarce 2011.11156 §4. The paper states non-negativity as the only constraint 2011.11156 §4.

ClassTTA: one weight per view and class

A learned, non-negative weight for each augmentation and each class, giving M×C weights in total 2011.11156 §4. It has more parameters than AugTTA and needs more labelled data per class to fit them 2011.11156 §5.1.

Random-view averaging (CT detection, per candidate)

Not learned. For each candidate from an existing detector, draw many random views: several physical scales, random shifts of up to 3 mm and random rotations. A ConvNet scores each view, and the candidate's score is the mean of the probabilities 1505.03046 §II-C 1505.03046 eq. 1. The paper uses N = 100 views per candidate for bone and lymph nodes and N = 40 for polyps 1505.03046 §III-A. Card: Improving Computer-aided Detection using Convolutional Neural Networks and Random View Aggregation.

The paper's comparisons are against the existing detector for each dataset, not against other ways of combining views 1505.03046 tab. II. Reported sensitivity at 3 false positives per patient rises from 57% to 70% for bone metastases, 43% to 77% for lymph nodes, and 58% to 75% for colonic polyps 1505.03046 tab. II. The lymph-node headline mixes subsets, so use the separate mediastinum and abdomen figures 1505.03046 §III-E. The bone candidate generator reaches at most 88.9% sensitivity, and the paper does not say whether its figures are relative to that ceiling 1505.03046 §III-D. Sensitivity saturates quickly with N for lymph nodes and bone, and the polyp curves level off after about 10 views 1505.03046 §III-D 1505.03046 §III-E 1505.03046 §III-H.

When to use which

The paper chooses between AugTTA and ClassTTA per dataset on a small held-out part of the labelled data, and does not name the selection criterion 2011.11156 §4 2011.11156 §5. Results are author-reported. On the standard policy the learned weights beat the original model in every listed comparison 2011.11156 tab. 1. On the expanded policy, the learned-weights column for CIFAR-100 with CNN-7 is below both the original model and the mean (73.18 against 74.37 and 73.41) 2011.11156 tab. 2. The two papers use different tasks, baselines and metrics, so their numbers are not comparable. Check the gain on your own model and data before adopting a variant.

Implementation sketch

Language-agnostic outline of the common core. The two learned variants differ only in the shape of theta. Shapes: N images, M views, C classes.

# Collect logits for M views of each of N labelled validation images
# [arxiv:2011.11156 §4]
# Input:  outputs (N, M, C), labels (N,)
# Output: theta, shape (M, C) for ClassTTA or (M, 1) for AugTTA

def fit_weights(outputs, labels, epochs):
    # Start from uniform weights  [arxiv:2011.11156 §4]
    theta = full(shape, 1 / M)
    for epoch in range(epochs):
        for batch in minibatches(outputs, labels):
            # Weighted sum over views, one score per class  [arxiv:2011.11156 eq. 1]
            scores = sum(theta * batch.outputs, axis=view_axis)
            # Cross-entropy on the aggregated scores  [arxiv:2011.11156 §4]
            loss = cross_entropy(scores, batch.labels)
            # Gradient step: SGD, lr 0.01, momentum 0.9, weight decay 1e-4
            # over 30 epochs  [arxiv:2011.11156 §5]
            theta = sgd_step(theta, loss)
            # Keep weights non-negative by clipping  [arxiv:2011.11156 §4] [arxiv:2011.11156 §5]
            theta = maximum(theta, 0)
    return theta

def aggregate(outputs, theta):
    # Apply the learned weights to the M views of a test image  [arxiv:2011.11156 eq. 2]
    return sum(theta * outputs, axis=view_axis)  # (C,) scores, then softmax

Notes on the sketch:

• The paper calls the classifier output probabilities but aggregates logits 2011.11156 §4. Use the logits and check the output type of any classifier you wrap.
• The paper gives no batch size and no weight-initialisation rule 2011.11156 §5. The sketch uses uniform initialisation as a placeholder.
• The paper does not describe any temperature scaling or any extra penalty term 2011.11156 §4 2011.11156 §5. The card records several places where the released code differs from the paper, including the epoch count for AugTTA and a fitted temperature. Read the card before implementing Better Aggregation in Test-Time Augmentation.
• Learned weights reduce the number of views you need only if you skip the views with zero weight. All M views are still computed by default 2011.11156 §5.4.
• Random-view averaging has no fitting step. It draws new views for each candidate and takes the mean of the probabilities, not of the logits 1505.03046 §II-C 1505.03046 eq. 1. The card lists the released sampling details, which differ from the paper's text Improving Computer-aided Detection using Convolutional Neural Networks and Random View Aggregation.

Contradictions

• Gain size in the introduction. The introduction claims an increase of up to 2.5% in top-1 accuracy 2011.11156 §1. Table 1 shows ImageNet InceptionV3 going from 69.66 (original) to 72.79 (learned weights), a gain of 3.13 points 2011.11156 tab. 1. The two statements do not match.
• Probabilities or logits. The method section describes the classifier as returning probabilities, but the aggregation and baselines operate on logits 2011.11156 §4 2011.11156 §5. The paper does not say which is intended.
• Number of views (within one paper). The paper gives N = 100 in total for bone and lymph nodes, which is 25 per scale 1505.03046 §III-A. It also says the views are generated "N = 100 times ... at each scale" 1505.03046 §III-E. The two readings differ by a factor of four.

Papers

• Improving Computer-aided Detection using Convolutional Neural Networks and Random View Aggregation: averages a ConvNet's probabilities over random scale, shift and rotation views of each CT candidate, as a second stage after an existing detector, and reports sensitivity gains on bone, lymph-node and polyp sets.
• Better Aggregation in Test-Time Augmentation: learns non-negative weights for augmented views (AugTTA) or view-class pairs (ClassTTA) on a labelled split, and reports gains over mean, max and GPS on most ImageNet and Flowers-102 models.