PaperlayerView exact revision on arXiv

Source-linked AI summary

Test-Time Training with Self-Supervision for Generalization under Distribution Shifts

Yu Sun, Xiaolong Wang, Zhuang Liu, John Miller, Alexei A. Efros, Moritz Hardt

arXiv:1909.13231v3cs.LGcs.CVstat.ML

TL;DR

Supervised models can generalize poorly when training and test distributions differ. Test-Time Training adapts using an unlabeled test sample and a self-supervised task before predicting. It improves performance under distribution shifts while maintaining performance on the original distribution.

  • Problem

    Supervised models can generalize poorly when training and test data come from different distributions.

    Source: §1 · p. 1

  • Method

    Test-Time Training updates model parameters on a self-supervised task built from an unlabeled test sample before making a prediction.

    Source: §1 · p. 1, §2 · p. 2

  • Results

    The method substantially improves performance under distribution shifts while maintaining performance on the original distribution.

    Source: §1 · p. 1

  • Takeaways & Limitations

    The method extends to online streams, where carrying parameter state forward yields further improvements under the assumption of the same or smoothly changing shifts.

    Source: §1 · p. 1

    The method assumes training examples are drawn independently and identically from a distribution P.

    Source: §2 · p. 2

Abstract

from arXiv · show

In this paper, we propose Test-Time Training, a general approach for improving the performance of predictive models when training and test data come from different distributions. We turn a single unlabeled test sample into a self-supervised learning problem, on which we update the model parameters before making a prediction. This also extends naturally to data in an online stream. Our simple approach leads to improvements on diverse image classification benchmarks aimed at evaluating robustness to distribution shifts.

1. Introduction

The paper addresses weak generalization when training and test distributions differ, proposing to learn from shifts at test time rather than anticipate them. Test-Time Training adapts from unlabeled test samples and improves robustness across several object-recognition benchmarks.

1. Introduction

  • The proposed approach lets model parameters depend on the unlabeled test sample, using it as a hint about its distribution.Designing a practical feedback mechanism from the sample to the parameters remains a challenge.Source: §1 · p. 1
  • Test-Time Training creates a self-supervised problem from one unlabeled test sample, updates model parameters, and then makes a prediction.The experiments use image rotation by multiples of 90 degrees as the auxiliary task.Source: §1 · p. 1
  • Online Test-Time Training carries parameter state forward and yields further improvements when test samples come from the same or smoothly changing shifts.The approach can also use a batch of test samples.Source: §1 · p. 1
  • The method substantially improves performance under distribution shifts while maintaining performance on the original distribution.Evaluations cover corruption benchmarks, video frames of moving objects, and a test set of unknown shifts.Source: §1 · p. 1
  • The paper compares against joint training, which combines supervised and self-supervised training but keeps the model fixed at test time.The baseline corresponds to an improved implementation of prior training-time self-supervision work.Source: §1 · p. 1
  • A sufficient condition in a convex model for Test-Time Training to help is, roughly, correlated gradients between the main and auxiliary losses.The paper presents this as a theoretical investigation alongside its empirical results.Source: §1 · p. 1

2. Method

Test-Time Training jointly trains a main classification task and a self-supervised auxiliary task, then adapts their shared feature extractor on each test sample. In the online variant, the updated parameters are carried forward to subsequent samples.

