Robustness and Scalability Of Machine Learning for Imbalanced Clinical Data in Emergency and Critical Care

TL;DR

The study evaluates robustness and scalability of ML methods for imbalanced tabular clinical data in ED/ICU settings by comparing classical tree-based ensembles (DT, RF, XGBoost) with deep tabular models (TabNet, TabResNet) across MIMIC-IV-ED and eICU-CRD. It leverages Bayesian hyperparameter optimization and three imbalance metrics (CVCF, IR, NECD) to assess predictive performance and scaling, finding that XGBoost offers the most stable and scalable performance under increasing imbalance, while deep tabular models incur higher costs and degrade more sharply. The results advocate for ensemble methods as the practical default for real-time clinical deployment, and provide an imbalance-aware framework for model selection and cross-institution validation. The work also discusses deployment, fairness, and governance considerations, emphasizing cross-institution generalizability and the potential for federated or collaborative AI in acute care.

Abstract

Emergency and intensive care environments require predictive models that are both accurate and computationally efficient, yet clinical data in these settings are often severely imbalanced. Such skewness undermines model reliability, particularly for rare but clinically crucial outcomes, making robustness and scalability essential for real-world usage. In this paper, we systematically evaluate the robustness and scalability of classical machine learning models on imbalanced tabular data from MIMIC-IV-ED and eICU. Class imbalance was quantified using complementary metrics, and we compared the performance of tree-based methods, the state-of-the-art TabNet deep learning model, and a custom lightweight residual network. TabResNet was designed as a computationally efficient alternative to TabNet, replacing its complex attention mechanisms with a streamlined residual architecture to maintain representational capacity for real-time clinical use. All models were optimized via a Bayesian hyperparameter search and assessed on predictive performance, robustness to increasing imbalance, and computational scalability. Our results, on seven clinically vital predictive tasks, show that tree-based methods, particularly XGBoost, consistently achieved the most stable performance across imbalance levels and scaled efficiently with sample size. Deep tabular models degraded more sharply under imbalance and incurred higher computational costs, while TabResNet provided a lighter alternative to TabNet but did not surpass ensemble benchmarks. These findings indicate that in emergency and critical care, robustness to imbalance and computational scalability could outweigh architectural complexity. Tree-based ensemble methods currently offer the most practical and clinically feasible choice, equipping practitioners with a framework for selecting models suited to high-stakes, time-sensitive environments.

Figures (19)

• Figure 1: Architecture of TabResNet. Sequential structure of the network architecture for tabular data. The Input Layer (Linear + BatchNorm + ReLU + Dropout) is followed by 1--3 Compact Residual Blocks, each containing two Linear layers with Batch Norm, ReLU, Dropout, and a skip connection. An optional Reduction Layer precedes the Output Layer, which produces class predictions.
• Figure 2: Class distribution of of target outcomes. Class distribution of target variables in the MIMIC-IV-ED (top row) and eICU (bottom row) datasets. The histograms illustrate the frequency of samples across various clinical prediction tasks.
• Figure 3: Correlation of class imbalance metrics (MIMIC-IV-ED). Correlation heatmaps of class imbalance metrics for different prediction tasks on the MIMIC-IV-ED dataset. Pairwise correlations are shown between the CVCF, IR, and NECD metrics across weighting strategies (inverse, effective, median) for three prediction tasks.
• Figure 4: Effect of class imbalance on discharge diagnosis. Weighted F1 performance across varying levels of class imbalance for primary diagnosis prediction. The performance curves for 20 classifiers are shown, with the weighted F1 value decreasing as the imbalance severity increases.
• Figure 5: Effect of class imbalance on ICD Code prediction. Weighted F1 performance across varying levels of class imbalance for ICD code group prediction. Compared with fine-grained diagnosis prediction, grouped ICD categories reduce label sparsity, and classifiers generally maintain greater stability.