Skip to content

ML4H merge: tranche 1 - #1185

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

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

Conversation

@will-pang

@will-pang will-pang commented Aug 11, 2026

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
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>
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>
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>
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

Copy link
Copy Markdown
Collaborator

Pushed four proven repairs onto this branch (fast-forward from 2b7b2c5, no rebase).

  1. 9782aca Stop emitting fake missing-event placeholders ([MISSING_TEXT] / pad visits / black frames). Empty modalities are zero events.
  2. 869ac8e Keep frozen text encoders in eval when Trainer calls train(), so BERT dropout cannot fire on a frozen encoder.
  3. 11beefc Cache frozen [CLS] embeddings keyed on real tokens, not padded rows (padded-row keys never hit under shuffle).
  4. 86fe73b Restore BaseDataset._scan_table / _scan_parquet. MEDS still calls _scan_parquet after a0f1422 deleted them. Proof: tests/test_p0_parquet_scan.py reads a real two-row parquet file.

Still local, not on this PR yet: pad-mask / fp16 fill, nested padding_idx=0, sunlab resized_images + cache write.

Rian354 and others added 4 commits August 18, 2026 09:52
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>
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>
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>
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

Copy link
Copy Markdown
Collaborator

Pushed four more proven repairs (fast-forward 86fe73bb6fb31e). Each commit has a unit test whose docstring records the measured numbers and the pytest repro.

  1. 965a87b Fill attention masks with finfo(dtype).min and use fused SDPA. -1e9 is outside fp16, so AMP overflowed on padded positions. Fused vs explicit max abs diff: 7.153e-07 (no padding), 4.768e-07 (with padding). A10 notes_labs transformer 128/2/4, 2 epochs, batch 8: bf16 5,275 s vs fp32 10,198 s (epoch-1 2.84×, mean 2.41×, VRAM -24.5%, final train loss within 0.63%). Proof: tests/test_p1_fp16_attention.py.
  2. 7c4c056 Record {field}__pad_mask and skip those slots in the unified sequence. The collator padded with 0.0 and nothing marked it, so padded slots looked like real measurements at admission time. Token-budget notes crashed tensor a (7) must match b (14); reusing event pad_mask as a BERT token mask crashed shape '[96, 512]'. RNN packed lengths now clamp at 1. Proof: tests/test_p1_pad_masks.py.
  3. ca9b63d Keep padding_idx=0 on nested code embeddings. Pad row is zeros and receives zero gradient. Proof: tests/test_p2_nested_padding_idx.py.
  4. b6fb31e Accept resized_images and write sunlab CXR metadata to cache. Complete resized cohort 377,110 / 377,110. Hardcoded images/ was FileNotFoundError; default CXR config was KeyError: 'studytime_normalized'; read-only root now writes the CSV under cache_dir. After the layout worked (one seed, 6 epochs, 18,542 / 2,285, prevalence 0.0565): cxr_only PR-AUC 0.0602, +labs 0.3082, +notes+labs 0.4096. Those three compare with each other, not with the primary notes_labs table. Proof: tests/test_p2_sunlab_cache.py.

Repro for all four:

PYTEST_DISABLE_PLUGIN_AUTOLOAD=1 PYTHONPATH=. python -m pytest \
  tests/test_p1_fp16_attention.py tests/test_p1_pad_masks.py \
  tests/test_p2_nested_padding_idx.py tests/test_p2_sunlab_cache.py -q

12 passed on this checkout.

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

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants