Multi-task Learning via Adaptation to Similar Tasks for Mortality Prediction of Diverse Rare Diseases

Abstract

The mortality prediction of diverse rare diseases using electronic health record (EHR) data is a crucial task for intelligent healthcare. However, data insufficiency and the clinical diversity of rare diseases make it hard for deep learning models to be trained. Mortality prediction for these patients with different diseases can be viewed as a multi-task learning problem with insufficient data but a large number of tasks. On the other hand, insufficient training data makes it difficult to train task-specific modules in multi-task learning models. To address the challenges of data insufficiency and task diversity, we propose an initialization-sharing multi-task learning method (Ada-SiT). Ada-Sit can learn the parameter initialization and dynamically measure the tasks’ similarities, used for fast adaptation. We use Ada-SiT to train long short-term memory networks (LSTM) based prediction models on longitudinal EHR data. The experimental results demonstrate that the proposed model is effective for mortality prediction of diverse rare diseases.

1. Introduction

Mortality prediction [1] of diseases plays a crucial role in clinical work and helps doctors to take early interventions based on timely alert of patients’ adverse health status. With the immense accumulation of Electronic Health Records (EHR) available [2, 3], deep learning models [4], typically requiring a large volume of data, have been developed for mortality prediction of common diseases, demonstrating state-of-the-art performance. However, mortality prediction of rare diseases is relatively unexplored in the domain of intelligent healthcare and personalized medicine.

Predicting mortality of rare diseases suffers from the problem of data insufficiency. A rare disease is the disease that affects a small percentage of the population and has a different disease mechanism from common ones. However, there are more than 300 million people worldwide living with one of the approximately 7,000 rare diseases (US organization Global Genes^‡). Therefore, there is always not enough data for a specific rare disease. In some real-world data sources [3], only tens of data samples on average could be collected for each rare disease.

Besides data insufficiency, the clinical behavior diversity of these diseases is another challenge for mortality prediction of rare diseases. Behaviors of different diseases vary a lot and make potential conflicts to the training of global deep learning models [5, 6, 7]. For example, the high heart rate raises the mortality risk of people with heart disease, but for patients having a cold, it is a common symptom that does not indicate danger. So the insufficiency of diverse rare diseases data cannot be simply resolved by training a global model using samples from all common diseases.

Multi-task learning models [8, 9] can be used to settle the problem of disease behavior diversity. Mortality prediction for each kind of rare disease is viewed as a task, and multi-task learning is supposed to capture the task-specific characteristics as well as the shared information of all tasks. However, data insufficiency of the rare disease tasks makes it hard to train the corresponding modules and further utilize the shared characteristics of similar tasks in a multi-task learning framework.

Meta-learning methods can learn meta-knowledge of training models, which makes it possible to learn fast with few samples, such as few-shot learning [10]. To build a better multi-task model suitable for tasks with little training data, we bring in the idea of fast adaptation in meta learning, which learns a shared initialization as meta-knowledge for adaptation to new tasks. However, since the fast adaptation method adapts shared initialization to each task independently, it cannot directly take into account the relationship of similar tasks, which is important and can provide useful information to enhance multi-task learning [11, 12].

To deal with the above-mentioned challenges (summarised in Table 1), we propose an initialization-shared multitask learning method, named as Ada-SiT (Adaptation to Similar Task), in which the task similarity is dynamically measured in the process of meta-training according to our new definition. Therefore, the task similarity, as a part of learned meta-knowledge, can enhance the fast adaptation procedure of genetic meta-learning models. Moreover, Ada-SiT is model-agnostic and can employ all existing deep learning based approaches as the basic predictive model. Experimental results on real medical datasets demonstrate that the proposed model is able to make similar tasks cooperate in initialization-shared multi-task learning, and it outperforms state-of-the-art global models as well as multi-task methods for mortality prediction of diverse rare diseases.

Table 1:

Comparing the three kinds of methods in terms of their ability to handle three main challenges in mortality prediction for rare diseases

