StageNet: Stage-Aware Neural Networks for Health Risk Prediction

1 INTRODUCTION

• Earlier works [29] implicitly assume smooth progression over time, but progression speed can vary by stage.The paper defines a stage as a period of consistent status progression across a patient’s comorbidities, rather than a disease-specific stage.Source: §1 · p. 1
• The challenge is to extract stages that are often unavailable or unclear in EHR data and use them to model patterns that differ across stages.Temporal models may gate historical information without explicitly specifying stages or change points; the paper proposes extracting and selecting informative stage-specific patterns for risk prediction.Source: §1 · pp. 1–2
• StageNet pairs unsupervised stage-aware LSTM extraction with stage-adaptive convolution that learns and re-calibrates progression patterns for risk prediction.Inter-visit time informs LSTM cell states, allowing dimensions to balance long-term progression and short-term health status.Source: §1 · p. 2
• StageNet consistently outperforms state-of-the-art models on risk prediction across two datasets, with up to 12% higher AUPRC than the best baseline.It also performs better than baseline models in patient subtyping.Source: §1 · p. 2

2 RELATED WORK

• Markov-based trajectory models include HMM approaches for Alzheimer’s disease and continuous-time models for disease progression and glaucoma or Alzheimer’s.These works assume the Markov property, which may not hold when long-term dependencies exist.Source: §2 · p. 2
• General health-status progression models support prediction tasks including subtyping [10] and risk prediction.T-LSTM incorporates elapsed time, while Health-ATM integrates timestamps into convolution.Source: §2 · p. 2
• Unlike those approaches, StageNet explicitly extracts disease-stage progression information and uses it for risk prediction.The authors describe this as the first work to do so.Source: §2 · p. 2
• StageNet re-calibrates progression patterns using the current stage to emphasize features indicative of health risks at different stages.Earlier EHR feature re-calibration mainly supported interpretability, while computer-vision work used it to emphasize informative features.Source: §2 · p. 2
• Unlike the ordered neuron mechanism for NLP parsing [33], StageNet uses clinically meaningful inter-visit intervals to summarize patients’ status progression.Source: §2 · p. 2

3.1 Overview

• StageNet processes patient records and elapsed-time intervals through stage-aware LSTM, stage-adaptive convolution, and prediction modules to produce dynamic health-risk predictions.The LSTM outputs hidden states and stage-variation factors; convolution extracts and re-calibrates progression patterns before risk prediction.Source: §3.1 · pp. 2–3

3.2 Stage-aware LSTM module

• The stage-aware LSTM uses separate cell-state dimensions for short-term health status and longer-term progression, with time-aware gates that support unsupervised stage-variation estimation.Its soft masks distinguish recent from older history; a larger stage-variation factor indicates a greater chance of entering a new stage.Source: §3.2 · pp. 3–4

3.3 Stage-adaptive convolutional module

• Stage-weighted convolution reweights historical hidden states by their distance from the current stage, with multiple kernels extracting progression patterns.Historical states farther from the current stage receive lower weights; setting the kernel size to the observation window represents the whole stage.Source: §3.3 · p. 5
• A weighted average of hidden states in the observation window forms the current-stage progression theme, which guides re-calibration of the extracted patterns.The theme is mapped to pattern-importance weights; the mechanism is described as channel-wise attention whose weights use the current-stage theme.Source: §3.3 · p. 5

3.4 Prediction module

• The prediction module uses the stage-adaptive convolution output to produce a binary current-risk label and includes residual connections.The paper uses cross-entropy loss and Adam optimization.Source: §3.4 · p. 5

4 EXPERIMENT

• StageNet was compared with other baselines on the public MIMIC-III and ESRD datasets.The passage says code is provided.Source: §4 · p. 5

4.1 Dataset description

