Repository navigation
[Survival MoE] Add personalized mixture-of-experts model for survival prediction - #1345
AanshDevpura wants to merge 3 commits into
Conversation
Logiquo
added
the
status: automatic
|
/review |
Test CoverageTotal: 73.36% (19693 / 26846 lines)
|
![github-actions[bot]](https://avatars.githubusercontent.com/in/15368?s=60&v=4){kind=small}
| pending.replace(path) | ||
|
|
||
|
|
||
| def load_data(root: str | None, synthetic: bool, seed: int) -> pd.DataFrame: |
There was a problem hiding this comment.
[medium] Re-implements existing SUPPORT2 loader and KM helper
load_data re-implements a SUPPORT2 loader on top of Support2Dataset, which examples/survival_mdn_data.py::read_support2 already provides. censoring_km (line 149) duplicates the KM logic in examples/survival_mdn_metrics.py::km_fit. The tie handling differs (tied events leave the risk set first here), so the two are not interchangeable. Import the shared loader and KM helper from the sibling MDN example modules, as that example already does, or add a documented reason for the differing tie convention. Keep only the new IPCW Brier/ECE and preprocessing here.
| raise ValueError(f"Expected exactly one SUPPORT2 row for {pid}") | ||
| rows.append( | ||
| { | ||
| key: getattr(events[0], key, None) |
There was a problem hiding this comment.
[low] Missing attributes silently become imputed values
getattr(events[0], key, None) returns None for any absent attribute. A missing or misspelled SUPPORT2 column would then be silently imputed as missing for every patient, with no error. Use events[0][key] or getattr(events[0], key) without the default so a bad column fails loudly.
| stale += 1 | ||
| if stale >= args.patience: | ||
| break | ||
| model.load_state_dict(state) |
There was a problem hiding this comment.
[low] Crash if no epoch yields a finite validation NLL
state starts as None and is set only when score < best. If every validation NLL is NaN, model.load_state_dict(None) raises a TypeError. Initialize state from the starting weights, or raise an explicit error when no checkpoint was selected.
| train, | ||
| batch_size=args.batch_size, | ||
| shuffle=True, | ||
| generator=torch.Generator().manual_seed(42), |
There was a problem hiding this comment.
[low] Batch order is identical across seeds; docs imply otherwise
The DataLoader generator is seeded with 42 for every seed, so shuffle order is the same in all runs and only initialization varies. The module docstring says repeated seeds measure optimization variability, which overstates the spread. State that batch order is fixed in the docstring.
| if not torch.isfinite(x).all(): | ||
| raise ValueError("features must be finite; impute missing values first") | ||
| if self.embeddings: | ||
| categories = torch.as_tensor(kwargs["categories"], device=self.device) |
There was a problem hiding this comment.
[low] Missing categories input raises KeyError, not ValueError
When the model is built with categorical cardinalities, kwargs["categories"] raises a bare KeyError if the caller omits it. Every other input problem raises ValueError. Check for the key and raise ValueError with a clear message.
There was a problem hiding this comment.
AI Review
suggestion: need change
decision: need change
Thank you for this contribution. It adds a personalized mixture-of-experts survival model with query/key routing, MTLR expert heads and a censoring-aware NLL. It also adds a SUPPORT2 dropout-ablation example, API docs, and 21 synthetic tests that cover gradients, a hand-computed censoring likelihood, input validation, dataloader integration and checkpointing. The model looks correct and well tested. It is a self-contained addition, so I don't think a maintainer needs to review it for core-code reasons.
Coverage meets the bar. personalized_survival_moe.py is at 98.97% and models/__init__.py at 100%, against the 80% per-file guideline. The whole repository is at 73.36%, above the 60% bar.
To get this merged:
- The example duplicates helpers that already exist in the repo.
examples/support2_survival_personalized_moe.pyre-implements the SUPPORT2 loading inload_dataand the reverse Kaplan-Meier helpercensoring_km.examples/survival_mdn_data.py(read_support2) andexamples/survival_mdn_metrics.py(km_fit) already cover these. Please reuse them where you can. - If the tie handling or the 0.01 floor in your KM really has to differ to match the reference implementation, keep a thin wrapper or option around the shared helper. Alternatively, add a short comment explaining why a separate version is needed.
Optional suggestions:
getattr(..., None)inload_datasilently turns a missing SUPPORT2 column into an imputed value. Raising an error would be safer.model.load_state_dict(state)fails withstate=Noneif no epoch gives a finite validation NLL. A clear error would help.- A missing
categorieskwarg raisesKeyError. AValueErrorwould match the other validation errors. - The headline results are hardcoded in the docstring, and no output file is committed. Consider committing the JSON output or noting how to reproduce it.
SurvivalEVALis an undeclared dependency. Please say so clearly in the example's docstring or in the dependency metadata.- The example's IPCW Brier/ECE code is untested. That is acceptable, but a small test would be nice.
Reviewed commit ff11d23
Contributors: Aansh Devpura (devpura2@illinois.edu), Sahana Hariharan (sahanah2@illinois.edu), Shrida Shivaram Bhat (shridab2@illinois.edu)
Contribution Type: Model
Description
Add the personalized survival mixture-of-experts model from Let the Experts Speak: Improving Survival Prediction & Calibration via Mixture-of-Experts Heads to PyHealth. Includes a SUPPORT2 example comparing dropout 0, 0.1, 0.3 and 0.5 using C-index, IPCW Brier score and IPCW ECE. The example follows the released SUPPORT2 preprocessing and architecture, runs five paired seeds with early stopping, and reports full-grid calibration and measured results. Across 20 real-data runs, dropout 0.1 had the best mean C-index (0.8105), IPCW Brier (0.1409), and IPCW ECE (0.0413).
Files to review
pyhealth/models/personalized_survival_moe.py- New model implementation including:pyhealth/models/__init__.py- Added import forPersonalizedSurvivalMoE.docs/api/models/pyhealth.models.personalized_survival_moe.rst- Sphinx API page using the standard PyHealth autoclass layout.docs/api/models.rst- Added the model documentation to the index.examples/support2_survival_personalized_moe.py- Downloads missing SUPPORT2 data from UCI, loads it through Support2Dataset, and runs the dropout ablation; includes synthetic mode, metrics and recorded results.tests/core/test_personalized_survival_moe.py- 21 synthetic model tests covering initialization, forward/output shapes, finite nonzero gradients, categorical embeddings, censoring, dropout, numerical stability, input validation, PyHealth dataloader integration, and checkpoint save/load with temporary-directory cleanup.