Automated Multi-Task Learning for Joint Disease Prediction on Electronic Health Records

Authors: suhan  cui
prasenjit  mitra
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-14 18:48:32

Abstract: In the realm of big data and digital healthcare, Electronic Health Records (EHR) have become a rich source of information with the potential to improve patient care and medical research. In recent years, machine learning models have proliferated for analyzing EHR data to predict patients' future health conditions. Among them, some studies advocate for multi-task learning (MTL) to jointly predict multiple target diseases for improving the prediction performance over single task learning. Nevertheless, current MTL frameworks for EHR data have significant limitations due to their heavy reliance on human experts to identify task groups for joint training and design model architectures. To reduce human intervention and improve the framework design, we propose an automated approach named AutoDP, which can search for the optimal configuration of task grouping and architectures simultaneously. To tackle the vast joint search space encompassing task combinations and architectures, we employ surrogate model-based optimization, enabling us to efficiently discover the optimal solution. Experimental results on real-world EHR data demonstrate the efficacy of the proposed AutoDP framework. It achieves significant performance improvements over both hand-crafted and automated state-of-the-art methods, also maintains a feasible search cost at the same time.

Main paper: 
 introduction: 
 In the era of big data and digital healthcare, the voluminous Electronic Health Records (EHR) can revolutionize patient care, medical research, and clinical decision-making. Using these, the machine learning (ML) community has been designing models to predict patients' future health conditions, e.g., models for mortality prediction [1] , diagnosis prediction [2, 3] and hospital readmission [4] . Although most existing machine learning based prediction models are designed to be single-task, i.e. predicting the risk of a single target disease, some works [5, 6, 7, 8, 9] designed multi-task learning (MTL) models to jointly predict multiple targets. The motivation lies in the fact that two or more diseases might be related to each other in terms of sharing common comorbidities, symptoms, risk factors, etc. Consequently, training on related diseases simultaneously can offer additional insights and potentially enhance prediction performance. While multi-task learning offers potential advantages, the existing MTL frameworks for EHR data still suffer from the following limitations.Limitations of the existing MTL frameworks for EHR data. To design an effective MTL framework, two fundamental challenges need to be addressed:(1) How can we determine which tasks should be trained together? The task grouping problem [10] involves finding groups of tasks to train jointly. Multi-task learning only provides advantages when the tasks are synergistic, i.e., training on the tasks together makes the model learn general knowledge that helps in performing the tasks better in the test set and prevents overfitting. Thus, given a large set of related tasks in a domain, we may need to group the tasks (allowing tasks to belong to multiple groups) together to create groups of tasks on each of which we will train a model. However, existing works usually rely on human expert discretion to select multiple tasks upfront and create a shared model for those tasks [5, 6, 7, 8, 9] . Hence, none of them has addressed the general problem of task grouping for EHR data. Moreover, due to the complexity of disease correlations, grouping synergistic tasks together is extremely challenging for human experts. It not only demands substantial effort (trying out every possible task combination) but also introduces the risk of task interference (putting disparate diseases together), potentially leading to performance degradation. Therefore, how to design the appropriate task grouping for MTL on EHR data presents a critical challenge.(2) How can we design model architectures for MTL? Existing works [5, 6, 7, 8, 9] typically rely on hand-crafted architectures for multi-task learning, which consist of a shared EHR encoder followed by several task-specific classifiers. However, due to the large number of possible operations as well as network topologies, manually tuning an optimal architecture for MTL is impossible. Furthermore, the optimal architectures for different task groups might also be distinct. Thus, things can even get worse when the number of tasks grows and different task combinations are involved for joint tuning. Therefore, we need a more efficient and effective approach to design the optimal MTL architectures for EHR data.Automating the MTL framework design for EHR data. To address the aforementioned challenges, we look to Automated Machine Learning (AutoML) [11] . Since AutoML relies on data-driven approaches to automate the design of machine learning algorithms, it has the potential to improve the design of an MTL framework for EHR data and reduce human interventions. Several attempts have been explored in other domains, e.g., computer vision, to improve the design of task grouping [12, 10, 13] and MTL architectures [14, 15, 16, 17, 18] . However, the exploration of AutoML in healthcare domain remains relatively limited [19] . To the best of our knowledge, there are no existing work that automates the finding of groups of tasks for MTL towards designing an optimal framework for classification tasks using EHR data, which is a notable gap in the field.Joint optimization over task grouping and architecture search. Morever, currently there exists no end-to-end optimization framework for automating MTL, even in other domains. Current approaches independently address the problems of task grouping and architecture design. First, a line of work [12, 10, 13] solves the task grouping problems by learning the task correlations. They operate under the underlying assumption that MTL architectures are the same across different task groups, which might not be practical nor optimal. Second, researchers also apply Neural Architecture Search (NAS) [20] to automatically design MTL architectures for a predefined set of tasks [14, 15, 16, 17, 18, 21] .No existing work has integrated these two approaches to address both problems simultaneously. However, combining them naively could lead to sub-optimal results, as sequential optimization might result in inaccurate estimations for both aspects. Therefore, we need a more generalized AutoML framework for the joint optimization of both task grouping and architecture search.Overview of the proposed approach. Therefore, in this paper, we show that an integrated approach for multi-task grouping and neural architecture search provides significant improvements. First, we extend existing single-task models like Retain [22] , Adacare [16] to MTL in an EHR setting. Second, we apply DARTS [23] , an NAS method used in MTL settings in different domains to the EHR domain. We use one shared model for predicting multiple tasks. These adaptations improve over the single-task setting. Second, we explore the impact of automated task grouping in the EHR setting by grouping tasks and finding an optimal NAS model for each task group. This further improves the performance. Finally, we propose an integrated framework an Automated multi-task learning framework, AutoDP, for joint Disease Prediction on electronic health records, which aims at jointly searching for the optimal task grouping and the corresponding neural architectures that maximize the multi-task performance gain. We show that this third method provides the maximal performance gain.Specifically, in AutoDP, we employ a surrogate model-based optimization approach [24] for efficient search. First, we define the joint search space of task combinations and architectures that includes all possible configurations for MTL. We want to find optimal solutions from this search space. To achieve that, the first question is how we can evaluate the performance of each configuration. Performing the ground truth evaluation for every configuration is infeasible, since it requires an entire multi-task learning procedure for each pair of architecture and task combination. Therefore, instead of exhaustively evaluating all the configurations, we build a surrogate model to predict the multi-task gains for any given configurations from the search space. In this way, we only need to evaluate the ground truth gains for a subset of samples from the search space, and use them to train the surrogate model for estimating the rest ones. The intuition is that there exists an underlying mapping from each configuration to the expected multi-task gains; thus it can be learned by a neural network. The remaining question is how we can effectively train the surrogate model using as few samples as possible. To this end, we further propose a progressive sampling strategy to guide the surrogate model training for improving sample efficiency. That is we train the surrogate model through multiple iterations. At every iteration, we select some points from the search space and update the surrogate model accordingly. The selection is conditioned on the current surrogate model and involves both exploitation and exploration. That is, we iteratively select the points that bring higher performance gains and also come from unexplored areas, which makes the training samples represent the whole search space. Eventually, after we obtain the trained surrogate model, we further use it to derive the final optimal task grouping and architectures. Because of the huge search space, it is not practical to use brute-force search. Hence, we develop a greedy search method to find the near-optimal solution. In summary, our contributions are as follows:• We are the first to propose an automated approach for multi-task learning on electronic health records AutoDP, which largely improves the design of task grouping and model architectures by reducing human interventions. Specifically, this work is the first to automate the design for the optimal task grouping and model architectures for MTL on EHR data.• We are the first to propose a surrogate model based optimization framework that jointly searches for the optimal task grouping and corresponding model architectures with high efficiency in any domain.• We propose a progressive sampling strategy to construct the training set for the surrogate model, which improves sample efficiency by reducing the required number of ground truth evaluations during searching. Importantly, we balance exploitation and exploration so that the sampled configurations can represent the whole search space and are highly accurate.• We propose a greedy search algorithm to derive the final MTL configuration using the trained surrogate model and find a near-optimal solution from the huge search space efficiently.• Experimental results on real world EHR data -MIMIC IV [25] demonstrate that AutoDP improves classification performance significantly over existing hand-crafted and automated methods under feasible computational costs.Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?Answer: [Yes] Justification: The introduction section presents all contributions and scope.Guidelines:• The answer NA means that the abstract and introduction do not include the claims made in the paper. • The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers. • The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings. • It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper. 
