Skip to content

Repository files navigation

PRIMED-AI

EchoJEPA + HuBERT-ECG: frozen-embedding, deployment-risk study of multimodal LVEF estimation

criticaldata/PRIMED-AI


Overview

PRIMED-AI studies whether fusing two cardiac foundation models — EchoJEPA-L (echocardiogram video) and HuBERT-ECG (12-lead ECG) — can estimate left ventricular ejection fraction (LVEF) in a way that is ready for real-world deployment, not just benchmark accuracy.

The pipeline uses frozen embeddings only: foundation models act as fixed feature extractors; all task-specific learning happens in lightweight probes on top of cached representations.

ECG is ubiquitous and cheap; echo requires a sonographer and a cart. If we train a fused model but only ECG is available at inference, how much accuracy do we lose — and does the model degrade gracefully or fail silently?


Motivation

LVEF is a core measure of heart pump function, used to diagnose heart failure and guide treatment. Echocardiography is accurate but resource-intensive; ECG is cheap, fast, and available almost everywhere.

Most multimodal ML papers ask whether fusion beats unimodal baselines on a leaderboard. This project asks the deployment-and-risk questions instead:

Question What we evaluate
Does it work? Missing-modality robustness when echo or ECG is unavailable at inference
Is it equitable? Fairness audit of EF error across sex, age, and race
Is it safe? Whether accuracy degrades gracefully or fails silently under shift

Fusion vs. unimodal accuracy is supporting evidence, not the thesis.


Pipeline

Stage Input Output
Cohort construction MIMIC-IV-Echo studies, MIMIC-IV-ECG records, MIMIC-IV demographics Paired rows on subject_id within 24–48h, joined to structured LVEF, split by subject
Embedding extraction Paired cohort only (~few K–tens of K studies, not all 525K echos) Frozen EchoJEPA-L (1024-d) and HuBERT-ECG (768-d) vectors, cached to Parquet
Probes Cached embeddings ECG-only, echo-only (attentive), concat-MLP, and cross-attention fusion heads
Deployment analyses Trained fused checkpoint + held-out test split Missing-modality degradation, fairness stratification, EF≤40% calibration

Expensive forward passes run once and are cached; probe training reads only from cache and completes in minutes. Foundation-model weights are never fine-tuned.

See TECHNICAL.md for full pipeline details.


Results

Held-out test split, n = 245 after dropping three non-finite HuBERT-ECG rows from the original 248-row test split. One fused cross-attention checkpoint (M09) trained on both modalities, scored under three inference-time conditions — no separate unimodal models are trained for the dropped conditions.

Condition Inference input LVEF MAE (95% CI) EF≤40% AUROC (95% CI)
full Echo + ECG present 10.41 (9.27–11.64) 0.771 (0.698–0.839)
ecg_dropped ECG branch masked, echo present 11.22 (10.00–12.40) 0.382 (0.279–0.493)
echo_dropped Echo branch masked, ECG present 20.54 (19.01–22.00) 0.689 (0.605–0.765)

Dropping echo costs much more MAE than dropping ECG. The ECG-dropped row is the important one: its AUROC is 0.382 with a 95% interval of 0.279–0.493, which excludes 0.5, so the EF≤40% ranking is not merely uninformative but inverted — the gate reads below chance. MAE barely moves in that condition (11.22 vs 10.41), so the regression output looks healthy while the clinical decision it feeds is wrong. That is a silent failure, not graceful degradation, and it is the pattern this repo exists to surface. This run is a pooled-manifest baseline: the echo branch still receives one mean-pooled vector per study, not retained clip tokens — and on pooled inputs the cross-attention block provably cannot attend across modalities, so these numbers describe a concat-equivalent model (see TECHNICAL.md §6.4).

