Add MixLSTM and BaselineLSTM model for clinical mortality prediction

[michellekankan]

Contributor

• Name: Temin Kan
• NetID: temink2
• Email: temink2@illinois.edu

Type of Contribution

Model (Option 2)

Paper

Jeeheh Oh, Jiaxuan Wang, Shengpu Tang, Michael Sjoding, Jenna Wiens
Relaxed Parameter Sharing: Effectively Modeling Time-Varying Relationships in Clinical Time-Series
https://arxiv.org/abs/1906.02898

Description

Implement the mixLSTM architecture and its BaselineLSTM counterpart for modeling irregular clinical time-series data.

• mixLSTM: Addresses temporal conditional shift by relaxing the standard LSTM parameter-sharing constraint. It utilizes a mixture of LSTM cells, dynamically combining their parameters at each timestep using learnable mixing coefficients.

• BaselineLSTM: A standard LSTM implementation that serves as the control group for ablation studies.

• Both models are fully integrated with the PyHealth BaseModel structure, supporting dynamic schema parsing for features and labels.

File Guide

• pyhealth/models/mix_lstm.py
• pyhealth/models/__init__.py
• tests/core/test_mix_lstm.py
• examples/mimic3_mortality_mixlstm.py
• docs/api/models.rst
• docs/api/models/pyhealth.models.mix_lstm.rst

Ablation Study Summary

We conducted a data-level and hyperparameter ablation study on the MIMIC-III In-hospital Mortality task comparing BaselineLSTM with mixLSTM (K=2, K=4) across standard and ablated (masks removed, zero-imputed) data regimes.

• Performance: mixLSTM (K=4) achieved the highest performance on the baseline data (AUROC: 0.8020, AUPR: 0.3651), outperforming the BaselineLSTM (AUROC: 0.7868, AUPR: 0.3522).

• Robustness: All models exhibited significant performance drops in the ablated regime, indicating a heavy reliance on explicit missingness masks for clinical prediction tasks.