2. Method

  • The method assumes training examples are drawn independently and identically from a distribution P.Source: §2 · p. 2
  • The auxiliary task is rotation prediction: classify whether an image was rotated by 0, 90, 180, or 270 degrees.The method requires an auxiliary self-supervised task; other self-supervised tasks may also be used.Source: §2 · p. 2
  • The model uses a shared feature extractor with separate main-task and self-supervised branches, forming a Y-shaped architecture.In experiments, the auxiliary branch matches the main branch architecture except for its output dimensionality.Source: §2 · p. 2
  • Joint training minimizes the summed main-task and auxiliary-task losses on training data drawn from P.The parameters are optimized jointly in a multi-task learning setup.Source: §2 · p. 2
  • Test-Time Training fine-tunes the shared feature extractor by minimizing the auxiliary loss on a test sample, then predicts with the updated extractor and the main-task branch.The paper reports negligible empirical difference between optimizing the shared extractor alone and optimizing it with the auxiliary branch.Source: §2 · p. 2
  • Test-time training applies the same random transformations used for training augmentation to create a batch of augmented copies of the test sample.These copies are used for the test-time optimization.Source: §2 · p. 2
  • For online samples, the next update starts from parameters adapted on the previous sample, letting the model use distribution information from prior samples.The standard version instead initializes each test-sample optimization from the jointly trained parameters.Source: §2 · p. 2

3. Empirical Results

The experiments evaluate Test-Time Training on corruption, video-frame, and unknown-shift benchmarks. Across these settings, TTT and TTT-Online generally improve on fixed-test-time baselines, with results depending in part on whether the self-supervised rotation task is meaningful.

3.1. Object Recognition on Corrupted Images

  • Across five severity levels and 15 CIFAR-10-C corruption types, TTT beats joint training, while TTT-Online often improves by >10% and never hurts by more than 0.2%.TTT-Online contributes >24% on the three noise types and 38% on pixelation.Source: §3.1 · p. 3, Figure 1 · p. 3
  • On ImageNet-C level-5 corruptions, TTT always improves over both baselines; TTT-Online improves many corruption types by more than a factor of three, while hurting original accuracy by 0.1%.With corruptions, TTT-Online accuracy is already significantly higher after 10,000 samples and continues rising toward the end of the 50,000-sample test stream; the original test set shows no visible trend.Source: §3.1 · p. 4, Figure 2 · p. 4
  • As noise severity increases through the test stream, errors rise for all methods, but TTT-Online has gentler error slopes than joint training and beats UDA-SS on the first two noise types.This experiment relaxes the fixed test-distribution assumption to allow gradually changing distributions.Source: Figure 3 · p. 5, §3.1 · p. 5

3.2. Object Recognition on Video Frames

  • On CIFAR classes in VID-Robust, TTT and TTT-Online improve average accuracy over both baselines, with large gains on ship and dog, where rotation prediction is more meaningful than for airplane.For airplane images, black margins can provide a trivial rotation cue, and some images are difficult even for humans to orient, illustrating the need for a well-defined, non-trivial self-supervised task.Source: Table 2 · p. 6, §3.2 · p. 6
  • On VID-Robust anchor frames, TTT and TTT-Online improve accuracy over the baselines without video-specific modifications to the ImageNet model.The results use anchor frames from 1,109 video-frame sets across 30 classes.Source: §3.2 · p. 6, Table 3 · p. 6

3.3. CIFAR-10.1: Unknown Distribution Shifts

  • On CIFAR-10.1, where prior work had not improved an existing model, TTT improves over joint training by 0.8% (4.8% relative), a small gain against the typical factor-of-two performance drop.The test set was created to investigate distribution shifts from CIFAR-10 despite the datasets having the same classes and image dimensionality.Source: §3.3 · p. 6, Table 4 · p. 6

4. Theoretical Results

Theoretical analysis links Test-Time Training’s effect on the main task to positive gradient correlation between the main and self-supervised losses. The paper also reports that stronger correlation tracks greater improvement over baseline on a deep model across corruptions.

4. Theoretical Results

  • For smooth convex models, positive gradient correlation between the main and self-supervised losses guarantees reduced main-task loss after one Test-Time Training step.The theorem assumes differentiable, convex, β-smooth main-task loss and bounded gradients.Source: §4 · pp. 7–8
  • In the toy regression model, Test-Time Training updates the shared feature extractor using a self-supervised task that automatically generates a label for the input.The example uses a two-layer linear network; the update is one gradient step on the self-supervised loss.Source: §4 · p. 7
  • A carefully chosen learning rate makes the toy model’s main-task loss zero after one step, while any sufficiently small positive rate suffices when the ideal rate is positive.The ideal rate depends on the unknown main-task label.Source: §4 · p. 7
  • Stronger gradient correlation corresponds to greater improvement over baseline on the deep model across the diverse corruptions considered.The paper reports this empirical relationship for non-convex loss functions as well.Source: §4 · p. 8