Provenance: real run on cached EchoJEPA (vjepa2.1-vitl-mimic-pt-100) + HuBERT-ECG embeddings, local Mac CPU, seed 42, fusion width 256, 1,000 bootstrap resamples. The canonical fused checkpoint is the one scripts/train_probes.py writes to probes/fused/cross_attn_fused.pt, selected by validation MAE from the all-probe M10 run; manifest SHA-256 81694c9b…, checkpoint SHA-256 7cae4f92…. results/ and probes/ are gitignored (per-example predictions and checkpoints stay local), but sanitized aggregate copies of every result JSON plus the full checksums are committed under docs/results/ — verify a reproduction with shasum -c docs/results/SHA256SUMS. See CONTRIBUTING.md for what reproducing takes.

Reproduce with:

python scripts/evaluate_missing_modality.py \
  --manifest data/processed/echo_hubert_manifest.parquet \
  --checkpoint probes/fused/cross_attn_fused.pt \
  --embed-dim 256 \
  --echo-dim 1024 \
  --ecg-dim 768

Defaults assume the standard cohort/embedding/checkpoint layout; --help lists the paths. --embed-dim, --echo-dim, and --ecg-dim must match the checkpoint architecture or the state_dict load fails on shape.

In-cohort baselines

All four probes on identical splits (seed 42, same 245-row test frame, scripts/diagnose_baseline_gap.py):

Probe LVEF MAE (95% CI) EF≤40% AUROC (95% CI)
Cross-attention fused 10.41 (9.27–11.64) 0.771 (0.698–0.839)
Concat-MLP 10.93 (9.77–12.22) 0.716 (0.634–0.795)
Echo-only (attentive) 11.09 (9.79–12.25) 0.771 (0.698–0.845)
ECG-only (ridge) 11.60 (10.36–12.84) 0.671 (0.582–0.756)

Fusion beats ECG-only on both metrics with paired-bootstrap significance (ΔMAE −1.18, CI −1.99 to −0.33; ΔAUROC +0.100, CI +0.03 to +0.17) and is never behind either solo probe — the gap to published single-modality baselines is a cohort and label-regime property, not a fusion failure (TECHNICAL.md §8).

Fairness

Test-split stratification of the fused checkpoint (scripts/evaluate_fairness.py, same run as above; strata under n=20 are flagged in the artifact and omitted here):

Stratum n MAE full MAE echo_dropped AUROC full
Sex F 144 10.13 21.18 0.770
Sex M 101 10.82 19.62 0.764
Age 40–54 37 8.94 20.43 0.742
Age 55–64 57 9.44 23.66 0.727
Age 65–74 57 13.10 21.70 0.752
Age 75–89 77 10.26 18.28 0.788
Race: Black 43 6.61 21.39 0.973
Race: White 175 11.52 20.12 0.761

The sex gap is small in both conditions. The largest full-condition gaps are age 65–74 (13.10 vs 10.41 overall) and the Black–White MAE difference (6.61 vs 11.52). Note that MIMIC-IV records administrative gender and admission-reported race, both with known curation bias, and the remaining race/age strata are too small for stable estimates (flagged, not silently included, in results/fairness/).

Calibration

EF≤40% Platt scaler fit on val predictions only, applied to test (scripts/evaluate_calibration.py), 10 bins:

Condition ECE
full 0.072
echo_dropped 0.044
ecg_dropped 0.017

Post-scaling ECE stays low even with a modality dropped, and ecg_dropped has the lowest ECE of the three. That number is an artifact of collapse, not of quality. Its reliability bins put 234 of 245 test cases in the single 0.1–0.2 bin, where mean confidence (0.167) and observed rate (0.171) nearly coincide — so ECE is small because almost every case gets the same near-constant probability, not because the score discriminates. Against that condition's 0.382 AUROC (CI 0.279–0.493, excluding 0.5) the reading is: well-calibrated probabilities attached to a below-chance gate, with what ranking survives pointing the wrong way. Low ECE is not evidence the degraded model is usable — here it is close to the opposite. Bins are in docs/results/calibration_ecg_dropped.json.


