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

configs/pretrain/base.yaml

@@ -0,0 +1,70 @@
+# Base config for SSL pretraining on unified multimodal sequences.
+# Use with scripts/pretrain_ssl.py.
+
+method: mae
+
+# Encoder backbone: transformer | jamba | mamba.
+arch: transformer
+
+# Model — standardized compute: 128-dim, 2 layers, 4 heads (head_dim 32),
+# matching the downstream e2e backbone so SSL encoders transfer 1:1.  Bumped up
+# from 64/1 (too narrow: V-JEPA drifted, downstream sat near prevalence); VRAM is
+# BERT-bound so the larger backbone is effectively free on 48 GB RTX 6000 Ada.
+embedding_dim: 128
+heads: 4
+num_layers: 2
+dropout: 0.1
+
+# Mamba / Jamba backbone knobs (ignored by the transformer backbone).
+state_size: 16
+conv_kernel: 4
+jamba_transformer_layers: 1
+jamba_mamba_layers: 1
+
+# MAE decoder (ignored for SimMIM / I-JEPA). Discarded after pretraining, so it
+# does NOT count toward the deployed-encoder 128/2 standard — kept at 2 layers so
+# the reconstruction objective remains learnable.
+decoder_layers: 2
+decoder_heads: 2
+decoder_dim: 128
+norm_pix_loss: false
+
+# I-JEPA / V-JEPA predictor (ignored for MAE / SimMIM). Also discarded after
+# pretraining; 2 layers keeps latent prediction non-degenerate at 128-dim.
+predictor_layers: 2
+predictor_heads: 2
+predictor_dim: 128
+ema_decay: 0.996
+ema_end: 1.0
+num_target_blocks: 4
+target_block_len: 4
+
+# Masking
+mask_ratio: 0.5
+mask_strategy: block
+lab_mask_ratio: null
+text_mask_ratio: null
+
+# Training
+epochs: 50
+batch_size: 32
+lr: 0.0001
+weight_decay: 0.05
+max_grad_norm: 1.0
+scheduler: cosine
+warmup_steps: 1000
+save_every_n_epochs: 5
+
+# Task
+observation_window_hours: 24
+note_extraction: regex
+note_source: discharge
+icd_codes: false
+include_vitals: false
+tokenizer_model: null
+freeze_encoder: false
+text_finetune_mode: full
+
+# Data / compute
+num_workers: 4
+seed: 42

configs/pretrain/ijepa_labs_only.yaml

@@ -0,0 +1,7 @@
+_inherit: base.yaml
+
+method: ijepa
+task: labs_only
+num_target_blocks: 4
+target_block_len: 4
+output_dir: ./output/pretrain/ijepa_labs_only

configs/pretrain/ijepa_notes_labs.yaml

@@ -0,0 +1,7 @@
+_inherit: base.yaml
+
+method: ijepa
+task: notes_labs
+num_target_blocks: 6
+target_block_len: 3
+output_dir: ./output/pretrain/ijepa_notes_labs

configs/pretrain/mae_labs_only.yaml

@@ -0,0 +1,7 @@
+_inherit: base.yaml
+
+method: mae
+task: labs_only
+mask_ratio: 0.5
+lab_mask_ratio: 0.5
+output_dir: ./output/pretrain/mae_labs_only

configs/pretrain/mae_notes_labs.yaml

@@ -0,0 +1,8 @@
+_inherit: base.yaml
+
+method: mae
+task: notes_labs
+mask_ratio: 0.5
+lab_mask_ratio: 0.5
+text_mask_ratio: 0.25
+output_dir: ./output/pretrain/mae_notes_labs

configs/pretrain/simmim_labs_only.yaml

@@ -0,0 +1,7 @@
+_inherit: base.yaml
+
+method: simmim
+task: labs_only
+mask_ratio: 0.5
+lab_mask_ratio: 0.5
+output_dir: ./output/pretrain/simmim_labs_only

configs/pretrain/simmim_notes_labs.yaml

