Skip to content

ML4H merge: tranche 1 - #1185

Open
will-pang wants to merge 27 commits into
sunlabuiuc:masterfrom
Multimodal-PyHealth:ml4h-merge-tranche-1
Open

ML4H merge: tranche 1#1185
will-pang wants to merge 27 commits into
sunlabuiuc:masterfrom
Multimodal-PyHealth:ml4h-merge-tranche-1

Conversation

@will-pang

@will-pang will-pang commented Aug 11, 2026 •
edited
Loading

Copy link
Copy Markdown
Contributor

A first pass at consolidating all the work we did so that we can merge to main. I'll probably need everyone's help in reviewing this (@joshuasteier, @Rian354, @joshuxchn, as well as @jhnwu3), but in short the updates here add in unified embedding support and tested across the different backbones (RNN, transformer, bottleneck
transformer, EHRMamba, JambaEHR) that we're interested in presenting results. In terms of modality, this focuses only on labs, labs + notes, labs + notes + cxr.

Three big updates made

  • Removed all sample re-weighting work
  • Removed observation windows
  • Removed ICD-10 codes

What to remove

  • All references in pyhealth/scripts/ which I've used to submit to condor and can be removed.

@will-pang will-pang changed the title Ml4h merge tranche 1 ML4H merge: tranche 1 Aug 11, 2026
Rian354 and others added 4 commits August 17, 2026 07:20
@Rian354 @cursoragent
Stop emitting fake missing-event placeholders.
9782aca 
Empty notes, labs, CXR, and ICD are now zero-length sequences instead of a
constant [MISSING_TEXT] row, black image, or pad visit. The tokenizer crash
on empty batches is handled by constructing empty tensors, so BERT cannot
treat note presence as a free mortality feature.

Co-authored-by: Cursor <cursoragent@cursor.com>
@Rian354 @cursoragent
Keep frozen text encoders in eval when Trainer calls train().
869ac8e 
nn.Module.train() re-enables dropout inside Bio_ClinicalBERT even when
every weight has requires_grad=False. Pin those encoders back to eval so
a frozen note embedding is deterministic across steps.

Co-authored-by: Cursor <cursoragent@cursor.com>
@Rian354 @cursoragent
Cache frozen [CLS] embeddings keyed on real tokens, not padded rows.
11beefc 
A frozen BERT forward is identical across epochs, but batch padding width
changes every shuffle, so a key over the full padded row never hits.
Hash only the attended tokens so the same note reuses its [CLS] vector.

Co-authored-by: Cursor <cursoragent@cursor.com>
@Rian354 @cursoragent
Restore BaseDataset parquet scanning that MEDS still calls.
86fe73b 
a0f1422 deleted _scan_table/_scan_parquet while MEDSDataset._subset_patient_ids still calls _scan_parquet, so MEDS loads crashed. load_table routes through _scan_table again, and resolve_table_path keeps absolute cache paths. The test reads a real two-row parquet file.

Co-authored-by: Cursor <cursoragent@cursor.com>
@Rian354

Rian354 commented Aug 18, 2026 •
edited
Loading

Copy link
Copy Markdown
Collaborator

Empty notes, labs, and CXR are now zero-length sequences instead of a fake [MISSING_TEXT] row, a pad visit, or a black frame.

Frozen BERT stays in eval when Trainer calls train(), so dropout cannot fire on a frozen encoder. The [CLS] cache keys on the real tokens, not the padded row, so the same note hits after a shuffle.

_scan_table / _scan_parquet are restored. MEDS still calls them.

Rian354 and others added 4 commits August 18, 2026 09:52
@Rian354 @cursoragent
Fill attention masks with dtype min and use fused SDPA.
965a87b 
The explicit path filled padded scores with -1e9, which is outside the fp16 range, so AMP overflowed. Ordinary forwards now use fused scaled_dot_product_attention; the explicit path stays behind register_hook for interpretability and fills with finfo(dtype).min. The test checks a padded fp16 pass stays finite and that -1e9 still overflows.

Co-authored-by: Cursor <cursoragent@cursor.com>
@Rian354 @cursoragent
Record batch padding and skip those slots in the unified sequence.
7c4c056 
The collator padded short samples with 0.0 and nothing marked the extra slots, so they looked like real measurements at admission time. It now emits {field}__pad_mask, the unified heads thread it through, and RNN packed lengths clamp at 1 so an all-pad sample does not crash. Tests check the collate mask, pad-last sort, and an all-pad RNN step.

Co-authored-by: Cursor <cursoragent@cursor.com>
@Rian354 @cursoragent
Keep padding_idx=0 on nested code embeddings.
ca9b63d 
NestedSequenceProcessor used padding_idx=None so a fake empty visit could have a non-zero vector. Index 0 then received gradients. Empty visits are now zero events, so the pad row stays frozen zeros. The test checks both the zeros and a zero gradient on that row.

Co-authored-by: Cursor <cursoragent@cursor.com>
@Rian354 @cursoragent
Accept resized_images and write sunlab CXR metadata to cache.
b6fb31e 
The sunlab loader required a directory named images and wrote the derived CSV into the PhysioNet root, which is read-only on the cluster. Both images and resized_images are accepted, cache is tried first, and the generated YAML points at the absolute CSV. The test chmods the root to 555 and checks the CSV lands in cache.

Co-authored-by: Cursor <cursoragent@cursor.com>
@Rian354

Rian354 commented Aug 18, 2026 •
edited
Loading

Copy link
Copy Markdown
Collaborator