Method	Examples	Challenges
Data Insufficiency	Task Similarity	Task Diversity
Global Models	LSTM [5], TCN [6]	✓	✓	×
Multi-task Models	Multi-SAnD [8], Multi-Dense [9]	×	✓
Meta-learning Models	MAML [10]	✓	×
Our Model	Ada-SiT	✓	✓	✓

It is worthwhile to highlight the contributions of the proposed model as follow:

• To the best of our knowledge, this is the first attempt to simultaneously tackle the challenge of disease diversity and data insufficiency in mortality prediction of rare diseases.

• We propose a novel initialization-shared multi-task learning method Ada-SiT, which can utilize information of task similarity for adaptation to each small sample-size task.

2. Related Works

2.1. Deep Learning for Healthcare

The accumulation of Electronic Health Records (EHR) has enabled research on deep learning methods for healthcare [13, 14, 15, 16, 17]. Multi-layer Perceptron (MLP) [18], Convolutional Neural Network (CNN) [19] and Recurrent Neural Network (RNN) [4, 20, 21] have been used in healthcare domain. Among these methods, there are many works on mortality prediction. The good performance of these models depends on a large volume of EHR data, which cannot be satisfied in our scenario of mortality prediction for diverse rare diseases. As a result, these models cannot make precise mortality predictions for patients with different rare diseases. Our work is suitable for these settings because it simultaneously tackles the challenges of disease diversity and data insufficiency. Furthermore, our method is a general framework and can be applied to train deep learning models to improve their performance.

2.2. Multi-task Learning

Multi-task learning is an efficient method to improve the performance by jointly learning multiple related tasks. In deep multi-task learning models, the information sharing mechanism is based on specific network structures, including shared layers [22], shared functions [23] and additional constraints [24]. However, task similarity cannot be directly interpreted in these models . We propose a model-agnostic multi-task learning method which can share the parameter initialization for fast adaptation to each task and task similarity can be dynamically measured in the training prossess.

In the healthcare domain, multi-task learning is used for prediction of various clinical events [25], mortality prediction of multiple patient cohorts [9] and patient-specific diagnosis [26], in which the ”tasks” have different definitions. Similar to the paper [9], our work also treats the mortality prediction of a certain patient cohort as a task. However, the method proposed by Suresh [9] is suitable for a small number of tasks with a large volume of data, meanwhile, our method on mortality prediction for rare diseases is designed to deal with hundreds of tasks with insufficient data.

2.3. Optimization-based Meta Learning

To solve the problem of data insufficiency and task diversity, our method borrows the idea behind optimization-based meta-learning [10], which can adapt to new environments with a few training samples by modifying the parameter optimization process. MAML (Model Agnostic Meta Learning) [10] uses fast adaptation to find a good initialization for the parameters of deep neural networks. And the idea of meta-learning is also applied in personalized dialog system for improving the diversity of text generation [27]. Our work is similar to MAML for using fast adaptation to get parameters of each task, but different from it in two ways: First, the objective of MAML is to learn a good parameter initialization for fast adapting to new tasks but our work is to find good model parameters for each given task. Second, our work measures task similarity in model space dynamically and uses samples in similar tasks to assist the adaptation to each task while MAML does fast adaptations to each task independently.

In clinical scenario, MetaPred [28] use MAML for clinical risk prediction with limited patient electronic health records. Its task is similar to ours but its method is different from ours in two ways: First, it trains a parameter initialization on source domain and simulated target domain via fast adaptation, but our method can learn from multiple small target domains without the source domain knowledge. Second, like MAML, MetaPred also doesn’t consider task similarity.

3. Data and Task Descriptions

We give the notations and data descriptions of the predictive tasks in the following.

3.1. Heterogeneous Temporal Events in EHR data

