What’s a good imputation to predict with missing values?

1 Introduction . Data with missing values are ubiquitous in many applications , as in health or business : some observations come with missing features . There is a rich statistical literature on imputation as well as inference with missing values [ Rubin , 1976 , Little and Rubin , 1987 , 2002 , 2019 ] . Most of the theory and practices build upon the Missing At Random ( MAR ) assumption that allows to maximize the likelihood of observed data while ignoring the missing-values mechanism , for instance using expectation maximization [ Dempster et al. , 1977 ] . On the contrary , Missing Not At Random settings , where missingness depends on the unobserved values , may not be identifiable and require dedicated methods with a model of the missing-values mechanism . Learning predictive models with missing values poses distinct challenges compared to inference tasks [ Josse et al. , 2019 ] . Indeed , when the input is an arbitrary subset of variables in dimension d , there are 2d potential missing data patterns and as many sub-models to learn . Consequently , even simple data-generating mechanisms lead to complex decision rules [ Le Morvan et al. , 2020b ] . To date , there are few supervised-learning models natively suited for partially-observed data . A notable 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . exception is found with tree-based models [ Twala et al. , 2008 , Chen and Guestrin , 2016 ] , widely used in data-science practice . The most common practice however remains by far to use off-the-shelf methods first for imputation of missing values and second for supervised-learning on the resulting completed data . Such a procedure may benefit from progress in missing-value imputation with machine learning [ van Buuren 2018 , Yoon et al . 2018 , Mattei and Frellsen 2019 ] . However , there is a lack of learning theory to support such Impute-then-Regress procedures : Under what conditions are they Bayes consistent ? Which aspects of the imputation are important ? There is empirical realization that the choice of imputation matters for predictive performance . The NADIA R package [ Borowski and Fic ] can select an imputation method to minimize a prediction error on a test set . Auto-ML is used to optimize full pipelines , including imputation [ eg Jarrett et al. , 2021 ] . Ipsen et al . [ 2020 ] optimize a constant imputation for supervised learning . However , the imputation is only weakly guided by the target in these approaches , it is set either from a family of black-box methods using gradient-free model selection , or from trivial imputation functions . In addition , there is a lack of insight on what drives a good imputation for prediction . We contribute a systematic analysis of Impute-the-Regress procedures in a general setting : non-linear response function and any missingness mechanism ( no MAR assumptions ) . We show that : • Impute-then-Regress procedures are Bayes optimal for all missing data mechanisms and for almost all imputation functions , whatever the number of variables that may be missing . This very general result gives theoretical grounding to such widespread procedures . • We study “ natural ” choices of imputation and regression functions : the oracle imputation by the conditional expectation and oracle regression function on the complete data . We show that chaining these oracles is not Bayes optimal in general and quantify its excess risk . We show that in both cases , choosing an oracle for one step , imputation or regression , imposes discontinuities on the other step , thus making it harder to learn . • As these results suggest that imputation and regression should be adapted to one another , we contribute a method that jointly optimizes imputation and regression , using NeuMiss networks [ Le Morvan et al. , 2020a ] as a differentiable imputation procedure . • We compare empirically a number of Impute-then-Regress procedures on simulated non-linear regression tasks . Joint optimization of both steps provides the best performance . 2 Problem setting . Notations We consider a dataset of i.i.d . realizations of the random variable ( X , M , Y ) 2 Rd ⇥ { 0 , 1 } d ⇥ R where X are the complete covariates , M a missingness indicator , and Y a response of interest . For each realization ( x , m , y ) , mj = 1 indicates that xj is missing , and mj = 0 that it is observed . We denote by mis ( m ) ⇢ J1 , dK the indices corresponding to the missing covariates ( and similarly obs ( m ) the observed indices ) , so that xobs ( m ) corresponds to the entries actually observed . We define the incomplete covariate vector eX 2 ( R [ ( NA ) ) d as eXj = Xj if Mj = 0 and eXj = NA otherwise , where NA represents a missing value . Assumptions We assume that X admits a density on Rd and that , for all j 2 J1 , dK , each component Xj has finite expectation and variance , that is E ⇥ X 2 j ⇤ < 1 . Moreover , we assume that the response Y is generated according to : Y = f ? ( X ) + ✏ , with E ⇥ ✏|Xobs ( M ) , M ⇤ = 0 and E ⇥ Y 2 ⇤ < 1 . ( 1 ) where f ? : Rd ! R is a function of the complete input data X , ✏ 2 R is a random noise variable . 2.1 Supervised learning with missing values . Optimization problem In practice , in the presence of missing values , we do not have access to the complete data ( X , M , Y ) but only to the subset of it that is observed , i.e , ( Xobs ( M ) , M , Y ) . Thus instead of learning a mapping from Rd to R , we need to learn a mapping from ( R [ ( NA ) ) d to R , where the set of observed covariates can be any subset of J1 , dK . It is this unusual input space , partly discrete , that makes supervised learning with missing values challenging and different from classical supervised learning problems . Formally , the optimization problem we wish to solve is : min f : ( R [ ( NA ) ) d 7 ! R R ( f ) : = E ⇣ Y f ( eX ) ⌘2 ( 2 ) Bayes predictor The function which minimizes ( 2 ) , called the Bayes predictor , is given by : f̃ ? ( eX ) = E ⇥ Y |Xobs ( M ) , M ⇤ = E ⇥ f ? ( X ) |Xobs ( M ) , M ⇤ . ( 3 ) As eX is a function of Xobs and M , we will sometimes slightly abuse notations and write f̃ ? ( eX ) = f̃ ? ( Xobs , M ) . The risk of the Bayes predictor is called the Bayes risk , which we denote as R ? . It is the lowest achievable risk for a given supervised learning problem . Definition 1 ( Bayes optimality ) . A Bayes optimal function f achieves the Bayes rate , i.e , R ( f ) = R ? . As can be seen from ( 3 ) , the Bayes predictor is a function of M , a discrete random variable that can take one of 2d values since M 2 { 0 , 1 } d. The function f̃ ? can thus be viewed as 2d different functions , one for each possible subset of variables . This view raises questions that are central to this paper : How should we parametrize functions on such input domains ? And which function families should we consider to approximate f̃ ? ? These questions have been studied in the case where f ? is assumed to be a linear function , and X follows a Gaussian distribution . Indeed , under these assumptions , Le Morvan et al . [ 2020b , a ] have derived analytical expressions for the Bayes predictor and deduced appropriate parametric estimators . However , aside from specific cases , it is impossible to derive an analytical expression for the Bayes predictor , and novel arguments are needed to understand which classes of functions should be considered in general . 3 Asymptotic analysis of Impute-then-regress procedures . 3.1 Impute-then-regress procedures . Let |mis ( m ) | ( resp . |obs ( m ) | ) be the number of missing entries ( resp . observed ) for any missing data pattern m. For each m 2 { 0 , 1 } d , we define an imputation function ( m ) : R|obs ( m ) | ! R|mis ( m ) | which outputs values for the missing entries based on the observed ones . We denote by ( m ) j : R|obs ( m ) | ! R the component function of ( m ) that imputes the j-th component in X if it is missing . Classical choices of imputation functions include constant functions or linear functions . Finally , we introduce the family of functions FI that transform an incomplete vector into a complete one , precisely : FI = ( : ( R [ { NA } ) d ! Rd : 8j 2 J1 , dK , j ( eX ) = ( Xj if Mj = 0 ( M ) j ( Xobs ( M ) ) if Mj = 1 ) . ( 4 ) Let us define FI1 in the exact same way but for imputation functions ( m ) 2 C1 , for all m 2 { 0 , 1 } d. Here we study Impute-then-regress procedures , which we define as two-step procedures where the data is first imputed using a function 2 FI , and then a regression is performed on the imputed data . Such a procedure is quite natural to deal with arbitrary subsets of inputs variables . It embeds the data into Rd to reduce the problem to a classical one . In practice , impute-then-regress procedures are widely used . However , the choice of function class is so far mostly ad-hoc and raises multiple questions : How close to the Bayes rate can functions obtained via such procedures be ? Should we prefer some choices of imputation functions over others ? What happens when the missing data mechanism is missing not at random ? In this section , we will give answers to these questions . Below , we write obs ( resp . mis ) instead of obs ( M ) ( resp . mis ( M ) ) to lighten notations . 3.2 Impute-then-regress procedures are Bayes optimal . Definition 2 ( Universal consistency ) . An estimator fn is Bayes consistent if limn ! 1 R ( fn ) = R ? . It is said to be universally consistent if the previous statement holds for all distributions of ( X , Y ) . The following theorem shows that Impute-then-regress procedures are Bayes optimal for almost all imputation functions . In other words , it means that a universal learner trained on imputed data provides optimal performances asymptotically for almost all imputation functions . Let us now define , for all imputation functions 2 FI , the function g ? 2 argmin g : Rd 7 ! R E ⇣ Y g ( eX ) ⌘2 . Theorem 3.1 ( Bayes consistency of Impute-then-regress procedures ) . Assume the data is generated according to ( 1 ) . Then , for almost all imputation function 2 FI1 , the function g ? is Bayes optimal . In other words , for almost all imputation functions 2 FI1 , a universally consistent algorithm trained on the imputed data ( eX ) is Bayes consistent . Appendix A.3 gives the proof . Theorem 3.1 states a very general result : Impute-then-regress procedures are Bayes consistent for all missing data mechanisms , almost all imputation functions , regardless of the distribution of ( X , Y ) and the number of missing covariates . Since Theorem 3.1 holds for almost all imputation functions , it implies that good imputations are not required to obtain good predictive performances , at least asymptotically . Note that here , the notion of almost all is to be understood in its topological sense , and not in its measure theory sense . Moreover , this theorem does not make any assumption on the missing data mechanism , and is therefore valid for Missing Not At Random ( MNAR ) data . This contrasts with most methods for inference and imputation with missing values , valid only for MAR data . Finally , the theorem remains valid for any configuration of variables that may contain missing values , including the case in which all variables may contain missing values . Bayes consistency of Impute-the-Regress procedures has already been studied , but in much more restricted settings . Josse et al . [ 2019 ] proved that such procedures are Bayes consistent under the MAR assumption , for constant imputations functions and for only one potentially missing variable . Bertsimas et al . [ 2021 ] refined this result to almost surely continuous imputation functions . While these two prior works build on very similar proofs , we use here very different arguments summarized in the next paragraph . The first key idea of the proof is that , after imputation , all data points with a given missing data pattern m are mapped to a manifold M ( m ) of dimension |obs ( m ) | . For example in 3D , data points are mapped to R3 when completely observed , to 2D manifolds when they have one value missing , to 1D manifolds when they have two values missing , and to one point when all values are missing ( see Figure 1 ) . Thus , Impute-then-Regress procedures first map data points to various manifolds depending on their missing data patterns and then apply a prediction function defined on the whole space including manifolds . The second key idea of the proof is to ensure that the original missing data patterns of imputed points can almost surely be identified . For this , the proof requires that all manifolds of the same dimension are pairwise transverse . This assumption is sufficient , though not necessary , to ensure that the intersection of two manifolds of dimension |obs ( m ) | can not itself be of dimension |obs ( m ) | . Transversality is a weak assumption . In fact , Thom ’ s transversality theorem , ( which we rely on in our proof ) says that it is a generic property : it holds for “ typical examples ” , i.e almost all imputation functions will lead to transverse manifolds . To clarify this concept , we provide a particular case in 2D where 1D manifolds are not transverse in Appendix A.4 . The proof is constructive and exhibits a function g ? which achieves the Bayes rate for a given set of imputation functions . For each manifold M ( m ) , ordered from smallest dimension to largest , we require that g ? on M ( m ) equals the Bayes predictor for missing data pattern m except on points for which g ? has already been defined , i.e , the points where M ( m ) intersects with the manifolds ranked before it . Thus , we obtain a function g ? that does not depend on m , and which for each manifold , equals the Bayes predictor except on subsets of measure zero under the assumption that manifolds of the same dimension are pairwise transverse . Refer to appendix A.3 for more details . While this theorem is a very general result , it does not say what the optimal function associated to a given imputation looks like . In fact , depending on the imputation function it may be non-continuous , vary widely , and require a very large number of samples to be approximated . Note on Impute-then-classify procedures - Theorem 3.1 applies to regression problems . However , it can easily be shown that a similar result holds in binary classification settings . Indeed , in a binary classification setting , the Bayes predictor predicts class 1 if P ( Y = 1|X ) > 0.5 and -1 otherwise . Thus , it suffices to consider that the function of interest f ? ( X ) is the posterior probability P ( Y = 1|X ) . Then the same arguments as those used to prove Theorem 3.1 can be used to show that Impute-then-classify procedures are Bayes optimal for almost all imputation functions .

Summary: The authors look at the impute and regress framework and provide arguments as to Bayes optimality of the process for missing data mechanisms and imputation functions. It gives a general theoretical framework for the impute and regress procedures. Unlike the missing at random setups, they look at more general missing value mechanisms. They also argue that the imputation and regression processes should be properly aligned with each other and also provide a method for optimizing the two processes jointly.