Test-time Batch Statistics Calibration for Covariate Shift

1 INTRODUCTION . Deep neural networks ( DNNs ) achieve impressive success across various applications , but heavily rely on the independent and identical distribution ( i.i.d . ) assumption . However , in real-world applications , the model is prone to encounter the novel instances . For examples , an automatic pilot should have robust performance under different weather conditions . Unfortunately , when applying DNNs to novel environment , the performance has a clear degradation due to the covariate shift ( Ben-David et al. , 2010 ) , i.e. , the test data distribution differs from the training distribution . Domain adaptation ( DA ) is a promising alternative , which transfers the knowledge learned on labeled source domain to unlabeled target domain , where the data distribution is distinct ( Long et al. , 2015 ) . However , domain adaptation needs the pre-collected target data , which is not applicable . Unlike DA , Domain generalization ( DG ) aims at training a general model from multiple source domains and generalizing to the unseen target domain . DG is more challenging since the target domain is totally unseen during training . Recently , another practical scenario named test-time adaptation ( TTA ) is proposed . In TTA , the target data are not pre-collected for iterative training , but used for adapting the source-trained model during inference . We show the comparison between different settings in Table 1 . To make the comparison clear , we introduce two indicators . Iterative training means the model is trained on the unlabeled target data iteratively . Online training means the model parameters are updated during inference . Noticing that a recent DG work ( Pandey et al. , 2021 ) adapts the model during inference , while most previous DG methods did not . We call it as optimization-free TTA . In this paper , we focus on the practical scenarios : DG and TTA . One of the main approaches to adapt during inference is test-time normalization ( T-BN ) ( Nado et al. , 2020 ) . T-BN re-calculates the target batch statistics to replace the source statistics in BN layers during inference . Motivated by T-BN , Wang et al . ( 2021 ) proposed to perform test time adaptation by re-calculating the target batch statistics and updating the affine parameters in BN layers with entropy minimization . However , this paradigm has critical restrictions . For representation learning , the discriminative representations are crucial for recognition task . Substituting the source statistics with target statistics in BN layers will inevitably lead to a mismatch with the source-trained model parameters . This mismatch will probably perturb the original discriminative structures . Another limitation is the estimation error on target batch statistics . The source statistics in BN layers are updated in a moving average manner during training time , while the target statistics are calculated in each batch during test time . The statistics estimated in a batch introduces more errors in reflecting the domain characteristics . A preliminary empirical investigation of the mentioned restrictions are shown in Fig . 1 and Table 2 . Motivated by the hidden restrictions of T-BN , we propose a more general method named α-BN to calibrate the batch statistics during inference . Specifically , we mix up the source and target statistics in BN layers to both alleviate the domain shift and preserve the discriminative structures . Equipped with α-BN , common DG models can be further improved without any training . Based on α-BN , we further propose an unified test-time adaptation framework named CORE with an online optimization , which exploits the pairwise Class correlation to facilitate robust and accurate test time adaptation . To sum up , we have following contributions : 1 . We investigate two practical yet challenging transfer learning scenarios : domain generalization and test time adaptation , which release the requirement of pre-collected target domain data . 2 . Motivated by the hidden restrictions of test time normalization , we present a general formulation α-BN for both alleviating domain shifts and preserving discriminative informations . It can be seamlessly incorporated into mainstream deep neural networks to improve the generalization on unseen domains without any training . 3 . Based on α-BN , we propose a unified framework CORE for test-time adaptation . CORE optimizes the pairwise class correlation in an unsupervised online learning manner . 4 . We conduct numerous experiments on total twelve datasets from three topics : robustness to corruptions , DG on image classification and DG on semantic segmentations . The empirical results show that both α-BN and CORE achieve the state-of-the-art ( SoTA ) performance in their respective communities . The result on GTA5→ Cityscapes , for instance , is improved from 28.4 % to 43.9 % without any training , which even outperforms the SoTA source-free DA method . 2 RELATED WORKS . Domain Adaptation and Generalization Domain adaptation ( DA ) enables transferring the knowledge from source domain to target domain by jointly optimizing both labeled source data and unlabeled target data . This paradigm has gained a lot of attention in the last decade and various methods are proposed , which can be roughly divided into three categories : metric learning ( Long et al. , 2015 ; Sun & Saenko , 2016 ; Kang et al. , 2019 ; Li et al. , 2020 ) , adversarial training ( Ganin et al. , 2016 ; Long et al. , 2018 ; Tsai et al. , 2018 ) and self-training Zou et al . ( 2019 ) ; Liang et al . ( 2020 ) ; Ge et al . ( 2020 ) . However , training with the large amounts source data is inefficient and impractical in many real-world applications . To address this issue , source-free DA is proposed , which adapts to the target domain with target data and the model pre-trained on source domain . Chidlovskii et al . ( 2016 ) are the first to investigate source-free DA and proposed a denoising auto-encoder for adaptation . SHOT ( Liang et al. , 2020 ) is another representative work , which proposed information maximization loss and clustering-based pseudo-labeling . However , despite getting rid of the source domain data , source-free DA also requires the pre-collected target domain data for iterative training , limiting its application scenarios . Therefore , domain generalization ( DG ) and fully test-time adaptation ( TTA ) are proposed ( Let us elaborate TTA in the next paragraph ) . Domain generalization aims at generalizing the model trained on a ( multiple ) source domain ( s ) to the unseen target domain directly . Recently , various DG methods are proposed including domain-invariant representation learning ( Zhao et al. , 2020 ; Matsuura & Harada , 2020 ) , proxy tasks ( Carlucci et al. , 2019 ; Huang et al. , 2020 ) , augmentations ( Volpi et al. , 2018 ; Zhou et al. , 2021 ) , meta-learning ( Li et al. , 2018 ; Balaji et al. , 2018 ) and so on . However , Gulrajani & Lopez-Paz ( 2020 ) provided a DG benchmark named DomainBed for fair comparison and found a well-implemented empirical risk minimization ( ERM ) model outperforms most DG methods . Recently , Pandey et al . ( 2021 ) proposed label-preserving target projections during inference time for DG . This work differs from most previous works , which focus on learning from source domains , while it performs optimization-free TTA during inference . Our proposed α-BN also belongs to it , but the methods differ . Fully Test-time Adaptation Fully test-time adaptation is proposed by Wang et al . ( 2021 ) , which adapts the model to target domain by online training . TTA can be seen as a compromise between source-free DA and DG . Different from source-free DA , TTA does not require the pre-collected target domain data but trains on the target data in an online manner . Also , different from DG , TTA allows optimization during test , which introduces additional test time cost but usually guarantees better performance . Therefore , TTA is a really practical scenario since it gets rid of iterative training and yields better performance compared to generalizing to the new environment directly . Wang et al . ( 2021 ) proposed TENT to achieve TTA by feature modulation . Feature modulation contains two steps : test-time normalization ( we will elaborate it in the next paragraph . ) and affine parameters optimization by entropy minimization , which is a widely-used regularization term on DA and semisupervised learning ( Grandvalet et al. , 2005 ) . Another similar scenario is test-time training , which optimizes the networks before making a prediction during inference Sun et al . ( 2019 ) . Since this setting is weaker than TTA , we mainly talk about TTA in this paper . Normalization and Adaptation Batch normalization ( BN ) is widely-used in DNNs nowadays for stable training and fast converge . BN is originally proposed to alleviate the internal covariate shift during training a very deep neural networks Ioffe & Szegedy ( 2015 ) . Recently , Schneider et al . ( 2020 ) and Nado et al . ( 2020 ) discovered that updating the batch statistics during testing improves the robustness to common corruptions . In this paper , we call it as test-time normalization . Similar to their works , Wang et al . ( 2021 ) proposed feature modulation , which also updates the batch statistics rather than freezes them . The key insight for these methods is that batch statistics are closely related to the domain characteristics ( Li et al. , 2016 ; Pan et al. , 2018 ) . Based on this finding , Jeon et al . ( 2021 ) and Zhou et al . ( 2021 ) proposed to synthesise the novel domain styles to facilitate generalization . Another kind of work focuses on instance normalization ( Ulyanov et al. , 2016 ) . Similarly , Huang & Belongie ( 2017 ) found the instance-wise statistics are related to instance characteristics ( i.e. , image styles ) . Motivated by this finding , Zhou et al . ( 2021 ) proposed to generate novel in- stances by mixing instance-level feature statistics to enhance out-of-distribution generalization . It is worth noticing that the aforementioned DG method can be summarized as an augmentation-based technique . Different from them , we propose a post-processing method to calibrate the batch statistics on target domain during test time . 3 UNDERSTANDING TEST-TIME NORMALIZATION . Test-time normalization ( T-BN ) re-calculates the batch statistics on target domain during inference . Since the batch statistics are closely related to the domain characteristics ( Li et al. , 2016 ; Pan et al. , 2018 ) , T-BN adapts the model to target domain explicitly . However , every coin has two sides , and T-BN is not a free lunch . During training on the source domain , the model parameters are associated with the source statistics . Therefore , substituting the source statistics with the target ones inevitably results in a mismatch with the model parameters , which leads to the degradation of discrim- inative structures . In short , T-BN alleviates the negative effects caused by domain shift , but perturbs the discriminative structures . We report the averaged accuracy of “ Source ” and “ T-BN ” on three DG classification benchmarks in Table 2 . We observe that the accuracy of “ T-BN ” is consistently lower than “ Source ” , revealing that substituting the source statistics by the estimated target batch statistics directly is not effective in generalizing to the new environment . To further understand T-BN , we begin with the following two perspectives . Error of ideal target hypothesis . Based on domain adaptation theory ( Ben-David et al. , 2010 ) , the domain shift can be reflected by the error of the ideal target hypothesis based on the target representations learned by source model . “ target representation ” means the representations are obtained on the target domain data . To obtain the ideal target hypothesis , we train a new classifier over the target representations with corresponding labels . Two methods are compared : “ Source ” and “ T-BN ” . “ Source ” obtains the target representations by the source model directly , while “ T-BN ” performs test-time normalization . The error of the ideal target hypothesis is shown in Fig . 1 ( b ) . As expected , the error of the ideal target hypothesis in “ T-BN ” is lower over all tasks . It is worth noticing that the only difference between them is that T-BN normalizes the BN layer inputs by the target statistics rather than the source statistics , and others remain consistent ( e.g. , the same network architecture and the same network parameters ) . Therefore , we reasonably postulate that T-BN alleviates the domain shift , which results in the lower error of the ideal target hypothesis . Representation visualization . Discriminative representation learning is essential for recognition task . The discriminative representation satisfies two basic principles : intra-class tightness and interclass separation . To qualitatively verify that how T-BN affects the learned representations , we visualize the target representations in Fig . 1 ( c ) . The variance of each category cluster in “ T-BN ” is significantly larger compared to “ Source ” , which indicates the discriminative structures are injured due to the mismatch between target statistics and source model parameters .

Summary: This paper presents a domain adaptation algorithm for the recently introduced test-time adaptation setting by adapting batch-norm parameters. The algorithm has two key components: first, it adapts batch-norm statistics using a linear combination of source and estimated target domain statistics; second, it uses a class correlation optimzation loss to optimize batch-norm parameters on the test batches. These design choices differ from previously described methods. The resulting algorithm is evaluated on a range of datasets and tasks including common corruptions and the DomainBed benchmark for image classification, as well as a benchmark for image segmentation.