Gradient Importance Learning for Incomplete Observations

1 INTRODUCTION . We consider learning from incomplete datasets . Components of the data could be missing for various reasons , for instance , missing responses from participants , data loss , and restricted access issues . This phenomenon is prevalent in the healthcare domain , mostly in the context of electronic health records ( EHRs ) , which are structured as patient-specific irregular timelines with attributes , e.g. , diagnosis , laboratory tests , vitals , thus resulting in high missingness across patients for any arbitrary time point . Such missingness introduces difficulties when developing models and performing inference in real-world applications such as Kam & Kim ( 2017 ) ; Scherpf et al . ( 2019 ) . Existing works tackle this problem by either proposing imputation algorithms ( IAs ) to explicitly produce estimates of the missing data , or by imposing imputation objectives during inference , e.g. , withholding observed values as being the ground-truth and learning to impute them . However , some of these either require additional modeling assumptions for the underlying distributions ( Bashir & Wei , 2018 ; Fortuin et al. , 2020 ) , or formatting of the data ( Yu et al. , 2016 ; Schnabel et al. , 2016 ) . Other applications depend on domain knowledge for pre-processing and modeling such as Yang et al . ( 2018 ) ; Kam & Kim ( 2017 ) , or introduce additional information-based losses ( Cao et al. , 2018 ; Lipton et al. , 2016 ) , which are usually intractable with real-world data . Moreover , generative methods ( Mattei & Frellsen , 2019 ; Yoon et al. , 2018 ) could result in imputations with high variation ( e.g. , low confidence of the output distribution ) , when data have high missingness rates or small sample sizes . Figure 1 illustrates imputations generated by two state-of-the-art deep generative models on the MNIST dataset with 50 % , 70 % and 90 % of the pixels set as missing ( i.e. , masked out ) . It is observed that both approaches suffer from inaccurate reconstruction of the original digits as the missingness rate increases , which is manifested as some of the imputed images not being recognizable as digits . Importantly , the error introduced by the imputation process can be further propagated into downstream inference and prediction stages . Imputing the missing data ( either explicitly or implicitly ) is , in most cases , not necessary , more so , considering that sometimes where and when the data are missing can be intrinsically informative . ∗Duke University , USA . Contact : { qitong.gao , miroslav.pajic } @ duke.edu . †King Abdullah University of Science and Technology , Saudi Arabia . Code available at https : //github.com/gaoqitong/gradient-importance-learning . Consider a scenario in which two patients , A and B , admitted to the intensive care unit ( ICU ) suffer from bacterial and viral infections , respectively , and assume that the healthcare provider monitors the status of the patients by ordering two ( slightly ) different blood tests periodically , namely , a culture test/panel specific to bacterial infections and a RT-PCR test/panel specific to viral infections . Hence , patient A is likely to have much fewer orders ( and results ) for viral tests . In both cases , the missingness patterns are indicative of the underlying condition of both patients . Moreover , such patterns are more commonly found in incomplete data caused by missing at random ( MAR ) or missing not at random ( MNAR ) mechanisms , which usually introduce additional difficulties for inference ( Little & Rubin , 2019 ) . Inspired by this , we propose the gradient importance learning ( GIL ) method , which facilitates an imputation-free learning framework for incomplete data by simultaneously leveraging the observed data and the information underlying missingness patterns . Specifically , we propose to use an importance matrix to weight the gradients that are used to update the model parameters in the backpropagation process , after the computational graph of the model is settled , which allows the model to exploit the information underlying missingness without imputation . However , these gradients can not be tracked by automated differentiation tools such as Tensorflow ( Abadi et al. , 2016 ) . So motivated , we propose to generate the importance matrix using a reinforcement learning ( RL ) policy on-the-fly . This is done by conditioning on the status of training procedure characterized by the model parameters , inputs and model performance at the current training step . Concurrently , RL algorithms are used to update the policy by first modeling the back-propagation process as an RL environment , or a Markov decision process ( MDP ) . Then , by interacting with the environment , the RL algorithm can learn the optimal policy so the importance matrix can aid the training of the prediction models to attain desirable performance . Moreover , we also show that our framework can be augmented with feature learning techniques , e.g. , contrastive learning as in Chen et al . ( 2020 ) , which can further improve the performance of the proposed imputation-free models . The technical contributions of this work can be summarized as follows : ( i ) to the best of our knowledge , this is the first work that trains deep learning ( DL ) models to perform accurate imputationfree predictions with missing inputs . This allows models to effectively handle high missing rates , while significantly reducing the prediction error compared to existing imputation-prediction frameworks . ( ii ) Unlike existing approaches that require additional modeling assumptions or that rely on methods ( or pre-existing knowledge ) that intrinsically have advantages over specific types of data , our method does not require any assumptions or domain expertise . Consequently , it can consistently achieve top performance over a variety of data and under different conditions . ( iii ) The proposed framework also facilitates feature learning from incomplete data , as the importance matrix can guide the hidden layers of the model to capture information underlying missingness patterns during training , which results in more expressive features , as illustrated in Figure 3 . 2 GRADIENT IMPORTANCE LEARNING ( GIL ) . In this section , we introduce GIL ; a method that facilitates training of imputation-free prediction models with incomplete data . In Section 2.2 , we show how the gradient descent directions can be adjusted using the importance matrix A . In Section 2.3 , we introduce an RL framework to generate the elements of A which can leverage the information underlying missingness during training , as illustrated in Figure 2 . 2.1 PROBLEM FORMULATION . We consider a dataset containing N observations X = ( X1 , X2 , . . . , XN ) where each Xj can be represented by a vector xj ∈ Rd ( for tabular data ) or a matrix Xj ∈ RTj×d ( for time-series ) with Tj denoting the time horizon of the sequence Xj . Since the focus of this work is principally on time-series data , recurrent neural networks will receive significant attention . We also define the set of missing indicatorsM = ( M1 , M2 , . . . , MN ) , where Mj = mj ∈ { 0 , 1 } d or Mj ∈ { 0 , 1 } Tj×d depending on the dimension of Xj , and 0 ’ s ( or 1 ’ s ) correspond to the entries in Xj that are missing ( or not missing ) , respectively . Finally , we assume that each Xj is associated with some label yj ∈ Y ; thus Y = ( y1 , y2 , . . . , yN ) denotes the set of labels . We define the problem as learning a model parameterized by W to directly predict yj by generating ŷj that maximizes the log-likelihood∑N j=1 log p ( ŷj |Xj , Mj , W ) , without imputing the missing values in Xj . 2.2 GRADIENT IMPORTANCE . Missingness is a common issue found in tabular data and time series , where multi-layered perceptron ( MLP ) ( Bishop , 2006 ) and long short-term memory ( LSTM ) ( Hochreiter & Schmidhuber , 1997 ) models , respectively , are often considered for predictions . Below we illustrate how the importance matrix can be applied to gradients , which are produced by taking regular stochastic gradient descent ( SGD ) steps , toward training MLP models with incomplete tabular inputs . Note that this idea can be easily extended to sequential inputs with LSTM models and the corresponding details can be found in Appendix A . First , the definitions of the dataset and missing indicators can be reduced to Xtab = ( x1 , x2 , . . . , xN ) and Mtab = ( m1 , m2 , . . . , mN ) , with xj ∈ Rd and mj ∈ Rd , respectively . Then we define an MLP with k hidden layers , for k ≥ 2 , as follows ŷ = φout ( Woutφk ( Wk . . . φ2 ( W2φ1 ( W1x ) ) ) ) , ( 1 ) where x ∈ Rd is the input , Wi and φi are the weight matrix and activation functions for the i-th hidden layer , respectively , and we have omitted the bias terms for notational convenience . The first layer can be interpreted as the encoding layer Wenc = W1 that maps inputs to features ζ = fenc ( x|Wenc ) , while the rest are the inference layers Winf = { W2 , . . . , Wk , Wout } that map the features to prediction ŷ = finf ( ζ|Winf ) . In the following proposition , we show that the gradients of some commonly used loss functions , E ( ŷ , y ) , such as cross entropy or mean squared errors , w.r.t . W1 , can be formulated in the form of an outer product . The proof and details can be found in Appendix E. Proposition 1 Given a MLP and a smooth loss function E ( ŷ , y ) , the gradients of E w.r.t . the encoding layer can be formulated as ∂E/∂W1 = ∆1x > , where ∆1 contains the gradients propagated from all the inference layers , x ∈ Xtab and y ∈ Y . From this Proposition it can be seen that the gradients that are used to train the encoding layers , following regular SGD solvers , can be formulated as the outer product between the gradients ∆ propagated from the inference layers and the input x as ( ∂E/∂Wenc ) SGD = ∆ · x > , ( 2 ) where Wenc ∈ Re×d , ∆ ∈ Re , x ∈ Rd , e is the dimension of the features , and d is the dimension of the input . To simplify notation , note that henceforth we use x to refer to an individual tabular data xj ∈ Xtab , regardless of its index . The corresponding SGD updates for training Wenc , using learning rate α , are Wenc ←Wenc − α∆ · x > . As a result , it is observed from ( 2 ) that the j-th ( j ∈ [ 1 , d ] ) column of ( ∂E/∂Wenc ) SGD is weighted by the j-th element in x given ∆ . However , given that some entries in x could be missing ( according to the missing indicator m ) and their values are usually replaced by a placeholder ξ ∈ R , it may not be ideal do directly back-propagate the gradients in the form of ( 2 ) , as the features ζ captured by the encoding layers may not be expressive enough for the inference layers to make accurate predictions . Instead , we consider introducing the importance matrix A ∈ [ 0 , 1 ] e×d to adjust the gradient descent direction as Wenc ←Wenc − α ( ∂E/∂Wenc ) SGD A , ( 3 ) which not only trains the model to capture information from the observed inputs that are most useful for making accurate predictions , but also to leverage the information underlying the missingness in the data . The elements of A can be generated using the RL framework introduced in the next section . Note that above we have omitted all the bias terms to simplify our presentation and note that : i ) gradients w.r.t . to biases do not conform to the outer product format , and ii ) more importantly , these gradients do not depend on the inputs – thus in practice , the importance matrix A is only applied to the gradients of the weights . Moreover , we do not consider convolutional neural network ( CNN ) models ( LeCun et al. , 1999 ) in this work because of i ) . Though the proposed framework could still be applied to CNNs , it may not be as efficient as for MLPs or LSTMs , where the search space for A is significantly reduced by taking advantages of the outer product format , as shown in ( 4 ) below . Importantly , our focus is to address the missingness in tabular data and time-series , in which case MLPs and LSTMs are appropriate , while the missingness in images is usually related to applications in other domains such as compressed sensing ( Wu et al. , 2019 ) or image super-resolution ( Dong et al. , 2015 ) , which we plan to explore in the future .

Summary: The paper proposes a method to train on data with missing features. It does so with a single model that can handle missing features, not with extra imputation methods. This is done by weighting the gradient update of the first weight matrix of the neural network with a vector a that is produced by an RL agent. This agent's reward in turn is the performance of the model after the update.