[Paper Review] Test-Time Training with Self-Supervision for Generalization under Distribution Shifts

Yu Sun, Xiaolong Wang+4 more|arXiv (Cornell University)|Sep 29, 2019

Intelligent Tutoring Systems and Adaptive LearningComputer Science199 citations

TL;DR

This paper introduces Test-Time Training (TTT), which updates the model with a self-supervised task on unlabeled test samples before predicting, to improve robustness to distribution shifts both in batch and online settings.

ABSTRACT

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.

Motivation & Objective

• Motivate robust generalization when training and test distributions differ.
• Propose a generic test-time adaptation framework using a self-supervised auxiliary task.
• Investigate both standard and online test-time training regimes across diverse benchmarks.

Proposed method

• Construct a multi-task network with a shared feature extractor and two branches: main task and self-supervised task (rotation prediction).
• Train jointly on labeled data for the main task and unlabeled data for the self-supervised task with a shared extractor.
• At test time, update only the shared feature extractor by minimizing the self-supervised loss on the test sample(s) before predicting the main task.
• Utilize data augmentation to form test-time batches analogous to training-time augmentation.
• Optionally extend to online settings by updating across a stream of test samples, carrying over state between samples.

Experimental results

Research questions

• RQ1 Can test-time updates driven by self-supervision improve performance under distribution shifts without labels at test time?
• RQ2 How does online test-time training compare to batch/test-time adaptation under evolving test distributions?
• RQ3 What is the role of gradient alignment between main-task and self-supervised losses in the effectiveness of test-time training?

Key findings

• TTT and its online variant (TTT-Online) consistently improve over a plain object-recognition baseline under multiple corruption types and levels.
• TTT-Online often reduces error by large margins compared with joint training baselines and can also improve on the original distribution.
• In CIFAR-10-C, TTT-Online surpasses unsupervised domain adaptation with rotation prediction on most corruptions, suggesting forgetting of the training distribution can be advantageous.
• On ImageNet-C, TTT-Online yields substantial improvements across many corruption types, with performance improving as more test samples are processed.
• TTT shows a small gain on CIFAR-10.1, a dataset designed to reveal unknown distribution shifts, marking a first step toward test-time adaptation for unseen shifts.
• A theoretical result links positive gradient correlation between main and self-supervised losses to improved main-task performance during test-time updates.