Test-Time Classifier Adjustment Module for Model-Agnostic Domain Generalization

1 Introduction . Deep neural networks often fail to generalize to out-of-distribution samples . Accuracy suffers when the model performs under conditions different to those of training , such as variations in light [ 8 ] , weather [ 51 ] , object poses [ 2 ] , textures [ 16 ] , or object backgrounds [ 5 ] . Nevertheless , the model may be deployed to different conditions in practical situations ; thus , some countermeasures are needed . Over the past decade , various studies have focused on training a generalizable model to unseen domains given a dataset consisting of several source domains . This setting is usually denoted as domain generalization ( DG ) [ 6 , 61 ] . Domain generalization operates under the assumption that one can improve robustness to domain shift by incorporating the structure common to multiple domains . For example , domain-invariant feature learning constrains the representation to be invariant to domain shifts [ 14 , 45 , 29 ] . Other methods use meta-learning [ 19 ] to learn how to regularize the model to improve the robustness [ 28 , 4 , 31 ] . However , despite significant work on this front , machine learning systems are still vulnerable to domain shifts even after using the above methods during training . Notably , recent large-scale benchmarks [ 17 ] show that many approaches do not provide significant improvement compared to simple supervised learning , i.e. , empirical risk minimization ( ERM ) , with a proper and practical experimental setup . It suggests that the setup in its current state may be too difficult , and a different approach might be needed . This paper proposes a method of using additional off-the-shelf data in the DG setup , the unsupervised data available at the test-time . Since no data about the target domain is available during training in 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) , virtual . a DG setup , the existing domain generalization algorithms focus on how to use labeled data from multiple-source domains . However , at test-time the model always has access to test data from the target domain . Although the available data is constrained to be ( 1 ) unlabeled and ( 2 ) only available online ( models can not know all test cases in advance ) , this data provides clue about the target distribution that is not available during training . It is natural to ask the question : How can we use the off-the-shelf unlabeled data available at test-time to increase performance on the target domain ? It is worth emphasizing that our setting is different from the transductive setting [ 49 , 22 , 57 ] where all test cases are known in advance , even though we use test data for adjustment . When testing , the model is usually deployed in some environment , and must work well on various samples that will appear continuously . Similarly , the deployed model usually needs to make correct predictions at that moment ; there is no point in going back in time and correcting the predictions . Therefore , it is desirable that adjustment and inference be performed at the same time , not offline after a large amount of data has been accumulated . Looking beyond domain generalization , some recent studies suggest optimizing the model during test time using objective function defined only by unsupervised data ( e.g , prediction entropy ) [ 34 , 52 ] . However , updating parameters using stochastic gradient descent ( SGD ) increases computational costs and harm inference throughput . In addition , data available at test time is limited , and stochastic optimization can lead to catastrophic failure . To this end , we present test-time templates adjuster ( T3A ) , which adjusts the linear classifier ( the last layer of deep neural networks ) at test-time . T3A adjusts the weights of the linear classifier as the following optimization-free procedure : ( 1 ) create a pseudo-prototype for each class using online unlabeled data and the classifier trained in the source domains , ( 2 ) and then classify each sample based on its distance to the pseudo-prototype . This procedure makes the adjusted decision boundary avoid the high-data density region on the target domain and reduce the ambiguity ( entropy ) of predictions , which is known to be connected to classification error [ 52 ] . Since T3A does not alter the training phase , it can be used together with existing DG algorithms . Moreover , it can be used together with any classification model since it only adjusts the linear classifier on top of the representations . Some readers may wonder how effective it is to modify only the linear classifier while freezing the representation itself . Later in this paper ( Section 3.2 ) , we empirically demonstrate that this modification is indeed beneficial . We evaluate our method on multiple standard domain generalization benchmarks , namely VLCS [ 12 ] , PACS [ 27 ] , OfficeHome [ 50 ] , and TerraIncognita [ 5 ] . We compare our method with ( 1 ) various DG algorithms reported in [ 17 ] and ( 2 ) Tent [ 52 ] that minimizes the prediction entropy at test time using SGD . With the standard ResNet50 backbone [ 18 ] , T3A improves ERM by 1.5 points on average accuracy over four dataset , and outperforms most existing DG algorithms . Furthermore , we evaluated our method with 10 different backbone networks , including residual networks ( resnet18 and resnet50 ) , big transfer ( BiT-M-R50x3 , BiT-M-R101x3 , and BiT-M-R152x2 [ 24 ] ) , vision transformers ( ViT-B16 , ViT-L16 , Hybrid ViT [ 9 ] , DeiT [ 48 ] ) , and MLP-Mixer ( Mixer-L16 ) [ 47 ] . The results show that T3A gives a statistically significant performance gain against ERM on all backbone networks . 2 Preliminary and Related Work . 2.1 Domain Generalization . Problem setup Following [ 6 ] , we assume multiple datasets Dd = { ( xdi , ydi ) } nd i=1 collected from several different domains d ∈ { 1 , · · · , dtr } . The dataset Dd from domain d contains identically and independently distributed samples characterized by some probability distribution P d ( X , Y ) , where X and Y are random variables of input and target , respectively . Then , our goal is to develop a predictor f ( X ) that performs well on some unseen test domain , which is characterized by a different probability distribution P ( X , Y ) 6= P d ( X , Y ) for all d ∈ { 1 , · · · , dtr } . Note that one can not assume the target distribution during training , e.g. , no data about the target distribution is available at the time . Therefore , the predictor is usually trained on datasets from several source domains . For example , the predictor can be trained by minimizing the empirical risk : argmin φ 1 d dtr∑ d=1 1 nd nd∑ i=1 ` ( f ( xdi ) , y d i ) , ( 1 ) where φ is the set of the parameters of the function , and ` is a loss function measuring prediction error . In the rest of this paper , optimizing the predictor with eq . 1 is called ERM . In a real application , the model will be deployed after training and expected to classify the data in an online manner . For benchmarking the algorithm , given a dataset containing nd domains , we usually use the leaveone-domain-out procedure , which uses a single domain as a test domain and the others as training domains . The procedure is repeated nd times , changing the test domain every time . Algorithms A central branch of DG algorithm is domain-invariant feature learning , which explicitly reduces the domain gaps on a space of latent representations . For example , [ 14 ] proposed domainadversarial networks ( DANN ) which measure the domain gaps via an external domain classifier . CORAL [ 45 ] align the second-order statistics of representations among different domains . [ 29 ] uses maximum mean discrepancy ( MMD ) to measure the domain gap . Many extension have been proposed [ 33 , 1 , 21 ] , but they are all the same in that they enhance domain invariance . Another branch are meta-learning-based methods [ 28 , 4 , 31 ] , which divide the available domains into meta-train-domains and meta-test-domain and regulate the model trained in meta-train-domains to be useful for the meta-test-domain . Invariant risks minimization [ 3 ] regularizes ERM with a gradient normalization penalty over a dummy classifier . Several studies propose to augment the data using mixup [ 59 ] between two source domains , which implicitly enhances invariance to domain shifts [ 56 , 58 , 53 ] . Key differences As briefly mentioned above , existing domain generalization algorithms focus on the training phase ; how to regularize the predictor using the knowledge from multiple source domains . Our work focuses on the test phase ; how to adjust the model using online and unlabeled data , which can characterize the target distribution . Note that proposed method works fully online ; It does not require access to offline unlabeled data , and therefore can be compared fairly with existing DG methods 2.2 Other Related Work . Unsupervised domain adaptation . Our work is related to unsupervised domain adaptation ( UDA ) [ 37 , 38 , 55 ] as both methods aim to adapt a model given unsupervised data . However , our work primarily differs from UDA in that UDA focuses on adapting during training , while we focus on adapting during testing . In other words , UDA assumes we can access labeled data from the source domain and ( unlabeled ) data from the target domain at the same time , which is not always possible . Source-free domain adaptation . Among them , recent “ source-free ” setups are particularly similar to our setting [ 30 , 25 , 34 ] . In these setups , source data is not needed during the adaptation phase , and the model is adapted using the unlabeled data solely from the target domain . However , this adaptation is usually made in an offline manner , i.e. , these source-free methods optimize offline with multiple losses for multiple epochs . Our method adjusts the classifier in an online-manner , and therefore is suitable for a domain generalization setup where the trained model is assumed to be deployed . Test-time adaptation . Regarding the problem setup , our work is most closely related to testtime adaptation [ 52 ] or test-time training [ 46 ] . Notably , [ 52 ] proposes fully test-time adaptation , which modulates the BN parameter by minimizing the prediction entropy using stochastic gradient descent . The concept of test-time adaptation is very similar to our work and can be used in domain generalization , however , it has not been fully investigated under the domain generalization setup . Moreover , recent architectures do not employ batch normalization either on pre-training ( mainly to avoid the large memory usage required by BN ) or fine-tuning phase ( for improving performance ) . Besides , minimizing the prediction entropy using SGD could lead to trivial solutions , such as being biased to predict only a particular class . In comparison , ( 1 ) our method can be used together with any classification models since it only adjusts the linear classifier on the top of the representations , ( 2 ) our method alleviates catastrophic forgetting due to not using SGD during test-time . Prototypical networks . Before deep learning become popular , the prototype-based classifier is well-investigated in the context of semi-supervised learning [ 10 ] , continual learning [ 35 , 39 ] , and few-shot learning [ 44 ] . Our work is most similar to prototypical networks [ 44 ] , which combine prototypical classifier and deep neural networks as with our method . However , the use-cases mentioned above of prototypical networks assume access to a few labeled data from the same domain , which differs in our case . To handle the difference , we combine prototypical networks with pseudolabeling techniques , which are often used in domain transfer literature [ 42 ] . Besides , connection to entropy minimization is a new perspective introduced in this paper . Algorithm 1 Algorithm of T3A for prediction . Input : Feature extractor fθ , the batch of input B , and support sets Sk available at this point . Output : Prediction for all x ∈ B , where x ∼ P ( X ) . # Step1 . Adjust the template for each class using the B. for x ∈ B do ŷ = argmax qω ( Y = yk|fθ ( x ) ) ( eq . 2 ) Sk = Sk ∪ { fθ ( x ) ‖fθ ( x ) ‖ } for y k = ŷ ( eq . 3 ) end for Filter support sets with eq . 6 # Step2 . Predict based on the distance between the adjusted template . return argmax γ ( Y = yk|fθ ( x ) ) for all x ∈ B ( eq . 4 )

Summary: This paper introduces one algorithm for domain generalization and test-time template adjuster (T3A) for developing a model that performs well under conditions that are different from training conditions. The designed T3A is composed with two procedures: (1) compute a pseudo prototype representation for each class based on unlabeled online dataset, then (2) classify the data based on its distance to the pseudo prototypes. The proposed T3A can improve the performance on the unseen domains, and also outperform the state-of-the-arts.