Test-time training

Overview

A model fixed at deployment can fail on inputs that differ from its training data, such as blurred, noisy or compressed photos. Test-time training (TTT) gives each incoming input its own short adaptation before the prediction is made. The adaptation uses a loss whose signal comes from the input itself, so no ground truth is needed 2209.07522 §1. The same per-input, unlabelled setting is the one the prompt-tuning work studies, where a model sees one unlabelled test image at a time 2308.06038 §1 2308.06038 §3.1.

The invariant shared by the variants is that the adaptation is per input and temporary: the model starts from the same pretrained state for every input, takes a few gradient steps on a loss computed from that input alone, predicts, and then discards the adapted state 2209.07522 §3 [arxiv:2308.06038 §3.3]. What differs is the loss, which parameters move, and the optimiser settings. The masked-autoencoder variant reconstructs hidden patches and moves the whole encoder and decoder. The prompt-tuning variant minimises prediction entropy over augmented views and moves only four context tokens 2308.06038 §4.1.

The two cards that cover this concept are Test-Time Training with Masked Autoencoders and Diverse Data Augmentation with Diffusions for Effective Test-time Prompt Tuning. The second is a neighbour rather than a pure test-time-training method: its loss is an unlabelled objective, not a pretext task. Its predecessor, the rotation-prediction TTT of Sun et al., and the prompt-tuning method it builds on, TPT (Shu et al., 2022), have no card here, so they are described only as the papers' baselines 2209.07522 §2 2308.06038 Lineage.

Variants

Rotation-prediction TTT (predecessor, described in the paper)

The earlier method attaches a rotation-prediction head to a shared network and adapts on the rotation task at test time 2209.07522 §2. The paper argues that rotation is a poor self-supervised task for some inputs: outdoor scenes often reveal orientation through the horizon, so the task is too easy, while top-down views make most rotations look alike, so it is too hard 2209.07522 §2. The paper also reports that the rotation task lowers accuracy on rotation-invariant classes 2209.07522 tab. 1. Only the paper's account of this method is available here, not a primary source.

Masked-autoencoder TTT (TTT-MAE)

Keep the shared-encoder layout of the earlier method and replace the rotation head with a masked-autoencoder decoder. Per test input, the encoder and decoder are reset to their pretrained weights, a random 75% of patches are hidden, and the decoder must redraw the hidden patches. The loss is taken on those patches 2209.07522 §3 2209.07522 §4.1. The paper reports 20 SGD steps with momentum 0.9 and a fixed learning rate of 5e-3 at batch 128 [arxiv:2209.07522 §4.1]. The main-task head stays frozen during the test-time steps and is trained beforehand with the encoder frozen (ViT probing) 2209.07522 §3.

Reported gains (author-reported): on ImageNet-C at severity 5, the method beats the same model without test-time training on all 15 corruptions [arxiv:2209.07522 tab. 2]. The paper also reports gains on ImageNet-A and ImageNet-R [arxiv:2209.07522 tab. 3] and on a portrait classification task 2209.07522 tab. 4.

Diffusion-view test-time prompt tuning (DiffTPT)

Keep the image model frozen and tune a text prompt instead. The prompt is a set of four learnable context tokens on CLIP, initialised from "a photo of a" 2308.06038 §4.1. For each test image the method builds 128 views: the original, 64 crop-and-flip views, and 63 views from Stable Diffusion V2 that are conditioned on the image embedding rather than a label 2308.06038 §3.3 2308.06038 §4.1. Two filters follow. A cosine filter keeps the 80% of generated views most similar to the original, then an entropy filter keeps the 30% most confident views 2308.06038 §3.4 2308.06038 §4.1. The loss is the entropy of the mean prediction over the kept views [arxiv:2308.06038 eq. 4] 2308.06038 eq. 7. The paper runs four optimisation steps at learning rate 5e-3 2308.06038 §4.1, and its ablation finds that going from one to four steps helps on cross-dataset generalisation, with further steps no better 2308.06038 §4.3 2308.06038 fig. 7.