related work: 
 Multi-Task Learning with EHR. To enhance prediction performance while forecasting patients' health conditions based on their historical data [26] , existing studies employ multi-task learning to simultaneously predict multiple related target diseases or conditions, resulting in improved performance compared to single-task training. For example, Wang, et al. [7] investigated the advantages of joint disease prediction using traditional machine learning models. More recently, researchers have applied recurrent neural network (RNN) based models to conduct multi-task learning on EHR data [27, 6, 5] , which is able to predict tasks like mortality, length of stay, ICD-91 diagnoses and etc. Additionally, Zhao, et al. [8] also utilized a transformer based method for multi-task clinical risk prediction on multi-modal EHR data. However, all these studies manually select the set of tasks for joint training without task grouping and utilize a hand-crafted MTL model architecture, which largely limits their performance.Multi-Task Grouping. Due to the limitation of manually selected task groups, some of the work focus on obtaining the optimal task grouping through searching. Specifically, Standley, et al. [10] is the first work that systematically analyze the task correlations. For improving the efficiency, they use pair-wise MTL gains to estimate the high-order MTL gains, and obtain the pair-wise gains by training one model for each task pair. Based on the estimated gains, they derive the optimal task grouping using brute-force search. Fifty, et al. [12] further improves the efficiency by training one model to derive all the pair-wise gains. They derive the task affinity based on the gradient information during training. Furthermore, Song, et al. [13] propose a more general method that employs a meta model to learn the task correlations and estimates the high-order MTL gains more effectively. These works normally assume that the model architecture is the same across different task groups. But in practice, we can maximize performance gains by applying different model architectures with respect to each task group. Thus, we need a more general framework that considers the model architectures during task grouping.Multi-Task NAS. Neural Architecture Search (NAS) [20] stands as a prominent research area in AutoML, focusing on the exploration of optimal deep network architectures through a data-driven approach. Although the main stream of NAS focuses on the setting of single task learning, some researchers also try to employ NAS in multi-task learning applications, predominantly for searching computer vision MTL architectures. Notably, studies done by Ahn, et al. [14] and Bragman, et al. [15] employ reinforcement learning and variational inference, respectively, to determine whether each filter in convolutions should be shared across tasks. Furthermore, other recent works [16, 17, 18, 21] leverage differentiable search algorithms [23] , to determine the optimal sharing patterns across multiple network layers for diverse tasks. Despite the demonstrated advancements, a common limitation is their reliance on human experts to pre-define a set of tasks for joint training. This constraint poses challenges in practical scenarios where task grouping is not readily available, thereby limiting their broader applicability. What is more, their frameworks often search for better MTL architectures on top of one or several backbone architectures such as ResNet [28] . However, such backbone architectures might not be available for EHR applications in medical domain. Therefore, a new multi-task NAS framework is needed for EHR data.3 MethodologyQuestion: Does the paper discuss the limitations of the work performed by the authors?Answer: [Yes]Justification: We have discuss the limitation and future work in Section 5.Guidelines:• The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper. • The authors are encouraged to create a separate "Limitations" section in their paper.• The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be. • The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated. • The authors should reflect on the factors that influence the performance of the approach.For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon. • The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size. • If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness. • While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations. 
