Test-Time Training with Masked Autoencoders

In plain words

A model trained on clean photos often gets worse on photos that look different: blurry, noisy, foggy or compressed. The usual fix is to train on many kinds of damaged photos in advance, but nobody can list every damage in advance.

This paper takes a different route. When a new photo arrives, the model gets a short private training session on that one photo before it answers. The training task needs no labels: hide most of the photo and ask the model to redraw the hidden parts. Doing that forces the model to learn how the picture is built, and that adjusted model is then used for the classification. After the answer is given, the model is reset to its original state and the next photo gets its own session.

The authors use a masked autoencoder for the hide-and-redraw task, because it works well on images and needs no labels. They report that this simple recipe helps on several benchmarks of shifted images.

TL;DR

Test-time training (TTT) with a masked autoencoder (MAE) instead of the rotation task of the original TTT method. Per test image: start from the pretrained weights, run 20 SGD steps that reconstruct masked patches of that image, predict on the unmasked image, then discard the adapted weights. The main evidence is ImageNet-C level 5, where every corruption improves over a fixed ViT-probing baseline. Results are author-reported; the method runs 20 SGD steps per test image. Relevant if you need robustness to unseen shifts and can afford per-input adaptation.

Problem

Models are usually fixed at deployment, so they must be robust to every shift that might occur §1. The paper argues that adapting to each test input is an alternative: each input is a one-sample learning problem where the label is missing but a self-supervised label can be made from the input itself §1.

The earlier TTT method §2 uses rotation prediction as its self-supervised task. The paper argues that rotation is sometimes too easy (outdoor scenes, where the horizon gives the answer) or too hard (top-down views where every rotation looks plausible) §2, and that on rotation-invariant classes it hurts accuracy tab. 1. The assumption behind the method is that the shift is unknown in advance and that the test input is the only information about it §1, §5.

Method

The method replaces the self-supervised head of the Y-shaped TTT network with MAE reconstruction §3. The network has one encoder f (the MAE encoder) that feeds two heads: a reconstruction decoder g (the MAE decoder) and a main-task head h for classification. Encoder and decoder are both ViTs, and the paper keeps them unchanged so the comparison with MAE stays clean §3, Architecture.

Training happens in three steps §3, Training setup:

1. Pretrain the MAE on ImageNet-1k (encoder f0, decoder g0); this is taken from the MAE authors' checkpoint §3, Training setup §4.1.
2. Train only the main head h0 on the labelled ImageNet training set while the encoder stays frozen ("ViT probing"), with a ViT-Base head instead of the linear head used by MAE eq. 1 §3.
3. Keep f0, g0 and h0 as the starting point for test time.

The paper compares fine-tuning, probing and joint training of the main head and selects ViT probing as default: fine-tuning makes the encoder specialise for recognition, which the test-time step then undoes; joint training performed worse on ImageNet validation §3, Training setup tab. 2.

At test time, each input x is handled independently §3, Test-time training:

1. Reset the encoder and decoder to f0 and g0.
2. Minimise the reconstruction loss of x alone, with a random mask per copy eq. 2 §4.1.
3. Predict with the adapted encoder and the fixed head h0, on the unmasked input §3.
4. Discard the adapted weights §3.

The reconstruction loss is the pixel-wise mean squared error between decoded and original patches §3, eq. 2. Only the masked patches enter the loss in the code models_mae_shared.py; the paper's equation does not spell that out.

# TTT-MAE for one test image  [§3] [eq. 2]
# Input:  crop of shape (224, 224, 3); pretrained encoder0, decoder0, head0
# Output: predicted class

def predict(crop, encoder0, decoder0, head0):
    # Start from the pretrained weights; discard them after this input  [§3]
    encoder = copy(encoder0)
    decoder = copy(decoder0)
    # Batch of 128 identical copies of the crop  [§4.1]
    # [code: tt_image_folder.py#L37-L39]
    batch = repeat(crop, 128)                 # (128, 224, 224, 3)
    for step in range(20):
        # Each copy gets its own random 75% patch mask  [§4.1]
        # [code: models_mae_shared.py#L155-L162]
        masked = random_mask(batch, ratio=0.75)
        # Reconstruct the removed patches; loss on those patches only  [§3]
        # [eq. 2]
        loss = mse_on_removed(decoder(encoder(masked)), batch)
        # SGD with momentum 0.9 and fixed lr 5e-3 at batch 128  [§4.1]
        # Paper also states weight decay 0.2; the code does not apply it
        # [code: engine_test_time.py#L64-L65]
        encoder, decoder = sgd_step(encoder, decoder, loss)
    # Predict on the unmasked crop with the adapted encoder and the fixed head  [§3]
    features = encoder(crop)                  # (1, 768)
    return argmax(head0(features))