@@ -0,0 +1,8 @@
+_inherit: base.yaml
+
+method: simmim
+task: notes_labs
+mask_ratio: 0.5
+lab_mask_ratio: 0.5
+text_mask_ratio: 0.25
+output_dir: ./output/pretrain/simmim_notes_labs

configs/pretrain/vjepa_labs_only.yaml

@@ -0,0 +1,9 @@
+_inherit: base.yaml
+
+method: vjepa
+task: labs_only
+num_target_blocks: 4
+target_block_scales: [2, 4, 8]
+require_multimodal_blocks: false
+normalize_targets: true
+output_dir: ./output/pretrain/vjepa_labs_only

configs/pretrain/vjepa_notes_labs.yaml

@@ -0,0 +1,11 @@
+_inherit: base.yaml
+
+method: vjepa
+task: notes_labs
+# More, multi-scale blocks for the longer multimodal sequence; prefer windows
+# that span both labs and notes so the model learns cross-modal prediction.
+num_target_blocks: 6
+target_block_scales: [2, 4, 8, 16]
+require_multimodal_blocks: true
+normalize_targets: true
+output_dir: ./output/pretrain/vjepa_notes_labs

pyhealth/_wandb.py

@@ -0,0 +1,116 @@
+"""Thin, opt-in Weights & Biases logging helper.
+
+Everything here is a no-op unless ``WANDB_PROJECT`` is set in the environment
+(and ``wandb`` is importable), so importing/using it never breaks a run that
+doesn't want tracking. Entity/project/mode come from the standard W&B env vars:
+
+    WANDB_PROJECT   e.g. pyhealth-multimodal   (required to enable logging)
+    WANDB_ENTITY    e.g. pyhealth-multimodal   (team; optional)
+    WANDB_MODE      online (default) | offline  (offline for no-internet nodes;
+                                                 sync later with `wandb sync`)
+
+Failures (no network, bad key, wandb missing) degrade to a warning + no-op —
+tracking must never take down training.
+"""
+from __future__ import annotations
+
+import os
+from typing import Any, Dict, Optional
+
+
+def enabled() -> bool:
+    return bool(os.environ.get("WANDB_PROJECT"))
+
+
+def sweeps_project() -> Optional[str]:
+    """Dedicated project for Optuna sweeps, kept separate from the real experiment
+    runs so the main project isn't cluttered by per-trial study runs. Override with
+    WANDB_PROJECT_SWEEPS; otherwise it's ``<WANDB_PROJECT>-sweeps``."""
+    p = os.environ.get("WANDB_PROJECT")
+    return os.environ.get("WANDB_PROJECT_SWEEPS") or (f"{p}-sweeps" if p else None)
+
+
+def init_run(config: Optional[Dict[str, Any]] = None, name: Optional[str] = None,
+             group: Optional[str] = None, job_type: Optional[str] = None,
+             project: Optional[str] = None, tags: Optional[list] = None):
+    """Start a W&B run, or return None if disabled/unavailable."""
+    if not enabled():
+        return None
+    try:
+        import wandb
+    except ImportError:
+        print("[wandb] WANDB_PROJECT set but wandb not installed — skipping tracking.")
+        return None
+    try:
+        return wandb.init(
+            project=project or os.environ["WANDB_PROJECT"],
+            entity=os.environ.get("WANDB_ENTITY"),
+            name=name,
+            group=group,
+            job_type=job_type,
+            config=config or {},
+            tags=[t for t in (tags or []) if t],
+            reinit=True,
+        )
+    except Exception as e:  # network/auth/etc. — never fail the run over telemetry
+        print(f"[wandb] init failed ({e}) — continuing without tracking.")
+        return None
+
+
+# Canonical metric namespaces so W&B groups panels into a few tidy sections
+# (val/, test/, best/, loss/, sys/) instead of dozens of flat, redundant keys.
+# Applied at log()/summary() so every caller stays consistent for free.  Only
+# affects W&B display names; on-disk metrics_history.json keeps its raw keys.
+_CANON = {
+    "pr_auc": "val/pr_auc", "val_pr_auc": "val/pr_auc", "val_roc_auc": "val/roc_auc",
+    "val_f1": "val/f1", "val_accuracy": "val/accuracy", "val_loss": "val/loss",
+    "test_pr_auc": "test/pr_auc", "test_roc_auc": "test/roc_auc", "test_f1": "test/f1",
+    "test_accuracy": "test/accuracy", "test_loss": "test/loss", "test_test_loss": "test/loss",
+    "test_n": "test/n", "test_pos": "test/pos",
+    "best_pr_auc": "best/pr_auc", "best_val_pr_auc": "best/pr_auc", "best_epoch": "best/epoch",
+    "total": "loss/total", "train_loss": "loss/train",
+    "train_vram_peak_mb": "sys/vram_peak_mb", "train_vram_allocated_mb": "sys/vram_allocated_mb",
+    "learning_rate": "sys/lr", "epoch_time_s": "sys/epoch_time_s", "global_step": "sys/global_step",
+}
+
+
+def _canon_key(k: str) -> str:
+    if k in _CANON:
+        return _CANON[k]
+    if k.startswith("modality_") or k.startswith("scale_"):  # per-modality / multi-scale SSL loss
+        return "loss/" + k
+    if k.startswith("train_loss_"):                          # per-component supervised loss
+        return "loss/" + k[len("train_loss_"):]
+    return k  # already-namespaced (hp/, arch/, sweep/, epoch, trial, ...) pass through
+
+
+def _canon(record: Dict[str, Any]) -> Dict[str, Any]:
+    return {_canon_key(k): v for k, v in record.items()}
+
+
+def log(run, record: Dict[str, Any], step: Optional[int] = None) -> None:
+    if run is None:
+        return
+    try:
+        run.log(_canon(dict(record)), step=step)
+    except Exception:
+        pass
+
+
+def summary(run, record: Dict[str, Any]) -> None:
+    if run is None:
+        return
+    try:
+        for k, v in _canon(record).items():
+            run.summary[k] = v
+    except Exception:
+        pass
+
+
+def finish(run) -> None:
+    if run is None:
+        return
+    try:
+        run.finish()
+    except Exception:
+        pass

