Test-Time Training for Out-of-Distribution Generalization

1 INTRODUCTION . Supervised learning remains notoriously weak at generalization under distribution shifts . Unless training and test data are drawn from the same distribution , even seemingly minor differences turn out to defeat state-of-the-art models . Transfer learning , domain adaptation and adversarial robustness are but a few existing paradigms that anticipate differences of sorts between the training and test distribution . In this work , we explore a new take on generalization , without any mathematical structure or data available at training-time to anticipate the distribution shifts . We start from a simple observation . The unlabeled test instance x presented at test-time gives us a hint about the distribution from which it was drawn . Our approach , called test-time training , allows the model parameters θ to depend on the test instance x , but not its unknown label y . This variable decision boundary θ ( x ) is powerful in theory ( see Appendix A ) , and raises new challenges in practice that we only begin to address here . Our algorithm creates a self-supervised learning problem based on this single test instance x , updating θ at test-time before making a prediction . Self-supervised learning uses an auxiliary task that automatically creates labels from unlabeled inputs . For the visual data we work with , the task rotates an image x by a multiple of 90 degrees , and assigns the angle as the label ( Gidaris et al. , 2018 ) . Our approach can also be easily modified to work outside the standard supervised learning setting . If several test instances arrive in a batch , we can use the entire batch for test-time training . If test instances arrive online sequentially , we obtain further improvements by keeping the state of the parameters . After all , prediction is rarely a single event . The online setting can be the natural mode of deployment in practice , and shows the strongest improvements . We experiment with test-time training for generalization under distribution shifts in the context of object recognition on several benchmarks . These include images with diverse types of corruption at various levels ( Hendrycks & Dietterich , 2019 ) , video frames of moving objects ( Shankar et al. , 2019 ) , and a new test set of unknown shifts collected by Recht et al . ( 2018 ) . Our algorithm does not hurt on the original distribution , but makes substantial improvements under distribution shifts . In all our experiments , we compare with a strong baseline ( labeled joint training ) that uses both supervised and self-supervised learning at training-time , but keeps the model fixed at test-time . Very recent work shows that additional training-time self-supervision improves robustness Hendrycks et al . ( 2019b ) . The joint training baseline we use corresponds to an improved implementation of their work . We are also inspired by Sun et al . ( 2019 ) , which performs joint training with a selfsupervised task in the context of unsupervised domain adaptation , simultaneously on the source and target domain . A comprehensive review of related work follows in section 5 . We complement the empirical results with theoretical investigations in section 4 , of when test-time training is expected to help on a convex model , and establish an intuitive sufficient condition , which , roughly speaking , is to have correlated gradients between the loss functions of the two tasks . 2 METHOD . Next we describe the algorithmic details of our method . To setup notations , consider a standard Klayer neural network with parameters θk for layer k. The stacked parameter vector θ = ( θ1 , . . . , θK ) specifies the entire model for a classification task with loss function lm ( x , y ; θ ) on the test instance ( x , y ) . We call this the main task , as indicated by the subscript of the loss function . We assume to have training data ( x1 , y1 ) , . . . , ( xn , yn ) drawn i.i.d . from a distribution P . Standard empirical risk minimization corresponds to solving the optimization problem : min θ 1 n n∑ i=1 lm ( xi , yi ; θ ) . ( 1 ) Our method requires a self-supervised auxiliary task with loss function ls ( x ) . In this paper , we choose the rotation prediction task ( Gidaris et al. , 2018 ) , which has been demonstrated to be simple and effective at feature learning for neural networks . The task simply rotates x on the image plane by one of 0 , 90 , 180 and 270 degrees and have the model predict the angle of rotation as a four-way classification problem . Other self-supervised tasks in section 5 might also be used for our method . The auxiliary task shares some of the model parameters θe = ( θ1 , . . . , θκ ) up to a certain κ ∈ { 1 , . . . , K } . We think of those κ layers as a shared feature extractor . The auxiliary task uses its own task-specific parameters θs = ( θ′κ+1 , . . . , θ ′ K ) . We call the unshared parameters θs the selfsupervised task branch , and θm = ( θκ+1 , . . . , θK ) the main task branch . Pictorially , the joint architecture is a Y -structure with a shared bottom and two branches . For our experiments , the selfsupervised task branch has the exact same architecture as the main branch , except for the output dimensionality of the last layer due to the different number of classes in the two tasks . Training is done in the fashion of multi-task learning ( Caruana , 1997 ) ; the model is trained on both tasks on the same data drawn from P . Losses for both tasks are added together , and gradients are taken for the collection of all parameters . The joint training problem is therefore min θe , θm , θs 1 n n∑ i=1 lm ( xi , yi ; θm , θe ) + ls ( xi ; θs , θe ) . ( 2 ) Now we describe the standard version of test-time training on a single test instance x . Simply put , test-time training finetunes the shared feature extractor θe by minimizing the auxiliary task loss on x . This can be formulated as min θe ls ( x ; θs , θe ) . ( 3 ) Denote θ∗e the ( approximate ) minimizer of Equation 3 . The model then makes a prediction using the updated parameters θ ( x ) = ( θ∗e , θm ) . Empirically , the difference is negligible between minimizing Equation 3 over θe versus over both θe and θs . Theoretically , there exists a difference only when optimization is done with more than one step of gradient descent . In the standard version of our method , the optimization problem in Equation 3 is always initialized with parameters θ = ( θe , θs ) obtained by minimizing Equation 2 on data from P . After making a prediction on x , θ∗e is discarded . Outside of the standard supervised learning setting , when the test instances arrive online sequentially , the online version of test-time training solves the same optimization problem as in Equation 3 to update the shared feature extractor θe . However , on test instance xt , θ is instead initialized with θ ( xt−1 ) updated on the previous instance xt−1 . This allows θ ( xt ) to take advantage of the distributional information available in x1 , . . . , xt−1 as well as xt . Test-time training naturally benefits from standard data augmentation techniques . On each test instance x , we perform the exact same set of random transformations as used for data augmentation during training to form a batch only containing these augmented copies of x for test-time training . 3 EMPIRICAL RESULTS . We experiment with both versions of our method ( standard and online ) on three kinds of benchmarks for distribution shifts , presented here in the order of visually low to high-level , which is roughly also the order of easy to hard . Our code will be released once the paper is accepted . Network details . Our architecture and hyper-parameters are consistent across all experiments . We use Residual Networks ( ResNets ) ( He et al. , 2016b ) , which are constructed differently for CIFAR10 1 ( 26-layer ) and ImageNet 2 ( 18-layer ) . ResNets on ImageNet have four groups , each containing convolutional layers with the same number of channels and size of feature maps ; our splitting point is the end of the third group . ResNets on CIFAR-10 have three groups ; our splitting point is the end of the second group . In addition , Batch Normalization ( BN ) , a popular module in deep networks , is ineffective when training on small batches , for which the estimated batch statistics are less accurate ( Ioffe & Szegedy , 2015 ) . This technicality hurts test-time training since each batch only contains ( augmented ) copies of a single image . Therefore our networks instead use Group Normalization ( GN ) ( Wu & He , 2018 ) , which achieves similar performance as BN on large batches without hurting on small ones . Results with BN are shown in Appendix E for completeness . Optimization details . For Equation 2 , optimization hyper-parameters are set to the default 3 in standard practice ( Huang et al. , 2016 ; He et al. , 2016a ) . For Equation 3 , we use stochastic gradient descent ( SGD ) with the learning rate set to that of the last epoch during training , which is 0.001 in all our experiments . Following standard practice ( He et al. , 2018 ; Liu et al. , 2018 ) known to improve performance when finetuning , we do not use weight decay or momentum . For the standard version , we take ten gradient steps , using batches independently generated by the same image . For online we take only one step . The computational aspects of our method are discussed in Appendix C. Following standard practice , the transformations used for data augmentation are random crop with padding and random horizontal flip for CIFAR-10 ( Guo et al. , 2017a ; Huang et al. , 2016 ) , and random resized crop and random horizontal flip for ImageNet ( Ioffe & Szegedy , 2015 ; He et al. , 2016a ) . Specifically , these transformations do not contain information about the distribution shifts . In all the tables and figures , baseline refers to the plain ResNet model ( using GN , unless otherwise specified ) ; joint training refers to the model jointly trained on both the main task and the selfsupervised task , fixed at test-time as in Hendrycks et al . ( 2019b ) ; test-time training refers to the standard version described section 2 ; and test-time training online refers to the online version that does not discard θ ( xt ) for xt arriving sequentially from the same distribution . Performance for test-time training online is calculated , just like the others , as the average over the entire test set ; we always shuffle the test set before test-time training online to avoid ordering artifacts . 1CIFAR-10 ( Krizhevsky & Hinton , 2009 ) is a standard object recognition dataset with 10 classes of objects in natural scenes . The standard train / test split has 50,000 / 10,000 images , each of size 32-by-32 pixels . 2The ImageNet 2012 classification dataset ( Russakovsky et al. , 2015 ) for object recognition has images from 1 , 000 classes , 1.2 million for training and 50,000 for validation . Following standard practice ( He et al. , 2016a ; b ; Huang et al. , 2016 ) , the validation set is used as the test set . 3Namely , we use stochastic gradient descent ( SGD ) with weight decay and momentum ; learning rate starts at 0.1 and is dropped by a factor of ten at two scheduled milestones , to 0.01 and 0.001 .

Summary: The motivation is to increase accuracy of CNNs with unseen (unknown) distribution shifts. To this end, instead of the recent approach of training-time self-supervision, they adopted it for test-time [limited novelty]. More precisely, for a classification task, they considered two headed neural network with one head for main classification task and another for an auxiliary task (e.g. predicting rotation degree of rotated images). The feature extractor (up to k layers) is shared between these tasks thus updated using the two tasks, while two heads are updated according to each task. In test-time, the samples (drawn from shifted data distributions) are being used for updating the shared feature extractor through the auxiliary task. They have investigated their test-time approach in a series of experiments in online and offline settings on synthetic shift in data distribution of image classification tasks. The distributional shifts are synthesized adding some non-adversarial perturbation to clean images. The authors proof that their test-time approach can lead to lower error rate on given test-time samples when the underlying learning model is a linear regression model.