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 gradient importance learning for predicting labels on incomplete data. The problem is: during both training and testing, certain dimensions of the samples are masked randomly and the masks (i.e. the "missing indicators") are known for the algorithm. Among existing methods in this problem, the paper claims that it is the first "imputation-free" way to solve this problem. The paper focuses on MLP and LSTM and adds an (element-wise) multiplicative parameter matrix on the gradient of the first layer (the one that is mostly closed to the data). The matrix is trained using reinforcement learning. The authors also connect its proposal to visual attention. Experiments are conducted on tabular, time series, and a simple image dataset.