feat(datasets): oversample and weighted resampling strategies

[Rian354]

PR: Additional Sampling Strategies for Class-Imbalanced EHR Data

Problem

MIMIC-IV in-hospital mortality is heavily imbalanced (~10% positive rate). PR #28 added sample_balanced() (undersampling) as the first strategy. Two additional strategies are worth comparing:

Strategy	Mechanism	Tradeoff
Undersample	Drop majority-class samples	Loses negative diversity, faster epochs
Oversample	Duplicate minority-class samples w/ replacement	Full negative diversity, risk of positive overfitting (identical BERT embeddings repeat)
Weighted resample	Dataset-level class-proportional resampling	Full diversity in both directions, no external sampler needed

Note: WeightedRandomSampler (PyTorch) is incompatible w/ litdata.StreamingDataset (IterableDataset). The weighted strategy uses dataset-level resampling via sample_weighted() instead of an external sampler.

Changes

pyhealth/datasets/splitter.py

1. sample_oversample(dataset, ratio=1.0, seed=None, label_key="label")

   • Keeps ALL negative samples (no majority-class information is discarded).
   • Duplicates positive samples w/ replacement until n_pos ~= n_neg / ratio.
   • If the current positive count already satisfies the target ratio, returns the dataset unmodified.
   • Rebuilds patient_to_index and record_to_index on the oversampled result.
   • Mirrors sample_balanced() API exactly.

2. sample_weighted(dataset, seed=None, label_key="label")

   • Single pass to collect labels and class frequencies.
   • Assigns p[i] = 1 / class_count[label[i]] to each sample, normalized to a probability distribution.
   • Resamples the dataset w/ replacement using numpy (compatible w/ IterableDataset).
   • Returns a SampleDataset.subset() of the same length as the original.
   • Raises ValueError if a class is missing.

pyhealth/datasets/__init__.py

Exports: sample_oversample, sample_weighted.

examples/mortality_prediction/unified_embedding_e2e_mimic4.py

--sampling-strategy {none, undersample, oversample, weighted} (default: none)

--balanced-sampling remains as a legacy alias. If set and --sampling-strategy is still none, the strategy is promoted to undersample.

--balanced-ratio applies to both undersample and oversample (negatives per positive).

scripts/condor/run_e2e_full_scale.sh

Two new condition cases added:

• oversampled -> --sampling-strategy oversample --balanced-ratio 1.0
• weighted_sampler -> --sampling-strategy weighted

scripts/condor/batch_e2e_smoke.sub

Adds oversampled and weighted_sampler smoke jobs (MLP, seed=42, 5 epochs, dev=1000). Total smoke jobs: 4 -> 9.

Full-Scale Results (full MIMIC-IV, 50 epochs, seed=42, A100 80GB)

Condition	Model	Epochs	Best Ep	PR-AUC	ROC-AUC	F1
balanced (undersample)	mlp	21	10	0.1847	0.7828	0.1798
balanced (undersample)	transformer	50	49	0.1694	0.7733	0.1750
oversampled	mlp	25	14	0.2524	0.8582	0.2753
oversampled	transformer	14	3	0.1839	0.7803	0.1861
weighted_sampler	mlp	25	14	0.2170	0.8489	0.2535
weighted_sampler	transformer	25	14	0.1840	0.7722	0.1761
baseline (no resampling)	mlp	39	28	0.3206	0.8690	0.1571

Actual ranking: none >> oversampled > weighted > undersample

This is the reverse of the original hypothesis. Why:

• At full MIMIC-IV scale (train ~130k samples, 9.1% positive), balanced undersamples to ~14k samples, a 9x size reduction. The model sees far fewer negatives overall.
• oversample duplicates positives ~10x from ~12k to ~120k. W/ frozen BERT, all duplicated positives have identical embeddings. The model memorizes a small set of fixed vectors rather than generalizing. Transformer peaks early (ep3) and stagnates.
• weighted_sampler (dataset-level resampling) is functionally identical to oversample in expectation, w/ a similar ceiling.
• The unbalanced pos_weight=10 baseline trains on the full negative distribution w/ a scalar gradient correction and continues improving through ep28.

Takeaway: for MIMIC-IV mortality at full scale, pos_weight correction w/o resampling outperforms all three resampling strategies. Resampling may be more beneficial at smaller dataset sizes or w/ non-frozen encoders.

Files Changed

File	Change
pyhealth/datasets/splitter.py	+sample_oversample, +sample_weighted (replaces make_weighted_sampler)
pyhealth/datasets/__init__.py	exports for 2 new functions
examples/mortality_prediction/unified_embedding_e2e_mimic4.py	--sampling-strategy flag, routing logic
scripts/condor/run_e2e_full_scale.sh	oversampled + weighted_sampler condition cases
scripts/condor/batch_e2e_smoke.sub	4 -> 9 smoke jobs
scripts/condor/batch_e2e_weighted_sampler.sub	New. Standalone submit for weighted_sampler condition.

[will-pang]

looks good - will test a bit more later