TTT++: When Does Self-Supervised Test-Time Training Fail or Thrive?

1 Introduction . Machine learning models often struggle to generalize under distribution shifts . Even a perceptually mild shift between training and test data , e.g. , JPEG compression , may cause severe prediction errors [ 1 ] . One popular family of methods to address this challenge is to learn an invariant representation across domains by making use of labelled training data and unlabelled test data simultaneously [ 2–5 ] . However , revisiting training data at test time can be impractical due to increasing privacy concerns , inflating sizes of datasets as well as many other real-world constraints . This shortcoming prompts a more challenging yet appealing test-time adaptation paradigm : given a trained model , how can we adapt it from one domain to another on the fly , without access to training data and human annotations ? One promising approach towards this goal is test-time training ( TTT ) through self-supervision [ 6 ] . The key idea of TTT is simple and straightforward : train the model on two tasks , namely a main task and a self-supervised learning ( SSL ) task , and update the model based only on the SSL task at test time . This technique implemented with self-supervised rotation prediction has shown encouraging results for improving the robustness of image classifiers under a variety of distributional shifts . Yet , its empirical performance is still inferior to other families of test-time algorithms [ 7 , 8 ] . 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . In this paper , we first take an in-depth look at TTT with emphasis on its limitations . Our analysis starts with a basic question : can TTT always mitigate the effects of distributional shifts ? Through an illustrative problem , we show that the TTT framework can lead to surprising failures , deteriorating the test accuracy rather than improving it . This problem is largely attributed to the unconstrained update from the SSL task that interfere with the main task . To address this issue , we introduce a test-time feature alignment strategy by means of offline feature summarization and online moment matching : once training completes , we compute the mean and covariance matrix of training features and store them as part of the model , referred to as offline feature summarization ; at test time , we encourage the test feature distribution to be close to the training one by matching the moments estimated online with those pre-computed offline , a process referred to as online moment matching . One practical challenge for online feature alignment lies in scaling the strategy to problems with a large number of classes , as obtaining a robust estimate of moments often requires at least a handful of samples per class . To mitigate this issue , we draw inspiration from recent literature [ 9 ] and decouple the sample size from the batch size for moment estimates . Specifically , we maintain a large dynamic queue of encoded features and progressively update it in a mini-batch manner . This modification enables effective feature alignment even with limited batch size , greatly improving its viability in the online test-time setting . Finally , we shed light on the strong potential of TTT through a theoretical analysis of the test accuracy after adaptation . In particular , we derive a lower bound of the test accuracy on the main task and show that it is expected to grow rapidly when the SSL task gets closer to the main task . These findings motivate our integration of contrastive representation learning [ 9–12 ] , as a strong instance of SSL , into the TTT framework in visual recognition problems . By combining the three proposed components , we devise an improved version of test-time training , termed TTT++ . Experimental results show that TTT++ significantly outperforms other recent methods by significant margins on various robustness benchmarks . Our results suggest that exploiting extra information , including both task-specific information in the form of strong self-supervision and domain-specific information in the form of feature summarization , can be a promising direction to enhance the effectiveness of test-time adaptation . 2 Background . 2.1 Related Work . Test-time Adaptation . Adapting machine learning models based on test samples has garnered growing interests in both generative problems such as super-resolution [ 13 ] , image synthesis [ 14 ] and image manipulation [ 15 ] , and discriminative problems like image classification [ 16 ] . Our work is focused on the latter one in the presence of distributional shifts . Several recent works [ 16 , 17 ] have shown the advantage of adapting the learned classifier to new test domains in the unsupervised manner , without revisiting the source data . One simplest form is to replace the batch-norm statistics estimated on the training set with those on test examples [ 17 ] . Another line of work proposed to adapt the model parameters by exploiting the predicted labels on test examples , such as entropy minimization [ 8 ] and pseudo-labeling [ 7 ] . While these methods yield promising results on some benchmarks , they are inherently restricted to classification problems and often vulnerable under large distribution shifts [ 18 ] . More closely related to ours , [ 6 ] proposed test-time training through self-supervised learning , e.g. , predicting the type of image rotations . This approach does not involve any assumptions about the output for the main task and is therefore more generic . It has been successfully applied to a variety of problems , such as instance tracking [ 19 ] and reinforcement learning [ 20 ] . Nevertheless , it was shown empirically inferior to other test-time algorithms [ 8 ] . Our work provides an in-depth analysis of its limitations , introduces simple yet effective remedies , and consolidates more theoretical grounds . Feature Alignment . Aligning the distributions of training and test samples in the feature space is commonly used for domain adaptation . Previous feature alignment methods fall into two main categories : minimizing a divergence measure , such as MMD [ 21 ] , Coral [ 22 ] and CMD [ 23 ] , or encouraging the domain confusion through adversarial training [ 5 , 24 ] . However , most of these methods rely on the co-existence of source and target data , and thus can not be readily applied to the test-time setting where source data is not available . Our work revisits the critical role of feature alignment for test-time training and proposes a simple and practical strategy that enables online feature alignment , even with a limited batch size . Self-Supervised Learning . Self-supervised learning is a powerful paradigm to learn rich representations from unlabeled samples . Stunning progress has been made in recent years by designing informative self-supervised tasks [ 11 , 12 , 25–28 ] and stabilizing the training process [ 9 , 29 ] . Prior works are mainly focused on unsupervised pre-training , whereas our work looks into the importance of incorporating strong self-supervised learning methods for test-time adaptation . 2.2 Preliminary : Test-Time Training . Test-time training ( TTT ) [ 6 ] is a general framework for adapting neural network models to a new test distribution based on unlabeled samples . Different from the conventional approach that trains the model only on the task of interest , TTT considers two tasks : a main task and an auxiliary SSL task . The model is trained on both tasks simultaneously with a multi-task architecture composed of one shared encoder g and two separate heads πm and πs respectively . Given a labeled training dataset D = { ( xi , yi ) } i∈ { 1 , ... , N } , the model is trained to minimize two losses jointly : Ltrain ( D ; g , πm , πs ) = 1 N N∑ i=1 Lm ( xi , yi ; g , πm ) + λLs ( xi ; g , πs ) , ( 1 ) where λ is a hyper-parameter to balance the two tasks . In the presence of distributional shifts , the learned model often struggles to directly generalize to a new test set D′ = { x′i } i∈ { 1 , ... , N ′ } . The core idea of TTT is to fine-tune the encoder g based on the self-supervised task with the test examples , LTTT ( D′ ; g′ ) = 1 N ′ N ′∑ i=1 Ls ( x′i ; g′ ) , ( 2 ) with the hope that the updated model πm ◦ g′ yields improved results on the main task . TTT instantiated with a self-supervised rotation prediction task has been demonstrated effective for improving the robustness of image classifiers under common distributional shifts [ 6 ] . However , it was shown inferior to other families of test-time adaptation methods [ 7 , 8 ] . We will next look into its strengths and limitations , and propose an improved version for better adaptation performance . 3 When Does Test-Time Training Fail ? . In this section , we first throw light on a caveat of test-time training under large distributional shifts , and subsequently introduce practical solutions tailored for the test-time setting . 3.1 Illustrative Example of Failures . One implicit assumption behind TTT is that the encoder update based on SSL can counter the effect of the underlying distributional shift on the main task . This assumption , however , is likely broken under large shifts and results in unexpected adaptation failures . To illustrate this limitation , we introduce a simple toy problem in Figure 1 , where the main task and the SSL task are defined as classifying colors and symbols of encoded features in the 2-dimensional latent space . We consider an ideal scenario where the two tasks are highly correlated such that a well-trained model can attain high accuracies on both of them in the training domain . Nevertheless , in the presence of a significant distributional shift , the model may still suffer from substantial prediction errors in the test domain . In this scenario , while TTT may restore the discriminative power of the learned representation for the main task to a certain degree , the unconstrained self-supervised adaptation may yield severe overfitting to the auxiliary SSL task . As a consequence , the performance on the main task can even deteriorate as opposed to improving . This phenomenon is not restricted to our illustration and also occurs in practice , as shown in Section 5.1 . 3.2 Online Feature Alignment . As illustrated in the toy example above , simply applying self-supervised adaptation at test time can lead to arbitrarily poor results . To address this issue , we introduce an online feature alignment strategy to ensure robust adaptation at test time . The core idea of our strategy is to impose a constraint over the feature space during TTT such that the feature distribution of test examples remains close to that in the training domain . While some feature alignment techniques such as MMD [ 2 ] and adversarial training [ 24 ] have been widely used for domain adaptation , they often rely on sampling from training and test domains concurrently , which is impractical in the test-time setting . We , therefore , turn to classical divergence measures that can be estimated independently for each distribution . More specifically , we use the square distance of the first and second moments between two feature distributions , inspired by DDC [ 30 ] and Coral [ 22 ] , to approximate the domain discrepancy . This design choice allows us to summarize the distribution of training features in a compact format and store it as part of the model , eliminating the need to revisit the training data during test-time adaptation . Concretely , once training completes , we perform an offline feature summarization step that characterizes the distribution of feature vectors in the training domain Z = { zT1 , . . . , zTN } through the empirical mean µz = 1N ∑N i zi and covariance matrix ΣZ = 1 N−1 ( Z TZ − ( ITZ ) T ( ITZ ) ) . The former is essentially equivalent to the channel-wise batch normalization statistics while the latter is more informative yet light-weight for computation and storage . At test time , we regularize the self-supervised adaptation by minimizing the distance between the feature statistics estimated from a mini-batch of test samples , i.e. , µ′z and Σ ′ z , and the pre-stored quantities about the training domain : Lf , z = ‖µz − µ′z‖ 2 2 + ‖Σz − Σ ′ z‖ 2 F , ( 3 ) where ‖·‖2 is the Euclidean norm and ‖·‖F is the Frobenius norm . The basic form of online moment matching can be limiting in that the low-order statistics may be insufficient to fully capture complex distributions in high dimensions , e.g. , 2048 for the standard ResNet-50 . To alleviate this issue , we align the feature distributions at both the output of the encoder and the output of the self-supervised head , which are of lower dimensions , e.g. , 128 in the case of contrastive learning described in Sec . 4.2 . Our final objective at test time is a weighted combination of the self-supervised loss Ls , the feature alignment loss at the encoder Lf , z and that at the self-supervised head Lf , s , LTTT++ = Ls + λzLf , z + λsLf , s , ( 4 ) where λz and λs are hyper-parameters controlling the emphasis on each term .

Summary: This paper focus on improving the generalization of deep neural networks in the presence of distributional shift at test time i.e., unsupervised source-free network adaptation on data distribution that is different from training. They follow the footsteps of a prior work by Sun et al. [8] and leverage an auxiliary task at test time. However, this work points out that test time training with auxiliary tasks like rotation prediction need not necessarily encourage the network to improve its generalization on the main task. Thus, the authors propose to transfer the benefits of self-supervised learning (SSL) to the main task by introducing an online feature alignment strategy that utilizes offline feature summarization statistics collected on the original training (source) data and also choosing more suitable SSL task like contrastive learning instead of rotation prediction. Tthey denote the resulting approach as TTT++. Results and comparison to existing techniques are shown on popular datasets like cifar10-c, cifar100-c, cifar10.1 and VisDA-C.