Uncertainty Modeling for Out-of-Distribution Generalization

1 INTRODUCTION . Deep neural networks have shown impressive success in computer vision , but with a severe reliance on the assumption that the training and testing domains follow an independent and identical distribution ( Ben-David et al. , 2010 ; Vapnik , 1992 ) . This assumption , however , does not hold in many real-world applications . For instance , when employing segmentation models trained on sunny days for rainy and foggy environments ( Choi et al. , 2021 ) , or recognizing art paintings with models that trained on photographs ( Li et al. , 2017 ) , inevitable performance drop can often be observed in such out-of-distribution deployment scenarios . Therefore , the problem of domain generalization , aiming to improve the robustness of the network on various unseen testing domains , becomes quite important . Previous works ( Huang & Belongie , 2017 ; Li et al. , 2021 ) demonstrate that feature statistics ( mean and standard deviation ) , as the moments of the learned features , carry informative domain characteristics of the training data . Domain characteristics primarily refer to the information that is more specific to the individual domains but less relevant to the task objectives , such as the photo style and capturing environment information in object recognition . Consequently , domains with different data distributions generally have inconsistent feature statistics ( Wang et al. , 2020b ; 2019a ; Gao et al. , 2021a ) . Most deep learning methods follow Empirical Risk Minimization principle ( Vapnik , 1999 ) to minimize their average error over the training data ( Shen et al. , 2021 ) . Despite the satisfactory performance on the training domain , these methods do not explicitly consider the uncertain statistics discrepancy caused by potential domain shifts during testing . As a result , the trained models tend to overfit the training domain and show vulnerability to the statistic changes at testing time , substantially limiting the generalization ability of the learned representations . Intuitively , the test domains may bring uncertain statistics shifts with different potential directions and intensities compared to the training domain ( as shown in Figure 1 ) , implying the uncertain nature of domain shifts . Considering such “ uncertainty ” of potential domain shifts , synthesizing novel feature statistics variants to model diverse domain shifts can improve the robustness of the trained network to different testing distributions . Towards this end , we introduce a novel probabilistic method to improve the network generalization ability by properly modeling Domain Shifts with Uncertainty ( DSU ) , i.e. , characterizing the feature statistics as uncertain distributions . In our method , instead of treating each feature statistic as a deterministic point measured from the feature , we hypothesize that the feature statistic , after considering potential uncertainties , follows a multi-variate Gaussian distribution . The distribution “ center ” is set as each feature ’ s original statistic value , and the distribution “ scope ” represents the variant intensity considering underlying domain shifts . Uncertainty estimation is adopted here to depict the distribution “ scope ” of probabilistic feature statistics . Specifically , we estimate the distribution “ scope ” based on the variances of the mini-batch statistics in an efficient non-parametric manner . Subsequently , feature statistics variants are randomly sampled from the estimated Gaussian distribution and then used to replace the original deterministic values for modeling diverse domain shifts , as illustrated in Figure 2 . Due to the generated feature statistics with diverse distribution possibilities , the models can be trained to properly alleviate the domain perturbations and encode better domain-invariant features . Our proposed method is simple yet fairly effective to alleviate performance drop caused by domain shifts , and can be readily integrated into existing networks without bringing additional model parameters or loss constraints . Comprehensive experiments on a wide range of vision tasks demonstrate the superiority of our proposed method , indicating that introducing uncertainty to feature statistics can well improve models ’ generalization against domain shifts . 2 RELATED WORK . 2.1 DOMAIN GENERALIZATION . Domain generalization ( DG ) has been attracting increasing attention in the past few years , which aims to achieve out-of-distribution generalization on unseen target domains using only single or multiple source domain data for training ( Blanchard et al . ( 2011 ) ) . Research on addressing this problem has been extensively conducted in the literature ( Zhou et al . ( 2021a ) ; Wang et al . ( 2021 ) ; Shen et al . ( 2021 ) ) . Here some studies that are more related to our work are introduced below . Data Augmentation : Data augmentation is an effective manner for improving generalization ability and relieving models from overfitting in training domains . Most augmentation methods adopt various transformations at the image level , such as AugMix ( Hendrycks et al . ( 2020 ) ) and CutMix ( Yun et al . ( 2019 ) ) . Besides using handcraft transformations , mixup ( Zhang et al . ( 2018 ) ) trains the model by using pair-wise linearly interpolated samples in both the image and label spaces . Manifold Mixup ( Verma et al . ( 2019 ) ) further adopts this linear interpolation from image level to feature level . Some recent works extend the above transformations to feature statistics for improving model generalization . MixStyle ( Zhou et al . ( 2021b ) ) adopts linear interpolation on feature statistics of two instances to generate synthesized samples . The pAdaIn ( Nuriel et al . ( 2021 ) ) swaps statistics between the samples applied with a random permutation of the batch . Invariant Representation Learning : The main idea of invariant representation learning is to enable models to learn features that are invariant to domain shifts . Domain alignment-based approaches ( Li et al . ( 2018c ; b ) ) learn invariant features by minimizing the distances between different distributions . Instead of enforcing the entire features to be invariant , disentangled feature learning approaches ( Chattopadhyay et al . ( 2020 ) ; Piratla et al . ( 2020 ) ) decouple the features into domain-specific and domain-invariant parts and learn their representations simultaneously . In addition , normalizationbased methods ( Xingang Pan & Tang ( 2018 ) ; Choi et al . ( 2021 ) ) can also be used to remove the style information to obtain invariant representations . Learning Strategies : There are also some effective learning strategies that can be leveraged to improve generalization ability . Ensemble learning is an effective technique in boosting model performance . The ensemble predictions using a collection of diverse models ( Zhou et al . ( 2020b ) ) or modules ( Seo et al . ( 2020 ) ) can be adopted to improve generalization and robustness . Meta-learningbased methods ( Finn et al . ( 2017 ) ; Li et al . ( 2018a ) ) learn to simulate the domain shifts following an episode training paradigm . Besides , self-challenging methods , such as RSC ( Huang et al . ( 2020 ) ) , force the model to learn a general representation by discarding dominant features activated on the training data . 2.2 UNCERTAINTY IN DEEP LEARNING . Uncertainty capturing the “ noise ” and “ randomness ” inherent in the data has received increasing attention in deep representation learning . Variational Auto-encoder ( Kingma & Welling ( 2013 ) ) , as an important method for learning generative models , can be regarded as a method to model the data uncertainty in the hidden space . Dropout ( Srivastava et al . ( 2014 ) ) , which is widely used in many deep learning models to avoid over-fitting , can be interpreted to represent model uncertainty as a Bayesian approximation ( Gal & Ghahramani ( 2016 ) ) . In some works , uncertainty is used to address the issues of low-quality training data . In person re-identification , DistributionNet ( Gal & Ghahramani ( 2016 ) ) adopts uncertainty to model the person images of noise-labels and outliers . In face recognition , DUL ( Chang et al . ( 2020 ) ) and PFE ( ( Shi & Jain , 2019 ) ) apply data uncertainty to simultaneously learn the feature embedding and its uncertainty , where the uncertainty is learned through a learnable subnetwork to describe the quality of the image . Different from the aforementioned works , our proposed method is used to model the feature statistics uncertainty under potential domain shifts and acts as a feature augmentation method for handling our-of-distribution generalization problem . 3 METHOD . 3.1 PRELIMINARIES . Given x ∈ RB×C×H×W to be the encoded features in the intermediate layers of the network , we denote µ ∈ RB×C and σ ∈ RB×C as the channel-wise feature mean and standard deviation of each instance in a mini-batch , respectively , which can be formulated as : µ ( x ) = 1 HW H∑ h=1 W∑ w=1 xb , c , h , w , ( 1 ) σ2 ( x ) = 1 HW H∑ h=1 W∑ w=1 ( xb , c , h , w − µ ( x ) ) 2 . ( 2 ) As the abstraction of features , feature statistics can capture informative characteristics of the corresponding domain ( such as color , texture , and contrast ) , according to previous works ( Huang & Belongie , 2017 ; Li et al. , 2021 ) . In out-of-distribution scenarios , the feature statistics often show inconsistency with training domain due to different domain characteristics ( Wang et al. , 2019a ; Gao et al. , 2021a ) , which is ill-suited to deep learning modules like nonlinearity layers and normalization layer and degenerates the model ’ s generalization ability ( Wang et al. , 2020b ) . However , most of the deep learning methods only treat feature statistics as deterministic values measured from the features while lacking explicit consideration of the potential uncertain statistical discrepancy . Owing to the model ’ s inherent vulnerability to such discrepancy , the generalization ability of the learned representations is limited . Some recent methods ( Nuriel et al. , 2021 ; Zhou et al. , 2021b ) utilize feature statistics to tackle the domain generalization problem . Despite the success , they typically adopt linear manipulation ( i.e. , exchange and interpolation ) on pairwise samples to generate new feature statistics , which limits the diversity of synthetic changes . Specifically , the direction of their variants is determined by the chosen reference sample and such internal operation restricts their variant intensity . Thus these methods are sub-optimal when handling the diverse and uncertain domain shifts in real world . 3.2 MODELING DOMAIN SHIFTS WITH UNCERTAINTY . Given the arbitrary testing domains with uncertain feature statistic shifts in both direction and intensity , properly modeling the domain shifts becomes an important task for tackling the challenge of domain generalization problem . Considering the uncertainty and randomness of domain shifts , it is promising to employ the methods of “ uncertainty ” to treat the “ uncertainty ” of domain shifts . In this paper , we propose a novel method by modeling Domain Shifts with Uncertainty ( DSU ) . Instead of treating each feature statistic as a deterministic value measured from the learned feature , we hypothesize that the distribution of each feature statistic , after considering potential uncertainties , follows a multi-variate Gaussian distribution . This means each feature statistic has a probabilistic representation drawn from a certain distribution , i.e. , the feature statistics mean and standard deviation follow N ( µ , Σµ ) and N ( σ , Σσ ) , respectively . Specifically , the corresponding Gaussian distribution ’ s center is set as each feature ’ s original statistics , while the Gaussian distribution ’ s standard deviation describes the uncertainty scope for different potential shifts . Through randomly sampling diverse synthesized feature statistics with the probabilistic approach , the models can be trained to improve the robustness of the network against statistics shifts .

Summary: This paper studies the problem of out-of-distribution generalization by modeling uncertainty in feature statistics. It improves the network generalization ability by modeling the uncertainty of domain shifts with synthesized feature statistics during training. Instead of being deterministic values, the feature statistics are hypothesized to follow a multivariate Gaussian distribution. The proposed method is tested on various tasks including image classification, semantic segmentation, and instance retrieval, and shows strong performance compared to state-of-the-art methods.