Prediction with Incomplete Data under Agnostic Mask Distribution Shift

1 Introduction

Data with missing values is ubiquitous in many applications due to sensor malfunction, incomplete sensing coverage, etc. Recent years have witnessed increasing attention on prediction with only incomplete feature, which consists of observed feature values and a mask that indicates which features are observed. Existing methods [Morvan et al., 2020; Morvan et al., 2021] assume that the training and testing distributions are the same. However, this assumption can be violated in real-world scenarios. In this paper, we study the problem of prediction with incomplete data in the presence of distribution shift. We focus on the scenario where the underlying joint distribution of complete features and label is invariant, but the missing pattern, i.e., mask distribution may be different between training and testing. Such mask distribution shift may result from different sensor deployment, data management, etc. For example, DiDi Traffic Speed [DiDiChuxing, 2018] is a naturally incomplete dataset throughout one year. The traffic network is almost unchanged, so we may reasonably assume the speed distribution is relatively stable. The average missing rate from Jan to Jun is 37% but drops to 23% from Jul to Dec, possibly because more and higher-quality sensors are deployed. Moreover, different from the transfer learning formulation [Pan and Yang, 2010], we assume the mask distribution shift is agnostic, since testing distribution is usually unavailable during training in practice.

∗Extended version: https://arxiv.org/abs/2305.11197.

†Corresponding author.

Many methods for missing data can be used to predict with incomplete feature. Morvan et al. [2020] propose NeuMiss for linear regression with incomplete Gaussian feature. Morvan et al. [2021] then extend NeuMiss to nonlinear case by training it jointly with an MLP. Some missing data imputation methods [Yoon et al., 2018; Li et al., 2019; Ma et al., 2019; Mattei and Frellsen, 2019; Li and Marlin, 2020] can also be used by treating the label as missing feature. However, none of these methods consider distribution shift. Their models learn the information of mask distribution and thus can hardly generalize under mask distribution shift.

Several methods have been proposed for prediction under agnostic feature distribution shift [Shen et al., 2018; Kuang et al., 2018; Kuang et al., 2020; Shen et al., 2020; Zhang et al., 2021; Xu et al., 2022]. It is assumed that the conditional label distribution given complete feature is invariant. They learn this invariant conditional distribution to achieve generalization. However, they are designed for complete data. In our setting, the conditional label distribution given trivially-imputed complete feature is not invariant, so these methods cannot be applied to incomplete data.

We observe that the conditional label distribution given observed feature values and mask is invariant between training and testing. As a result, there is an invariant optimal predictor for each mask respectively. Learning these optimal predictors can generalize under agnostic mask distribution shift. Since the number of optimal predictors increases exponentially with feature dimension [Morvan et al., 2020], we approximate the optimal predictors jointly using a double parameterization technique, i.e., we first parameterize the optimal predictors and then further parameterize the mapping from mask to the parameters of the optimal predictors. However, such parameterization has the undesirable side effect that the learned model may depend on the intra-mask correlation and the correlation between features and mask, since the training loss depends on these correlations. This may defy the generalizability of the learned model, as such correlations will change under mask distribution shift. Inspired by Xu et al. [2022], we decorrelate features and mask to make the learned model independent of their correlations and help it well approximate the optimal predictors. Combining all the techniques above, we have our StableMiss that can achieve generalization under agnostic mask distribution shift.

The contributions of this paper are summarized as follows.

• This paper proposes StableMiss, a novel method for prediction with incomplete data that is robust to agnostic mask distribution shift. To the best of our knowledge, this is the first method that considers agnostic mask distribution shift.
• Extensive experiments are conducted on both synthetic and real-world datasets. The results show that StableMiss is robust and outperforms the state-of-the-art methods under agnostic mask distribution shift.

2 RelatedWork