• StageNet’s training loop applies its stage-aware LSTM, progression-pattern and prediction modules at each visit, then updates parameters using the loss.The algorithm also takes patient records, elapsed-time intervals, and observation-window length K as inputs.Source: §4.1 · p. 6
• The MIMIC-III cohort contains 33,678 patients and 2,202,114 records, with 17 physiologic variables transformed into 76-dimensional visit vectors.The ESRD dataset contains 656 patients and 13,091 visits, with 17 numeric variables and 17.5% positive labels.Source: §4.1 · pp. 5–6
• The ESRD dataset uses patients’ previous records to fill missing data to prevent leakage of future information.The cohort was cleaned and patient identities could not be identified during data collection and analysis.Source: §4.1 · p. 6

4.2 Baselines

• The evaluation also includes attention-based clinical prediction models, whose contribution the authors describe as orthogonal and combinable with StageNet.The passage identifies attention models [4] [34].Source: §4.2 · p. 6
• LSTM predicts from its hidden state, while ON-LSTM allocates hidden dimensions to long- and short-term information.T-LSTM incorporates irregular visit intervals through time decay in the recurrent cell; the original unsupervised clustering model is modified for supervised learning.Source: §4.2 · p. 6
• Decay-LSTM applies feature-level time intervals to decay memory at the input gate; Health-ATM uses irregular intervals in a hybrid CRNN.Health-ATM’s target-aware attention is removed because the task lacks a specific target embedding.Source: §4.2 · p. 6
• StageNet-I pairs regular LSTM with stage-adaptive convolution, replacing weighted convolution with regular convolution and averaging hidden states to calculate z_t.StageNet-II retains stage-aware LSTM and predicts directly from its hidden state.Source: §4.2 · p. 6

4.3 Health Risk Prediction

• Decompensation prediction on MIMIC-III detects whether death is expected within 24 hours; ESRD mortality prediction targets death within 12 months.MIMIC-III labels are assigned hourly from four hours after ICU admission until death or discharge; ESRD predictions are made at each timestep.Source: §4.3 · p. 6
• 10% higher AUPRC and min(Re,P+) on MIMIC-III versus the best baselines ON-LSTM and Health-ATM; on ESRD, 12% higher AUPRC and 6% higher min(Re, P+) versus ON-LSTM.StageNet consistently outperforms all state-of-the-art models on both datasets.Source: Table 2 · p. 7, §4.3.1 · p. 7
• Performance is assessed with AUROC, AUPRC, and min(Re,P+), the maximum of the minimum precision and recall on the precision-recall curve.The passage identifies these as metrics used in prior work.Source: §4.3 · p. 7

4.3.1 Results.

• The reduced StageNet-I and StageNet-II models outperform state-of-the-art models in most cases on both datasets, supporting the usefulness of higher-level temporal features and stage information.The authors report that ON-LSTM and Health-ATM perform better than other baselines in most cases.Source: §4.3.1 · p. 7

4.3.2 Observations.

• Peritonitis, cardiovascular disease, and cerebrovascular disease have the highest average stage variation, whereas cancer patients have the most stable status.The analysis interprets higher stage variation as greater change from historical status and recommends clinician attention to heart-disease or peritonitis history.Source: §4.3.2 · p. 7, Figure 5 · p. 7
• Patients in the high-risk group have more unstable health status, while those in the low-risk group have the most stable status.Visits are grouped by predicted risk score into low (≤0.4), medium (>0.4 and ≤0.7), and high (≥0.7) categories.Source: §4.3.2 · pp. 7–8, Table 3 · p. 8

4.3.3 Case study.

• In the case study, stage variation and predicted mortality risk rise around 2012 as the patient experiences acute myocardial infarction, before death from heart disease.Before May 2011, both stage variation and risk remain relatively low; the predicted risk stays high afterward.Source: §4.3.3 · p. 8, Figure 6 · p. 8

4.4 Patient Subtyping

• The experiment clusters patients using learned representations and k-means, evaluating cluster quality with the Calinski-Harabasz score and reporting results for k = 4.The ESRD evaluation uses 80% of patients for training and a 20% test set.Source: §4.4 · p. 8
• For clustering with k = 4, StageNet achieved over 58% higher Calinski-Harabasz score than the best baseline, Health-ATM−; it divided patients into two high-risk groups (Cluster III and IV) and two low-risk groups (Cluster I and II).Source: Table 4 · p. 8, §4.4 · p. 8
• StageNet divides patients into two high-risk and two low-risk groups, while ON-LSTM, Decay-LSTM, and Health-ATM− produce different risk-group patterns.T-LSTM identifies one high-risk and three low-risk groups; baseline cluster risk scores increase more evenly.Source: §4.4 · p. 8