A few more measurement fixes:

The attention mask fill is finfo(dtype).min, and ordinary forwards use fused SDPA. -1e9 is outside the fp16 range, so AMP overflowed on padded positions.

The collator now emits {field}__pad_mask. Padded zeros were looking like real measurements at admission time.

Nested code embeddings keep padding_idx=0, so the pad row stays frozen zeros.

The sunlab CXR loader accepts resized_images and writes the derived CSV under cache_dir. The PhysioNet root is read-only on the cluster.

Rian354 and others added 11 commits August 18, 2026 13:45
@Rian354 @cursoragent
Refuse unknown amp_dtype instead of silently selecting fp16.
7494787 
"bfloat16" and any other spelling previously fell through to float16 and a
GradScaler, so a typo changed both precision and gradient scaling with no
message.

Co-authored-by: Cursor <cursoragent@cursor.com>
@Rian354 @cursoragent
Thread collate pad_mask through the unified MLP path.
f3689cb 
Without this the sixth backbone is missing from the table, and a unified MLP
would score padded slots as real events the same way the other heads used to.

Co-authored-by: Cursor <cursoragent@cursor.com>
@Rian354 @cursoragent
Fit lab z-scores on observed train rows via region_of_interest.
ee17eb6 
Iterating a StreamingDataset under torchrun fitted 1/N of the split (len()=5
of 20 at WORLD_SIZE=4). Padded zeros also pulled sodium's mean from 140 to
105. patient_to_index after subset() still holds parent indices and raised on
index 237.

Co-authored-by: Cursor <cursoragent@cursor.com>
@Rian354 @cursoragent
Write run_config.json next to the metrics of a finished run.
224e7f4 
metrics_history.json stored the score but not the conditions, so a frozen-
encoder run and a fine-tuned run were indistinguishable after the job log was
gone.

Co-authored-by: Cursor <cursoragent@cursor.com>
@Rian354 @cursoragent
Stop dropping later stays against the first admission's CXR window.
8ae3c54 
NotesLabsCXRMIMIC4 and CXRMIMIC4 skipped any stay with admit time >=
first_admit + window_hours, so a later admission contributed no images. Lab
collection was already per-admission. Bump emitted_data_version to 2 so those
caches cannot be reused.

Co-authored-by: Cursor <cursoragent@cursor.com>
@Rian354 @cursoragent
Keep paired runs from overwriting each other or reporting train as test.
5def4ea 
The directory was {model}_seed{seed}, so labs and notes_labs at one seed
destroyed the first arm. split_by_patient fell back to split_by_sample with
no warning, and predictions came from test or val or train. Wire MLP and the
lab standardiser; leave the Jamba library default at 6.

Co-authored-by: Cursor <cursoragent@cursor.com>
@Rian354 @cursoragent
Give NotesLabs a 24h window and put concatenated stays on one timeline.
5be22e4 
The class default was None (labs through discharge) while the docstring, LabsMIMIC4, and the runner all used 24. Event times were hours from each stay's own admit, so stay 2 at +6h sorted with stay 1 at +6h. Collection is still per stay; times are hours from the first stay in the sample. Bump emitted_data_version to 3.

Co-authored-by: Cursor <cursoragent@cursor.com>
@Rian354 @cursoragent
Give time embeddings a 10-year span and drop ICDLabsMIMIC4.
9f50b76 
The old sinusoid wrapped every 720 hours, so later stays aliased with
the first. MIMIC timestamps ICD at discharge, which leaks the
in-hospital mortality label.

Co-authored-by: Cursor <cursoragent@cursor.com>
@Rian354 @cursoragent
Point the old unified_embedding import at the package copy.
89f7bc7 
Tests and docs imported pyhealth.models.unified_embedding while runtime
used embedding.unified, so the two files could drift. embedding.py next
to the embedding/ package was unused.

Co-authored-by: Cursor <cursoragent@cursor.com>
@Rian354 @cursoragent
Point AMP autocast at the trainer device instead of hardcoding CUDA.
b50bf2f 
use_amp=True on CPU or MPS warned that CUDA was unavailable and skipped
mixed precision. Unknown amp_dtype spellings were already rejected.

Co-authored-by: Cursor <cursoragent@cursor.com>
@Rian354 @cursoragent
Default to a full stay and stamp admission-context notes at admit.
8d4a4c9 
Will's protocol is through discharge, so the 24h class/runner default is
gone. Discharge-section text is still written at the end of the stay;
stamping it at charttime leaked length of stay. Radiology stays at exam
time. Cache version 4.

Co-authored-by: Cursor <cursoragent@cursor.com>
@Rian354

Rian354 commented Aug 19, 2026

Copy link
Copy Markdown
Collaborator

Pushed the rest of the measurement fixes onto this branch (fast-forward to 8d4a4c9).

Protocol matches the opening post: full stay (no 24h cap), no ICD, no re-weighting. Event times are hours from the first stay in the sample, so later admissions do not land on top of the first. Admission-context discharge sections are stamped at that stay's admit; radiology stays at exam time. Empty notes are still empty sequences, not a fake [MISSING_TEXT] token.

Also: AMP dtype is validated and autocast follows the trainer device; unified MLP gets the collate pad mask; lab z-scores fit on observed train rows; run directories are {task}_{model}_seed{seed}; run_config.json sits next to metrics; time embeddings no longer wrap every 30 days.

Task cache version is 4. Rebuild caches before comparing numbers.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

No reviews

Assignees

No one assigned

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

2 participants

@will-pang @Rian354