The input X of each mortality prediction task is a given episode of patient EHR, which could be represented as T heterogeneous temporal events [4]: $X={\left\{e_t\right\}}_{1\le t\le T}$. $e_t$ in this sequence is a tuple with four element: e = (type, value_c, value_n, time) , where type is the clinical event type, value_c and value_n are the categorical and numerical attributes of e_t, and time is the record time of e_t.

3.2. Multi-Task Mortality Prediction

The mortality prediction of each rare disease is defined as a task. Specifically, assuming that there are M diseases, we refer task_i as the corpus of the i-th mortality prediction task for patients of rare disease d_i with N_i samples:

$$tas{k}_i={\left\{\left(X_k^{(i)},y_k^{(\widehat{i})}\right)\right\}}_{k=1}^{N_i}$$ (1)

where $X_k^{(i)}$ and $y_k^{(\widehat{i})}$ denote the k-th sample and its label respectively in the i-th task.

Specifically, X is a given episode of a patient’s EHR data, represented as heterogeneous temporal events, and $\widehat{y}$ is the binary label indicating whether a patient will die in 24 hours.

3.3. Patient Cohort Setup

3.3.1. Heterogeneous Temporal Event Datasets

We set up two heterogeneous temporal event datasets based on MIMIC-III [3] database and eICU [2] database. MIMIC-III is a large, freely-available database comprising health data of patients in critical care units from Beth Israel Deaconess Medical Center between 2001 and 2012, and eICU is populated with data from a combination of many critical care units throughout the continental United States between 2014 and 2015.

The two datasets have the same data preprocessing framework: For each patient, we select all the events with their features from the original database and arrange the events in the temporal order. The descriptions of the selected events in eICU are listed in Table 2. The details of the events in MIMIC-III can be found in [13]. Then we annotate the mortality label for each patient event sequence.

Table 2:

The tables in elCU to construct heterogeneous temporal events

Table Name in eICU

Description

lab intakeoutput medication infusiondrug careplan admissiondrug nursecharting physicalexam diagnosis respiratorycare allergy

Laboratory tests mapped to a standard set of measurements. E.g. labTypeID, labResult Intake and output recorded for patients. E.g. intakeTotal, outputTotal, dialysisTotal Active medication orders for patients. E.g. drugHiclSeqno, dosage Details of drug infusions. E.g. drugRate, infusionRate, drugAmount Documentation relating to care planning. E.g. cplGeneralID, cplItemValue Details of medications that a patient was taking prior to admission to the ICU. E.g. drugUnit Information entered in a semi-structured form by the nurse. E.g. nursingchartcelltypecat Patients’ results of the physical exam. E.g. physicalExamText, physicalExamValue Diagnosis information of Patients. E.g. ICD9Code, diagnosisPriority Information related to respiratory care for patient. E.g. airwayType, airwaySize, cuffPressure Details of patient allergies. E.g. allergyType

3.3.2. Rare Disease Selection

We reorganize the heterogeneous temporal event datasets to get two rare disease datasets, MiniMIMIC and MiniEICU.

For each ICD (International Classification of Diagnose ^§) code in MIMIC-III and eICU, we calculate its sample size (i.e. the number of patients with this code). We select 858 ICD codes with less than 40 samples in MIMIC-III and 70 ICD codes with less than 100 samples in eICU as rare diseases. The heterogeneous temporal event sequences of patients with selected rare disease d_i form task_i. So MiniMIMIC is the task list {task₁, ..., task₈₅₈} and MiniEICU is the task list {task₁, ..., task₇₀}.

The statistics of MiniMIMIC and MiniEICU are summarized in Table 3.

Table 3:

Characteristics of Datasets

Name	MiniMIMIC	MiniEICU
# of tasks	858	70
# of samples	16610	7000
positive sample rate (mortality rate)	7%	13%
max # of samples per task	40	100
min # of samples per task	10	100
mean # of samples per task	19.36	100

Each task_i is split into 3 parts with fixed proportions, namely Train_i(70%), Valid_i(10%) and Test_i(20%). The validation set is used for conducting “early stop” and selecting hyper-parameters. The results of the evaluation metrics on the test set and their stander variations are used to compare different models.