4.4.1 High-risk patient subtypes.

• Baseline models identify similar features across high- and medium-risk clusters, whereas StageNet separates two high-risk groups and identifies more discriminative features.The paper compares features across risk groups using t-tests and reports the top five significant features.Source: §4.4.1 · p. 8, Table 5 · p. 9

ON-LSTM

• StageNet identifies hemoglobin as significant in Cluster IV, unlike the baseline models, alongside indicators associated with heart disease and severe infection.The paper links potassium and diastolic blood pressure to heart disease, and high white blood cell count and C-reactive protein to severe infection.Source: ON-LSTM · p. 9
• StageNet’s Cluster III and IV have different cause-of-death distributions, while ON-LSTM’s clusters show no significant distribution difference.StageNet’s clusters differ in chronic versus acute causes; Figure 8 shows the comparatively similar ON-LSTM distributions.Source: Figure 7 · p. 9, ON-LSTM · p. 9, Figure 8 · p. 9

4.4.2 Low-risk patient subtypes.

• Cluster I has higher serum creatinine, glucose, and blood chlorine and lower albumin than Cluster II, indicating higher potential health risk.The passage relates albumin and creatinine to nutritional status and ESRD progression, chlorine to kidney damage, and glucose to possible diabetes.Source: §4.4.2 · pp. 9–10, Table 6 · p. 10

4.4.3 Health status stability vs. Patient subtypes.

• Stage variation is positively correlated with cluster health risk; Cluster IV is most unstable, while high-risk Cluster III has lower variation and more chronic progression.The authors report these observations as consistent with their earlier results and medical findings.Source: §4.4.3 · p. 10, Table 7 · p. 10

5 CONCLUSIONS

• StageNet achieves 10% higher AUPRC and min(Re,P+) on MIMIC-III, and 12% higher AURPC and 6% higher min(Re,P+) on ESRD, compared with the best baseline.The conclusion reports consistent improvements over state-of-the-art methods across both datasets.Source: §5 · p. 10
• StageNet combines a stage-aware LSTM that extracts health stages unsupervisedly with a stage-adaptive convolutional module that uses and re-calibrates stage-related variation patterns.The conclusion describes these components as the model’s approach to health-risk prediction.Source: §5 · p. 10
• The patient-subtyping experiment finds that StageNet learns more discriminative patient representations than baseline models; the authors hope this can help physicians identify unstable patients and prevent or delay adverse outcomes.The clinical use is framed as a hope, not a demonstrated outcome.Source: §5 · p. 10

A DATASET DETAILS

• Table 8 presents basic statistics for the ESRD and MIMIC datasets.The supplied passage does not state the table’s individual values.Source: A DATASET DETAILS · p. 10, Table 8 · p. 10

• Baseline hyperparameters use recommended original-paper settings when available; otherwise, values are selected by grid search on the validation set.Source: B IMPLEMENTATION DETAILS · p. 10
• LSTM, T-LSTM, and Decay-LSTM use 64/128 hidden units, while ON-LSTM uses 72/384 hidden units and chunk factors 36/128 for ESRD/MIMIC-III, respectively.Source: B IMPLEMENTATION DETAILS · p. 10
• Health-ATM− uses 64/128 hidden units and 32/64 convolutional filters of size 3/5; StageNet uses K=10, 72/384 hidden units, and chunk factors 36/128 for ESRD/MIMIC-III, respectively.StageNet also applies dropout before the output layer and dropconnect in the LSTM layer, with dropout rates 0.5/0.3 for ESRD/MIMIC-III, respectively.Source: B IMPLEMENTATION DETAILS · p. 10
• Models train for 50 epochs on MIMIC-III and 200 epochs on ESRD, with a learning rate of 0.001.Source: B IMPLEMENTATION DETAILS · p. 10