pyhealth/models/pretrain/__init__.py

@@ -0,0 +1,35 @@
+"""Self-supervised pretraining models for PyHealth multimodal sequences.
+
+Models:
+- :class:`MultimodalMaskedAutoencoder` — true MAE with a transformer decoder.
+- :class:`MultimodalSimMIM` — SimMIM-style linear-head reconstruction.
+- :class:`MultimodalIJEPA` — I-JEPA latent predictive architecture.
+- :class:`MultimodalVJEPA` — V-JEPA: multi-scale spans + location/scale-aware
+  per-position latent prediction (extends I-JEPA).
+
+Utilities:
+- :class:`UnifiedMaskGenerator` — masking strategies for unified event seqs.
+- :func:`apply_mask_token` — replace masked positions with a learnable token.
+"""
+
+from .jepa import MultimodalIJEPA, MultimodalVJEPA
+from .mae import MultimodalMaskedAutoencoder, PerModalityMAEDecoder
+from .masking import UnifiedMaskGenerator, apply_mask_token, random_mask_like
+from .rope import RoPEMultiHeadedAttention, RoPETransformerLayer, RotaryPositionEmbedding
+from .simmim import MultimodalSimMIM
+from .trainer import PretrainTrainer
+
+__all__ = [
+    "MultimodalMaskedAutoencoder",
+    "MultimodalSimMIM",
+    "MultimodalIJEPA",
+    "MultimodalVJEPA",
+    "PerModalityMAEDecoder",
+    "UnifiedMaskGenerator",
+    "apply_mask_token",
+    "random_mask_like",
+    "RotaryPositionEmbedding",
+    "RoPEMultiHeadedAttention",
+    "RoPETransformerLayer",
+    "PretrainTrainer",
+]