PaperlayerView exact revision on arXiv

Source-linked AI summary

Graph-Enhanced Disease Trajectory Discovery and Temporal Endpoint Forecasting from Longitudinal Electronic Health Records

Jin-Jia Zheng, Vincent S. Tseng

arXiv:2610.20502v1cs.LG

TL;DR

Longitudinal EHR models can forecast disease, but often leave explicit progression evidence difficult to inspect, while trajectory catalogs are not routinely integrated into endpoint and event-time forecasting. This paper combines leakage-controlled graph discovery, patient-compatible branching motifs, and a time-aware Transformer. It reports modest overall gains, with timing improvements concentrated in longer-range events.

  • Problem

    Sequence models can predict future disease but often encode progression implicitly, while trajectory studies are not systematically integrated into endpoint ranking and event-time forecasting.

    Source: §1 · p. 1

  • Method

    The framework builds train-only diagnosis transitions, freezes validation-selected branching-motif catalogs, and combines retrieved patient-compatible motifs with a time-aware Transformer.

    Source: §1 · p. 2, §3.5 · p. 6

  • Results

    Event-time MAE decreased by 19.2 days versus a matched history-only Transformer, with timing gains driven by events beyond 730 days.

    Source: §6 · p. 13, §8 · p. 14

  • Takeaways & Limitations

    Branching motifs provide structured trajectory context whose reported timing benefit is concentrated in long-range events rather than uniformly improving forecasting.

    Source: §6 · p. 13, §8 · p. 14

    MIMIC-IV represents one hospital system, and incomplete post-discharge ascertainment remains despite censor-aware masking.

    Source: §7 · p. 13

Abstract

from arXiv · show

Longitudinal electronic health records (EHRs) can reveal disease-progression structure, while useful forecasting requires leakage-controlled discovery, censor-aware outcomes, and inspectable trajectory evidence. We developed a graph-enhanced framework in MIMIC-IV v3.1 that splits patients before graph construction, learns train-only consecutive-admission transitions, and freezes endpoint-specific branching-motif policies selected on validation data. The final catalog contained 99,527 motifs, and a time-aware Transformer retrieved at most 24 patient-compatible motifs per endpoint through endpoint-specific cross-attention. Across five training seeds, micro-AUPRC increased from 0.9330 to 0.9357, while event-time MAE decreased from 676.9 to 657.7 days (reduction 19.2 days; 95% paired patient-cluster bootstrap CI, 13.8-24.8). Timing effects were heterogeneous: MAE worsened for events within 730 days but improved by 58.9 days for events occurring after 730 days. In an exploratory internal temporal split, branch conditioning reduced conditional event-time MAE by 26.9 days (95% CI, 23.0-30.8) without evidence of better micro-AUPRC. Corrected survival evaluation over the supported 0-2,555-day interval did not establish macro C-index or integrated Brier score superiority over a matched history-only discrete-hazard model. A directed, degree-matched GNN repair substantially improved unseen-edge discovery, yet replacing GNN-derived scores within fixed motif identities did not improve downstream forecasting. These results position branching motifs as structured trajectory context whose main value is long-range timing support rather than uniformly stronger endpoint discrimination, survival prediction, or causal explanation.

1. Introduction

Longitudinal EHR models can forecast future disease, but often leave progression evidence implicit, while trajectory studies have not routinely connected discovered patterns to endpoint and event-time forecasting. This paper combines train-only graph discovery with patient-matched branching motifs and censor-aware forecasting.

Research gap

  • Existing neural sequence models can encode long-range disease dependencies, but their implicit representations make explicit disease paths supporting endpoint forecasts difficult to identify.Trajectory research reveals ordered disease associations, yet often remains descriptive and is not systematically integrated into endpoint ranking or event-time forecasting.Source: §1 · pp. 1–2

Framework

  • The framework splits patients before discovery, builds graph evidence and motif candidates from training data, and freezes validation-selected catalogs before testing.Its final forecaster combines patient-compatible branching motifs with a time-aware Transformer for risk, endpoint ranking, and time-to-event prediction.Source: §1 · p. 2