preliminaries: 
 Problem definition. Assume we have the input EHR data for multiple patients where each patient is represented as X ∈ R L×de , where L is the time sequence length and d e is the hidden dimension of the input features. We have N prediction tasks using the EHR data, denoted asT = {T 1 , T 2 , • • • , T N }.We seek to maximize the overall MTL performance gain for all these prediction tasks compared to single task training. First, we define MTL gain. Conduct a single task training on each task independently using a specific backbone model (such as RNN), and obtain the single-task performance for all tasks in terms of a predefined metric (such as average precision), denoted as{s 1 , s 2 , • • • , s N }.Then, the MTL gain is defined as:EQUATIONwhere m i is the multi-task performance for T i . Therefore, our objective is to maximize the overall gain for all tasks: G = 1 N N i=1 g i . To achieve that, our proposed method solves two searching problems at the same time using AutoML. First, we search for a list of task combinations that defines which tasks should be trained together. Second, we determine the optimal model architecture for each task combination. We aim at finding the optimal configuration for both, such that the highest overall gain G is attained.Task grouping search space. For N tasks, there are2 N -1 task combinations, C = {C 1 , C 2 , • • • , C 2 N -1 }, where every C is a subset of T . Given a budget B, we aim at searching for maximally B task combinations from C to determine which tasks should be trained together. The task combinations should cover all N tasks so that we are able to obtain {m 1 , m 2 , • • • , m N }. If one task T n appears in multiple task combinations, we simply choose the highest performance for it as m n .Architecture search space. For every task combination, we also need to search for an MTL architecture to model the EHR data. We adopt the hard sharing mechanism as in most existing works [27, 5] , which consists of a shared encoder for extracting the latent representation of the input EHR and multiple task specific classifiers to generate the output for every task. Specifically, we enable the search for the optimal shared encoder. For the search space of the encoder, we adopt the setting of directed acyclic graph (DAG) [23] . The architecture is represented as a DAG that consists of P ordered computation nodes, and each node is a latent feature that has connections to all previous nodes. For each connection (also called edge), we can choose one operation from a EQUATIONwhere node features E i ∈ R L×de 's all have the same dimension as X, and o (i,p) is the operation that transform E i to E p . Essentially, sampling one architecture from the search space is equivalent to sampling one operation for every edge in the DAG. In this way, we can get the set of all possible architectures denoted as A.Finally, to predict, we take the last node representation E P as the encoded feature for the input EHR, and use task-specific classifiers to output final predictions, which are all fixed fully connected network layers.MTL procedure. To evaluate a specific sample from the joint search space C × A, we need to conduct an MTL experiment to obtain the multi-task performances. Specifically, given an architecture A ∈ A and a task combination C ∈ C, we train the model A to predict for all tasks in C and get the multi-task performances for those tasks. Then, we can compute their gains by Eq.( 1). In this way, we are able to evaluate how much gains that this sample (C, A) could achieve.We propose a surrogate model based AutoML framework to search for the optimal task grouping and corresponding architectures simultaneously. To achieve that, we need to first evaluate the MTL gains for all the samples in the joint search space C × A, and then select the best B samples (pairs of task combinations and architectures) that maximize G. However, it is not practical to obtain the ground-truth gains for every sample, since the whole search space is normally very huge and every MTL procedure is also considerably expensive. Therefore, we build a neural network (called surrogate model) to learn the mapping from a pair of task combination and architecture to the multi-task gains:EQUATIONwhere g (C,A) ∈ R |C| is the per-task gains for task combination C ∈ C if using A as the model, and F (•) is the surrogate model. In this way, we only need to evaluate the ground truth gains for a small subset of samples from the search space, and use them to train the surrogate model for estimating all other unseen samples. The assumption is that the multi-task gains are essentially determined by the configuration of the task combination and the architecture, so there exists an underlying mapping that could be learned by a neural network. We set universal hyperparameters and optimization settings for all MTL procedures, hence the influence of other factors can be ignored.Specifically, we introduce the model architecture of the surrogate model in Section 3.3. Then, we outline the training procedure of the surrogate model in Section 3.4, where we propose an active learning strategy to collect training samples. Eventually, we use greedy search to derive the final configuration of task grouping and architectures by utilizing the trained surrogate model, as discussed in Section 3.5. The framework overview is shown in Figure 1 .For learning the mapping from an input configuration to the multi-task gains, the surrogate model is required to encode both architectures and task combinations. Also, the model needs to output multi-task gains. Therefore, we design a new surrogate model that consists of two encoders that respectively transform the input architecture and task combination into latent representations. Then, two representations are fused together to predict the multi-task gains.Architecture Encoding. For encoding a given architecture A, we apply a graph encoder [29] that is specifically designed for modeling DAGs, which is suitable for encoding the architectures in our search space. It can sequentially update the hidden states for the P computation nodes in preceding order by aggregating information from all predecessors. For node p, we have:EQUATIONwhere h 0 ∈ R ds is the input node representation which contains trainable parameters, and W ∈ R ds×ds 's are learnable transition matrices constructed for each operation in O. For every operation in the architecture, we also apply the corresponding W in our graph encoder. For aggregating all incoming representations, we apply average pooling to obtain the node representation h p . Finally, we use the node representation for the last node h P as the overall encoding for the input architecture.Task Combination Encoding. For encoding a given task combination C, we use the self attention mechanism [30] to model the high order interactions among the selected tasks in C. Specifically, we randomly initialize the embedding for all N tasks, and for task combination C, we have:EQUATIONwhere u ∈ R ds 's are corresponding embeddings for the selected tasks and z ∈ R ds is the final representation for task combination C. Additionally, we also use average pooling on top of the self attention layers to obatin z.Prediction. Eventually, we apply a two layer MLP to fuse both architecture encoding h P and task combination encoding z, and output the predicted gains for all selected tasks ĝ(C,A) ∈ R |C| . We use the mean absolute error to supervise the surrogate model as follows:EQUATIONwhere g (C,A) ∈ R |C| is the ground truth gains generated by conducting an MTL procedure for (C, A).In order to efficiently train the surrogate model defined in previous section, we develop a progressive sampling method to collect training samples. Start with an empty training set and a random initialized surrogate model, we progressively sample more points from the search space C × A, and then use them to train the surrogate model. Specifically, we include two stages for the surrogate model training:Warm start. Firstly, we warmup the surrogate model by selecting a small number of samples from the search space. Specifically, we use the task combination that contains allN tasks C 0 = {T 1 , • • • , T n }and randomly sample Q 0 architectures from A. Then we conduct Q 0 MTL procedures to evaluate their gains by training on C 0 . In this way, we collect Q 0 training samples as the initial training set denoted as D. We further train the surrogate model on D, and denote the model parameters as Θ.Progressive selection. Then, we progressively select more points and train the surrogate model as introduced in Algorithm 1. Totally, we conduct K 1 rounds of sampling. For each round, we iterate through all N tasks. With respect to one task T n , we build the acquisition function Γ over the set of task combinations that contains T n based on the predicted gains for T n . Then, we select one task combination C * that have highest value. We apply Upper Confidence Bound [31] as the acquisition function that considers both exploration and exploitation by explicitly estimating the mean and variance of predicted gains (line 11 marked by blue). Besides that, we would also like to see the effect of exploration vs exploitation. so we try out different settings of Γ. Specifically, we propose three variants of AutoDP, namely AutoDP µ+σ , AutoDP µ and AutoDP σ , which corresponds to the original setting, including only µ or only σ in Γ. In this way, we can compare the results with pure exploration and pure exploitation during sampling, and find out the optimal strategy for AutoDP.Algorithm 1: Progressive Selection Input: Training set D, surrogate model parameter Θ, Q 1 , Q 2 , K 1 ; Q 1 > Q 2 . Output: Updated D and Θ for k = 1, 2, • • • , K 1 do for n = 1, 2, • • • , N doCollect all task combinations that contains T n :C Tn = {C j |∀C j ∈ C, T n ∈ C j };for ∀C j ∈ C Tn do 5Randomly sample Q 1 architectures from A, denote the set as A Cj ;6Forward the surrogate model to collect gains for T n with every architecture in A Cj : G = {g[T n ]|∀A ∈ A Cj , g = F (C j ,We derive the final results using the trained surrogate model. Due to the huge search space, it is still not practical to use brute force search to get the global optimum. Therefore, we propose to apply a greedy method to search for near-optimal solutions. We introduce the detailed procedure in Algorithm 2. The high level idea is that we first randomly initialize the configuration, and then gradually improve its multi-task gain by random mutation and greedy selection. Specifically, given the budget B, we aim at searching for B samples from the search space C × A such that the overall gain G is maximized. We first randomly initialize the population P that contains B pairs of task combinations and architectures. Then, at every iteration, we randomly mutate one pair (C, A) from the population and see whether the overall multi-task gain will increase. If so, we Estimate P ′ and P using F (•);Choose the better one: P ← Select(P ′ , P) ; end update P accordingly. After K 2 iterations, we can obtain a near-optimal solution. In practice, we also apply multiple initial populations to avoid getting stuck on local optima. Although we only get an approximate solution, our method can already achieve significant improvements over baselines as shown in Section 4.2.Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?Answer: [NA]Answer: [Yes] Justification: We have uploaded the code in supplementary materials, and we use public data.Guidelines:• The answer NA means that paper does not include experiments requiring code. • Providing as much information as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted. 