5. Related Work

Related work spans adaptation at test time, self-supervised learning, and robustness to shifts. The paper distinguishes its approach by updating on each unlabeled test instance without anticipating the test distribution.

Related work

  • Prior test-instance adaptation includes learning image super-resolution from each input [Shocher et al. (2018)] and adapting a pre-trained GAN to input-image statistics [Bau et al. (2019)].The paper also cites online-update work in video segmentation, tracking, detection, and echocardiography.Source: §5 · p. 8
  • Self-supervised learning can learn semantic information without human annotations through context prediction, jigsaw puzzles, colorization, noise prediction, and feature clustering; this paper uses rotation prediction (Gidaris et al., 2018).Jointly training a main task and a self-supervised task can improve robustness on the main task (Hendrycks et al., 2019a).Source: §5 · p. 8
  • Adversarial robustness methods rely on a modeled perturbation set, and the paper notes an accuracy–robustness tradeoff for approaches that smooth a decision boundary [Tsipras et al. (2018); Zhang et al. (2019)].Corruption-specific training may not transfer to other corruption types [Geirhos et al. (2018)].Source: §5 · p. 8
  • Unlike unsupervised domain adaptation, which has target-distribution data available during training, Test-Time Training does not anticipate a test distribution and can be viewed as one-sample adaptation on the fly.The paper cautions that standard domain-adaptation techniques may be ill-defined with only one target sample.Source: §5 · p. 9
  • Domain generalization uses some environments for training and others for testing, but its fixed decision boundary is limited by a specificity–generality tradeoff and the meta-distribution scope.Source: §5 · p. 9
  • The update rule can also be viewed as one-shot self-supervised learning and may benefit from progress in one-shot learning.The paper contrasts its test-time setting with continual learning, where avoiding forgetting earlier tasks is central.Source: §5 · p. 9
  • Unlike standard online learning, Test-Time Training receives no test labels or post-prediction environmental feedback, so it updates before prediction using self-supervision.The paper notes that “online learning” is used differently in some cited applications.Source: §5 · p. 9

6. Discussion

The discussion points to applications beyond image classification and suggests that test-time training can also serve as a benchmark for general-purpose self-supervised tasks.

6. Discussion

  • Test-time training may apply to segmentation, detection, speech recognition, and natural language processing, with field experts able to design special-purpose self-supervised tasks.Source: §6 · p. 9
  • The authors propose test-time training as an evaluation benchmark for general-purpose self-supervised tasks alongside pre-training and fine-tuning.Source: §6 · p. 9
  • The paper frames its contribution as a small step toward a paradigm in which more learning occurs after model deployment.Source: §6 · p. 9

Appendix: Test-Time Training with Self-Supervision

The supplied passage names generalization under distribution shifts.

A1. Informal Discussion on Our Variable Decision Boundary

The discussion distinguishes fixed decision boundaries from variable ones by whether the model parameters can depend on the input. It argues that variable boundaries avoid relying on an enormous model class whose generalization is impossible with finite data absent prior knowledge, while noting that the relationship to traditional learning theory is not formally developed here.

A1. Informal Discussion on Our Variable Decision Boundary

  • A fixed boundary uses one parameter setting for every input, whereas a variable boundary allows a setting to be selected for each input.The discussion defines both in terms of whether a mapping can be represented by the model class.Source: A1. Informal Discussion on Our… · p. 14
  • Without prior knowledge, generalization over an enormous class capable of representing every mapping is impossible with any finite amount of training data.The authors argue that selecting parameters based on an input instead amounts to training only on the points the model actually encounters.Source: A1. Informal Discussion on Our… · p. 14
  • The distinction between variable boundaries and traditional learning theory is raised, but a formal discussion is left beyond the paper’s scope.Source: A1. Informal Discussion on Our… · p. 14