Study aims

  • The study tests whether train-only discovery can yield a high-coverage branching catalog without a fixed top-N cap and whether branch tokens improve censor-aware prediction.It also evaluates stratified effects, matched ablations, survival comparisons, and whether better unseen-edge discovery improves forecasting when motif identities are fixed.Source: §1 · p. 2

Contributions

  • The contributions include a patient-disjoint discovery protocol, a validation-selected catalog of 99,527 motifs, and a five-seed censor-aware evaluation.The evaluation also includes candidate-count and lead-time analyses, matched ablations, a censor-aware hazard comparator, and graph-repair and motif-deletion studies.Source: §1 · p. 2

2. Related work

Prior work offers powerful sequential EHR prediction, population trajectory maps, and graph-based representations, but leaves open how to connect patient-specific trajectory evidence to endpoint forecasting. The paper frames this as a joint problem of catalog design, motif retrieval, and progression-sensitive evaluation.

2.1 Longitudinal EHR forecasting

  • Sequence models evolved from recurrent predictors to Transformers, yet latent states and attention do not necessarily expose explicit directional progression evidence for a forecast.The proposed motifs complement rather than replace the sequence encoder, allowing tests of whether explicit progression structure adds information beyond a history-only model.Source: §2.1 · pp. 2–3

2.2 Disease trajectory mining and transition graphs

  • Population trajectory studies often describe disease sequences, leaving open how large catalogs should be matched to individual histories and incorporated into endpoint-specific risk and timing models.The framework treats catalog policy as validation-selected and uses branching motifs to connect shared prefixes with alternative outcomes.Source: §2.2 · p. 3

2.3 Graph-enhanced and knowledge-based clinical prediction

  • Graph-based clinical methods vary in graph type and representation; latent embeddings may not expose active paths, while explicit paths can be unreliable with sparse edge evidence.The paper uses graph-derived branching motifs and separately tests whether improved held-out edge prediction helps forecasts when motif identities are fixed.Source: §2.3 · pp. 3–4

2.4 Time-to-event and progression-sensitive evaluation

  • Classification, endpoint ranking, and event-time metrics capture complementary forecasting dimensions, so gains in one may coexist with unchanged or worse performance in another.The evaluation accounts for repeated patient indices and horizon heterogeneity using censor-aware masks, patient-cluster resampling, and lead-time stratification.Source: §2.4 · p. 4

3. Proposed framework

The framework discovers directional diagnosis transitions using training patients, selects and freezes endpoint-specific branching catalogs using validation data, and retrieves patient-compatible motifs for censor-aware forecasting. A time-aware Transformer combines visit histories with motif context through separate endpoint-specific risk and timing pathways.

3. Proposed framework

  • Patients are split before graph discovery; training data supply graph evidence and motif candidates, validation selects catalog policies, and the frozen catalog is applied to censor-aware test examples.The forecaster retrieves at most 24 motifs per endpoint and uses endpoint-specific cross-attention with separate risk and timing gates.Source: Figure 1 · p. 4, §3 · p. 4

3.1 Problem setting, index examples, and endpoint outcomes

  • Each nonterminal admission can be an index; prevalent endpoints are removed, events within 30 days are excluded, and insufficiently observed event-free rows are unknown.Unknown rows are excluded from losses and metrics; onset means the first recorded endpoint occurrence in MIMIC-IV, not biological onset.Source: §3.1 · p. 5

3.2 Consecutive-admission transition graph

  • The transition graph counts directed diagnosis pairs across consecutive admissions separated by 30 to 3,650 days, using training patients only.Transition confidence is edge support divided by the number of qualifying source-admission opportunities containing its source diagnosis.Source: §3.2 · p. 5

3.3 GraphSAGE-style directed link predictor

  • The directed link predictor uses diagnosis structural features and a directed edge scorer, while its message-passing adjacency is symmetrized and row-normalized.The post-primary directed, degree-matched repair was supporting evidence and did not replace the frozen primary catalog.Source: §3.3 · p. 5

3.4 Explicit paths, branching motifs, and adaptive thresholds

  • Discovery produces explicit paths and branching motifs, but the final forecaster uses branching motifs selected under validation-chosen stable or broad policies.Both policies disable fixed top-N export; the frozen catalog contains 99,527 motifs, with stable selected for COPD and dementia and broad for stroke and CKD.Source: §3.4 · pp. 5–6

