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

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. Source code can be found via the link: \url{https://github.com/SH-Src/AutoDP}.

Paper

References (38)

Scroll for more · 26 remaining

Similar papers

Peer review

Reviewer azWt8/10 · confidence 4/52024-06-18

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.

Strengths

- Novel AutoML approach for EHR tasks - Superior performance when compared to the traditional approach - Evaluated with the adequate dataset - Good range of baselines - Ablation study included - Open access data and code - Clear methodology presentation - Results contribute to AutoML and EHR fields

Weaknesses

- It claims "feasible search cost", but there is no runtime comparison or evaluation

Questions

- How much longer does it take to run compared to DARTS and other baselines?

Rating

8

Confidence

4

Soundness

4

Presentation

3

Contribution

3

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.

Area Chair yLee2024-08-10

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

Reviewer WEv35/10 · confidence 4/52024-07-11

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.

Strengths

1. The proposed AutoDP automates the search for optimal task groupings and model architectures simultaneously, reducing the reliance on human experts and improving the efficiency of designing multi-task learning frameworks. 2. The experimental results demonstrate that AutoDP outperforms existing hand-crafted and automated methods on real-world EHR datasets, achieving higher multi-task learning gains.

Weaknesses

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

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?

Rating

5

Confidence

4

Soundness

3

Presentation

3

Contribution

3

Limitations

1. The paper lacks a discussion on model interpretability, which is a critical aspect for trust and clinical decision-making in healthcare.

Reviewer 23Pe1/10 · confidence 5/52024-07-13

Summary

N/A

Strengths

N/A

Weaknesses

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

N/A

Rating

1

Confidence

5

Soundness

1

Presentation

1

Contribution

1

Limitations

N/A

Reviewer Zp9P4/10 · confidence 3/52024-07-13

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

Strengths

The paper introduces a novel approach, AutoDP, that addresses limitations in current MTL frameworks for EHR data by automating task grouping and model architecture design. The idea of using a surrogate model-based optimization for joint search is innovative and holds promise for efficient exploration of the search space. The paper is well-written and the overall presentation is clear.

Weaknesses

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

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.

Rating

4

Confidence

3

Soundness

2

Presentation

2

Contribution

3

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.

Area Chair yLee2024-08-10

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

Reviewer S7e45/10 · confidence 2/52024-08-04

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.

Strengths

S1. By proposing an automated approach, AutoDP, the paper contributes to reducing human intervention in the configuration process, thereby streamlining the identification of optimal task combinations and neural architectures tailored for EHR data. S2. The employment of surrogate model-based optimization and a progressive sampling strategy demonstrates a novel and efficient methodology for navigating the vast search space, leading to improved performance with a feasible computational cost. S3. The experimental results on the real-world EHR dataset, MIMIC-IV, validate the efficacy of the proposed AutoDP framework, showcasing substantial performance improvements over both hand-crafted and existing automated methods. S4. This work not only advances the state-of-the-art in automated machine learning for healthcare but also provides valuable insights and tools for the broader research community working on multi-task learning problems.

Weaknesses

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

-

Rating

5

Confidence

2

Soundness

3

Presentation

2

Contribution

3

Limitations

-

Area Chair yLee2024-08-10

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

Reviewer WEv32024-08-09

Reply to authors

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.

Reviewer Zp9P2024-08-12

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.

Program Chairsdecision2024-09-25

Decision

Accept (poster)

© 2026 NYSGPT2525 LLC