Parameter scope: the paper updates encoder and decoder, plus the class token and mask token §4.1. The released command sets --finetune_mode encoder, which freezes the decoder engine_test_time.py README.md. The appendix reports the difference as negligible app. tab:encoder_decoder.

Method details used in the experiments §4.1:

• Only non-masked tokens go to the encoder during test time; the decoder receives all tokens §4.1.
• Test-time augmentation: none beyond random masking. Every TTT step uses the same 224×224 centre crop used for the prediction §4.1.
• Pixel targets with or without per-patch normalisation; normalised targets are slightly better and are the default in the released command §4.1 [app. tab:pix_norm] README.md.

Figures in words: fig. 1 shows one test image being reconstructed over 500 SGD steps, with the recognition loss still falling while reconstruction continues fig. 1. fig. 2 compares AdamW and SGD for test-time steps: AdamW improves then needs early stopping, while SGD keeps improving over 20 steps; the paper therefore uses SGD fig. 2 §3, Optimizer for TTT.

Theory §5: in a linear PCA setting with known covariance and an orthogonal corruption, the paper shows that a test-time update that mixes the training covariance with the single test sample strictly reduces expected risk for some mixing weight above zero (Theorem 1, proof in the appendix). This is a result about linear PCA, not about the MAE used in the experiments.

Hyperparameters

Test-time training §4.1 main_test_time_training.py README.md. Values in the README command are what the released code runs.

Results

Author-reported. Setup: ImageNet-C, severity level 5, 15 corruptions, each with 50,000 images §4.2 engine_test_time.py. Model: MAE ViT-L encoder, ViT-B head, ViT probing trained on the original ImageNet training set, no augmentation from the corruptions §4.1 §4.2. Baseline: the same model with no test-time training §4.2.

Main ImageNet-C numbers. fig. 3 is the bar graph; the table below gives the same numbers as the paper's own table tab. 2.

Source: tab. 2. Accuracy in percent, ImageNet-C level 5, 20 TTT steps.

Other reported numbers (author-reported, same hyperparameters as above):

• Level 5 is the main result. The appendix gives all five severities with 10 steps app.. Most entries improve, but not all: at level 2 elastic falls from 51.0 to 49.7 and at level 4 brightness from 73.1 to 72.7 [app. tab:acc_2] app. tab:acc_4. The level-5 10-step numbers sit below the 20-step main table app. tab:acc_5.
• Rotation-invariant classes (5 classes chosen with a ResNet-18 filter [§4.2 footnote]): TTT with rotation hurts on these classes, while TTT-MAE still helps on the four corruptions shown tab. 1.
• ImageNet-A and ImageNet-R, with the same hyperparameters as ImageNet-C: ImageNet-A accuracy 15.3 (baseline) → 21.3 (TTT-MAE); ImageNet-R 31.3 → 38.9 tab. 3.
• Portraits (gender classification, trained on three year-splits of a yearbook dataset and tested on the fourth; the paper reports accuracy per training split): baseline 76.1, 76.5, 78.2 → TTT-MAE 76.4, 76.7, 79.4 tab. 4 §4.4.

Paper vs code