4. Methodology

In this section, we begin with some basic notations for the adaptation to a single task and the backbone deep network for mortality prediction based on EHR data. Then we introduce the framework of our method Ada-Sit (Adaptation to similar tasks), which could be applied to train multiple mortality predictive models for different rare disease.

4.1. Adaptation to a Single Task

For a task task_i of a given disease and given initial parameters θ (either random or learned), we formulate the learning process Learn(task_i; θ) of model parameters θ_i as minimizing the loss function of model parameters on the data of the given task from the initialization θ.

$${\theta }_i=\text{Learn}(tas{k}_i;\theta )=\arg \underset{\theta }{\min }{L}_{tas{k}_i}(\theta )$$ (2)

We assume p(y|X, θ) is the mortality rate predicted by model with parameters θ. The loss function $L_{tas{k}_i}(\theta )$ of model parameters θ on task corpus task_i is defined using the cross entropy CE(·) between the model output p(y|X, θ) and the true label $\widehat{y}$ of the outcome:

$$L_{tas{k}_i}(\theta )={\displaystyle \sum _{(X,\widehat{y})\in tas{k}_i}\text{CE}(p(y|X,\theta ),\widehat{y})}$$ (3)

In this work, the mortality prediction model p(y|X, θ) based on heterogeneous temporal events in EHR data is mainly composed of attributed event embedding [4] and long short-term memory (LSTM). Firstly, clincial events with attributes are embedded into vectors via information fusion of their type, categorical attributes and numerical attributes. After the attributed event embedding module, the temporal dependencies encoded in the sequence of embedded vectors are then captured by long short-term memory (LSTM) [5], which outputs the prediction results for mortality with a sigmoid layer at the last LSTM cell. The mortality prediction model is illustrated in the middle of Figure 2.

Figure 2:

Training Mortality Prediction Model with Ada-Sit

4.2. Adaptation to Similar Tasks

The architecture of our proposed multi-task learning method Ada-SiT is represented in Figure 1, where Ada-Sit is compared with generic multi-task methods and generic meta-learning based multi-task methods. In the following, we first introduce the architecture of Ada-SiT, and then the task similarity measurement module, which is a key component of Ada-SiT. The overall framework of how to train mortality prediction model with Ada-Sit is illustrated in Figure 2.

Figure 1:

(A) Generic Architecture of Multi-Task Learning. The shared module (in blue), whose output will be taken as the input of specific modules, is shared between different tasks. (B) Generic Architecture of Meta-Learning based Multi-Task Learning. The parameter initialization θ in blue, which will be adapted to each specific task task_i , is shared between all tasks (C) Architecture of our Ada-SiT (Adaptation to Similar Task). The shared initialization θ will be adapted to similar tasks N(task_i) of each task task_i, resulting in a predictive model θ_i for each task.

4.2.1. Architecture of Ada-SiT

Fast adaptation, one of the meta-learning methods, is applied into Ada-SiT for the multi-task learning scenario. Ada-SiT differs from the original fast adaptation [10] which separately adapts to each task. Ada-SiT can dynamically measure task similarity and learn to adapt to multiple similar tasks.

The idea of fast adaptation is to learn a good parameter initialization for fast adaptation to new tasks. Specifically, in our scenario of multi-task learning, it means using patient data from all the disease tasks to find a good initial parameter initialization θ and mortality prediction model parameters θ_i of each disease which is adapted from the found initial parameters θ.

The initialization θ is learned by repeatedly simulating scenarios of mortality prediction of each disease with its similar diseases. We achieve this goal by defining the meta-objective function as:

$$\underset{\theta }{\min }{\displaystyle \sum _{tas{k}_i}{L}_N(tas{k}_i)}(\text{Learn}(N(tas{k}_i);\theta ))$$ (4)

where N(task_i) is the extended sample set composed of tasks similar with task_i. Details of N(task_i) will be described in the next subsection.