set up: 
 Dataset & Tasks. We adopt MIMIC -IV dataset [25] for our experiments, which is a publicly available database sourced from the electronic health record of the Beth Israel Deaconess Medical Center. Specifically, we extract the clinical time series data for the 56,908 ICU stays from the database as our input EHR data, with an average sequence length of 72.9. With respect to each ICU stay, we also extract 25 prediction tasks (listed in Table 3 ), including chronic, mixed, and acute care conditions. Each condition is associated with a binary label indicating whether the patient has the corresponding condition during the ICU stay.Baselines. To compare the proposed method with existing work, we choose several state-of-artbaselines, including both hand-crafted and automated methods. Specifically, as described below, we include several human-designed EHR encoders to compare with the searched architecture we defined in Eq. ( 2). Also, we include one NAS method and one multi-task grouping method as the automated baselines. More importantly, we combine the multi-task grouping method with the NAS method and hand-crafted encoders to show the superiority of our joint optimization method.• EHR encoders: We choose four models that are widely utilized for analyzing EHR time series, including LSTM [32] , Transformer [30] , Retain [22] and Adacare [1] .• NAS: We choose DARTS [23] as the NAS baseline, which is a differentiable search method for efficient architecture search. We apply it to our search space A to find better EHR encoders. Several state-of-the-art works in other domains have also used it to find MTL architectures [16, 17, 18, 21] .• Multi-task grouping: MTG-Net [13] is the current state-of-the-art multi-task grouping algorithm, which uses a meta learning approach to learn the high-order relationships among different tasks. We refer to this method as MTG in latter sections. Evaluation Metric. We use two widely used metrics for binary classification to evaluate our method and baselines: ROC (Area Under the Receiver Operating Characteristic curve) and AVP (Averaged Precision). During surrogate model training, we use AVP as the metric to compute multi-task gains as in Eq. ( 1), since it is a more suitable choice for considering the class imbalance.Please also refer to Appendix A for the implementation details.We show our results in Table 1 . Each experiment is run five times and the average of the runs are reported. We run three settings: Task @ 5, Task @ 10 and Task @ 25, which refers to using the first 5 tasks, 10 tasks and 25 tasks respectively. Since grouping all 25 tasks takes a long time to run, we include two small settings that only have the first 5 or 10 tasks in Table 3 for grouping. Our results demonstrate our hypotheses: (a) applying Retain, Adacare, and DARTS improves over the single-task setting, (b) applying different NAS models for each group further improves the performance, and finally, (c) AutoDP provides the best results in terms of averaged per-task gain for ROC and AVP, a significant improvement over existing MTL frameworks for EHR data.First, without considering task grouping, we train one shared model to predict for all tasks in three settings and compute the multi-task gains for them. Results show that this setting only provides minimal improvement over single task training. Note that the automated method (DARTS) performs better than other hand-crafted methods. We also see that sequential optimization over task grouping and architecture search (MTG+DARTS) performs better than MTG + other hand-crafted encoders.Moreover, we see that the three variants of AutoDP performed better than the other methods. Among them, AutoDP µ+σ performs the best, which means the balance of exploration and exploitation is the most effective strategy for training the surrogate model. For the last method, we also report the standard deviations and p-values of statistical tests (compared to MTG+DARTS), which justifies that Beside the overall performance gain, we also look at the distribution of performance gains for each individual task as shown in Figure 2 . We can observe that the proposed method does not have the issue of negative transfer, since all tasks have a positive gain. Also, for some of the tasks, it can achieve over 20% improvement, which further shows the effectiveness of We analyze the effect of two vital hyperparameters of our method: K 1 and B, since they are the crucial parameters that largely define the complexity of our method during searching and inference respectively. We choose the setting of Task @ 25 for a comprehensive analysis of all tasks. We try out different values and report the corresponding performance gain (AVP) in Figure 3 .First, K 1 determines the number of training samples collected during searching. Given that each sample invokes an MTL procedure, it constitutes the major portion of the search cost. Therefore, our goal is to find an optimal value for K 1 , striking a balance between cost-effectiveness and achieving commendable performance. We notice that the performance change plateaus after K 1 reaches 25.That is, the surrogate model effectively learns the distribution of the search space after consuming 25 × 25 training samples during active selection (25 samples per round). Consequently, we can empirically decide to halt the iteration at this point.Second, B determines the number of task groups for the final configuration, which indicates the number of MTL models required for achieving the expected performance gain after searching. We also observe similar phenomenon that the performance becomes stable after B reaches 12. We could also choose the optimal value for B accordingly. We further analyze the effect of several components within AutoDP, including progressive sampling (Section 3.4), greedy search (Section 3.5), and task grouping as a whole. We replace these components with naive or human intuitioninspired baselines and report the performances in Table 2 . Removing any of the components from the original framework leads to noticeable performance decreases, demonstrating the effectiveness of the designed components.Specifically, we replace progressive sampling and greedy search with purely random methods, referred to as Random Sampling and Random Search. In all three settings, performance generally decreases, highlighting the contributions of these components of AutoDP.Additionally, we use disease-based grouping (Appendix B) to first assign tasks into different groups based on their medical relevance and then employ DARTS to search for the model architecture for each group. This allows us to analyze the effectiveness of automated task grouping. By comparing disease-based grouping with the searched configurations (Appendix C), we observe that AutoDP does not strictly follow medical classifications for task grouping but achieves significant performance improvements over disease-based grouping. This indicates the necessity of using an automated search algorithm to find the optimal task grouping, which surpasses human intuition. 