Prediction with Incomplete Data. Prediction with incomplete feature has attracted increasing attention recently. Morvan et al. [2020] derive the analytical expression of the optimal predictor for linear regression with incomplete Gaussian feature and proposes the NeuMiss network to approximate the optimal predictor. Morvan et al. [2021] then extend NeuMiss to nonlinear case by training it jointly with a Multi-Layer Perceptron (MLP). Besides these methods natively proposed for prediction, some missing data imputation methods can also be used by treating the label as missing feature. GAIN [Yoon et al., 2018] is an adaptation of GAN [Goodfellow et al., 2014], where the generator imputes missing entries, which the discriminator tries to distinguish from observed entries with partial information about the mask. MisGAN [Li et al., 2019] learns the complete data distribution from incomplete data and uses it to supervise imputation. Partial VAE [Ma et al., 2019], MIWAE [Mattei and Frellsen, 2019] and P-BiGAN [Li and Marlin, 2020] extend VAE [Kingma and Welling, 2014], IWAE [Burda et al., 2016] and BiGAN [Donahue et al., 2017] respectively to learn the prior and posterior distributions of incomplete feature given its latent representation and mask. However, all these methods assume that the training and testing distributions are the same. They can hardly generalize under distribution shift.

Prediction under Agnostic Feature Distribution Shift. Several methods have been proposed for prediction under agnostic feature distribution shift. They typically assume that the conditional label distribution given complete feature is invariant. Then they learn this invariant conditional distribution to achieve generalization by decorrelating the features. CRLR [Shen et al., 2018] learns a weight for each sample respectively by minimizing the so-called confounder balancing loss, which is zero when the features are decorrelated. Then weighted Logistic regression is carried out with the learned weights. Based on CRLR, DGBR [Kuang et al., 2018] also extracts nonlinear representations of features with deep autoencoder. DWR [Kuang et al., 2020] learns a set of sample weights by minimizing the sum of covariance between pairs of features and then carries out weighted linear regression. SRDO [Shen et al., 2020] first constructs a new dataset by sampling each feature independently from the training set to decorrelate among features and then learns sample weights by density ratio estimation [Sugiyama et al., 2012]. Weighted linear regression is also used with the learned weights. StableNet [Zhang et al., 2021] adopts Random Fourier Features to measure nonlinear correlations among features. It iteratively optimizes a weighted regression model and a set of sample weights by minimizing the prediction error and the nonlinear correlations, respectively. Xu et al. [2022] propose a general framework with DWR and SRDO as specific implementations and gives theoretical analysis on the framework. However, all these methods are designed for complete data and cannot be applied to incomplete data.

Compared to the existing methods, StableMiss can not only predict with incomplete feature but also generalize under agnostic mask distribution shift.

3 ProblemFormulation

3.1 Preliminary

We use capital and lowercase letters, e.g., X and x, to denote random variable and its realization, respectively. We use subscripts to index the entries of a vector, e.g., xi is the i-th entry of x. Let x ∈ Rn and y ∈ Rd denote the feature and label, respectively. In the presence of missing data, we consider the case where x is partially observed and y is fully observed during training. A binary mask m ∈ {0, 1}n indicates which entries of x are observed: mi = 1 if xi is observed, and mi = 0 if xi is missing. The complementary mask m is defined by mi = 1 - mi, ∀i. With a slight abuse of notation, we regard m and m as the index sets of the observed and missing entries, so that the observed and missing feature values are xm = {xi | i ∈ m} and xm = {xi | i ∈ m}, respectively. We consider the case where the mask m is known, since it is common to know which features are observed within incomplete feature. The incomplete feature is given by (xm,m). We consider the case where label generation process depends on the feature but not the mask, i.e., p(y | x,m) = p(y | x). We do not make assumptions on p(x).

Following Little and Rubin [1986], we model the generative process of incomplete feature as follows. A complete feature sample x is first drawn from the complete feature distribution p(x). Given x, a mask sample m is then drawn from the conditional mask distribution p(m | x). The resulted incomplete feature (xm,m) follows the distribution

We focus on the Missing Completely At Random (MCAR) case and Missing At Random (MAR) case [Little and Rubin,

1986]. Under MCAR, the mask M is independent of the underlying complete feature X, i.e., p(m | x) = p(m), ∀m, x; under MAR, M only depends on the observed feature values XM, i.e., p(m | x) = p(m | xm), ∀m, x.

3.2 Problem Statement