We maximize the meta-objective function using stochastic approximation with gradient descent. For each epoch, we find similar tasks N(task_i) for each task, and then independently sample two samples subsets (D_tr and D_val) from the training set of the task N(task_i). The former D_tr is used to simulate the learning process of mortality prediction models, and the latter D_val is used to evaluate the precision of the learned models for the updating of the shared initialization. Here, a single-step gradient descend is applied in the training simulation:

$${{\theta }^{\prime }}_i=\text{Learn}(N(tas{k}_i);\theta )=\theta -\alpha {\nabla }_{\theta }{L}_{N(tas{k}_i)}^{tr}(\theta )$$ (5)

where α is the learning rate and $L_{N(tas{k}_i)}^{tr}(\cdot )$ represents the loss function calculated from sample set D_tr from N(task_i).

Next, we evaluate the updated task parameters ${{\theta }^{\prime }}_i$ on D_val. The gradient computed from the evaluation, referred to as meta-gradient, is used to update the initial parameters θ. Gradients from a batch of tasks are aggregated to updating θ as follow:

$$\theta \leftarrow \theta -\beta {\nabla }_{\theta }{\displaystyle \sum _{tas{k}_i}{L}_{N(tas{k}_i)}^{val}}({{\theta }^{\prime }}_i)$$ (6)

where β is the meta learning rate, and $L_{N(tas{k}_i)}^{val}(\cdot )$ represents the loss function calculated from sample set D_val from N(task_i).

In the process of the calculation and approximation of the meta-gradient by the chain rule, the second-order derivative term can be ignored without much accuracy loss [10]. And the meta-gradient can be approximated as the following simplified gradient:

$${\nabla }_{\theta }{L}_{N(tas{k}_i)}^{val}({\theta }^{\prime })={\nabla }_{{\theta }^{\prime }}{L}_{N(tas{k}_i)}^{val}({\theta }^{\prime })\left(1-\alpha H_{\theta }\left(L_{N(tas{k}_i)}^{tr}\left(\theta \right)\right)\right)\approx {\nabla }_{{\theta }^{\prime }}{L}_{N(tas{k}_i)}^{val}({\theta }^{\prime })$$ (7)

where the term including $H_{\theta }\left(L_{N(tas{k}_i)}^{tr}\left(\theta \right)\right)$, the Hessian matrix, a square matrix of second-order partial derivatives of the loss function $L_{N(tas{k}_i)}^{tr}(\cdot )$ at θ, is ignored.

As our goal is multi-task learning via adaptation to similar tasks, model parameters θ_i of each task are adapted from the newly updated initialization parameters θ at the end of each iteration. These parameters θ_i of tasks are comparable in the model space, and will be used for calculating the task similarity in model space in the next iteration.

4.2.2. Task Similarity Measurement in Model Space

To measure the similarity of all tasks in terms of clinical behavior, we define the similarity of tasks in the model space, where each predictive model for its corresponding task is represented as a vector composed of all the parameters.

Formally, the similar tasks N(task_i) of task_i is defined as N_cos(·):

$$N_{cos}(tas{k}_i)=\left\{s\right|s\in tas{k}_j\&\cos ({\theta }_i-\theta ,{\theta }_j-\theta )>\eta \}$$ (8)

where θ_i and θ_j are the parameters of corresponding model of task_i and task_j, θ is the initial parameters. cos(θ_i − θ, θ_j − θ) reflects cos function of the angle between gradient directions of task_i and task_j, and η is threshold of the cos function.

The models of similar diseases also have similar gradient directions when they are adapted from the initialization. And the big cos (close to 1) value of included angle between two gradient directions indicates the tasks are similar.

Notice that a natural alternative way to get similar tasks is selecting k nearest neighbors in the model space. However, the absolute distance of models is more meaningful than relative distance, because the distance of models is generated by the gradient descent of the adaptation process from a common initialization. So selecting models in the neighborhood of a certain model as its similar models are more suitable in our Ada-SiT method, which is demonstrated by the experimental results in Section 5.4.

