Maximum Mean Discrepancy for Generalization in the Presence of Distribution and Missingness Shift

1 INTRODUCTION . Machine learning models often work on the assumption that the training and test datasets follow the same distribution . In practice , this might not be the case , yet the difference in distribution , called dataset shift , is often ignored . This might be acceptable when the dataset shift is small , but small scale shift is not always guaranteed . When it is not , it can lead to poor performance on shifted test sets . Dataset shift can be divided into three main categories : covariate shift , prior probability shift , and concept shift . Covariate shift occurs when the distribution of input features changes , prior probability shift occurs when the distribution of target variables changes , and concept shift occurs when the relationship between input features and target variables changes ( Quinonero-Candela et al. , 2008 ) . In this paper , we focus on covariate shifts . There have been some attempts to address covariate shift in previous literature . One straightforward approach is to identify the shifted covariates and drop them from the feature set . However , this can be challenging for high-dimensional data , and one may also lose useful information in the process . A softer method is to use importance re-weighting by assigning higher weights to training samples that are more similar to the samples in the test dataset . The importance weights can be calculated through various techniques , including density estimation ( Shimodaira , 2000 ; Kanamori & Shimodaira , 2000 ; Zadrozny , 2004 ; Dudı́k et al. , 2005 ) , kernel mean matching ( Huang et al. , 2007 ) , Kullback-Leibler importance estimation ( Sugiyama et al. , 2008 ) , and discriminative learning ( Bickel et al. , 2009 ) . These methods work particularly well for sample selection bias , but they treat samples at a row/sample level which is not flexible enough . In this paper , we propose optimizing an auxiliary maximum mean discrepancy ( MMD ) loss with a mixture of kernels to treat the covariate shift problem . MMD is already widely used in unsupervised domain adaptation ( UDA ) ( Tzeng et al. , 2014 ; Long et al. , 2015 ; 2017 ; Yan et al. , 2017 ; Chen et al. , 2020 ) . Here we extend it to the more general covariate shift problem . Conceptually , we can think of domain adaptation as an extreme case of covariate shift as the joint distribution of input variables shifts to a different domain entirely . Moreover , previous UDA works has never considered an important special case of covariate shift that involves data missingness . For example , in an internal project of carbon emission prediction in Section 4.3 , we try to train a model with the companies who reported carbon emission to infer carbon emission for the unreported companies . However , we find that the unreported companies usually have higher missing rates in various sets of features compared to the reported companies . A similar situation can happen in other use cases . For instance , various physical measurements and clinical exams need to be taken before a diagnostic decision of a disease can be made , but when estimating the probability of risk of the disease on a new group of people , different measurements and exams can be missing for different subgroups . Or , in financial fraud detection , the unlabeled set may also have higher feature missing rates compared to the labeled set , especially in the self-reported features . If we simply ignore the data missingness shift , the model performance may be degraded on the unlabeled set . In this paper , we propose building a novel masker model to mask training samples using MMD and thereby align the missingness patterns in the training and test datasets . We find that this can help improve the test performance by preventing the model from depending on relationships that will be unavailable at test time . Also , note the MMD masker model is especially relevant when categorical features are involved , as new categories in the test set correspond to missing features in the training set , see Section 4.2 for an example . Expanding on this , we combine a missingness shift treatment and a general feature distribution shift treatment together in a hybrid model , and show that the hybrid model is superior through various experiments done with both synthetic data and real data . 2 RELATED WORK . MMD was first introduced in the two-sample tests ( Gretton et al. , 2007 ) . It is a non-parametric metric to measure the distance between two distributions . Therefore , to compare two distributions , one does not need to know their probability density functions , or PDFs . This is useful as the PDFs can be difficult to calculate or can even be unknown . Because of its simplicity , MMD has been widely applied to various problems , such as goodness-of-fit testing ( Jitkrittum et al. , 2018 ) , MMD GAN ( Li et al. , 2015 ; Sutherland et al. , 2017 ; Li et al. , 2018 ) , and MMD VAE ( Zhao et al. , 2019 ) . One technique related to using MMD for covariate shift is using kernel mean matching to re-weight training samples ( Huang et al. , 2007 ) . In that paper , the authors formulated a quadratic problem between empirical means to find suitable weights for training samples , after which they could apply weighted ordinary least squares or support vector machines . The difference is that they tackled the problem from the importance resampling angle and used the results for more traditional machine learning methods . We incorporate MMD directly into a neural network model in one step instead of two by using the representation learning ability of neural networks , and are able to treat covariate shift in a more granular level than the row/sample level if necessary . The most related application of MMD to our work is in UDA . Deep Domain Confusion ( DDC ) ( Tzeng et al. , 2014 ) is one of the earliest attempts to use MMD for domain invariant learning . Then , Deep Adaptation Networks ( DAN ) ( Long et al. , 2015 ) were proposed to match an optimal multi-kernel MMD to reduce domain discrepancy . This was followed by Joint Adaptation Networks ( JAN ) ( Long et al. , 2017 ) , which were proposed to apply joint MMD in multiple domain-specific layers across domains . Yan et al ( 2017 ) proposed adjusting plain MMD by class-specific auxiliary weights to account for the class weight bias across domains . Chen et al ( 2020 ) proposed matching higher-order moments to perform fine-grained domain alignment . Note that these applications all focused on domain shifts in classification problems , while we extend the application to more general distribution shift and to missingness shift that has never been considered before . 3 METHOD . 3.1 PRELIMINARY : MAXIMUM MEAN DISCREPANCY . Assume we have two sets of samples X = { xi } Ni=1 and Y = { yj } Mj=1 drawn from two distributions P ( X ) and P ( Y ) . MMD is a non-parametric measure to compute the distance between the two sets of samples in mean embedding space . Let k be a measurable and bounded kernel of a reproducing kernel Hilbert space ( RKHS ) Hk of functions , then the empirical estimation of MMD between the two distributions inHk can be written as LMMD2 ( X , Y ) = 1 N2 N∑ i=1 N∑ i′=1 k ( xi , xi′ ) + 1 M2 M∑ j=1 M∑ j′=1 k ( yj , yj′ ) − 1 NM N∑ i=1 M∑ j=1 k ( xi , yj ) ( 3.1 ) When the underlying kernel is characteristic ( Fukumizu et al. , 2008 ; Sriperumbudur et al. , 2010b ) , MMD is zero if and only if P ( X ) = P ( Y ) ( Gretton et al. , 2012a ) . For example , the popular Gaussian RBF kernel , k ( x , x′ ) = exp ( − 12σ2 |x − x ′|2 ) , is a characteristic kernel and was widely used in previous literature ( Li et al. , 2015 ; Sutherland et al. , 2017 ; Long et al. , 2017 ; Li et al. , 2018 ) . Dziugaite et al ( 2015 ) compared the RBF kernel with the rational quadratic ( RQ ) kernel and the Laplacian kernel and found that the RBF kernel performed best . Therefore , in this work , we also used the RBF kernel . What remains is how to select the bandwidth parameter σ in the Gaussian RBF kernel , which can affect the hypothesis testing power of MMD measurement . There were some heuristics in previous literature ( Sriperumbudur et al. , 2010a ; Gretton et al. , 2012b ) , but later research ( Li et al. , 2015 ; Sutherland et al. , 2017 ; Li et al. , 2018 ) found that using a mixture of 5 or more bandwidths performed well , so we followed this procedure as well . 3.2 MMD REPRESENTATION . Let us first define the context of our predictive problem . Assume that we have training and test features Xtr , Xte ∈ Rm , where m is the number of features . We also know the target value ytr for the training set , and we want to build a model to predict yte for the test set . When we train the model with ( Xtr , ytr ) , a covariate shift between Xtr and Xte can cause the trained model to perform poorly on the test set . Assume we build a neural network model for the task , then similar to what has been widely used in domain adaptation , we can match the MMD statistics between embtr and embte at some intermediate embedding layer so that the learned embeddings have similar distributions between the training and test sets and will therefore reduce the amount of extrapolation . We call this approach MMD Representation . See Figure 1 for the architectural details . Note that models like DAN ( Long et al. , 2015 ) and JAN ( Long et al. , 2017 ) can be considered as special forms of MMD Representation , as they also apply MMD to multiple specific intermediate layers . The final loss we will optimize is the original task loss plus the MMD loss in the representation space : L = Ltask + λLMMD2 ( embtr , embte ) ( 3.2 ) where λ is a hyperparameter to adjust the ratio between the original loss and the MMD loss . Generally , we found that setting λ to match the magnitudes of the two loss terms leads to the best results . Compared to the approach of dropping features with covariate shift before training , MMD Representation is able to retain as much information as possible from each individual feature and find a common feature space that has as little shift as possible between the training and test sets . embtr , embte ← fF ( X ′tr ) , fF ( Xte ) ŷtr ← fP ( embtr ) Update fF , fP ← minLtask ( ŷtr ) + λLMMD2 ( embtr , embte ) end for end for 3.3 MMD MASK . Despite the advantages of MMD Representation , we found that it is insufficient when a missingness shift is involved , as was the case for the carbon estimation problem mentioned in Section 1 . In general , our MMD Mask can be effective when the missingness rates for some set of features are low or zero in the training set but much higher in the test set . The architecture of MMD Mask is shown in Figure 1 . Denote Itr and Ite as the original missingness indicators for the training and test features Xtr and Xte , respectively ; we want to learn a new missingness indicator Îtr conditioned on Xtr and Itr that can minimize the MMD loss between the joint distributions P ( X ′tr , Îtr ) and P ( Xte , Ite ) : L = LMMD2 ( ( X ′tr , Îtr ) , ( Xte , Ite ) ) ( 3.3 ) where X ′tr is generated by masking the original Xtr using learned Îtr . The downstream model is then trained using X ′tr instead of Xtr . In this way , the downstream model can focus on the right features for each training sample . Importance resampling using kernel mean matching also operates in the original feature space , but it removes or downweights each sample as a whole ( at the row level ) , while we treat each sample feature by feature . Compared to MMD Representation , another advantage of MMD Mask is that it can still be applied when the downstream model can not learn an intermediate representation ( e.g . a tree-based model ) .