The problem is to predict with incomplete feature under agnostic mask distribution shift. Given a training set D = {(x(i) m(i) ,m(i) , y(i) )}Ni=1 , consisting of N samples from the training distribution ptr (x,m, y), the goal is to learn a prediction function g(xm,m) for agnostic testing distribution pte (x,m, y), where the input to g is only incomplete feature. We seek the prediction function g to minimize the mean squared loss

Ideally, the optimal g is given by the conditional expectation:

How we learn it approximately from training data will be introduced in the next section.

Note that the testing error can be arbitrarily large without any prior knowledge about the testing distribution. We consider the case where only the mask distribution may change between training and testing. More specifically, we make the following assumption on the testing distribution.

Assumption 1. The joint distribution of complete feature and label is invariant between training and testing:

The mask distribution shift still remains agnostic under the above assumption. Different from the transfer learning formulation [Pan and Yang, 2010], the testing distribution is unavailable during the training process.

4 Methodology

4.1 Prediction Framework

Our method relies critically on the following simple result, the proof of which is given in the extended version.

Theorem 1. Under Assumption 1, the conditional label distribution given observed feature values and mask is invariant between training and testing in MCAR or MAR:

As a consequence,

Theorem 1 holds since MCAR and MAR guarantee that the missing feature XM is independent of of mask M given the observed feature XM. It shows that the ideal optimal predictor is invariant between training and testing. If we can learn it under the training distribution, it will automatically generalize to the agnostic testing distribution. However, as noted in Morvan et al. [2021], it is essentially an aggregation of 2n optimal predictors, one for each specific mask m. Since the number of optimal predictors increases exponentially with feature dimension, it is infeasible to learn them separately.

Note that the optimal predictor can be thought of as a function of the observed feature xm parameterized by ϕ, where ϕ is a function of m, i.e., g(xm,m) = gϕ(m)(xm). Learning all the 2n optimal predictors then corresponds to learning the 2n different values of ϕ(m). The function ϕ(m) corresponding to the optimal g is very complicated in general. To address the exponential explosion problem, we adopt the common technique of approximating it by a simpler function ϕθ(m) parameterized by θ, which can be implemented by a neural network with parameter θ. Since we will also implement g by a neural network, we use the zero-imputed feature instead of xm to uniformize the input size. The final form of our predictor is thus

where ⊙ is element-wise multiplication and θ is the parameter we need to learn. The framework is shown in Figure 1.

We illustrate the above framework by a simple example. Suppose Y = ∑ni=1 αiXi, where Xi's are mutually independent except that X1 = X2. The optimal predictor is

where

Note that, if xi is missing, the value of ϕi actually does not matter, as it is multiplied by (x ⊙ m)i = ximi = 0. Here the function g is linear in x ⊙ m with the parameter ϕ(m) = (ϕ0(m), . . . , ϕn(m)) being quadratic in m. We can easily parameterize ϕ(m) by θ ∈ R(n+1)×(n+1)×(n+1) , where ϕk(m) = ∑i,j θijkmimj. The corresponding g is then

where x0 = m0 = 1. Here approximation is not necessary as the optimal ϕ(m) is simple enough. In general, however, ϕ(m) may be very complex and approximation is necessary. For example, when the Xi's are not independent, the ϕ(m) for the linear model can be an n-th order polynomial in m with 2n parameters.

Note that the architecture of NeuMiss [Morvan et al., 2020] fits into our framework, although it does not explicitly consider mask distribution shift. NeuMiss uses a specific g and ϕθ that are carefully designed to approximate the analytical form of the optimal predictor for a linear label generation process with Gaussian feature. Our framework is more flexible and applies to more general settings, as shown in Section 5.

4.2 Decorrelation

The use of ϕθ(m) solves the exponential explosion problem, but it introduces another. As noted in prior work [Shen et al., 2018; Kuang et al., 2018; Kuang et al., 2020; Shen et al., 2020; Zhang et al., 2021; Xu et al., 2022], the learned value of θ can be affected by the correlation between the variables that we condition on in the optimal predictor, i.e., X and M in our case. This defies our original goal, as a learned value of θ that performs well on the training distribution may not do so on the testing distribution, where the correlation changes. We explain this problem in details below.

