Research map / Architectures and efficiency / Neural architecture search
Automated Multi-Task Learning for Joint Disease Prediction on Electronic Health Records
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes AutoDP, a surrogate-model-based AutoML framework that jointly searches task groupings and DAG-based encoder architectures for multi-task disease prediction on MIMIC-IV EHR data.
Problem
Multi-task learning on EHR data relies on human-chosen task groups and hand-crafted shared architectures, and sequential or separate task-grouping and NAS approaches are sub-optimal.
Approach
Jointly search task combinations and encoder architectures using a surrogate model that predicts multi-task gains, trained via progressive sampling with an upper-confidence-bound acquisition, and a greedy mutation search to pick final groups and architectures.
Key results
- AutoDP (mu+sigma variant) reaches +13.43% AVP gain over single-task LSTM on Task @ 25, versus +11.76% for MTG+DARTS.
- On Task @ 5, AutoDP achieves +1.69% ROC and +7.74% AVP gain, with p-values below 0.05 against MTG+DARTS.
- Disease-based grouping (human intuition) yields only +8.61% AVP on Task @ 25 versus +13.43% for AutoDP.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use joint task-grouping and architecture search (AutoDP) | EHR multi-task disease prediction · MIMIC-IV Tasks @ 5 | ROC gain (%): +1.47 -> +1.69 | Helped |
| Use joint task-grouping and architecture search (AutoDP) | EHR multi-task disease prediction · MIMIC-IV Tasks @ 5 | AVP gain (%): +6.41 -> +7.74 | Helped |
| Use joint task-grouping and architecture search (AutoDP) | EHR multi-task disease prediction · MIMIC-IV Tasks @ 10 | ROC gain (%): +2.02 -> +2.55 | Helped |
| Use joint task-grouping and architecture search (AutoDP) | EHR multi-task disease prediction · MIMIC-IV Tasks @ 10 | AVP gain (%): +6.87 -> +8.81 | Helped |
| Use joint task-grouping and architecture search (AutoDP) | EHR multi-task disease prediction · MIMIC-IV Tasks @ 25 | ROC gain (%): +2.41 -> +2.80 | Helped |
| Use joint task-grouping and architecture search (AutoDP) | EHR multi-task disease prediction · MIMIC-IV Tasks @ 25 | AVP gain (%): +11.76 -> +13.43 | Helped |
| Use only μ instead of μ+σ in the acquisition function (ablation) | AutoDP · MIMIC-IV Tasks @ 5, Tasks @ 10, and Tasks @ 25 | ROC and AVP gain (%): ROC: +1.69 -> +1.49, +2.55 -> +2.08, +2.80 -> +2.68; AVP: +7.74 -> +7.12, +8.81 -> +7.53, +13.43 -> +12.70 | Mixed |
| Use only σ instead of μ+σ in the acquisition function (ablation) | AutoDP · MIMIC-IV Tasks @ 5, Tasks @ 10, and Tasks @ 25 | ROC and AVP gain (%): ROC: +1.69 -> +1.95, +2.55 -> +2.49, +2.80 -> +2.62; AVP: +7.74 -> +7.68, +8.81 -> +8.45, +13.43 -> +13.37 | Mixed |
| Add progressive sampling rather than random sampling (ablation) | AutoDP · MIMIC-IV Tasks @ 5, Tasks @ 10, and Tasks @ 25 | AVP gain (%): +6.75 -> +7.74; +7.04 -> +8.81; +11.30 -> +13.43 | Helped |
| Add greedy search rather than random search (ablation) | AutoDP · MIMIC-IV Tasks @ 5, Tasks @ 10, and Tasks @ 25 | AVP gain (%): +6.89 -> +7.74; +7.15 -> +8.81; +12.04 -> +13.43 | Helped |
| Add automated task grouping rather than disease-based grouping (ablation) | AutoDP · MIMIC-IV Tasks @ 5, Tasks @ 10, and Tasks @ 25 | AVP gain (%): +6.29 -> +7.74; +6.99 -> +8.81; +8.61 -> +13.43 | Helped |
Surprising findings
- Plain single shared model for all tasks gives only minimal gains over single-task training in the ROC metric, though AVP gains are larger.
- Searched task groupings do not follow medical disease classifications, yet outperform disease-based grouping.
Limitations
- Assumes all tasks share the same input EHR data.
- Only evaluated on one dataset (MIMIC-IV) with a single LSTM-based single-task backbone for gains.
- Search requires many MTL training runs (sampling rounds), with notable GPU-hour cost.
Experimental setup
Datasets: MIMIC-IV
Benchmarks: MIMIC-IV (25 CCS-based disease prediction tasks)
Metrics: ROC AUC average precision averaged per-task MTL gain
Baselines: LSTM Transformer Retain AdaCare DARTS MTG-Net MTG+DARTS Disease-based grouping
Builds on: DARTS MTG-Net RETAIN AdaCare Upper Confidence Bound
Closest papers by meaning
- Context Clues: Evaluating Long Context Models for Clinical Prediction Tasks on EHR Data (ICLR 2025 · similarity 0.84)
- Toward Better EHR Reasoning in LLMs: Reinforcement Learning with Expert Attention Guidance (AAAI 2026 · similarity 0.84)
- One Loss to Rule Them All: Marked Time-to-Event for Structured EHR Foundation Models (NeurIPS 2026 · similarity 0.83)
- ACES: Automatic Cohort Extraction System for Event-Stream Datasets (ICLR 2025 · similarity 0.82)
- Collaborative Synthesis of Patient Records through Multi-Visit Health State Inference (AAAI 2024 · similarity 0.82)
- MedAgentBoard: Benchmarking Multi-Agent Collaboration with Conventional Methods for Diverse Medical Tasks (NeurIPS 2025 · similarity 0.82)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.