3.5 Soft temporal-prefix retrieval and branch-token representation

  • Soft temporal-prefix matching scores motifs against current and earlier admissions, then diverse-prefix retrieval passes at most 24 tokens to each endpoint query.The runtime token budget does not limit the size of the discovery catalog.Source: §3.5 · p. 6

3.6 Time-aware Transformer and endpoint-conditioned context gates

  • A time-aware Transformer encodes visit diagnoses and inter-admission gaps, while endpoint queries attend to retrieved branch tokens through separate risk and timing modules.Each pathway has its own scalar context gate.Source: §3.6 · p. 6

3.7 Multi-task objectives and optimization

  • Training combines at-risk endpoint classification, observed-positive timing regression, and pairwise endpoint ranking, with validation micro-AUPRC selecting checkpoints.Forecasting runs use five seeds, 42–46.Source: §3.7 · p. 7

4. Experimental setup

The study used a patient-disjoint MIMIC-IV v3.1 cohort and censor-aware outcomes, comparing a branch-enhanced model with matched alternatives under a leakage-controlled evaluation protocol.

4.1 Data source, coding, endpoints, and cohort

  • The cohort contained 100,018 patients and 422,096 admissions, split before discovery into 70,012 training, 10,001 validation, and 20,005 test patients.Diagnoses were mapped to a hybrid CCS/CCSR vocabulary, with COPD, stroke, CKD, and dementia as endpoints.Source: §4.1 · p. 7

4.2 Censor-aware outcome construction and follow-up

  • Prevalent endpoints were masked, events within 30 days excluded, and event-free rows with under 3,650 days of follow-up excluded as unknown.Follow-up ended at last hospital discharge or earlier recorded death; timing errors were calculated only for observed positive at-risk events.Source: §4.2 · p. 8

4.3 Compared models

  • The principal comparison matched a history-only time-aware Transformer against the branch-enhanced model; additional controls included pattern-token and Mamba-based comparators.Controls included mean-token aggregation, null and training-patient donor tokens, and removal of online numeric channels; survival extensions used a separate cohort and estimand.Source: §4.3 · p. 8

4.4 Metrics and reporting

  • Risk discrimination, at-risk endpoint ranking, observed-event timing, and survival outcomes were assessed with patient-cluster resampling and separate uncertainty estimates.An exploratory temporal evaluation used a 1,095-day horizon after the planned 3,650-day horizon lacked enough late-era negatives; it was not fully testblind external validation.Source: §4.3 · p. 8, §4.4 · p. 9

5. Results

The final model combined a high-coverage branching-motif catalog with modest overall forecasting gains. Timing benefits were concentrated in long-range events, while discrimination and temporal-split findings were qualified.

5.1 Discovery scale and branch coverage

  • The validation-selected catalog contained 99,527 endpoint-specific branching motifs, with motifs retrieved for 99.48% of frozen test examples and branch evidence for 98.87% of positive endpoint rows.Coverage does not establish clinical validity or causal explanation.Source: §5.1 · p. 9

5.2 Main forecasting results

  • Micro-AUPRC improved by 0.002705 and gap MAE decreased by 19.23 days across five seeds; the overall gains were modest and discrimination gains endpoint-specific.Branch-versus-history AUPRC, MAE, and sMAPE comparisons remained supported in Holm multiplicity sensitivity.Source: Table 3 · p. 9, §5.2 · pp. 9–10

5.2.1 Candidate-count-aware endpoint ranking

  • Among 3,604 examples with at least two eligible endpoints, MRR rose from 0.9373 to 0.9432 and Hit@1 from 0.8893 to 0.9015, but eight candidate-count-by-metric comparisons did not survive Holm adjustment.73.56% of positive ranking examples had only one eligible endpoint, creating a structural ceiling.Source: §5.2.1 · p. 10

5.2.2 Timing heterogeneity by observed lead time

  • MAE worsened by 37.78 days for events within 180 days and 36.85 days for events at 181-730 days, but improved by 58.87 days for events after 730 days.The aggregate timing benefit was driven by long-range progression timing rather than uniform near-term improvement.Source: §5.2.2 · p. 10

