Temporal Probabilistic Asymmetric Multi-task Learning

When performing multi-task predictions with time-series data , knowledge learned for one task at a specific time step may be useful in learning for another task at a later time step ( e.g . prediction of sepsis may be useful for prediction of mortality for risk prediction at intensive care units ) . To capture such dynamically changing asymmetric relationships between tasks and long-range temporal dependencies in time-series data , we propose a novel temporal asymmetric multi-task learning model , which learns to combine features from other tasks at diverse timesteps for the prediction of each task . One crucial challenge here is deciding on the direction and the amount of knowledge transfer , since loss-based knowledge transfer ( Lee et al. , 2016 ; 2017 ) does not apply in our case where we do not have loss at each timestep . We propose to tackle this challenge by proposing a novel uncertaintybased probabilistic knowledge transfer mechanism , such that we perform knowledge transfer from more certain tasks with lower variance to uncertain ones with higher variance . We validate our Temporal Probabilistic Asymmetric Multi-task Learning ( TP-AMTL ) model on two clinical risk prediction tasks against recent deep learning models for time-series analysis , which our model significantly outperforms by successfully preventing negative transfer . Further qualitative analysis of our model by clinicians suggests that the learned knowledge transfer graphs are helpful in analyzing the model ’ s predictions . 1 INTRODUCTION . Multi-task learning ( MTL ) ( Caruana , 1997 ) is a method to train a model , or multiple models jointly for multiple tasks to obtain improved generalization , by sharing knowledge among them . One of the most critical problems in multi-task learning is the problem known as negative transfer , where unreliable knowledge from other tasks adversely affects the target task . To prevent negative transfer , researchers have sought ways to allow knowledge transfer only among closely related tasks , by either identifying the task groups or learning optimal sharing structure among task ( Duong et al. , 2015 ; Misra et al. , 2016 ) . However , not only the task relatedness but the relative reliability of the task-specific knowledge also matters , and recent asymmetric multi-task learning models Lee et al . ( 2016 ; 2017 ) tackle this challenge by allowing tasks with low loss to transfer more . While the asymmetric knowledge transfer between tasks is useful , it does not fully exploit the asymmetry in the case of time-series analysis , which has an additional dimension of the time axis . With time-series data , knowledge transfer direction may need to be different depending on the timestep . For instance , suppose that we predict infection and mortality for patients in intensive care units based on their medical records . At earlier timesteps , prediction of infection may be more reliable than mortality , thus we may want knowledge transfer to happen from task infection to mortality ; at later timesteps , we may want the opposite situation to happen . Moreover , knowledge transfer may happen across timesteps . For example , a high risk of infection in early timestep will alarm high risk of mortality at later timesteps . To exploit such temporal relationships between tasks , we need a model that does not perform static knowledge transfer between two tasks ( Figure 1a ) , but dynamically changes the knowledge transfer amount and direction at each timestep , and also transfers knowledge across timesteps ( Figure 1b ) . Toward this objective , we propose a multi-task learning framework for time-series data , where each task not only learns its own latent features at each timestep but also aggregates the latent features from the other tasks at the same or different timesteps via attention allocation . Yet this brings in another challenge . On what basis should we promote asymmetric knowledge transfer ? For asymmetric knowledge transfer between tasks , we could use task loss as a proxy of knowledge reliability . However , loss is not a direct measure of reliability , since a model trained with few instances may have a small loss , while the knowledge from the model could be highly biased and unreliable . Further , for prediction with time-series data , loss may not be available at every timestep . Thus , instead of task loss , we focus on the uncertainty , which can be obtained with probabilistic Bayesian models . Basically , if a latent feature learned at a certain timestep has large uncertainty ( variance ) , we can consider its knowledge as unreliable . In such a case , the model may allocate small attention values for the feature ; that is , the attention will be attenuated based on the uncertainty . To validate the superiority of our model , we experiment with it on three clinical risk prediction datasets against multiple baselines . The results show that our model obtains significant improvements over strong multi-task learning baselines . Further , both the asymmetric knowledge transfer between tasks at two different timesteps as well as the uncertainty-based attenuation of attention weights are found to be useful in improving the generalization performance . Our contribution in this work is threefold : • We propose a novel asymmetric multi-task learning framework for time-series analysis , which utilizes feature-level uncertainty to perform knowledge transfer among tasks and across time-steps , thereby exploiting both the task-relatedness and temporal dependencies . • We use a probabilistic Bayesian formulation for asymmetric knowledge transfer , where the amount of knowledge transfer depends on the uncertainty at the feature level . • We validate our model on clinical risk prediction tasks , on which it achieves significant improvements over baselines and provides meaningful interpretations , including temporal relationships between tasks . 2 RELATED WORK . Multi-task Learning While the literature on multi-task learning ( Caruana , 1997 ; Argyriou et al. , 2008 ) is vast , we selectively mention the prior works that are closely related to ours . Historically , multi-task learning models have focused on what to share ( Yang & Hospedales , 2016a ; b ; Ruder12 et al . ) , as the jointly learned models could share instances , parameters , or features ( Kang et al. , 2011 ; Kumar & Daume III , 2012 ; Maurer et al. , 2013 ) . With deep learning , multi-task learning can be implemented rather straightforwardly by making multiple tasks to share the same deep network . However , since solving different tasks will require heterogeneous knowledge , complete sharing of the underlying network may be suboptimal and brings in a problem known as negative transfer , where certain tasks are negatively affected by knowledge sharing . To prevent this , researchers are exploring more effective knowledge sharing structures . Duong et al . ( 2015 ) proposed a soft parameter sharing method that uses a regularizer to enforce the network parameters for each task to be similar , while Misra et al . ( 2016 ) proposed to learn the optimal combination of shared and taskspecific representations in computer vision . Kendall et al . ( 2018 ) proposed a multi-task framework when the losses are weighed based on tasks ’ uncertainty , thereby reducing negative transfer from uncertain tasks . While finding a good sharing structure can alleviate negative transfer , negative transfer will still persist if we perform symmetric knowledge transfer among tasks . To resolve this symmetry issue , Lee et al . ( 2016 ) proposed an asymmetric MTL model with inter-task knowledge transfer that allows task-specific parameters for tasks with smaller loss to transfer more . Lee et al . ( 2017 ) proposed a model for asymmetric task-to-feature transfer that allows reconstructing features with task-specific features while considering their loss , which is more suitable for deep neural networks and scalable . Our model is also targetting asymmetric multi-task learning , but is different from these previous works in that it utilizes uncertainty rather than loss as the measure of task reliability , and performs asymmetric knowledge transfer at each timestep , and across timesteps . Clinical time-series analysis While our method is generic and applicable to any time-series prediction tasks , we mostly focus on clinical time-series analysis in this paper . Recently , there has been some progress on this topic , mostly focusing on interpretability and reliability of the model . Choi et al . ( 2016 ) proposed an attention-based model that generates attention for both the timesteps ( hospital visits ) and features ( medical examination results ) , to provide interpretations of the predictions . However , attentions are often unreliable since they are learned in a weakly-supervised manner , and Heo et al . ( 2018 ) proposed to obtain reliable interpretation and prediction by proposing a probabilistic attention mechanism that considers uncertainty as to how to trust the input . Our work shares the motivation with these prior works as we target interpretability and reliability . Recently , Song et al . ( 2018 ) proposed a model that uses the Transformer architecture to perform time-series prediction with multi-task learning experiments . However , their model takes a straightforward approach where all tasks share a single base network , which is susceptible to negative transfer . 3 APPROACH . We now formally describe the problem setting and our probabilistic asymmetric multi-task learning framework for time-series prediction . 3.1 TIME-SERIES PREDICTION WITH FEATURE-LEVEL UNCERTAINTY . Our goal is to jointly train time-series prediction models for multiple tasks at once . Suppose that we are given training data for D tasks , D = { ( X1 , Y1 ) , . . . , ( XD , YD ) } . Further suppose that each data instance x ( x ∈ Xd for some d ) consists of T timesteps . That is , x = ( x ( 1 ) , x ( 2 ) , . . . , x ( T ) ) , where x ( t ) ∈ R1×m denote the data instance for the timestep t. Additional , yd is the label for task d ; for binary classification task , yd ∈ { 0 , 1 } , and for regression case , yd ∈ R. Given time-series data and tasks , we want to learn the task-specific latent features for each task and timestep , and then perform asymmetric knowledge transfer between them . Our multi-task learning framework is comprised of the following components : Shared Low-Level Layers ( Figure 2a ) We allow our model to share low-level layers for all the tasks in order to learn a common data representation before learning task-specific features . At the lowest layer , we have a shared linear data embedding layer to embed the data instance for each timestep into a continuous shared feature space . Given a time-series data instance x , we first linearly transform the data point for each timestep t , x ( t ) ∈ Rm , which contains m variables . ( v ( 1 ) , v ( 2 ) , ... , v ( T ) ) = v = xWemb ∈ RT×k ( 1 ) where Wemb ∈ Rm×k and k is the number of hidden units . After embedding the data instances , we input them into a shared RNN layer for pre-processing : r = ( r ( 1 ) , r ( 2 ) , ... , r ( T ) ) = RNN ( v ( 1 ) , v ( 2 ) , ... , v ( T ) ) ( 2 ) Task- and Timestep Embedding Layers After embedding and pre-processing the input into a continuous space , we further encode them into task- and timestep-specific features . Since hardsharing layers may result in negative transfer between tasks , we use separate embedding layers for each task to encode task-specific knowledge . For each task d , the separate network consists of L layers ( Figure 2b ) of feed-forward networks , to learn disentangled knowledge for each timestep . The L feed-forward layers for task embedding ( Figure 2b ) can be formalized as : hd = σ ( ( ... σ ( σ ( rW 1 d + b 1 d ) W 2 d + b 2 d ) ... ) W L d + b L d ) ∈ RT×k ( 3 ) where Wid ∈ Rk×k , bid ∈ Rk and σ is a non-linear activation function ( e.g . leaky relu ) . Modeling feature-level uncertainty While the above embedding can capture knowledge for each task and timestep , we want to further model their uncertainties as well , to measure the reliability of the knowledge captured . Towards this objective , we model the latent variables as probabilistic random variables , with two types of uncertainty ( Kendall & Gal , 2017 ) : 1 ) epistemic uncertainty , which comes from the model ’ s unreliability from the lack of training data , and 2 ) aleatoric uncertainty , that comes from the inherent ambiguity in the data . We capture the former by using dropout variational inference ( Gal & Ghahramani , 2016 ) , and the latter by explicitly learning the model variance as a function of the input ( Figure 2c ) . Suppose that our generative model is : zd ∼ pθ ( zd|x , ω ) , yd ∼ pθ ( yd|x , zd , ω ) where ω is the set of all parameters , pθ ( zd|x , ω ) and pθ ( yd|x , zd , ω ) are networks parameterized by θ . We denote all joint , marginal and conditional distributions arising from this generative model by p. zd|x , ω ∼ N ( zd ; µd , diag ( σ 2 d ) ) ( 4 ) µd = σ ( hdW µ d + b µ d ) ( 5 ) σd = softplus ( hdW σ d + b σ d ) ( 6 ) The posterior distributions p ( zd|x , yd , ω ) and p ( ω|D ) are intractable so we use variational inference to approximate the joint distribution . The variational distribution is defined by : q ( zd , ω|x , yd , D ) = qM ( ω|D ) q ( zd|x , yd , ω ) ( 7 ) We set qM ( ω|D ) by dropout approximation ( Gal & Ghahramani , 2016 ) with parameter M and simply set q ( zd|x , yd , ω ) to the prior network pθ ( zd|x , ω ) ( which works well in practice ( Sohn et al. , 2015 ; Heo et al. , 2018 ) ) for the model consistency at both training and test time .

Summary: The submission proposes a probabilistic method for multitask learning for sequence data focusing on the use of uncertainty to control the amount of transfer between tasks. In comparison to prior work, the approach directly models uncertainty by parametrising Gaussian distributions and dropout variational inference. The approach additionally models transfer between features from different time-steps. The evaluation shows that the method outperforms recent state-of-the-art models on two medical benchmarks.