Fix post-outcome data leakage in StageNet mortality/LOS tasks - #1205
Merged
Conversation
lehendo
commented
Aug 25, 2026
lehendo
commented
Collaborator
- MortalityPredictionStageNetMIMIC4: exclude diagnosis/procedure codes for the admission that ends in death, and cap that admission's labs to the first 48 hours after admission instead of through discharge. Earlier, already-resolved admissions are unaffected
- LengthOfStayStageNetMIMIC4: same treatment for the target admission (the one whose LOS is the label)
diagnoses_icd/procedures_icd events are timestamped at dischtime (per the MIMIC-IV config), so for the admission whose own outcome is being predicted, those codes are only known at-or-after the outcome. Labs were also pulled through discharge/death for that same admission. - MortalityPredictionStageNetMIMIC4: exclude diagnosis/procedure codes for the admission that ends in death, and cap that admission's labs to the first 48 hours after admission instead of through discharge. Earlier, already-resolved admissions are unaffected. - LengthOfStayStageNetMIMIC4: same treatment for the target admission (the one whose LOS is the label). Both changes only affect the single outcome-adjacent admission per patient; prior history is untouched. Reported metrics on these two tasks are expected to drop after this fix -- the previous numbers were inflated by the leak, not a legitimate baseline. Verified end-to-end against real MIMIC-IV demo data: terminal/target admission codes are excluded, late labs near the outcome are dropped, early labs and all prior-admission data are preserved, and patients whose only admission is the terminal/target one are correctly excluded (no leak-free signal available), matching the existing convention in dka.py.
Contributor
|
I think this still leaks the label through preprocessing. |
A reviewer (DarylOkeke) on PR sunlabuiuc#1205 pointed out that the leakage fix only restricted the death class: is_terminal_admission gated on hospital_expire_flag==1, so death cases had their terminal admission's ICD codes excluded and labs capped to 48h, while survivors' final admission kept full codes and labs through natural discharge. Two otherwise-identical patients would get systematically different feature richness based on the label itself -- a model could learn 'richer features -> survived' as a shortcut without any real clinical signal, and the existing test explicitly locked this asymmetry in as intended behavior. Fixed by determining a target admission the same way for both classes: the hospital_expire_flag==1 admission if the patient died, otherwise the chronologically last valid admission (mirroring how LengthOfStayStageNetMIMIC4 already treats its own target admission unconditionally). That admission's codes are excluded and its labs capped to TARGET_ADMISSION_INPUT_WINDOW_HOURS (renamed from TERMINAL_ADMISSION_INPUT_WINDOW_HOURS) regardless of the eventual label. Updated the existing survivor test, which asserted the old asymmetric behavior (all 3 admissions retained), to assert the corrected symmetric behavior instead, and added a direct cross-task symmetry check comparing mortality's and LOS's target-admission exclusion on the same survivor patient. Confirmed both new/updated tests fail against the pre-fix code (3 admissions instead of 2; leaked target-admission codes E1011/N179/ 5A1D70Z) and pass against the fix.
Collaborator
Author
|
@DarylOkeke , fixed. Pls review. |
Contributor
|
Looks good now. Both classes use the same target-admission window, so the label-dependent input difference is gone. |
jhnwu3
approved these changes
Sep 2, 2026
jhnwu3
left a comment
jhnwu3
left a comment
Collaborator
There was a problem hiding this comment.
Thanks for the clean catch.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.