5. Results and Discussions

5.1. Comparing Methods and Experimental Settings

We compare Ada-SiT to both global single-task learning methods and multi-task learning methods for mortality prediction. The data size of each task is too small to train separate single-task models, so these baselines have not been included.

The global single-task learningmodels are trained by all the patients in the training set.

• LSTM LSTM [5] is used to learn the representation of the heterogeneous event sequence for each patient. Binary predictions for mortality based on the learned representations can be generated with a logistic regression layer.

• TCN The architect for prediction is the same as LSTM, except that the Temporal Convolutional Network (TCN) [6] is used to learn patient representation vectors from the heterogeneous event sequences instead of LSTM.

The multi-task learningmethods use Train_i to train the model for each task_i and get prediction results such as predicted label and probability on Test_i.

• Multi-SAnD Multi-SAnD [8] is a Transformer-based multi-task learning method which uses the weighted sum of loss functions on all tasks for the loss function.

• MMoE MMoE [29] is a Multi-gate Mixture-of-Experts model which shares the expert submodel across all tasks and has a gating network trained to optimize each task.

• Multi-Dense Multi-Dense is proposed by Suresh [9]. It has a shared LSTM layer for representation learning, followed by task-specific dense layers and logistic regression output layers.

The models in this section are implemented with Tensorflow [30] and trained with Adam. In Ada-SiT, we set N(task_i) = N_cos(task_i), where the cos function threshold η is 0.7, for the task similarity measurement. α and β are 0.0005 and 0.001 respectively.

5.2. Evaluation Metrics

The data for target prediction tasks are imbalanced-labeled. So metrics for binary labels such as accuracy are not suitable for measuring the performance. Similar to the work [20], we adopt AUC (the area under ROC curves (Receiver Operating Characteristic curves)) and AP (the area under PRC (Precision-Recall curves)) for evaluation. They both reflect the overall quality of predicted scores at each decision time.

5.3. Quantitative Results

Table 4 shows the AUC and AP of Ada-SiT, global single-task models, and multi-task models on MiniMIMIC and MiniEICU. From the results in Table 4, we draw the following conclusions.

Table 4:

performance of different models on MiniMIMIC and MiniEICU

Model Class	Model	MiniMIMIC		MiniEICU
AUC	AP	AUC	AP
Global Single-task	LSTM [5] TCN [6]	0.8162 (0.0026) 0.8008 (0.0024)	0.3830 (0.0055) 0.4120 (0.0011)	0.6642 (0.0227) 0.6107 (0.0055)	0.2692 (0.0193) 0.1945 (0.0052)
Multi-task	Multi-SAnD [8] MMoE [29] Multi-Dense [9]	0.8036 (0.0161) 0.7181 (0.0117) 0.8325 (0.0036)	0.2754 (0.0063) 0.2195 (0.0097) 0.3997 (0.0096)	0.6215 (0.0075) 0.6300 (0.0023) 0.6730 (0.0071)	0.1592 (0.0016) 0.1364 (0.0030) 0.1147 (0.0039)
Ours	Ada-SiT	0.8729 (0.0112)	0.4543 (0.0241)	0.6746 (0.0090)	0.2961 (0.0103)

First, Ada-SiT can significantly improve the performance of global single-task learning methods. On both datasets, Ada-SiT performs better than LSTM and TCN. For example, on MiniMIMIC, Ada-SiT improves AUC and AP by around 6.9% and 19.1% respectively compared to LSTM. We can conclude that Ada-SiT can capture specific characteristics of diverse tasks without being interfered by data conflicts.

Second, Ada-SiT outperforms the compared multi-task learning methods. For example, on MiniMIMIC, Ada-SiT improves AUC and AP of Multi-SAnD by 8.6% and 65.0% respectively. On MiniEICU, it improves the AUC of Multi-SAnD and MMoE by 8.5% and 7.1% respectively. It should be noted that most of the multi-task baselines do not perform better than the global single-task baselines. The possible cause is that task-specific parameters of multitask models cannot be trained well because of data insufficiency of each task. We can conclude that Ada-SiT has a more robust information-sharing mechanism among tasks on small-size dataset compared to the traditional multi-task learning baselines.

