Learning to Generalize across Domains on Single Test Samples

1 INTRODUCTION . Despite their widespread adoption and success in academia and industry alike , deep convolutional neural networks suffer from a fundamental flaw : they have insufficient generalizability to test data that is out of their training distributions ( Recht et al. , 2019 ) . Improving the generalization of machine learning methods therefore remains a challenging problem ( Moreno-Torres et al. , 2012 ; Recht et al. , 2019 ; Krueger et al. , 2021 ) . To deal with the distribution shift , domain adaptation ( Saenko et al. , 2010 ; Long et al. , 2015 ; Lu et al. , 2020 ; Li et al. , 2021 ) and domain generalization ( Muandet et al. , 2013 ; Li et al. , 2017 ; 2020 ) have been extensively investigated . In domain adaptation , the key assumption is that the target data , either labeled or unlabeled , is accessible during training , which allows the model to be adapted to the target domains . However , this assumption does not necessarily hold in realistic scenarios where the target domain data is not available . By contrast , domain generalization strives to learn a model on source domains that can generalize well to unseen target domains without any access to the target domain data during training . We aim for domain generalization . Previous domain generalization approaches have successfully explored domain-invariant learning ( Muandet et al. , 2013 ; Chattopadhyay et al. , 2020 ) or domain augmentation ( Shankar et al. , 2018 ; Zhou et al. , 2020a ) to handle the domain shift between the source and target domains . However , since those models are trained on source domains , there will always be an “ adaptivity gap ” when applying them to target domains without further adaptation ( Dubey et al. , 2021 ) . Thus , it is necessary to train models that are able to further adapt to target domains without ever using target data during training . Recently , further adapting a pre-trained model using target domain data by fine-tuning or selfsupervised training at test time has shown effectiveness in improving performance on domain generalization ( Sun et al. , 2020 ; Wang et al. , 2021 ) . To achieve proper adaptation to the target domain , these methods typically rely on extra fine-tuning operations on target data or extra networks ( Dubey et al. , 2021 ) at test time . Although some of these methods can achieve adaptation with a single test 1 Code available : https : //github.com/zzzx1224/SingleSampleGeneralization-ICLR2022 sample ( D ’ Innocente et al. , 2019 ; Sun et al. , 2020 ; Banerjee et al. , 2021 ) , batches of target data are required for good performance . This prevents these methods from effectively adapting to the target domain when provided with very few target samples or exposed to test samples from multiple target domains without domain identifiers . Different from these works , we aim to learn the ability to adapt to a single sample such that the model can further adapt to each test sample from any target domain . In this paper , we propose learning to generalize on single test samples for domain generalization . Since a single sample from the target distribution can not inform much about the whole distribution , we consider each target sample a domain by itself . Correspondingly , we propose that each target sample adjusts the trained model in its own way . This avoids the difficulty of using the limited information available to adapt the model to the entire target domain distribution . An illustration of the method is shown in Figure 1 . Using meta-learning ( Li et al. , 2017 ) , we train our model to acquire the ability to adapt with each single sample on source domains such that it can then adapt to each individual sample in the target domain to handle the domain shift at test time . Unlike previous methods , our model does not need to be fine-tuned using target data or train extra networks when generalized to the unseen target domains . Once it is trained on source domains , our model is able to adapt to each single target sample with just a feed-forward computation at test time , without the need for any extra target data . Note that not a single target sample is seen during training , which does not forgo the strict source and target separation setting of domain generalization . Moreover , in contrast to previous test-time adaptation methods ( Sun et al. , 2020 ; Wang et al. , 2021 ; Dubey et al. , 2021 ) , our method does not need to re-adapt the model to each target domain . Our model effortlessly handles test samples from multiple domains without knowing their domain identifiers . To be specific , we build our model under the meta-learning framework and formulate the single sample generalization as a variational inference problem . In the training stage , we divide the source domains into several meta-source domains and a meta-target domain and explore the adaptive model by incorporating the information of the meta-target sample into the generation of the model parameters . For any given meta-target sample , we propose a variational distribution generated by this sample and the meta-source data to approximate the model distribution obtained by the meta-target data . In the test phase , the adapted models for the samples from the real target domains are generated on the fly by the variational model distribution . The random splits of the meta-source and meta-target domains expose the model to domain shifts and mimic the real generalization process from source domains to the target domain . Thus , the model is endowed with the ability to adapt to any unseen sample by end-to-end training with the source data only . By doing so , our method does not need to introduce any extra target data or fine-tuning operations on the target domain . A comparison of the training/test-time settings among different domain generalization ( and domain adaptation ) methods is summarized in Table 1 . We demonstrate the effectiveness of our single sample generalization by conducting experiments on commonly-used benchmarks for domain generalization . Our ablation studies under different domain generalization settings demonstrate the advantages of our method in comparison to other recent adaptive methods . 2 RELATED WORK . Domain adaptation To handle the domain shift between source and target domains , considerable efforts have been devoted to developing methods for domain adaptation ( Long et al. , 2015 ; Lu et al. , 2020 ; Hoffman et al. , 2018 ; Kumar et al. , 2010 ; Tzeng et al. , 2017 ; Luo et al. , 2019 ) . However , the assumption that target data is accessible during training is often invalid . Recently , some even more challenging settings were proposed in domain adaptation , e.g. , few-shot ( Motiian et al. , 2017a ) and one-shot domain adaptation ( Dong & Xing , 2018 ; Luo et al. , 2020 ) , but there are still a few or one target sample available during training . In this work , we focus on the more strict setting , domain generalization , that assumes not a single target sample is accessible during training . We construct our model under a meta-learning framework to learn the ability of adaptation to single samples across source domains . At test time , no more learning is needed and we generate an adapted model for each individual target sample by only a forward pass , without any extra fine-tuning operation . Domain generalization The problem of domain generalization was introduced by Blanchard et al . ( 2011 ) , and formally posed by Muandet et al . ( 2013 ) . Muandet et al . ( 2013 ) and Ghifary et al . ( 2016 ) both considered domain-invariant learning by minimizing the dissimilarity of features across source domains . Following these works , many methods explored domain invariant learning for domain generalization ( Motiian et al. , 2017b ; Rojas-Carulla et al. , 2018 ; Li et al. , 2018b ; Seo et al. , 2020 ; Zhao et al. , 2020 ) . Alternatively , augmenting the source domain data has been explored to simulate more domain shift during training ( Shankar et al. , 2018 ; Volpi et al. , 2018 ; Zhou et al. , 2020a ; Qiao et al. , 2020 ; Zhou et al. , 2020b ) . Our method tackles domain shift by learning to adapt to each unseen target sample , which is orthogonal to domain-invariant and domain augmentation methods . Domain meta-learning Meta-learning methods have also been studied for domain generalization ( Balaji et al. , 2018 ; Du et al. , 2020 ; Xiao et al. , 2021 ; Du et al. , 2021 ) . Li et al . ( 2018a ) introduced the model agnostic meta-learning ( Finn et al. , 2017 ) into domain generalization . Dou et al . ( 2019 ) further enhanced the method with global and local regularization . We follow the experimental settings of previous methods ( Du et al. , 2020 ; Xiao et al. , 2021 ) that split source domains into meta-source and meta-target domains . By mimicking the adaptation process at training time , our model learns to learn the adaptable classifiers from meta-source domains and the meta-target sample . Test-time adaptation Recently , several methods have proposed to adapt a model trained on source data to unseen target data at test time ( D ’ Innocente et al. , 2019 ; Sun et al. , 2020 ; Wang et al. , 2021 ; Dubey et al. , 2021 ) . Sun et al . ( 2020 ) trained their model on source domains with both supervised and self-supervised objective functions , and then fine-tuned it on target data with only the self-supervised one . Wang et al . ( 2021 ) further improved this method by removing the self-supervised module , while fine-tuning the model with an entropy minimization loss . Differently , Dubey et al . ( 2021 ) trained a domain embedding network with unsupervised source data , which is utilized to generate the target domain embedding by hundreds of unlabeled target samples to adapt the model at test time , without fine-tuning or further training . These adaptive methods all need extra steps , e.g. , fine-tuning the model or training embedding networks , and batches of target samples for good adaptation . In contrast , our method learns to adapt to each target sample using just its own information without any extra target samples and fine-tuning operations . The adaptation is achieved by generating the model for each target sample with only one forward pass . Another recent work utilized test-time augmentation to improve robustness at test time ( Chai et al. , 2021 ) . While it is of interest to adapt a test-time sample to source domains with such test-time augmentations , our method aims to adapt the model trained by source domains to each target sample . 3 METHODOLOGY . In domain generalization , several domains are defined as different data distributions in the joint space X × Y , where X and Y denote the input space and the label space . The domains are divided into non-overlapping source domains , collectively denoted as S and the target domains T . During training , only data from source domains is accessible . Data from target domains is never seen . Our goal is to learn a model on the source domains only that generalizes well to the target domains . As we focus on homogeneous domain generalization ( Zhou et al. , 2021 ) in this work , all domains share the same label space . Since the target data is inaccessible during training , we design a meta-learning framework with episodic training ( Li et al. , 2019 ; Balaji et al. , 2018 ) that endows the model with the capability of learning to generalize on single test samples from any unseen domain . To this end , according to the precise domain annotations ( if available ) or image clusters , we divide the source domain ( s ) S into several meta-source domains S ′ and a meta-target domain T ′ in each iteration to expose the model to domain shifts . After training , the learned model is applied to the target domain T for evaluation . In this case , the model is exposed to different domain shifts to learn the ability to handle it . After training , the learned model is applied to the target domain T for evaluation . Variational learning We develop our model under the probabilistic framework . Assume we have a single sample xt′ from the meta-target domain T ′ , and we would like to predict its class label yt′ . We consider the conditional predictive log-likelihood log p ( yt′ |xt′ , T ′ ) . By incorporating the model θt′ into the prediction distribution , we have log p ( yt′ |xt′ , T ′ ) = log ∫ p ( yt′ |xt′ , θt′ ) p ( θt′ |T ′ ) dθt′ , ( 1 ) where we condition the prediction on T ′ to leverage information from the meta-target data and ( xt′ , yt′ ) denotes the input-output pair of the meta-target sample . We refer to p ( θt′ |T ′ ) as the meta-prior distribution since it is conditioned on the meta-target data . However , in practice , the labeled data from the target domain is inaccessible , leading to an intractable distribution p ( θt|T ) during inference . To address the problem , we propose to approximate the model function p ( θt|T ) with the accessible source data by the variational inference approach , which provides a framework to approximate intractable distributions via optimization . By training under the meta-learning framework , our method will learn the ability to approximate the model functions of unseen target domains by source domain data . Variational domain generalization We introduce a variational distribution q ( θt′ |S ′ ) to approximate the true posterior over the model parameters θ . By incorporating q ( θt′ |S ′ ) into eq . ( 1 ) , we derive a lower bound of the conditional predictive log-likelihood : log p ( yt′ |xt′ , T ′ ) = log ∫ p ( yt′ |xt′ , θt′ ) p ( θt′ |T ′ ) dθt′ ≥ Eq ( θt′ ) [ log p ( yt′ |xt′ , θt′ ) ] − DKL [ q ( θt′ |S ′ ) ||p ( θt′ |T ′ ) ] . ( 2 ) The introduced Kullback-Leibler ( KL ) divergence term is minimized to reduce the distance between q ( θt′ |S ′ ) and p ( θt′ |T ′ ) , which encourages the model θ inferred by q ( θt′ |S ′ ) from meta-source domains to be generalizable to the meta-target domain . Although eq . ( 2 ) reduces the difference between the variational distributions q ( θt′ |S ′ ) and p ( θt′ |T ′ ) , q ( θt′ |S ′ ) still has limited capacity of handling the meta-target samples due to “ adaptivity gap ” ( Dubey et al. , 2021 ) . Since only little information about the meta-target domain T ′ and the meta-target sample xt′ is available in S ′ , there is no guarantee that the variational distribution q ( θt′ |S ′ ) will be adapted to the meta-target domain T ′ . To reduce the “ adaptivity gap ” , the model needs to acquire the ability to effectively use the target information . Thus , we further propose to learn to adapt the model θt′ by taking more information of the meta-target data into account . Single sample generalization In domain generalization , especially for real applications , when given data from an unseen target domain , the only accessible information of the target data is from the target sample , without any annotations . We therefore propose to utilize the information contained in the given target sample to construct the adapted model . Instead of utilizing q ( θt′ |S ′ ) as the variational distribution , we incorporate the meta-target sample xt′ into the variational distribution q ( θt′ |xt′ , S ′ ) to approximate the true posterior distribution . Given an unlabeled sample xt′ from the meta-target domain T ′ , we have a new evidence lower bound ( ELBO ) as follows : log p ( yt′ |xt′ , T ′ ) = log ∫ p ( yt′ |xt′ , θt′ ) p ( θt′ |T ′ ) dθt′ ≥ Eq ( θt′ ) [ log p ( yt′ |xt′ , θt′ ) ] − DKL [ q ( θt′ |xt′ , S ′ ) ||p ( θt′ |T ′ ) ] . ( 3 ) By conditioning on the given sample xt′ and the source data S ′ , the adapted variational distribution q ( θt′ |xt′ , S ′ ) contains both the categorical information provided by the meta-source data and the domain information given by the meta-target sample . This to a large extent guarantees the inferred model to be discriminative and adapted to the given meta-target domain sample . Moreover , the KL divergence in eq . ( 3 ) further acts as a regularizer to guide the variational distribution to be similar to the parameter distribution p ( θt′ |T ′ ) obtained by the meta-target data , which is assumed to be the most suitable parameter for the meta-target sample xt′ . This makes the model better adapted to xt′ Based on the meta-learning framework and eq . ( 3 ) , our model learns the ability to adapt the metasource model to each meta-target instance across different domain shifts . The learned ability is exploited to handle domain shifts at test time by adapting the source model to each target instance . During inference , our method directly uses the variational distributions q ( θt|xt , S ) as the adapted model generated by the information of the given target sample xt and the data from all source domains S . Thus , the method is able to generalize to each unseen target sample without the requirement of any extra data from the target domains or fine-tuning operations . The derivations of eq . ( 2 ) and eq . ( 3 ) , as well as more discussions on their tightness are provided in Appendix A . Learning to adapt classifiers The variational inference framework in eq . ( 3 ) indicates that we can learn to adapt all model parameters θ to the target sample with its own information . This could , however , be computationally infeasible due to the large number of parameters in the network . Thus , for computational efficiency and feasibility we utilize the features from the feature extractor as the sample-specific information to construct the adapted classifier , which contains much fewer parameters . To do so , we divide the model θ into a feature extractor ϕ and a classifier w. The feature extractor is shared across domains , while the classifier is trained to be adapted to each single sample by eq . ( 3 ) . Both the meta-prior distribution p ( wt′ |T ′ ) and the variational posterior distribution q ( wt′ |xt′ , S ′ ) of the classifier wt′ are generated by amortized inference using the amortization technique ( Kingma & Welling , 2013 ; Gershman & Goodman , 2014 ; Gordon et al. , 2019 ; Shen et al. , 2021 ) . The amortization technique provides a natural way to generate model parameters by feature representations , which is one appropriate solution for our method to incorporate the target feature into the construction of the adapted variational posterior distribution . Moreover , the amortization networks can be trained and evaluated together with the model , without introducing extra training steps and fine-tuning operations to achieve the adaptation ability . Specifically , we use the center features of samples in each class from the meta-target domain and meta-source domains as the representatives of T ′ and S ′ . We generate the meta-prior distribution p ( wt′ |T ′ ) by taking the center features of T ′ as input to the amortized inference network θa . For the variational distribution q ( wt′ |xt′ , S ′ ) , the input to the amortized inference network is the concatenation of the center features of the meta-source domain S ′ with the features of the target sample xt′ for adaptation . The amortized inference network θa is shared for the inference of both the meta-prior and posterior distributions . However , the direct concatenation operation is not necessarily optimal as the input for the amortized inference network . We propose to generate the parameter distribution using the meta-source domains and then adapt it with the given meta-target sample . To do so , we first generate a classifier distribution p ( ws′ |S ′ ) using the center features of the meta-source data S ′ . After that , we take p ( ws′ |S ′ ) as a prior and combine it with the meta-target sample to estimate the variational posterior , as follows : q ( wt′ |xt′ , S ′ ) = ∫ p ( wt′ |xt′ , ws′ ) p ( ws′ |S ′ ) dws′ . ( 4 ) This establishes a hierarchical variational inference of the posterior distribution . We show in our experiments that the hierarchical inference yields better results than the direct concatenation . In practice , we approximate the integral by : q ( wt′ |xt′ , S ′ ) = L∑ ℓ=1 p ( wt′ |xt′ , w ( ℓ ) s′ ) , ( 5 ) where w ( ℓ ) s′ ∼ p ( ws′ |S ′ ) and L is the number of Monte Carlo samples . The negative KL divergence term in eq . ( 3 ) tries to close the gap between the variational adaptive distribution q ( wt′ |xt′ , S ′ ) and the meta-prior distribution p ( wt′ |T ′ ) , aiming to make the variational distribution adapt better to the target data . However , there is no guarantee that the classifier produced by the meta-prior is discriminative enough for classification . To address this issue , we introduce an intermediate supervision based on the meta-prior distribution to make it as predictive as possible , by maximizing : Ep ( wt′ |T ′ ) [ log p ( yt′ |xt′ , wt′ ) ] . ( 6 ) This is also implemented as a cross-entropy loss using Monte Carlo sampling from the meta-prior . An illustration of our architecture is provided in Figure 2 . We also provide an algorithm in Appendix B. Empirical objective function By incorporating eq . ( 5 ) and eq . ( 6 ) into eq . ( 3 ) , we obtain the final empirical objective function : L= M∑ m=1 log p ( yt′ |xt′ , w ( m ) t′ ) + N∑ n=1 log p ( yt′ |xt′ , w ( n ) t′ ) − DKL ( L∑ ℓ=1 p ( wt′ |xt′ , w ( ℓ ) s′ ) ||p ( wt′ |T ′ ) ) ( 7 ) where we use Monte Carlo sampling : w ( m ) t′ ∼ ∑L ℓ=1 q ( wt′ |xt′ , w ( ℓ ) s′ ) and w ( n ) t′ ∼ p ( wt′ |T ′ ) . As illustrated in Figure 3 , at training time ( a ) , we train the model episodically on meta-source domains S ′ and a meta-target domain T ′ . At inference time ( b ) , given a single sample xt from the target domain T , we use this sample to adapt the ws generated by the source data . The adaptation is achieved by generating the adapted classifiers wt for each target sample with only one forward pass using the learned amortization inference network θa .

Summary: - The paper describes a method for domain generalization that performs test-time adaptation using a single test example at a time (as opposed to a transductive setting used in other works where a whole batch of test examples are used). - The method is cast as a meta learning task. The training data is split into meta-training/test domains to mimic the adaptation process during training.