We assume E[Y | xm,m] is realizable, i.e., there exists a parameter θ∗ corresponding to the parameter of optimal predictor ϕ. In the above example, this assumption is satisfied by the form of g in Equation (2), and θ∗ corresponds to the function ϕ as given by Equation (1). In practice, we use neural networks for g to approximately satisfy this assumption, which shows good empirical results in Section 5.

Consider the population loss. θ∗ minimizes training loss

Since parameter θ∗ corresponds to the parameter of optimal predictor, it also minimizes the above loss under testing distribution pte , i.e., θ∗ generalizes under mask distribution shift.

Suppose ˆ θ is the parameter learned from training distribution, i.e. ℓtr(ˆ θ) = ℓtr(θ∗ ). ˆ θ can be different from θ∗ , but they perform equally well on the training distribution. In fact,

However, ˆ θ and θ∗ may not perform equally well on testing distribution ptr , as formally stated in Theorem 2, the proof of which is given in the extended version.

Theorem 2. In the realizable case, there exists a problem instance, in which there exists a ˆ θ other than θ∗ such that Equation (3) holds ptr -a.s. but not pte -a.s..

Such ˆ θ is undesirable, since it cannot generalize as well as θ∗ . To address this problem, decorrelation among X and M may help. We explain the reason with an example below.

Consider the example in Section 4.1. If the entries of X and the entries of M are mutually independent, the only solution for ˆ θ is the θ∗ corresponding to Equation (1). The proof is given in the extended version. Thus we can perform decorrelation to help obtain the desired parameter θ∗ .

Under MCAR, we decorrelate the entries of X and those of M respectively; under MAR, we further decorrelate between X and M. In this paper, we follow Zhang et al. [2021] to decorrelate by sample reweighting. Specifically, we learn a weighting function w(xm,m) by minimizing the correlation under distribution ˜ pw, where

The learned weighting function will be normalized to make ˜ pw well-defined. Under ˜ pw, the conditional label distribution p(y | xm,m) does not change. Inspired by Zhang et al. [2021], we measure the correlation empirically by partial cross-covariance matrix with Random Fourier Feature [Strobl et al., 2019]. We adapt it to the case of incomplete data as follows.

The partial cross-covariance matrix of Xk and Xl, denoted by Σ′Xk,Xl;w , is computed with only samples in which Xk and Xl are both observed:

where N kl is the number of samples in which Xk and Xl are both observed, N k and N l are the number of samples with observed Xk and Xl, respectively, wi = w(x(i) m(i) ,m(i) ), and u, v are function vectors with elements drawn from the space of Random Fourier Feature [Strobl et al., 2019], the definition of which is given in the extended version. Similarly, Σ′Xk,Ml;w is computed with only samples in which Xk is observed, and Σ′Mk,Ml;w is computed with all the N samples. See the extended version for their detailed forms.

When decorrelating the entries of X and those of M, we optimize sample weight w ∈ RN+ by minimizing the correlation of all pairs of Xi and Xj and all pairs of Mi and Mj :

where Std(w) is standard deviation of the wi's. The last regularization term is used to prevent assigning very large weights to a small proportion of samples. Hyper-parameter γ is coefficient for the regularization term. Under MAR, we further add ∑1≤k,l≤n ∥Σ′Xk,Ml;w ∥2F into the objective to decorrelate between X and M.

With the learned weighting function, we can conduct regression under the weighted distribution ˜ pw:

In the ideal case, the entries of X and those of M are mutually independent under ˜ pw. Then the learned θ will not be affected by the correlation and can generalize under mask distribution shift. In practice, we can equivalently conduct weighted regression under the training distribution:

How decorrelation is combined with the prediction framework is shown in Figure 1.

5 Experiment

5.1 Datasets

We evaluate StableMiss on synthetic and real-world datasets.

Gaussian. Following Morvan et al. [2021], we generate feature X from multivariate Gaussian distribution. The mean values are drawn from standard Gaussian distribution, and the covariance matrix is generated by Σ = BB⊤ +D, where the entries of B ∈ Rn×0.7n are drawn from standard Gaussian distribution and D ∈ Rn×n is diagonal with values uniformly drawn from [10-2 , 10-1 ]. The feature dimension n = 50.