conclusions and future work: 
 In this paper, we propose AutoDP, an automated multi-task learning framework for joint disease prediction on EHR data. Compared to existing work, our method largely improves the design of task grouping and model architectures by reducing human interventions. Experimental results on real-world EHR data demonstrate that the proposed framework achieves significant improvement over existing state-of-the-art methods, while maintaining a feasible search cost. There are also some valuable future directions based on the current version of AutoDP.First, from the application perspective, if we aim at deploying AutoDP to real-world healthcare systems, it would be advantageous to apply it to more complex problem settings. For example, the incorporation of diverse clinical data sources beyond EHR such as claims, drugs, medical images and texts will significantly enhance the practical utility of AutoDP.Additionally, considering the dynamic nature of healthcare environments with continuously updated input data and evolving tasks, adapting the surrogate model to accommodate new data and tasks would be imperative.Moreover, addressing privacy concerns within healthcare systems is a promising direction. Therefore, extending AutoDP with data processing pipelines for automatic feature engineering could offer enhanced privacy safeguards and further improve its applicability in sensitive healthcare contexts.Finally, we assume all the tasks have the same input EHR data in our problem setting, which might not always be the case in practical scenarios. Chances are that, for some diseases, there are large and well-annotated data, while for the others, there are limited data available. How we should extend the current framework to handle more heterogeneous diseases/tasks remains a challenge. 