5.4. Ablation Experiments of Task Similarity Measurement

To evaluate the effect of task similarity measuring in Ada-SiT, we vary this module while remaining other parts of the model identified in this section. We implement the following variants of Ada-SiT:

• Ada-SiT− Ada-SiT− is Ada-SiT without similar task measurement (i.e. N(task_i) = task_i), nearly the same as MAML [10]

• Ada-SiT (Static) According to the work [31], many static features can be used to measure task similarity. In Ada-SiT (Static), we choose the mortality rate as the static feature in the clinical scenarios and use it to measure task similarity instead of the proposed similarity measurement.

• Ada-SiT (KNN) Ada-SiT (KNN) selects k nearest neighbors instead of neighbors within a certain distance for N_cos(task_i) while finding similar tasks.

Table ?? shows the results of the ablation experiments of task similarity measurement. We can draw the following conclusions. First, the information in similar tasks can improve the performance of fast adaptation. Ada-SiT (KNN) and Ada-SiT improve the AP of Ada-SiT−. Ada-SiT also improves the AUC of Ada-SiT−. Second, the task similarity measurement in model space outperforms the compared static measurement. Ada-SiT improves the AUC and AP of Ada-SiT (Static) by 3.0% and 9.7% respectively. It is because traditional task similarity measurements via static features only leverage the metadata of tasks, but our measurement in model space can find potential information from the mapping of samples and labels. Third, finding neighbors within a certain distance as N_cos(task_i) is the most suitable way to get similar tasks. For example, Ada-SiT improves the AUC of Ada-SiT (KNN) by around 5.6%. It is noteworthy that the AUC of Ada-SiT (KNN) is even lower than Ada-SiT−. It is possibly because some tasks in the k nearest neighbors of task_i may be far from task_i in model space and they interfere with the fast adaptation process of task_i.

5.5. Relationship between Task Similarity and Mortality Rate

There is strong correlation between task similarity and mortality rate. By treating each model’s parameter vector θ_i as a point in the model space, we use t-SNE [32] to visualize similar task clusters in Figure 3. In Figure 3, the two task clusters represent two types of rare diseases. The average mortality rate of diseases in the blue cluster is 0.6% and that in the yellow cluster is 32.1%. Meanwhile, the total average mortality rate of MiniMIMIC is 7%. We can see that the mortality rate is the main factor to determine task similarity. It suggests that our task similarity measurement module is reasonable and consistent with the clinical knowledge because diseases with a lower mortality rate and life-threatening diseases with a high mortality rate have different clinical behavior.

Figure 3:

Visualization of Tasks in the Model Space

6. Conclusion

In this paper, we propose a novel method Ada-SiT for learning predictive models for diverse tasks with insufficient data. Ada-SiT has a new task similarity measurement method and a new knowledge-sharing schema, where the shared initialization is learned to fast adapt to similar tasks. Experiment results show that our method is suitable for mortality prediction of diverse rare diseases, and can improve the performance compared to global single-task models and genetic multi-task models.

Acknowledge

This paper is supported by National Key Research and Development Program of China with Grant No. 2018AAA0101900 / 2018AAA0101902 as well as the National Natural Science Foundation of China (NSFC Grant No. 61772039 and No. 91646202).

Figures & Table

Table 5:

Ablation study of task similarity measurement

Methods	AUC	AP
Ada-SiT− (MAML)	0.8577 (0.0015)	0.3936 (0.0025)
Ada-SiT (Static)	0.8474 (0.0123)	0.4143 (0.0144)
Ada-SiT (KNN)	0.8264 (0.0110)	0.4059 (0.0112)
Ada-SiT	0.8729 (0.0112)	0.4543 (0.0241)