Gaussian-Ind. The feature generation is the same as Gaussian, except that we make the entries of Xmutually independent by using a diagonal covariance matrix.

Gaussian-Mix. We generate feature X from the more general Gaussian mixture model with 3 components, each generated in the same way as Gaussian. The proportion of the i-th component πi is uniformly drawn from [0, 1) and normalized by πi/ ∑3i=1 πi.

For the above 3 synthetic features, the label generation process is linear: Y = α0 + ∑ni=1 αiXi + ϵ, where ϵ is a Gaussian noise such that the signal-to-noise ratio is 10. Different fromMorvan et al. [2020; 2021] where all αi's except α0 are equal, we draw more general αi's from Gaussian distribution.

House Sales. Following Shen et al. [2020], we use dataset of house sales in King County, USA, which contains n = 16 features and a scalar house price as label.

MNIST [Lecun et al., 1998]. The MNIST dataset of handwritten digit images. Given incomplete image, we aim to predict the complete image.

Traffic [DiDiChuxing, 2018]. Average traffic speed within every hour from 1343 roads in the city of Chengdu, China, in 2018. We build a graph for the dataset, where nodes represent the roads and edges indicate the adjacency of roads. Note that this dataset is naturally incomplete. Given incomplete history, we aim to predict the future traffic speed.

All the datasets except the Traffic dataset are complete. We generate incomplete datasets by imposing mask on the complete samples according to the missing patterns in Section 5.2.

5.2 Missing Patterns

We design the following 3 missing patterns. There are 9 missing levels, denoted by r, from 10% to 90% at a step of 10%.

Missing patterns examples are given in the extended version.

MCAR-Ind. We generate mask M that is independent of feature X, and the entries of mask are mutually independent. For each sample, its sample missing rate rs has 80% to be r and 2.5% to be one of the other 8 levels respectively. Each entry is independently missing with probability rs.

MCAR. We generate mask M that is independent of feature X, but the entries of mask can be dependent. The sample missing rate rs is determined in the same way as MCAR-Ind. In each sample, following Li and Marlin [2020], we generate a window of length ⌊n · rs⌋ at a random position, where the ⌊n · rs⌋ consecutive features in the window are missing.

MAR. Following Morvan et al. [2020; 2021], we generate feature-dependent mask, and the entries of mask can also be dependent. First, randomly selected 10% features are set to be observed in all the samples. The mask on the other features are generated according to a model whose parameters depend on the selected features. The sample missing rate rs means the missing proportion of the other 90% features, which is determined in the same way as MCAR-Ind.

5.3 Baselines

We compare StableMiss with two categories of baselines. (1) state-of-the-art methods on prediction with incomplete data: NeuMiss [Morvan et al., 2020; Morvan et al., 2021], Partial VAE [Ma et al., 2019], MIWAE [Mattei and Frellsen, 2019] and P-BiGAN [Li and Marlin, 2020]; (2) state-of-theart methods on generalization under agnostic feature distribution shift: DWR [Kuang et al., 2020], SRDO [Shen et al., 2020] and StableNet [Zhang et al., 2021]. For Partial VAE, MIWAE and P-BiGAN, we treat label as missing feature. For DWR, SRDO and StableNet, we use mean-imputed feature as input. See the extended version for implementation details.

We use the commonly adopted Root Mean Square Error (RMSE) as evaluation metric throughout.

5.4 Performance on Synthetic Dataset

We compare StableMiss with baselines under different missing patterns and missing levels. The models are trained under a specific missing level and tested under all the missing levels. The difference in missing level represents the mask distribution shift. Due to space limit, we only present the results of missing rate shift on Gaussian-Mix feature with MAR mask; see the extended version for the missing rate shift on other settings and missing pattern shift. We also present the performance of optimal predictor E[Y | xm,m]. It can be derived from known feature distribution and label generation process, without which the performance of optimal predictor is not reachable.

We will show the results from two views: results with fixed training missing level that reflects the generalizability of methods, and results with fixed testing missing level that reflects the robustness to the mask distribution of training data.