Summary: Suppose we are given training data  $(\mathbf{X}_{tr}, \mathbf{I}\_{tr},  \mathbf{Y}\_{tr})$, where $\mathbf{Y}\_{tr} \in \mathbb R^{n\_{tr}}$  is the response of interest, $\mathbf{X}\_{tr} \mathbb R^{n\_{tr} \times p}$ is the feature design matrix ($p$ features, $n\_{tr}$ samples) and  $\mathbf{I}\_{tr} \in  \\{0,1\\}^{n\_{tr} \times p}$ is an indicator matrix that marks which features are missing in $\mathbf{X}\_{tr}$. In the setting studied in this paper, the researcher also has access to the test samples  $(\mathbf{X}\_{te}, \mathbf{I}\_{te})$ where $\mathbf{X}\_{te} \in \mathbb R^{n\_{te} \times p}$ and $\mathbf{I}\_{te} \in  \\{0,1\\}^{n\_{te} \times p}$ are the features of the test samples and their missingness indicator. The goal is to predict the response of the test samples while accounting for the potential complications of covariate-shift and a shift in the missingness pattern between training and test set covariates.  The authors propose three approaches to this general problem.   1) MMD-Representation to account for covariate-shift (ignoring missingness): here one learns a feature representation $f_F$ such that the distribution of $f_F(x\_{tr})$ and  $f_F(x\_{te})$ are similar, where $x\_{tr}$ is a feature vector drawn from the training distribution and $x\_{te}$ from the test distribution. Closeness here is measured in terms of the maximum mean discrepancy of Gretton and colleagues (2007). Then one learns a predictive model $f_P$ so that the predictions take the form $f_P(f_F(x))$.  The two models $f_P$ and $f_F$ are learned simultaneously by minimizing a weighted combination of the predictive training loss and the MMD loss between test and train featurization.  2) MMD-mask to deal with varying missingness patterns: here on learns a masking model $f_M$ and applies it as $\hat{\mathbf{I}} = f_M(\mathbf{X}_{tr}, \mathbf{I}\_{tr})$ to learn a missingness pattern that matches more closely the missingness pattern of the test dataset. Then one applies the learned missingness pattern on the training samples to get modified training samples $\mathbf{X}'\_{tr}$ with missingness pattern $\hat{\mathbf{I}}$. The masking model is learned by minimizing the maximum mean discrepancy between the distributions of  $(x'\_{tr}, \hat{i}\_{tr})$ and $(x\_{te}, i\_{te})$. Finally one learns the predictive model based on  $(\mathbf{X}'\_{tr}, \hat{\mathbf{I}}, \mathbf{Y}\_{tr})$.  3) MMD-hybrid: a combined approach of MMD-mask and MMD-representation to account for both shifts in the missingness pattern and covariate-shift.   These three approaches are benchmarked and compared in one synthetic dataset and two applications (Bike Sharing and IEEE-CIS Fraud Detection).