experimental setting/details: 
 Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?Answer: [Yes]Justification: We have introduced all implementation details.Guidelines:• The answer NA means that the paper does not include experiments.• The experimental setting should be presented in the core of the paper to a level of detail that is necessary to appreciate the results and make sense of them. • The full details can be provided either with the code, in appendix, or as supplemental material. 
experiment statistical significance: 
 Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?Answer: [Yes]Justification: We include error bars and statistical test for the proposed method.Guidelines:• The answer NA means that the paper does not include experiments.• The authors should answer "Yes" if the results are accompanied by error bars, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper. • The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions). • The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.) • The assumptions made should be given (e.g., Normally distributed errors).• It should be clear whether the error bar is the standard deviation or the standard error of the mean. • It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have a 96% CI, if the hypothesis of Normality of errors is not verified. • For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates). • If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text. 
experiments compute resources: 
 Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?Answer: [Yes]Justification: We provide full information for the computing resources.Guidelines:• The answer NA means that the paper does not include experiments.• The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage. • The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute. • The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper). 
code of ethics: 
 Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines? 
broader impacts: 
 Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed? 
safeguards: 
 Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)? 
 
 Appendix: We include the following sections in Appendix, providing the additional details of the proposed framework, additional experimental results, and a discussion about limitation and future work.To prepare our dataset, we adopt the data pre-processing pipeline outlined in Harutyunyan, et al. [27] . Given that the original implementation2 is designed for MIMIC-III [33] , we make specific modifications to tailor it for MIMIC-IV. The 25 labels are defined using the Clinical Classifications Software (CCS) for ICD-9 code 3 . Consequently, we first map the ICD-10 codes4 in the MIMIC-IV database to ICD-9 codes before generating the labels. After processing, we have the feature dimension d e as 76. We partition the dataset as train, validation and test sets with a ratio of 0.7 : 0.15 : 0.15.We implement the framework using the PyTorch framework and run it on an NVIDIA A100 GPU.Given the dataset we have, we first train a vanilla LSTM for every task independently, and report the backbone performance in Table 3 , which can be further used to compute multi-task gains. For the proposed method, we run three settings of experiments: Task @ 5, Task @ 10 and Task @ 25, which refers to using the first 5 tasks, 10 tasks and 25 tasks respectively. For different settings, we use specific hyperparameters as shown in Table 4 . Besides that, we define the candidate operation set O as {Identity, Zero, FFN, RNN , Attention}, which includes widely used operations for processing EHR time series. Among them, Identity means maintaining the output identical to the input. Zero means setting all the values of the input feature to 0. Attention and FFN represents one self-attention layer and one feed-forward layer respectively, which are the same as in Transformer [30] . RNN is one recurrent layer, and we adopt LSTM [32] in our framework. For all the MTL procedures and baseline training, we apply the batch size of 64 and learning rate of 3e -4. For training the surrogate model, we use the batch size of 5 and learning rate of 5e -5. During searching, we compute all multi-task gains on the validation set for guiding the surrogate model training. After we obtain the optimal configuration, we train the searched models and report their multi-task gains on the test set.To show the effectiveness of automated task grouping, we conduct experiments using a predefined task grouping based on disease categories. We asked GPT-4 [34] to classify the 25 prediction tasks into different groups based on their medical meaning. The result is shown in Table 5 . Using this grouping, we further apply DARTS to each group and report the multi-task gains as shown in Table 2 . Compared to AutoDP, there is a notable performance drop for the disease based grouping. This means human intuition dose not provide the optimal task grouping, which underscores the necessity of employing search algorithm to automatically discover better task grouping for MTL.Here, we show two example of the final configuration for setting Task @ 10 in Figure 4 and for Task @ 25 in Figure 5 . The proposed AutoDP identifies 5 and 10 different task groups respectively and also searches for the corresponding architectures. We can observe that some of the tasks tend to be trained independently, while others are grouped together for joint training. This supports our claim that fine-grained task grouping is necessary to bring the optimal performance gain. Also, the optimal architecture is also different for each task group, which further justifies the necessity of joint optimization over task grouping and architecture search. NeurIPS Paper ChecklistJustification: We have read the code of ethics and make sure the paper conform with it.Guidelines:• The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.• If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics. • The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction).Justification: There are negligible societal impacts of the work.Guidelines:• The answer NA means that there is no societal impact of the work performed.• If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact. • Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.• The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster. • The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology. • If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).Justification: There are no such risks.Guidelines:• The answer NA means that the paper poses no such risks.• Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters. • Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images. • We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.12. Licenses for existing assets Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?Answer: [Yes] Justification: We have followed the standard procedure for accessing MIMIC-IV datset.Guidelines:• The answer NA means that the paper does not use existing assets.• The authors should cite the original paper that produced the code package or dataset.• The authors should state which version of the asset is used and, if possible, include a URL. • The name of the license (e.g., CC-BY 4.0) should be included for each asset.• For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided. • If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset. • For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.38th Conference on Neural Information Processing Systems (NeurIPS 2024).https://www.cdc.gov/nchs/icd/icd9.htmhttps://github.com/YerevaNN/mimic3-benchmarkshttps://www.cdc.gov/nchs/icd/icd9.htmhttps://www.cms.gov/medicare/coding-billing/icd-10-codes/ 2018-icd-10-cm-gem 