Fixing Training Missing Level. We show the results when trained under 50% missing level and tested under all the missing levels in Table 1; the other settings are similar. Since the amount of observed data is different between missing levels, which influences the prediction error, we use the gap to optimal to reflect the generalization performance.

As the missing level increases, the error of the optimal predictor also increases, as there are less observed data. When the mask distribution shifts, StableMiss has the best generalization performance, reducing the gap to optimal of the second best, NeuMiss in this case, by 31%-44%. When there is no mask distribution shift, StableMiss still has the best performance in this case, slightly outperforming NeuMiss by 4%

in terms of the gap to optimal. As shown in the extended version, for Gaussian features with MCAR or MAR mask, NeuMiss can slightly outperform StableMiss, by at most 2%, in the absence of mask distribution shift. This is not too surprising, as these are the cases that NeuMiss is tailored for.

Fixing Testing Missing Level. We show the results when tested under 50% missing level and trained under all the missing levels in Table 2; the other settings are similar. When the mask distribution shifts, StableMiss achieves the best performance, reducing the gap to optimal of the second best NeuMiss by 30%-50%. Reduction is especially large when trained under 10% missing level, since the optimal predictors are more learnable under lower missing level. The results show that StableMiss is robust to the quality, i.e., mask distribution, of training data and can generalize from various training mask distributions.

5.5 Performance on Real-World Dataset

House Sales. We show the results on House Sales dataset when trained under 50% missing level in Table 3. Without prior knowledge about label generation process, the optimal predictor cannot be derived. StableMiss has the best performance when trained in distribution, i.e., trained under testing distribution. Instead of Optimal, we show our in-distribution performance, which is denoted by 'StableMiss-ID'.

Error of StableMiss-ID also increases with missing level.

When the mask distribution shifts, StableMiss achieves the best performance, reducing the gap to StableMiss-ID of the second best, usually StableNet, by 26%-50%, with an average of 38%, which demonstrates our generalizability on realworld data.

MNIST. We show the results on MNIST dataset when trained under 50% missing level in Table 4. When the mask distribution shifts, StableMiss achieves the best performance, reducing the gap to StableMiss-ID of the second best, usually StableNet, by 28%-50%, with an average of 41%.

Traffic. Table 5 shows the results on Traffic dataset, which is naturally incomplete. The training and testing missing rates are 37% and 23%. Since ground truth of missing entries is unavailable, RMSE is only computed on observed entries. The value for StableMiss-ID is the training error, which only serves as a reference for our comparison. Note that StableMiss achieves the best performance, reducing the gap to StableMiss-ID of the second best, StableNet, by 38%, which shows that StableMiss can be applied to real-world datasets with complex missingness.

5.6 Ablation Study

We study the efficacy of decorrelation by comparing with ablated variants that only decorrelates the entries of X and M respectively (only intra), only decorrelates between X and M (only inter) and does not decorrelate (w/o). The result on Gaussian-Mix feature with MAR mask when trained under 10% missing level is given in Figure 2. Our method with full decorrelation achieves the best performance especially when mask distribution shift becomes stronger. The gap of the second best to optimal is reduced by 55% on average and by 83% under 90%missing level.

We also study the influence of decorrelation when it is not needed. The result on Gaussian-Ind feature with MCAR-Ind mask when trained under 10% missing level is given in Figure 3. The performance with unneeded decorrelation is close to that without decorrelation. For real-world data with unknown correlation, we can always conduct decorrelation.

6 Conclusion

In this paper, we propose a novel method StableMiss for the problem of prediction with incomplete data under agnostic mask distribution shift. We leverage the observation that for each mask, there is an invariant optimal predictor. We approximate the optimal predictors jointly using a double parameterization technique. We also perform decorrelation to minimize the side effect caused by the intra-mask correlation and the correlation between features and mask. Extensive experiments are conducted on both synthetic and real-world datasets. The results show that StableMiss is robust and outperforms state-of-the-art methods under agnostic mask distribution shift.

Acknowledgments

This work was supported by NSF China (No. 62072302, U20A20181, U21A20519, 42050105, 62262018) and the Open Research Project of the State Key Laboratory of Media Convergence and Communication, Communication University of China, China (No.SKLMCC2021KF011).