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 is about test-time adaptation (TTA), in particular about improving the recently published test-time training (TTT) approach. It provides an analysis about conditions that can lead to TTT failures, and overcome them with a method that relies on i) feature alignment and ii) more sophisticated self-supervised learning (SSL) objectives (when compared with the TTT's seminal work). The method is compared against other recent TTA approaches, showing good performance on standardized benchmarks.