Maximum Mean Discrepancy for Generalization in the Presence of Distribution and Missingness Shift
Abstract
Covariate shifts are a common problem in predictive modeling on real-world problems.
This paper proposes addressing the covariate shift problem by minimizing
Maximum Mean Discrepancy (MMD) statistics between the training and test sets in either feature input space, feature representation space, or both.
We designed three techniques that we call MMD Representation, MMD Mask, and MMD Hybrid to deal with the scenarios where only a distribution shift exists, only a missingness shift exists, or both types of shift exist, respectively.
We ﬁnd that integrating an
MMD loss component helps models use the best features for generalization and avoid dangerous extrapolation as much as possible for each test sample.
Models treated with this MMD approach show better performance, calibration, and extrapolation on the test set.
More details such as related work, additional experiences and a discussion of future work are in the appendix.
And an extended version of this paper can be found at https://arxiv.org/abs/2111.10344 .
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.
Here we focus on covariate shifts that occurs when the distribution of input features changes [20].
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 [21, 14, 29, 5], kernel mean matching (KMM) [11],
Kullback-Leibler importance estimation [24], and discriminative learning [3].
These methods work particularly well for sample selection bias, but they treat samples at a row/sample level which is not ﬂexible 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) [26, 18, 19, 28, 4].
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 3, we try to train a model with the companies who reported carbon emission to infer carbon emission for the unreported companies.
However, we ﬁnd that the unreported companies usually have higher missing rates in various sets of features compared to the reported companies.
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 ﬁnd that this can help improve the test performance by preventing the model from depending on relationships that will be unavailable at test time.
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.
Method 2.1 Preliminary: Maximum Mean Discrepancy
MMD 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 difﬁcult to calculate or can even be unknown.
Assume we have two sets of samples X
= { x i } Ni =1 and Y = {
y j } Mj =1 drawn from two distributions P ( X ) and P ( Y ) .
Let k be a measurable and bounded kernel of a reproducing kernel Hilbert space (RKHS) H k of functions, then the empirical estimation of MMD between the two distributions in H k can be written as
[Equation]
=
M (cid:88) j (cid:48)
=
k ( y j , y j (cid:48) ) −
NM
N (cid:88) i =
M (cid:88)
j = k ( x i , y j ) (1)
When the underlying kernel is characteristic [7, 23], MMD is zero if and only if P ( X )
=
P ( Y ) [9].
For example, the popular Gaussian RBF kernel, k ( x, x (cid:48) )
=
exp( − 12 σ 2 | x − x (cid:48) | 2 ) , is a characteristic kernel and was widely used in previous literature [17, 25, 19, 16].
The bandwidth parameter σ in the Gaussian RBF kernel can affect the hypothesis testing power of MMD measurement.
There were some heuristics in previous literature [22, 10], but later research [17, 25, 16] found that using a mixture of 5 or more bandwidths performed well, so we followed this procedure as well.
2.2 MMD Representation, MMD Mask, MMD Hybrid
Let us ﬁrst deﬁne the context of our predictive problem.
Assume that we have training and test features X tr , X te ∈ R m , where m is the number of features.
We also know the target value y tr for the training set, and we want to build a model to predict y te for the test set.
When we train the model with ( X tr , y tr ) , a covariate shift between X tr and X te 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 emb tr and emb te 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 .
Note that models like DAN [18] and JAN [19] can be considered as special forms of MMD Representation, as they also apply MMD to multiple speciﬁc intermediate layers.
The ﬁnal loss we will optimize is the original task loss plus the MMD loss in the representation space:
[Equation]
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 ﬁnd a common feature space that has as little shift as possible between the training and test sets.
Despite the advantages of MMD Representation, we found that it is insufﬁcient when a missingness shift is involved, as was the case for the carbon estimation problem mentioned in Section 1.
Denote
I tr and I te as the original missingness indicators for the training and test features X tr and X te , respectively; we want to learn a new missingness indicator ˆ I tr conditioned on X tr and I tr that can minimize the MMD loss between the joint distributions P ( X (cid:48)
tr , ˆ I tr ) and P ( X te , I te ) :
[Equation]
where X (cid:48) tr is generated by masking the original X tr using learned ˆ I tr .
The downstream model is then trained using X (cid:48) tr instead of X tr .
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.
Furthermore, we can combine MMD Representation and MMD Mask together as a more comprehensive shift treatment.
We call this MMD Hybrid .
Denote the masker model as f M , the feature extractor as f F , and the prediction network as f P .
We use an alternating update method to train the
MMD Hybrid model, see Algorithm 1 in Appendix for more details.
See Figure 3 in Appendix for the architectural details for all three MMD methods.
Evaluation
Due to the space limit, we leave all the data and experiments details in Appendix C.
We ﬁrst used synthetic data to gain a general understanding of which MMD models work best in which situations .
The synthetic data has two features X 1 and X 2 , both having some predictive power for the target y .
Although X 1 can predict y more accurately compared to X 2 , we designed a distribution shift in the positive range of X 1 , and a missingness shift in the negative range of X 1 between the training and test datasets.
See Appendix C.1 for more details on how to generate the synthetic data.
Figure 1 shows how each model performs in different ranges of X 1 and X 2 on the test data.
In the region where data shift does not exist, all models learn the relationship very well and have very small residuals.
In the region where missingness shift exists, the MMD Representation reduced residuals signiﬁcantly compared to the baseline model, but not as well as the MMD Mask and the MMD
Hybrid.
In the region where distribution shift exists, we see that the baseline model is worst, followed by MMD Mask, MMD Representation and MMD Hybrid, in order.
We further tested how our approaches compare to exisitng methods generally .
Table 1 shows the performance metrics for each model from 10 runs on the test set for our synthetic data, a Bike Sharing dataset [1] and an IEEE-CIS Fraud Detection dataset [2].
Our MMD Hybrid almost always performs best in all three datasets.
Moreover, we applied MMD Mask in an internal carbon emission estimation project where there is a missingness shift between labeled and unlabeled data.
The companies that reported carbon emission (labeled data) usually have more complete feature information, while the unreported companies (unlabeled data) have many missing values in different sets of features.
In this case, the predictive
[Table]
X- - -   r e
s i du a l
Baseline
Distribution Shift Area
Missingness Shift Area
Residuals mean w/ std bar - -
X- - -
MMD REPR - -
X- - -
MMD MASK - -
X- - -
MMD HYBRID (a) residuals v.s.
X- -
X- - - -    r e
s i du a l
Baseline - -
X- - - -
MMD REPR - -
X- - - -
MMD MASK - -
X- - - -
MMD HYBRID (
[Figure]
v
[Figure]
s
[Figure] model is a tree-based model without representation learning ability, so only MMD Mask can be applied here .
As a baseline to compare with our MMD Mask model, we also tried a Simple Model that randomly sampled masks from unlabeled data conditioned on industry. [Figure] [Figure]
[Figure] Figure 2 shows the percentage error of prediction from the MMD Model and the Simple Model on the masked dataset by carrying over masks for companies from a year in which they did not report carbon emissions to a following year in which they did.
Note that this is the most realistic set of masks we can get for the labeled data.
As we can see, the MMD Model has a large performance gain over the Simple Model.
To conclude, we proposed a novel MMD masker model to match the joint distribution of input features and missingness indicators in input space between the training and test datasets, an approach designed speciﬁcally for treating a missingness shift, which has never been considered before.
We have shown that MMD Representation works better for distribution shifts while MMD Mask works better for missingness shifts on both synthetic data and real data.
Another advantage of MMD Mask is that it can be applied even when the downstream model cannot learn an intermediate representation such as in a tree-based model.
Furthermore, we combined MMD matching in the feature embedding space and in the raw input space to yield superior results.
A Related Work
MMD was ﬁrst introduced in the two-sample tests [8].
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 difﬁcult to calculate or can even be unknown.
Because of its simplicity, MMD has been widely applied to various problems, such as goodness-of-ﬁt testing [13], MMD GAN [17, 25, 16], and MMD VAE [30].
One technique related to using MMD for covariate shift is using kernel mean matching to re-weight training samples [11].
In that paper, the authors formulated a quadratic problem between empirical means to ﬁnd 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) [26] is one of the earliest attempts to use MMD for domain invariant learning.
Then, Deep Adaptation
Networks (DAN) [18] were proposed to match an optimal multi-kernel MMD to reduce domain discrepancy.
This was followed by Joint Adaptation Networks (JAN) [19], which were proposed to apply joint MMD in multiple domain-speciﬁc layers across domains.
Yan et al [2017] proposed adjusting plain MMD by class-speciﬁc auxiliary weights to account for the class weight bias across domains.
Chen et al [2020] proposed matching higher-order moments to perform ﬁne-grained domain alignment.
Note that these applications all focused on domain shifts in classiﬁcation problems, while we extend the application to more general distribution shift and to missingness shift that has never been considered before.
B Model Details
We show model architectures for MMD Representation, MMD Mask and MMD Hybrid all in
Figure /refﬁg:architecture.
Algorithm for MMD Hybrid is detailed in Algorithm 1.
Input: X tr , I tr , X te , I te , y tr , λ
Initialize f M , f F , f P for each epoch do for each batch do ˆ I tr ← f M ( X tr , I tr )
X (cid:48)
tr ← Mask X tr by ˆ I tr
Update f M ← min L MMD 2 (( X tr , ˆ I tr ) , ( X te , I te )) emb tr , emb te ← f F ( X (cid:48) tr ) , f F ( X te )
ˆ y tr ← f P ( emb tr )
Update f F , f P ← min
L task (ˆ y tr ) + λ L MMD 2 ( emb tr , emb te ) end for end for
C Experiments Details
We tested our approaches on both synthetic data and real data.
The real data is comprised of a Bike
Sharing dataset from the UCI Machine Learning Repository, an IEEE-CIS Fraud Detection dataset from Kaggle, an internal project of carbon estimation, and the image dataset MNIST [15].
For all experiments except carbon estimation and MNIST, the baseline model is a multilayer perceptron (MLP) model.
MMD Representation always shares the same architecture, but also matches
MMD statistics at the last hidden layer.
MMD Mask outputs masks for each feature per sample; we then mask out training samples based on the learned masks, and use that to train the same downstream baseline model.
MMD Hybrid combines MMD Representation and MMD Mask.
For synthetic data, the Bike Sharing data and the IEEE-CIS Fraud Detection data, the baseline models take features and missing indicators as inputs, and the missing values were imputed using the means from the training set.
We didn’t impute missing values for carbon estimation because it used a tree-based model that can handle missing values inherently.
For all datasets, we use MMD statistics under a mixture of Gaussian RBF kernels with bandwidths of { 1 , 2 , 4 , 8 , 16 , 32 } .
The MMD Mask model uses a Relaxed Bernoulli distribution parametrized by a temperature τ to generate a mask probability for each feature.
The Relaxed Bernoulli is a continuous approximation to a Bernoulli distribution over (0 , 1) , similar to the Gumbel-Softmax trick [12] for categorical distributions.
In all experiments, we set τ to 0.1 at the start of training and then gradually decrease it to 0.01.
We implemented all models in PyTorch.
Neural network models were all trained on one single Nvidia GPU with driver version 440.82 and CUDA version 10.2.
For synthetic data, the baseline model is an MLP with 3 hidden layers with a hidden size offor each.
The MMD Mask model has 3 hidden layers with hidden sizes of { 32 , 32 , 20 } .
For the
MMD Representation and MMD Hybrid models, we also tested different hyperparameter values λ ∈ { 1 , 5 , 10 } and picked the model that performed best.
In all models, we used ReLU as our activation function and optimized network parameters using RMSProp with a learning rate of 0.01.
All models were trained for 5,000 epochs with full batch.
For the Bike Sharing, the baseline model is an MLP model with 4 hidden layers with hidden size of { 64 , 64 , 64 , 64 } , and the MMD Mask model has 3 hidden layers with hidden sizes of { 512 , 512 , 128 } .
For the IEEE-CIS Fraud Detection, the baseline model is an MLP model with 4 hidden layers with hidden size of { 1024 , 512 , 512 , 256 } , and the MMD Mask model has 3 hidden layers with hidden sizes of { 512 , 512 , 256 } .
For both datasets, we used ReLU as our activation function and optimized network parameters using RMSProp.
For the MMD Representation and MMD Hybrid models, we also tested different hyperparameter values λ ∈ { 1 , 10 , 100 } .
The labeled data is split into 3:1 as train v.s. validation.
The hyperparameters are tuned based on the performance on the validation set.
For carbon emission estimtion, the MMD Mask model is an MLP model with 4 hidden layers with hidden size of { 512 , 512 , 512 , 512 } , with ReLU activation.
We trained the model for 300 epochs with a batch size of 5000, using RMSProp with a learning rate of 0.001.
For all experiments, whenever possible, we also compared our methods with kernel mean matching (KMM) [11], and UDA baselines like DAN [18] and JAN [19].
As for KMM, we followed the procedure in their Empirical KMM optimization Section, and found the best weights for training samples by solving the quadratic problem using cvxopt package 2 , then used these weights to reweight samples in the original loss function.
Whenever DAN and JAN are used, we apply MMD or Joint
MMD to the last two hidden layers as MLPs in this paper are only three to four hidden layers.
C.1 Synthetic Data
In this experiment, we used synthetic data to gain a general understanding of which MMD models work best in which situations.
We constructed a synthetic dataset as shown in Figure 4.
There are two input features X 1 and X 2 , both having some predictive power for the target y .
The relationship between feature X 1 and y is linear while the relationship between X 2 and y is parabolic. We injected noise (cid:15) 1 ∼ N (0 , 0 . 1) into feature X 1 and larger-scale noise (cid:15) 2 ∼ N (0 , 0 . 5) into feature X 2 . Based on this alone, the model should prefer using feature X 1 to feature X 2 . However, we also engineered two types of data shifts into feature X 1 between the training and test datasets.
We ﬁrst resampled from the training dataset with replacement to get our test dataset such that both the training and test https://cvxopt.
X- - -
DAN
Distribution Shift Area
Missingness Shift Area
Residuals mean w/ std bar - -
X- - -
JAN - -
X- - -
KMM (a) residuals v.s. X- -
X- - - -
DAN
Distribution Shift Area
Missingness Shift Area
Residuals mean w/ std bar - -
X- - - -
JAN - -
X- - - -
KMM
[Figure]
datasets had 5,000 rows.
Then in the test dataset, we designed a missingness shift for X 1 in the range [ − 3 .
5 , 0] by masking it out, and a distribution shift for X 1 in the range (0 , 3 . 5] by subtracting it by (cid:15) ∼ N (1 , 0 . 1) - see Figure 4.
Therefore, the test dataset has both missingness and distribution shifts compared to the training dataset.
We hope that a properly designed model can learn to predict using feature X 2 instead of X 1 in the regions of X 1 affected by the shifts. - -   4 X
[Figure]
In main sections, we took a scrutiny of MMD
Representation, MMD Mask and MMD Hybrid in regions with different types of shift.
Here we also include similar analysis for KMM, DAN and JAN.
From Figure 5, we can see that KMM,
DAN and JAN also have very small residuals in the region where data shift does not exist.
When missingness shift or distribution shift exists, performance of all methods deteriorates to different degrees.
KMM has a similar performance compared to the Baseline model in both missingness shift and distribution shift regions.
In the region where distribution shift exists, DAN and JAN work better than MMD Mask, which is expected.
DAN and JAN also work better than MMD Representation as they try to match more intermediate layers, and work similarly with MMD Hybrid.
However, in the region where missingness shift exists, DAN and JAN are not as good as MMD Mask and MMD Hybrid, but similar to MMD Representation.
Furthermore, we also checked the embeddings in the last hidden layer from the training and test phases for each model.
Figure 6 shows 2-D tSNE
[27] plots of embeddings from the baseline model,
MMD Representation, MMD Mask and MMD Hybrid.
As expected, in all MMD approaches, the
[Figure]
Finally, we examined the masks generated by MMD Mask and MMD Hybrid.
Figure 7a shows histograms of features X 1 in the original training set, the test set, and the masked training set.
We focused on X 1 here because there is no shift in X 2 .
As we can see, the masker in MMD Mask and
MMD Hybrid learned to mask out the shifted range on both the positive and negative sides to some degree.
Since the negative side has purely a missingness shift, it learned better by masking out more samples for X 1 ∈ [ − 3 . 5 , 0] , while it masked out less samples for X 1 ∈ (0 , 3 . 5] .
The masking ability is consistent with the downstream performance we observed.
C.2 Bike Sharing and IEEE-CIS Fraud Detection
[Figure]
We use the Bike Sharing as the regression example and the IEEE-CIS Fraud Detection as the classiﬁcation example to show how our MMD models perform in general.
The Bike Sharing dataset [6] from UCI contains the hourly and daily bike rental counts in the Capital bikeshare system in Washington, D.C., USA for 2011 and 2012.
We focused on hourly data as it provides more samples.
The dataset contains features like hour, month, season, workday, holiday and some weather information.
We log-transformed the target variable (bike rental counts) due to their long tail, and deliberately selected 2011-03 to 2011-11 as the training data (6,567 rows) and 2011-12 to 2012-03 as the test data (2,917 rows) so that the features have shifts between the training and test data due to the different times of year.
Also, since there are some new months in the test set (December, January and February), after ordinal feature processing, they are treated as missing in the prediction stage, which can be viewed as a source of missingness shift.
The IEEE-CIS Fraud Detection dataset is from Vesta’s real-world-e-commerce transactions and contains a wide range of features from device type to product features.
The goal is to estimate fraud probability of the unlabeled test set.
We followed the champion model from Chris Deotte and
Konstantin Yakovlev 3 , and used all the features except 219 Vesta engineered rich features that were determined redundant by their correlation analysis, but did not conduct further feature engineering like they did, as our goal is not to outperform the best model but use this dataset to show how our proposed MMD models can handle distribution and missingness shift in the features and outperform the baseline models.
In the end, we used 212 features before one hot encoding the categorical features and adding missing indicators.
Missingness shift is observed in some of the 212 features to some degree, but not severe.
We also examined if the masks learned for the month feature in the Bike Sharing dataset meet our expectation.
As we can see from Figure 7b, the maskers learn to leave more of March 2011 data in [Figure] while masking out data from other months more, which matches with our expectation, as March is the only common month in both training and test data.
Note that the maskers try to match the joint distribution of all features, not just the marginal distribution of months.
C.3 Carbon Emission Estimation
Carbon Emission Estimation is an internal project that seeks to estimate companies’ carbon emissions using various sets of features including environmental, social and governance (ESG) data, industry segmentation data, fundamental ﬁnancial data, and country and regional data.
In total, there are about 1,000 features.
One challenge of the project is the missingness shift between labeled and unlabeled data.
The labeled data usually has more complete feature information, while unlabeled data has many missing values in different sets of features.
As a result, a model trained on the labeled data would generate poor results on the unlabeled data due to the much larger missing rates in various features.
We could not take the approach of dropping all features that are sometimes missing as nearly all features in this dataset can be missing.
To overcome this challenge, we propose using data augmentation where we generate masks for the labeled data to mimic the missingness patterns in unlabeled data.
In this case, the downstream predictive model is a tree-based model without representation learning ability, so only MMD Mask can be applied here.
We trained a masker model to generate masks for labeled data based on all available features by matching the MMD statistics between the labeled and unlabeled data.
As a baseline, we also tried using a simple mask model that randomly sampled masks from unlabeled data conditioned on industry.
The downstream carbon estimation model was then trained on the masked labeled data with 10 masks sampled from the mask model plus the original training data.
We denote MMD Model as the downstream model trained on data masked by masks from the MMD Mask model, and Simple Model as the downstream model trained on data masked by masks from the simple mask model.
To evaluate the results, we designed a series of tests.
Besides the one mentioned in the main text, we also tested the performance of the MMD Model and Simple Model on data masked by both MMD masks and simple masks.
We sampled 10 masks from the corresponding mask model and applied those to the original labeled data, then compared predictions from the downstream model trained with two different versions of masked data.
Figure 8a compares the percentage error of prediction for each combination.
As we can see, the MMD Model always beats the Simple Model whether the data is masked by MMD mask or the simple mask.
And when the original labeled data is masked by
MMD mask, the performance gain is even larger.
The fact that the MMD Model performs better than the Simple Model even on the data masked by the simple mask is interesting.
It indicates that the missingness pattern from the simple mask is in a sense contained in the missingness pattern from
MMD mask.
In addition, we tested the MMD Model and Simple Model in the speciﬁc scenarios that an entire set of features is missing, as this can be the case for certain sub-universes of companies.
Figure 8b and
Figure 8c plots the scenarios on unmasked data and masked data respectively.
Since ESG data is much more useful for estimating carbon emissions than industry segmentation data is, the performance decreases more when ESG data is missing.
However, no matter which set of features is entirely missing, the MMD Model always performs better than the Simple Model.
Also, if we compare the curves for unmasked data with those for masked data, we can see that most of them are very similar except for the Simple Model with Segmentation data missing.
The Simple Model performance on masked data is a bit worse compared to unmasked data when Segmentation data is completely missing.
However, MMD Model performance on masked data remains similar to that on unmasked data, indicating that MMD masks are more aligned with the scenarios when an entire set of features is missing.
D MNIST Experiments
In this section [Table] we investigate how our MMD methods perform on more unstructured image data.
Our proposed methods can be useful when occlusion is correlated with features/labels on image data.
For example, say you have intact images of groundhogs, squirrels and pangolins with labels in training set.
But in the unlabeled set, the images of groundhogs and pangolins are incomplete as those are taken when the animals are partially in ground.
So we need to learn appropriated masks to mask out right images in training set in order to learn a more accurate and robust classiﬁcation model
[Table]
[Table] [Table]
To demonstrate our approaches also work on high dimensional image data, we used MNIST which is a handwritten digits dataset where each image is 28 × 28 pixels.
It has 60,000 training images and 10,000 test images.
To demonstrate that our methods work when there is a shift, particularly a missingness shift, we created a patterned mask on the original test data.
That is, for digits in { 1 , 4 , 5 , 7 } , we mask out the last 12 rows of pixels, and for digits in { 0 , 2 , 3 , 6 , 8 , 9 } , we mask out the last 12 columns of pixels.
We compared MMD Representation, MMD Mask, and MMD Hybrid approaches with a baseline model trained on raw training data, a golden model trained with the training samples masked by true masks, as well as traditional data augmentation.
All models in this section adopt CNN architectures.
The baseline model in MNIST experiments uses the architecture shown in Table 2.
The golden model and two variants of data augmentation models also use the exactly same architecture.
The
MMD Representation uses the same architecture except that it also matches the embeddings between training and test set before entering the ﬁnal dense layer.
The MMD Mask model generates masks using a conv-deconv architecture shown in Table 3.
It generates masks to mask original training data, then the masked training data goes through the same downstream model.
The MMD Hybrid model combines MMD Representation and MMD Mask.
The missing values were imputed as zero (black pixel).
[Table] [Figure] For traditional data augmentation, we tested two versions of RandomErasing [31], with one version always erasing the same area ( 12 × 28 ) of images at a random place and the other erasing a random area ranging from 0 to the same area ( 12 × 28 ) at a random place.
We followed the standard procedure to split the training samples into a training set of 54,000 rows and a validation set of 6,000 rows, and trained all models for 20 epochs with a batch size of 64 (except for the MMD Mask model, which was trained for 600 epochs with a larger batch size of 10,000).
The hyperparameter λ in the MMD
Representation and MMD Hybrid model was set to 1.
The ﬁnal model in all methods was the one that gave the best validation loss.
The results are summarized in Table 4.
The performance reported is the average test accuracy from 10 runs for each model.
Note that this baseline model can achieve around 98% test accuracy on the original test dataset, but only around 87% accuracy on the corrupted test dataset.
However, MMD
Representation, MMD Mask, and MMD Hybrid can all improve the performance by a large margin, to 95.9%, 97.2%, and 97.2% accuracy, respectively.
Because we corrupted the test set by applying patterned masks, MMD Mask is a more natural treatment compared to MMD Representation.
Since
MMD Hybrid combines MMD Representation and MMD Masks, it works as well as MMD Mask in this case.
All three methods performed better than the two traditional data augmentation methods: these achieve accuracies of 93.4%and 95.0%.
We believe that this is simply because the learned transformation is more effective and efﬁcient compared to more random transformations.
The golden model achieves 98.7% accuracy, which is not too far away from performances from the best MMD approaches.
Figure 9 shows some samples of generated masks applied to the original training set, which are quite close to the true masks we imposed on the test set.
E Limitations and Future Work
With the current approaches introduced in this paper, the learned model is calibrated to a speciﬁc test set.
If there is a new test set with a different distribution shift and/or missingness shift, we need to retrain the model with this new test set.
We think that MMD has the potential to learn a more generalized model that is robust for different sets of test data if we combine this technique with meta-learning approaches.
This is out of scope for this paper and an interesting area for future work to explore.
However, in scenarios where one has a target test set to label, our proposed MMD approaches help the model learn to use the best features in both the input and embedding spaces and avoid dangerous extrapolation as much as possible, and as a result, achieve better performance on the test set.