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 method to calibrate batch normalization statistics at test time to improve a model’s cross-domain generalization ability under covariate shifts. The authors validate the effectiveness of the method on several datasets and tasks. The authors also conduct some interesting analyses to help better understand the method.