Decision: The paper under review introduces AutoDP, an automated multi-task learning (MTL) framework aimed at enhancing disease prediction using electronic health record (EHR) data. The paper introduces a novel approach to automating task grouping and model architecture design in MTL, addressing significant challenges in the field.

The primary strengths of this work lie in its innovative automation of MTL framework design and the use of surrogate model-based optimization for efficient search space exploration. The experimental results on the MIMIC-IV dataset demonstrate substantial performance improvements over both hand-crafted and existing automated methods, validating the effectiveness of AutoDP.

While some reviewers expressed concerns about interpretability and the use of a single MIMICT-IV dataset, I believe these issues do not significantly detract from the paper's contribution. The MIMIC-IV dataset is widely recognized as a benchmark in the EHR domain, and its use is sufficient to demonstrate the method's effectiveness. The authors have provided case studies in the appendix that offer insights into the searched configurations, which addresses interpretability to a reasonable extent.

The authors have shown willingness to address reviewer concerns in the camera-ready version, including improving Figure 1 for better clarity and adding a discussion on data imbalance issues. They have also explained their rationale for focusing on the MIMIC-IV dataset and provided details on their architecture search space and progressive sampling strategy.

Given the innovative nature of the work and its potential impact on the field of healthcare machine learning, I recommend accepting this paper. The authors have made a significant contribution by introducing an automated approach that reduces human intervention in MTL framework design for EHR data analysis. While there is always room for further evaluation and interpretation, the current work presents a solid foundation for future research in this area.

Reviews: review: paper_summary: paper_summary: The research work presents an exploration into the potential of multi-task learning (MTL) within the context of electronic health record (EHR) data analysis and clinical prediction tasks. This study innovatively addresses the critical challenges of task grouping and model architecture design, which are essential for optimizing MTL frameworks.
strength_weakness: strength_weakness: Q1. Figure Clarity and Comprehensibility: The overview of the proposed AutoDP in Figure 1 lacks self-evidency, making it challenging for readers, particularly those who are not specialists in electronic health record (EHR) modeling and analysis, to grasp the main insights at a glance. The figure is dense with symbols and elements that are not immediately interpretable without additional context. It is recommended that the authors enhance the figure's presentation to improve readability and comprehension. This could involve breaking down complex elements, using clearer labeling, and providing a more detailed legend or accompanying text that guides the reader through the visualization.

Q2. Evaluation Metric Consideration: For many clinical prediction tasks on datasets like MIMIC, where label balance may not be achievable, the authors might consider the use of AUPRC (Area Under the Precision-Recall Curve) as an alternative or complementary evaluation metric to AUC-ROC (Area Under the Receiver Operating Characteristic Curve). AUPRC can be a more informative measure when dealing with imbalanced classes, as it evaluates the model's ability to rank samples correctly, independent of the classification threshold. Precision-recall metrics are less sensitive to the choice of threshold compared to accuracy or precision alone, which is a critical consideration in clinical settings where the cost of false positives and negatives can vary significantly.

