Multimodal SSL pretraining (MAE, SimMIM, I-JEPA, V-JEPA) with multi-arch backbones

[Rian354]

Summary

• New: pyhealth/models/pretrain/, a self-supervised pretraining stack for the multimodal clinical-sequence encoder, with four SSL objectives sharing one unified embedding and mask generator
• New: build_backbone(arch) selects transformer / jamba / mamba behind a single (x, mask) -> (emb, cls) contract, so every objective works with every backbone
• New: PretrainTrainer, DDP-native with atomic checkpoint resume, so preempted cluster jobs continue instead of restarting from zero
• New: scripts/pretrain_ssl.py plus SLURM/Condor runners for full-scale 50-epoch runs at tuned hyperparameters
• Encoders are standardized at 128-dim / 2-layer / 4-head so weights transfer 1:1 into the downstream backbone

Backbones

arch	implementation	notes
transformer	TransformerLayer	optional RoPE (use_rope, rope_scaling for NTK-aware extrapolation)
mamba	MambaLayer (new)	stacks MambaBlock + RMSNorm; pads masked positions to 0 so the causal conv cannot leak across the mask
jamba	JambaLayer	interleaves attention and SSM layers; the mix is set by num_transformer_layers / num_mamba_layers

SSL objectives

method	mechanism
MAE	TransformerLayer decoder reconstructs masked tokens; target='token' (content-only, per-embedding normalized) or 'unified'
SimMIM	full-sequence encode + learnable mask tokens + linear head, no decoder, so cheaper than MAE
I-JEPA	context encoder + EMA target encoder (frozen), location-aware predictor, cosine EMA schedule
V-JEPA	I-JEPA plus multi-scale target blocks (2,4,8) with a scale_embed, and per-scale loss breakdown

UnifiedMaskGenerator supports random and block masking with per_modality_ratio overrides; block masking preserves contiguity and enforces a minimum block length.

PretrainTrainer

• DDP-native, auto-detecting RANK / WORLD_SIZE from torchrun; AMP and gradient accumulation; EMA hook every ema_update_every optimizer steps for the JEPA variants
• IterableDataset-aware: skips DistributedSampler for litdata streams, which already shard internally, avoiding double-sharding
• Atomic resume: writes last.ckpt (weights) and _resume.pt (optimizer/scheduler/epoch/step) atomically, and resumes when both exist. If only last.ckpt survives, it warm-resumes from the weights and infers the epoch from metrics_history.json rather than restarting at zero
• Per-modality loss tracking (modality_0/1/..., total, scale_*) to metrics_history.json and W&B, main process only
• Optional val_dataloader and epoch_callback for held-out scoring and pruning, used by the pretraining Optuna sweep

Scripts and runners

• scripts/pretrain_ssl.py: --arch / --method / --task, standardized 128/2/4 defaults, exp_name = {arch}_{method}_{task}_seed{seed}. --task notes_only routes through NotesLabsMIMIC4(include_labs=False) for a leakage-free notes-only variant
• scripts/run_full_pretrain.py: replays tuned hyperparameters from a best_params JSON, launches env-resolved torchrun rather than a bare PATH lookup, and accepts --extra to override tuned values (for example --batch-size 32 --grad-accumulation-steps 2 to fit a smaller GPU at the same effective batch size)
• scripts/run_full_pretrain_local.sh / run_fullpt_condor.sh / scripts/slurm/full_pretrain_cc.sh: skip an encoder only when metrics_history.json shows a full 50 epochs, so partial runs resume instead of being skipped. Sets PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True to avoid fragmentation OOM
• scripts/kill_ddp.sh: top-down kill cascade, needed because torchrun restarts killed workers

W&B

pyhealth/_wandb.py is env-gated on WANDB_PROJECT and swallows all exceptions, so tracking never takes down training. Runs are named {arch}_{method}_{task}_seed{N} and tagged kind:pretrain, stage:pretrain, bb:{arch}, mod:{task}, method:{method}. Metrics are namespaced into val/, test/, best/, loss/, sys/ so the dashboard groups into a few sections instead of dozens of flat keys.

Depends on

The foundation edits (UnifiedMultimodalEmbeddingModel with text_finetune_mode, TransformerLayer, MambaBlock / JambaLayer). Land those first, or fold them in. The Optuna PR sits on top of this one.

Testing

tests/test_pretrain.py and tests/test_pretrain_backbones.py, roughly 680 lines covering masking (random and block, contiguity, floor enforcement), MAE / SimMIM / I-JEPA / V-JEPA forward, backward and loss, RoPE extrapolation, multi-scale sampling, EMA schedule, token vs unified targets, checkpoint loading into the downstream Transformer, and per-modality decoding, plus regressions for degenerate-prediction and empty-batch edge cases. Backbone contract (shape, backward, padding isolation) is validated for all three architectures.

Full 50-epoch pretraining has been run end to end on MIMIC-IV across the architecture x method matrix for three modality combinations (notes_labs, notes_only, labs_only), 26 of 27 encoders complete at the time of writing, with healthy per-modality validation losses and no representation collapse.