Reported gains (author-reported): on the printed averages, the method beats TPT by about 2.1 points on the five-dataset natural-shift average with RN50 and by 0.7 with ViT-B/16 2308.06038 tab. 1. On the ten-dataset cross-dataset average the gaps are 2.9 (RN50) and 1.5 (ViT-B/16) 2308.06038 tab. 2. These differences are computed here from the printed averages; the paper's headline figure is not reproduced by them (see Contradictions).

When to use which

The table compares the test-time options against not adapting at all, and against each other where both apply.

Implementation sketch

Language-agnostic pseudocode for the masked-autoencoder variant. Steps are the paper's; the batching of copies is a detail of the released code and is noted where it matters 2209.07522 §4.1. The prompt-tuning variant has the same outer loop with a different loss and parameter set (see the notes).

# Test-time training for one input  [arxiv:2209.07522 §3]
# Inputs: pretrained encoder0, decoder0, head0; one test image x

def predict(x, encoder0, decoder0, head0):
    # Start from the pretrained weights every time  [§3]
    enc, dec = copy(encoder0), copy(decoder0)
    for step in range(20):                    # 20 SGD steps  [§4.1]
        # Hide a random 75% of patches  [§4.1]
        masked = random_mask(x, ratio=0.75)
        # Reconstruct only the hidden patches  [§3]
        loss = mse_on_hidden(dec(enc(masked)), x)
        enc, dec = sgd_momentum_step(enc, dec, loss, lr=5e-3)  # [§4.1]
    # Predict on the unmasked input; keep head0 fixed  [§3]
    return argmax(head0(enc(x)))
    # Adapted enc and dec are dropped here  [§3]

Notes the paper leaves open or where the released code differs:

• Batching: the paper's setup averages the update over many copies of one image with different masks 2209.07522 §4.1. The exact copy count and how the learning rate scales with batch size are implementation details that the card records from the code, not from the paper.
• Which parameters move: the paper says encoder and decoder both adapt 2209.07522 §4.1. The appendix reports that freezing the decoder makes little difference 2209.07522 app.. Choose one and state it.
• Prompt-tuning variant: the parameters are the four context tokens, the loss is the entropy of the mean prediction over the kept views [arxiv:2308.06038 eq. 7], and the paper's objective applies the selection in every step. The released code fixes the kept set at the first step, which the DiffTPT card records as a difference.
• Prompt-tuning optimiser: the paper says Adam 2308.06038 §4.1. The card records that the released code uses AdamW with no weight decay, so an implementer must pick one.

Contradictions

Across cards there are no conflicting results, since the two papers use different backbones, benchmarks and metrics. Their numbers are not comparable. TTT-MAE evaluates a single MAE ViT on ImageNet-C and related sets 2209.07522 tab. 2. DiffTPT evaluates CLIP on ImageNet variants and fine-grained datasets 2308.06038 tab. 1 2308.06038 tab. 2.

Within each paper, points differ between text, tables and the released code. The cards record these; this page does not resolve them:

• TTT-MAE weight decay: the paper gives 0.2 for test-time training [arxiv:2209.07522 §4.1]. The card's code check finds the test-time optimiser does not apply it.
• TTT-MAE test-time steps: the main results use 20 steps, and the appendix's per-severity tables use 10 2209.07522 §4.1 2209.07522 app.. Compare numbers only at the same step count.
• DiffTPT headline gain: the abstract and conclusion state an average gain of 5.13% over TPT. The paper gives 5.1% (RN50) and 2.3% (ViT-B/16) for the cross-dataset scenario in §4.2, while its printed averages give 2.9 and 1.5 points over TPT 2308.06038 §4.2 2308.06038 tab. 2. The paper does not say how 5.13% is computed, so treat it as unverified.
• DiffTPT printed average: the RN50 CLIP average in tab. 1 is printed as 42.26, but its five per-dataset values average to 44.26 2308.06038 tab. 1. Recompute averages before comparing.

Papers

• Test-Time Training with Masked Autoencoders: replaces the rotation task of test-time training with masked-patch reconstruction and reports gains on ImageNet-C level 5 (author-reported).
• Diverse Data Augmentation with Diffusions for Effective Test-time Prompt Tuning: tunes CLIP prompts per test image on entropy over cosine-filtered diffusion views, with author-reported gains over TPT on natural-shift and cross-dataset averages.