MIRACLE: Causally-Aware Imputation via Learning Missing Data Mechanisms

1 Introduction . Missing data is an unavoidable byproduct of collecting data in most practical domains . In medicine , for example , doctors may choose to omit what they deem to be irrelevant information ( e.g. , some patients may be asked to get comprehensive blood tests while others don ’ t ) , data may be explicitly omitted by the patient ( e.g. , avoiding questions on smoking status precisely because of their smoking habit ) or simply misrecorded in electronic health systems ( see e.g. , [ 3 , 12 , 28 ] ) . Imputation algorithms can be used to estimate missing values based on data that was recorded , but their correctness depends on the type of missingness . For instance , expanding on the example above , younger patients may also be more likely to omit their smoking status . As illustrated in Figure 1 , the challenge is that implicitly conditioning inference on observed data introduces a spurious path of correlation between age and the prevalence of smoking that wouldn ’ t exist with complete data . Figure 1 : Missingness may introduce spurious dependencies . Missing data creates a shift between the available missing data distribution and the target complete data distribution . It is a shift that may be explicitly modeled as missingness indicators in an underlying causal model ( i.e. , a missingness graph as proposed by Mohan et al . [ 19 ] ) as shown in Figure 1 . The learning problem is one of extrapolation , learning with access to a missing data distribution for ⇤Equal contribution 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . prediction and inference on the complete data distribution – that is , generated from a model where all missingness indicators have been intervened on ( interventions interpreted in the sense of Pearl [ 22 ] ) thus graphically removing the dependence between missingness and its causes , and any spurious correlations among its ancestors . With this causal interpretation , imputation of missing data on a given variable Y from other observed variables X is formulated as a problem of robust optimization , minimize ✓2⇥ sup P2P E ( X , Y ) ⇠P ⇥ ( f✓ ( X ) Y ) 2 ⇤ , ( 1 ) simultaneously optimizing over the set of distributions P arising from interventions on missingness indicators . Causal solutions – i.e . imputation using functions of causal parents of each missing variable in the underlying causal graph – are a closely-related version of this problem with an uncertainty set Q defined as any distribution arising from interventions on observed variables and variable indicators ( see e.g . sections 3.2 and 3.3 in [ 16 ] ) , sup P2P E ( X , Y ) ⇠P ⇥ ( f✓ ( X ) Y ) 2 ⇤ sup P2Q E ( X , Y ) ⇠P ⇥ ( f✓ ( X ) Y ) 2 ⇤ , ( 2 ) since P ⇢ Q . Our premise is that causal solutions , i.e . minimizing the right-hand-side of ( 2 ) , are expected to correct for spurious correlations introduced by distribution shift due to missing data and preserve the dependencies of the complete data for downstream analysis . 1.1 Contributions . In this paper , we propose to impute while preserving the causal structure of the data . Missing values in a given variable are replaced with their conditional expectation given the realization of its causal parents instead of the more common conditional expectation given all other observed variables , which absorbs spurious correlations . We propose a novel imputation method called Missing data Imputation Refinement And Causal LEarning ( MIRACLE ) . MIRACLE is a general framework for imputation that operates on any baseline ( existing ) imputation method . A visual description is given in Figure 2 : given some initial imputation from a baseline method , MIRACLE refines its imputations iteratively by learning a missingness graph ( m-graph ) [ 19 ] and regularizing the imputation function such that it is consistent with the causal graph generating the data , substantially improving performance . In experiments , we apply MIRACLE to improve six popular imputation methods as baselines . We present detailed simulations to demonstrate on synthetic and a variety of publicly available datasets from the UCI Machine Learning Repository [ 8 ] that MIRACLE can improve imputation in almost every scenario and never degrades performance across all imputation methods . 1.2 Related work . The literature on imputation is large and varied . Still , most imputation algorithms work with the prior assumption that the missing data mechanism is ignorable , in the sense that imputation does not require explicitly modeling the distribution of missing values for imputation [ 25 ] . Accordingly , classical imputation methods impute using a joint distribution over the incomplete data that is either explicit or implicitly defined through a set of conditional distributions . For example , explicit joint modeling methods include generative methods based on Generative Adversarial Networks [ 33 , 34 , 1 ] , matrix completion methods [ 15 ] , and parametric models for the joint distribution of the data . Missing values are then imputed by drawing from their predictive distribution . The conditional modeling approach [ 31 ] consists of specifying one model for each variable and iteratively imputing using estimated conditional distributions . Examples of discriminative methods are random forests [ 27 ] , autoencoders [ 11 , 10 , 14 ] , graph neural networks [ 35 ] , distribution matching via optimal transport [ 20 ] , and multiple imputation using chained equations [ 6 ] . In a different line of research , Mohan et al. , in a series of papers , see e.g . [ 19 , 18 ] , explicitly considered missing data within the underlying causal mechanisms of the data . Subsequently , a range of related problems has been studied , including identifiability of distributions and causal effects in the presence of missing data , see e.g . [ 5 , 26 , 21 ] , testable implications relating to the causal structure using missing data [ 18 ] , and causal discovery in the presence of missing data [ 9 , 30 ] . Our focus , in contrast , is algorithmic in nature . We aim to develop an algorithm that improves imputation quality by leveraging causal insights represented as an estimated missingness graph learned from data . 2 Background . The basic semantic framework of our analysis rests on structural causal models ( SCMs ) ( see e.g . Chapter 7 in [ 22 ] for more details ) explicitly introducing missingness indicators and their functional relationship with other variables , using in part the notation of [ 19 ] . We define an SCM M as a tuple ( X , R , U , F , P ) where X is a vector of d endogenous variables and U is a vector of exogenous variables.2 R is the vector of missingness indicators that represent the status of missingness of the endogenous variables X . Precisely , Rj is responsible for the value of a proxy variable Zj of Xj , i.e. , the observed version of Xj . For example , Zj is equal to Xj if the corresponding record is observed ( Rj = 1 ) , otherwise Zj is missing ( Rj = 0 ) . F is a set of functions where each fX , fR 2 F decide the values of an endogenous variable X and a missingness indicator variable R , respectively . The function fX takes two separate arguments as parts of X ( except X itself ) and U , termed as PaX and UX . That is , X fX ( PaX , UX ) and R fR ( PaR , UR ) . The randomness in SCMs comes from the exogenous distribution PU ( u ) where the exogenous variables in U are generated independently and are mutually independent . Naturally , through the functions in F , the SCM M induces a joint distribution PX ( x ) over the endogenous variables X , called the endogenous distribution . An intervention on some arbitrary random variables V in X and R , denoted by do ( v ) , is an operation which sets the value of V to be v , regardless of how they are ordinarily determined . For an SCM M , let Mv denote a submodel of M induced by intervention do ( v ) . The interventional distribution PX ( x|do ( v ) ) induced by do ( v ) is defined as the distribution over X in the submodel Mv , namely , PX , Mv ( x ) = PX ( x|do ( v ) ) . Each SCM in the context of missingness is associated with a m-graph G ( e.g. , Fig . 1a ) , which is a directed acyclic graph ( DAG ) where nodes represent endogenous variables X and missingness indicators R , and arrows represent the arguments PaX and PaR of each function fX and fR respectively . By convention , exogenous variables U are often not shown explicitly in the graph . Assumption 1 ( Missingness indicators are not causes ) No missingness indicator in R can be the cause of the endogenous variables X , i.e. , the arguments of the functions generating X . Assumption 2 ( Causal sufficiency ) Exogeneous variables U are mutually independent , i.e. , all common parents of the endogenous variables are included in X . Assumption 3 ( No self-masking missingness ) Self-masking missingness refers to missingness in a variable that is caused by itself . In the m-graph this is depicted by an edge from Xj to Rj ( as shown in Figure 3 ( d ) ) . We assume that there is no such edges in the m-graph . Assumption 4 ( Observed root nodes ) The endogenous variables Xj such that PaXj = ; ( i.e. , the root nodes ) in the m-graph are always observed ( Rj = 1 with probability 1 ) . 2Essentially , X is the ground-truth features ; U is the random noise in the data generating process . We make the four assumptions above throughout the following sections . Assumption 1 and 2 are employed in most related works using m-graphs ( see e.g . [ 18 , 19 ] ) . Assumption 1 is valid , for example , if R is generated in the data collection process after the variable values are assigned . Consequently , under this assumption , if two endogenous variables of interest X1 and X2 are not d-separated by some variable X3 , they are not d-separated by X3 together with their missingness indicators R1 and R2 . We denote an independent relation in a data distribution by “ |= ” and dseparation in a m-graph by “ |= d ” . We assume the data distribution is faithful to a m-graph , meaning that the two independencies are equivalent . As shown in Figure 3 , data is missing completely at random ( MCAR ) if X |= dR holds in the m-graph , missing at random ( MAR ) if for any endogenous variable Xj , Rj |= dXj | X j holds , and missing not at random ( MNAR ) otherwise , as stated in [ 19 ] . If Assumption 3 is violated , we are unable to learn the missingness for self-masked variables . Assumption 4 is necessary for imputing all the missing variables from their causal parents . These assumptions are imperative for MIRACLE to provide improved imputations by leveraging the causal structure of the underlying data generating process . In our experiments ( Section 4 ) , we apply MIRACLE to real-world datasets where these assumptions are not guaranteed .

Summary: The authors present MIRACLE, a neural network that learns to impute well thanks a specific regularization scheme. The proposed regularisation is motivated by the idea that when imputation is based on all observed data, spurious correlations can arise while it is not the case when it is based solely on the causal parents. Technically, the weights of the imputation network are regularised so that the dependencies between all imputation functions form a DAG. This is possible thanks to prior work which introduced a score based on adjacency matrices that is equal to zero if and only if the corresponding graph is a DAG.    The neural network proposed by the authors also includes a second regulariser (moment regularizer), this time aimed at matching two different estimators of the first moments of the data. Overall, it minimises a reconstruction loss combined with these two regularisers.    Extensive experiments are conducted on simulated and real data that compare this neural network to various baselines in terms of imputation quality ( as well as performance on downstream tasks and congeniality in Supplementary materials). On simulated data, it is able to improve imputations over all baselines across all missing data mechanisms. On real data, it does not provide improvements in MCAR, but does so for half of the datasets in MAR and MNAR.