A2. Computational Aspects of Our Method

Test-Time Training increases test-time computation, and the authors outline two potential ways to reduce that cost while retaining performance, though these solutions are not thoroughly verified.

A2. Computational Aspects of Our Method

  • Test-time computation is 2 × batch size × number of iterations times slower than regular testing, which uses one forward pass per sample.Source: A2. Computational Aspects of Our Method · p. 14
  • Thresholding could skip updates for roughly 80% of test instances with little effect on performance, since most wrong predictions occur among the other 20%.The authors describe this as a potential solution that has not been thoroughly verified.Source: A2. Computational Aspects of Our Method · p. 14
  • Preliminary experiments found results almost as good with one update iteration and learning rate 0.01 instead of 10 iterations and 0.001.The authors leave deeper efficiency analysis to future work.Source: A2. Computational Aspects of Our Method · p. 14

A3. Proofs

This section provides proofs for the theoretical results in the main paper.

A3. Proofs

  • The appendix section presents proofs of the main paper’s theoretical results.Source: A3. Proofs · p. 14

A3.1. The Toy Problem

The toy problem analyzes Test-Time Training in a two-layer linear network with a shared parameter matrix and fixed task-specific heads. Under a stated condition on η*, a sufficiently small update rate is guaranteed to improve the main loss.

A3.1. The Toy Problem

A3.2. Proof of Theorem 1

The appendix gives additional corruption-benchmark results, including Batch Normalization performance, baseline comparisons, and results across corruption levels. It also documents limitations of Batch Normalization and architectural differences relevant to a direct comparison with Hendrycks et al. (2019a).

A4. Additional Results on the Common Corruptions Dataset

  • The appendix labels baseline, joint training, and the two Test-Time Training variants as B, JT, TTT, and TTT-Online, and lists the corruption types covered.The listed set includes the original test set and 15 corruption types, from Gaussian noise to JPEG compression.Source: A4. Additional Results on the Common… · p. 16

A4.1. Results Using Batch Normalization

  • With Batch Normalization, Test-Time Training reports improvements of 34% on contrast and 18% on defocus blue, and 22% on Gaussian noise versus joint training or 16% versus baseline.Original-distribution performance remains almost the same; the Batch Normalization model has slightly lower original error than the Group Normalization model and converges in half as many epochs.Source: A4.1. Results Using Batch Normalization · p. 16
  • Batch Normalization is unsuitable for the online version because inaccurate batch statistics worsen over long training streams, while its baseline error is higher than Group Normalization for almost every corruption.The appendix also notes that Group Normalization improves robustness, although this is described as unrelated to the paper’s main idea.Source: A4.1. Results Using Batch Normalization · p. 16

A4.2. Additional Baseline: Adversarial Logit Pairing

  • The appendix includes Adversarial Logit Pairing as an additional baseline on CIFAR-10-based benchmarks using Group Normalization, but not on ImageNet.The authors followed the method’s original training procedure and recommended hyperparameters; they omit ImageNet because the cited papers did not evaluate or recommend it there.Source: A4.2. Additional Baseline: Adversarial… · p. 16

A4.4. Results on CIFAR-10-C, Levels 1-4

A4.5. Direct Comparison with Hendrycks et al. (2019a)

  • 22.8% average error across corruptions and levels for the joint-training baseline, versus 28.6% for Hendrycks et al. (2019a), despite architectural changes favoring the baseline.The authors report Group Normalization instead of Batch Normalization, a different residual-group split, and a ResNet-26 rather than a 40-2 Wide ResNet; the baseline is four times smaller.Source: A4.5. Direct Comparison with Hendrycks… · p. 16

A4.3. Results on CIFAR-10-C and ImageNet-C, Level 5

Loading 1909.13231v3…