Official repository, linked from the project page [https://yossigandelsman.github.io/ttt_mae/index.html] and the paper's own README. The GitHub owner name matches an author (Yossi Gandelsman); the project page does not state the owner explicitly. Cloned at code_commit. Code paths below are relative to the repository root.

Discrepancies, each with both citations:

• Weight decay 0.2 is stated for test-time training §4.1. The test-time SGD optimizer is built without a weight-decay argument [code: engine_test_time.py#L64-L65]. The flag exists with default 0.05 but is not used there main_test_time_training.py.
• The paper says encoder and decoder are both optimised §4.1. The released command passes --finetune_mode encoder, which freezes the decoder [code: README.md] engine_test_time.py.
• The paper's main results use 20 steps §4.1. The appendix tables use 10 steps and say so app.. The README uses 20 steps README.md.
• Probing for the main head: the paper says ViT-Base head, probing on the training set, but the epochs and learning rate are not stated §4.1. The README gives 20 epochs, base lr 1e-3 and weight decay 0.2 and loads a checkpoint described as using "slightly different parameters" [code: README.md]. The probing optimiser has an option between LARS and AdamW [code: main_prob.py#L288-L291].
• Learning rate: the paper states 5e-3 fixed at batch 128 §4.1. The code derives lr from a base rate and the batch size, so a different batch changes the step size main_test_time_training.py.
• Reconstruction loss: the paper's equation averages over all patches [§3, eq. 2]. The code averages over the removed patches only [code: models_mae_shared.py#L265-L282].
• No gradient clipping is set in the released command [code: main_test_time_training.py] misc.py.
• The main head and the ViT-Base head share one name in the paper's figures: "Joint Train" in fig. 3 is the ResNet-16 TTT-Rot baseline from the prior paper, while "Joint Train" in tab. 2 is this paper's ViT joint training. Do not mix them.

Other checks: the paper's ViT-Base dimensions (12 layers, 768 width, 12 heads) match engine_test_time.py. Masking ratio 75% and centre-crop handling match main_test_time_training.py [code: main_test_time_training.py#L171-L176]. Per-sample random masks match [code: models_mae_shared.py#L155-L162]. Reset to the pretrained weights before every test input matches engine_test_time.py.

Licence of the code: no LICENSE file was found in the cloned repository's root or in the README, so code_license is "not stated". No code is quoted here.

Pitfalls

• Per test image the code runs 20 steps × 128 copies of the crop through the encoder and decoder; the paper itself says test time is slower than the fixed baseline and gives no per-image wall-clock figure [Limitations] [code: engine_test_time.py#L121]. Budget for that, not for the single forward pass of the baseline.
• Weight decay 0.2 is in the text but absent from the test-time optimiser engine_test_time.py. An implementer who follows the text will use a different optimiser from the one that produced the numbers.
• Freezing vs training the decoder: the appendix measures the gap as small app. tab:encoder_decoder, but the text says both are trained §4.1. Pick one and state it.
• Learning rate depends on batch size in the code [code: main_test_time_training.py#L195-L197]. The paper gives only the product result (5e-3 at batch 128).
• Steps: 20 in the main table, 10 in the appendix. Compare like with like.
• Early stopping is what the paper avoids: with AdamW, accuracy peaks and then falls, which a test-time setting cannot tune for §3, Optimizer for TTT fig. 2. The SGD choice rests on this single comparison.
• Test images are not paired with labels, so the 20 SGD steps cannot be checked against accuracy during deployment. The paper's accuracy curves use labels only in evaluation §4.1 engine_test_time.py.
• Normalised pixel targets: the choice changes the baseline too, so the gain over baseline is what matters, not the absolute number app. tab:pix_norm.
• The paper says the corruptions must not be used for training or for corruption-specific choices §4.2. The rotation-invariant class list was chosen with a ResNet-18 on the validation set §4.2 footnote, which is a selection step outside the test-time procedure.
• Theorem 1 is for linear PCA with a known covariance and an orthogonal corruption §5. It does not transfer to the deep model by itself; the paper says so §5, Remark on the assumptions.
• Authors list limitations: no guarantee that reconstruction produces useful features for every main task, only one main task (object recognition) and a handful of benchmarks were tested [Limitations].
• Portraits data was accessed only by the Berkeley authors §4.4. The yearbook dataset itself is not redistributed by this card.

Lineage

• Builds on: §1 the TTT method of Sun et al. (rotation-prediction test-time training; no card in this wiki). This paper keeps the Y-shaped network and the per-input reset, and replaces the self-supervised task with MAE reconstruction §3.
• Builds on: masked autoencoders (MAE) of He et al., used as the pretrained encoder and decoder without modification §3 §4.1. No card in this wiki.
• Compared against: fixed models trained with ViT probing, fine-tuning and joint training tab. 2; the rotation-prediction TTT baseline from Sun et al. fig. 3 tab. 1. Also ImageNet-A, ImageNet-R and Portraits tab. 3 tab. 4.
• Superseded by: not stated in the paper; not known from content already in this wiki.