5.2.3 Internal temporal validation

  • In the exploratory 1,095-day late-era evaluation, conditional event-time MAE fell by 26.901 days, while micro-AUPRC showed no evidence of improvement.The analysis used an anchor-year split and was not fully testblind external validation.Source: §5.2.3 · p. 10

5.3 Final-setting ablation

Matched ablations supported the branch-enhanced model on several overall outcomes and favored the adaptive catalog over fixed top-300 for timing metrics. Controls qualify claims that motif content or every engineered component is necessary.

5.3 Final-setting ablation

  • The adaptive catalog remained supported over fixed top-300 for overall MAE and sMAPE after Holm adjustment.GNN scores and learned gates were supported mainly by timing, while entropy retained small discrimination effects.Source: §5.3 · p. 10
  • The full model remained supported over the no-branch variant for overall AUROC, AUPRC, MAE, and sMAPE after Holm adjustment within the exploratory ablation family.The eight-variant ablation used the same split, masks, training budget, and evaluation protocol.Source: §5.3 · p. 10, Table 4 · p. 11

5.3.1 Matched comparators and token controls

  • The proposed model achieved micro-AUPRC 0.935698 and conditional event-time MAE 657.712 days, versus 0.932283 and 668.254 days for the SPT-inspired adapter.Against the matched Mamba history comparator, MAE was 11.809 days lower, while the AUPRC advantage was only +0.000307 with a CI spanning zero.Source: §5.3.1 · p. 11
  • Against content-free null tokens, the proposed-minus-null MAE difference was -3.512 days with a 95% CI spanning zero, so an advantage over preserved token-count and mask structure was not established.Removing ranking loss slightly improved sMAPE; this ablation does not establish that ranking loss is required.Source: §5.3 · p. 10, §5.3.1 · p. 11

5.4 Calibration and survival sensitivity

Calibration and survival analyses did not establish macro survival superiority, while graph repair improved unseen-edge discovery without improving fixed-catalog forecasting. Motif deletion changed predictions, but signed effects were not uniformly positive.

5.4 Calibration and survival sensitivity

  • Branch conditioning changed macro C-index by -0.002405 and IBS by -0.000398 versus matched history hazard, establishing neither macro superiority nor equivalence over 0-2,555 days.Endpoint effects were heterogeneous, including worse CKD IBS and lower dementia concordance.Source: §5.4 · p. 12

5.5 Held-out graph and explanation evaluation

  • The directed, degree-matched GNN repair increased unseen-edge AUPRC from 0.1724 to 0.3825 and AUROC from 0.5034 to 0.7538.Its MRR remained lower than the support-confidence heuristic.Source: §5.5 · p. 12
  • Replacing only GNN-derived scores within fixed motif identities produced no clear discrimination or ranking gain and increased MAE by 2.10 days.This separates improved unseen-edge prediction from incremental fixed-catalog forecasting utility.Source: §5.5 · p. 12
  • Deleting the top-attention motif changed endpoint probability by 0.0070 on average versus 0.0011 after random deletion, and predicted event time by 21.6 versus 7.4 days.Signed effects were not uniformly positive, so attended motifs may be inhibitory as well as risk-increasing.Source: §5.5 · p. 13

5.6 Multiplicity sensitivity

Multiplicity sensitivity analyses preserved the primary AUPRC and timing findings for branch conditioning and the adaptive catalog while narrowing endpoint- and ablation-level claims. Reported percentile confidence intervals remain unadjusted and are not simultaneous intervals.

5.6 Multiplicity sensitivity

  • Primary AUPRC and timing findings for branch conditioning and the adaptive catalog persisted after Holm and Benjamini-Hochberg sensitivity analyses.The adjustments narrowed endpoint- and ablation-level claims.Source: §5.6 · p. 13
  • Holm and Benjamini-Hochberg adjustments applied to paired-bootstrap tail probabilities within prespecified families.Source: §5.6 · p. 13
  • Percentile confidence intervals were unadjusted and are not simultaneous intervals.Source: §5.6 · p. 13

6. Discussion

Branch conditioning yielded modest overall gains, with timing effects varying by lead time; comparator analyses also narrowed what can be attributed specifically to motif content. Improved graph link prediction did not translate monotonically into forecasting utility, and survival superiority was not established.