Q3. Generalizability and Extensibility: The current experiments are limited to the MIMIC dataset, and the baseline models have been fine-tuned for this specific dataset in the level of model-design. While the results on MIMIC are promising, readers may be interested in the general applicability of the proposed method to other datasets and clinical scenarios. It would be beneficial for the authors to provide experiments on additional datasets to demonstrate the robustness and generalizability of their approach. Furthermore, an exploration of how the method performs when integrated with different models or when adapted to various clinical prediction tasks would strengthen the paper's contribution and credibility. Given the focus on EHR data, it would be insightful if the authors could discuss the transferability of their model to other EHR datasets, potentially from different geographical regions or healthcare systems. This discussion could include challenges related to data heterogeneity, variations in clinical practices, and how these factors might affect the performance of the proposed AutoDP framework.
questions: questions: -
limitations: limitations: -
paper_summary: paper_summary: The paper discusses an automated approach for multi-task learning on electronic health records, called AutoDP. This approach aims to improve the design of task grouping and model architectures by reducing human intervention. Specifically, AutoDP searches for the optimal configuration of task grouping and architectures simultaneously. The document mentions that existing MTL frameworks for EHR data suffer from limitations such as requiring human experts to select tasks and design model architectures. AutoDP proposes a surrogate model-based optimization approach to address these limitations. The document also details related work in multi-task learning with EHR data, multi-task grouping, and multi-task neural architecture search
strength_weakness: strength_weakness: The paper could benefit from a more detailed explanation of the architecture search space, particularly the types of operations allowed in the directed acyclic graph (DAG).
Additional details regarding the progressive sampling strategy for the surrogate model would be helpful to understand how it balances exploration and exploitation during search.
The paper mentions experiments on real-world EHR data but lacks specific results or comparisons with other MTL approaches to demonstrate the effectiveness of AutoDP.
questions: questions: Can the authors elaborate on the specific operations allowed in the DAG-based architecture search space? Are there any constraints on the number or types of operations that can be included?
The paper mentions a progressive sampling strategy for the surrogate model. Can the authors provide more details on how this strategy works, particularly how it balances selecting informative data points for the surrogate model and exploring new areas of the search space?
While the paper mentions using AutoDP on real-world EHR data, it would be beneficial to see concrete results and comparisons with other MTL approaches to assess the performance gains achieved by AutoDP.
The use of EHR data raises ethical concerns about patient privacy. The authors should ensure proper anonymization techniques are used throughout the development and application of AutoDP.
limitations: limitations: Limited Evaluation: The experimental validation of AutoDP is confined to a single dataset (MIMIC-IV). Evaluating the framework on additional EHR datasets would bolster the generalizability of its findings.
Lack of Interpretability: The paper could provide deeper insights into the specific task groupings and architectures discovered by AutoDP. Visualizations or case studies could help elucidate the underlying reasons for the observed performance improvements.
Privacy Concerns: The paper does not explicitly address privacy concerns related to EHR data. Extending AutoDP with data processing pipelines for automatic feature engineering could offer enhanced privacy safeguards.
Other concerns: This paper did not discuss the data imbalance issues as some diseases might have large and well-annotated data, while others might have small and many missing data. For disease prediction, there are similarities among some diseases, and also training on more disease can help to predict other similar diseases.
paper_summary: paper_summary: N/A
strength_weakness: strength_weakness: The authors' identity can be easily inferred by googling the linux username ("sxc6192") found in the .idea/deployment.xml file in the supplementary material.
questions: questions: N/A
limitations: limitations: N/A
paper_summary: paper_summary: The paper proposes an automated multi-task learning framework, AutoDP, for disease prediction using EHR data. It optimizes task grouping and model architecture to enhance prediction performance. AutoDP efficiently searches a vast space of task combinations and architectures by employing a surrogate model-based optimization approach. Experiments show it outperforms existing methods.
strength_weakness: strength_weakness: 1. Generalization Concern: The paper shows AutoDP's effectiveness beyond disease prediction on EHR data. However, its effectiveness and applicability in other domains remain uncertain.
2. Computational Scalability: The proposed method involves a vast search space and significant computational costs. While feasible within the scope of the study, the proposed approach may hinder scalability in larger real-world scenarios.
questions: questions: 1. Mapping Assumption Validity: The paper assumes a learnable mapping for multi-task gains, which needs more explanation.
2. Comments on Disease-Based Grouping: As mentioned in Appendix D, Grouping similar diseases might make it hard to discriminate between them, reducing MTL effectiveness.
3. Reporting individual task gains is necessary to show the effectiveness of the proposed method on MTL.
4. Are those gains robust? Will the proposed method result in overfitting?
limitations: limitations: 1. The paper lacks a discussion on model interpretability, which is a critical aspect for trust and clinical decision-making in healthcare.
paper_summary: paper_summary: The paper introduces AutoDP, a novel AutoML framework that performs combined task grouping and architecture search with a novel surrogate model. The proposed model achieved higher performance in multiple tasks on MIMIC VI dataset compared to common single-task, multi-task, and AutoML approaches.
strength_weakness: strength_weakness: - It claims "feasible search cost", but there is no runtime comparison or evaluation
questions: questions: - How much longer does it take to run compared to DARTS and other baselines?
limitations: limitations: The traditional limitation in AutoML work is the computational budget. The paper does not discuss the runtime of the proposed model and alternatives.
Ethics limitations are discussed in the supplementary material.

Comments: title: Reply to authors
comment: Thanks for your rebuttal. After carefully reading your comments, i still want to maintain my original rating. Using only one public available EHR data is limited for the generalization purpose.
Dear Reviewer azWt, 
I am a NeurIPS 2024 Area Chair of the paper that you reviewed.

This is a reminder that authors left rebuttals for your review. 
We need your follow up answers on that. Please leave comment for any un-answered questions you had, or how you think about the author's rebuttal. 
The author-reviewer discussion is closed on Aug 13 11:59pm AoE.

Best regards, AC
Dear Reviewer Zp9P,

I am a NeurIPS 2024 Area Chair of the paper that you reviewed.

This is a reminder that authors left rebuttals for your review. We need your follow up answers on that. Please leave comment for any un-answered questions you had, or how you think about the author's rebuttal. The author-reviewer discussion is closed on Aug 13 11:59pm AoE.

Best regards, AC
Dear Reviewer S7e4,

I am a NeurIPS 2024 Area Chair of the paper that you reviewed.

This is a reminder that authors left rebuttals for your review. We need your follow up answers on that. Please leave comment for any un-answered questions you had, or how you think about the author's rebuttal. The author-reviewer discussion is closed on Aug 13 11:59pm AoE.

Best regards, AC
Thank you for your response and the additional information provided. After carefully considering the paper and the rebuttal content, I have decided to maintain my initial evaluation, based on a comprehensive assessment of the work's contribution to the field.

Additional clarity to my point Q2: I want to suggest that in the disease prediction domain, it is less likely to provide many benefits (in terms of performance) to group similar diseases into one learning task, as typically differential diagnosis is even harder. Therefore, claiming the proposed method performs better than such a method is less attractive.

Thank you for your efforts and good luck.