Add Survival MDN model with SUPPORT reproduction and ablation

[nhajela1]

Contributor

Neil Hajela (nhajela2@illinois.edu)

Contribution type: New model

Original paper: Han, Goldstein, and Ranganath, Survival Mixture Density Networks, MLHC 2022.

Description

Adds a continuous-time Survival Mixture Density Network (Survival MDN) for right-censored outcomes to PyHealth. The model inherits from BaseModel, predicts a covariate-dependent Gaussian mixture, and maps latent time onto positive survival time. The censored negative log-likelihood combines event-density and censoring-survival terms and includes the inverse-transform Jacobian. The reference transformation is softplus; an exponential transformation is included for a controlled ablation.

The runnable example supports synthetic inputs, PyHealth's existing Support2Dataset, and the recovered exact SUPPORT benchmark split memberships. It compares softplus against exponential and optionally demonstrates a hidden-width variation (32 versus 64). Feature scaling is fitted on training patients only.

Reproduction and ablation

The recorded historical study contains 10 SUPPORT splits × 5 paired seeds × 2 transformations = 100 runs. Its results were produced using the project's standalone model core, not newly generated by the contributed native PyHealth example. The committed example uses the native PyHealth model and has separately passed short integration checks (below).

Recorded metric	Softplus	Exponential	Interpretation
Test censored NLL (lower is better)	0.517746	0.502901	Exponential lower; all 10 split-level means favored exponential
Concordance (higher is better)	0.621659	0.624975	Small improvement for exponential
Integrated Brier Score (lower is better)	0.189817	0.189404	Difference inconclusive
Integrated binomial log-likelihood (higher is better)	-0.560497	-0.557844	Small improvement for exponential

The exponential transformation produced heavier extrapolated survival tails; this is not evidence of improved calibration beyond observed follow-up. The paper already identifies exponential as an alternative mapping: the contribution is the controlled quantitative SUPPORT comparison, not a claim to have invented that transformation. The split-level intervals are descriptive because splits overlap. See examples/survival_mdn_results/README.md and PHASE9_FINAL_RESULTS.md for methodology, figures, caveats, and recorded results.

Main files for review

• pyhealth/models/survival_mdn.py — model, likelihood, transformations, and input validation.
• pyhealth/models/__init__.py — public model export.
• tests/core/test_survival_mdn.py — model construction, distribution outputs, gradients, analytical likelihood, invalid/nonfinite inputs, serialization, and embeddings.
• docs/api/models/pyhealth.models.survival_mdn.rst and docs/api/models.rst — documentation and index entry.
• examples/support_survival_survival_mdn.py — paired transformation ablation and optional width variation.
• examples/survival_mdn_data.py — synthetic, native SUPPORT2, and exact-membership benchmark input routes.
• examples/survival_mdn_metrics.py — example-local benchmark metric conventions.
• examples/survival_mdn_results/ — historical run table, aggregate summaries, figures, and interpretation, explicitly identified as standalone-core results.

Raw patient-level data, source archives, checkpoints, and historical reproduction scripts remain outside the upstream PR.

Validation performed

Environment: macOS Apple Silicon M3; Python 3.12.15; PyTorch 2.14.1. The repository was installed in editable mode (python -m pip install -e .).

• Import: from pyhealth.models import SurvivalMDN — passed.
• Unit tests: python -m pytest tests/core/test_survival_mdn.py -v — 21 passed in 7.46 seconds, no failures.
• Synthetic smoke test: python examples/support_survival_survival_mdn.py --source synthetic --samples 200 --epochs 3 --metrics nll — both transformation arms completed; no failed runs or nonfinite loss/gradient batches.
• Exact-membership SUPPORT smoke test: --source benchmark, recovered SUPPORT feature matrix and split matrix, --splits 1 --seeds 1 --epochs 2 --metrics nll — both transformation arms completed with finite NLL and no numerical failures. Test NLL: softplus 0.667014, exponential 0.564340.
• Native Support2Dataset integration: --source support2, supplied raw SUPPORT2 CSV, --epochs 8 --vary-hidden — three configurations completed with no failures or nonfinite loss/gradient batches. Test NLL: softplus/32 0.568475, exponential/32 0.562249, softplus/64 0.570332.
• Staged whitespace validation: git diff --cached --check — passed with no output.

These short runs are functional smoke tests, not reruns of the 100-run study or demonstrations of paper-level convergence.

Limitations

Reproduction uses a fixed validated configuration rather than the paper's full hyperparameter search and covers SUPPORT only. The original experimental splits overlap; IBS differences were inconclusive. The synthetic example is not clinical replication evidence. Survival training and example-local metrics do not use PyHealth's classification metric dispatcher. Historical benchmark metric conventions and extrapolation limitations are documented alongside the results.