6. Discussion

  • Events beyond 730 days showed a 58.9-day MAE reduction, whereas events within 730 days were estimated less accurately.Exploratory internal temporal validation reproduced a conditional timing benefit without a discrimination gain, while timing priors or loss reweighting shifted error across horizons rather than eliminating the trade-off.Source: §6 · p. 13
  • The model improved conditional timing over SPT-inspired, Mamba-based, mean-token, and donor-token comparators, but not overall over content-free null tokens.The no-online-numeric-channel variant did not degrade, so extra capacity, token presence/count structure, and engineered numeric channels cannot be ruled out or claimed uniquely necessary.Source: §6 · p. 13
  • Directed, degree-matched GNN training repaired unseen-edge prediction, but substituting improved scores within fixed motif identities did not improve forecasts.A taxonomy-feature graph further improved unseen-link AUPRC while worsening downstream MAE; branch-token value therefore did not increase monotonically with graph link-prediction quality.Source: §6 · p. 13
  • Corrected survival analysis did not establish macro C-index or integrated Brier score superiority over a matched history-only hazard model.Conditional event-time gains therefore should not be generalized to survival prediction or clinical utility; motif-deletion analyses support computational sensitivity, not causality, clinical correctness, or clinician-validated reasoning.Source: §6 · p. 13

7. Limitations and practical considerations

The study's conclusions are bounded by single-system ascertainment, observational diagnosis coding, horizon-specific trade-offs, and limits on survival and comparator interpretation. The model is intended for retrospective methodological research and hypothesis prioritization, not direct clinical decision-making.

7. Limitations and practical considerations

  • MIMIC-IV represents one hospital system, and the last recorded discharge does not guarantee continuous endpoint-free follow-up elsewhere.The censor-aware mask avoids treating known insufficient follow-up as negative, but incomplete ascertainment remains.Source: §7 · p. 13
  • CCS/CCSR aggregation can merge heterogeneous diagnoses, and the transition graph represents observational association rather than causal disease mechanisms or biological onset.Source: §7 · p. 13
  • Timing worsened within 730 days and improved beyond 730 days, and timing reweighting did not remove this short-versus-long trade-off.Ranking is ceiling-limited, and candidate-count-stratified signals were not familywise significant after Holm adjustment.Source: §7 · p. 13
  • Null-token controls did not establish that motif content itself is necessary, and removing online numeric channels did not impair the frozen setting.Repaired or taxonomy-augmented GNNs improved link prediction without improving downstream forecasting.Source: §7 · p. 14
  • The corrected survival extension relies on an independent-censoring interpretation and a supported IBS interval of 0-2,555 days, uses a different eligible test universe, and did not establish macro superiority over history hazard.Mamba and DeepHit-style comparators are task-matched adapters rather than complete reproductions; explanation analyses assess computational sensitivity, not causal effect, clinical correctness, or clinician agreement.Source: §7 · p. 14
  • The model is intended for retrospective methodological research and hypothesis prioritization, not direct clinical decision-making.Source: §7 · p. 14

8. Conclusion

The study presents a leakage-controlled graph-enhanced forecasting framework whose branching-motif context produced modest overall gains and lower aggregate event-time MAE than a matched history-only Transformer. Its findings support conditional long-range timing value, while bounding claims about motif necessity, survival superiority, and the translation of graph link-prediction gains into forecasts.

8. Conclusion

  • Event-time MAE decreased by 19.2 days versus a matched history-only Transformer, alongside modest overall gains from the frozen 99,527-motif catalog and endpoint-specific branch-token fusion.Source: §8 · p. 14
  • Lead-time stratification showed that timing gains were driven by events beyond 730 days, while candidate-count analysis found a directionally consistent but multiplicity-limited ranking signal.Source: §8 · p. 14
  • Same-task SPT-inspired and Mamba-based comparators and an internal temporal split supported a conditional timing advantage, but null-token and no-online-numeric controls narrowed claims about which motif components are necessary.Source: §8 · p. 14
  • Directed GNN improvements did not translate monotonically to forecasting, and corrected survival analysis did not establish macro superiority over matched history hazard.Source: §8 · p. 14
  • Branching motifs are best viewed as inspectable structured context for retrospective conditional forecasting, not causal explanations or guaranteed survival improvements.Source: §8 · p. 14
Loading 2610.20502v1…