Models & data

Component Source
EchoJEPA-L Video JEPA for echo · arXiv:2602.02603 · bowang-lab/EchoJEPA
Echo embeddings Pre-extracted, all 6 checkpoint variants · MITCriticalData/mimic-iv-echo-jepa-embeddings (gated)
HuBERT-ECG Self-supervised 12-lead ECG model · Edoardo-BS/hubert-ecg-baseproduced the reported ECG embeddings
ECG-FM Alternative ECG foundation model · arXiv:2408.05178 · wanglab/ecg-fm — wrapper shipped, not used for reported results
MIMIC-IV-Echo 0.1 Echo DICOM + structured LVEF · PhysioNet
MIMIC-IV-ECG 1.0 12-lead waveforms · PhysioNet
MIMIC-IV 3.1 Demographics for fairness audit · PhysioNet

All MIMIC datasets and the embedding dataset require PhysioNet credentialing and signed Data Use Agreements. docs/embeddings.md covers loading and re-extraction.

Which ECG encoder? The reported results use pre-extracted HuBERT-ECG embeddings. The repo also ships a frozen ECG-FM wrapper (encoders/ecg_fm.py, scripts/extract_ecg_embeddings.py) as an alternative extraction path — it is implemented and unit-tested but did not produce any reported number.

Scope: Task A — continuous LVEF regression + EF≤40% binary clinical gate (HFrEF threshold). Valvular disease and HFpEF phenotyping are out of scope.


Metrics

Metric Purpose
LVEF MAE Continuous regression quality
EF≤40% AUROC Clinical gate for reduced ejection fraction
Missing-modality degradation Δ MAE / Δ AUROC when echo or ECG is dropped at inference
Fairness gap Per-stratum error across sex, age band, and race

Repository guide

Document Contents
TECHNICAL.md Pipeline architecture, probes, deployment analyses, expected artifacts
docs/embeddings.md Embedding sources, loading code, model registry, re-extraction
CONTRIBUTING.md Local setup, tests, linting, data access, reproducibility

Open work lives on the issue tracker.


Status

The work pivoted from a fusion benchmark to a model-agnostic modality-failure analysis framework — failure taxonomy, complementarity matrix, and loud-vs-silent dropout profile. That framework is the main contribution and lives in src/primed_ai/failure/; the fusion pipeline described above is the substrate it is instantiated on.

The canonical pooled rerun is done — probe training and every deployment analysis run on a laptop CPU in about 90 seconds total, no GPU (see TECHNICAL.md §9). Still open: rebuilding the manifest with clip-level echo tokens (needs the gated HF shards), token-level ECG re-extraction (#72), and the cohort-expansion decision (#62) — the first two gate real cross-modal attention, the last gates statistical power.


Related work

Work Relationship
EchoJEPA Echo foundation model used in this pipeline
HuBERT-ECG ECG foundation model behind the reported ECG embeddings
ECG-FM Alternative ECG foundation model — wrapper shipped, not used for reported results
EchoingECG Closest prior cross-modal echo+ECG work — differentiated on frozen-embedding fusion + deployment-risk framing

Citation

If you use this work, please cite the repository:

@misc{primedai2026,
  title        = {PRIMED-AI: Deployment-Risk Study of Multimodal LVEF Estimation with EchoJEPA and HuBERT-ECG},
  author       = {Critical Data},
  year         = {2026},
  howpublished = {\url{https://github.com/criticaldata/PRIMED-AI}}
}

License

MIT — see LICENSE.

scripts/embedding_extraction/src/ is vendored V-JEPA code from Meta Platforms and stays under its own MIT license (see scripts/embedding_extraction/src/LICENSE).


PRIMED-AI · Critical Data

About

A framework for per-modality failure analysis and missing-data evaluation in multimodal clinical AI.

Topics

Resources

Contributing

Stars

8 stars

Watchers

0 watching

Forks

Releases

Packages

Contributors

Languages