Imitation learning for clinical decision support in pediatric ECMO

TL;DR

This paper explores imitation learning using transformer-based models to improve clinical decision support in pediatric ECMO, outperforming traditional methods on real-world data.

Contribution

It introduces a transformer-based imitation learning approach for pediatric ECMO decision support, demonstrating superior performance over classical baselines.

Findings

TabPFN outperforms XGBoost and MLP baselines

Transformer-based models provide better decision support

Supports use of imitation learning in high-stakes pediatric care

Abstract

Pediatric critical care is a dynamic, high-stakes process involving constant monitoring and adjustments in life-saving treatments. Modeling these interventions is crucial for effective decision support. To address the challenges of high complexity and data scarcity in pediatric Extracorporeal Membrane Oxygenation (ECMO), we frame clinical decision-making as learning to act from trajectories, i.e., imitation learning that learns action models from observational data, with a key feature that actions are not directly observed. We consider TabPFN, a recent transformer-based approach for tabular data, and traditional baselines including XGBoost and Multi-Layer Perceptrons(MLPs) on real-world pediatric ECMO data to learn the action models. We find that the TabPFN-based approach consistently outperforms these classical baselines, supporting its use as a strong clinician-behavior baseline for pediatric ECMO decision support.