Siddharth63/medical_JEPA

Medical JEPA — Joint Embedding Predictive Architectures for Physiological Signals

A collection of self-supervised encoders pretrained on large-scale ICU, clinical, and wearable-sensor data. Each modality ships its own frozen, downstream-ready encoder. The waveform encoders (respiratory, PPG, ABP, ECG, EEG) follow an I-JEPA-style recipe (multi-block masking, EMA target encoder, smooth-L1 / cosine loss on normalized predictor and target embeddings). The motion encoders (accelerometer, 6-axis IMU) are trained with raw-vs-spectrogram / RelCon-style contrastive learning.

Modalities

Directory	Signal	Window	Encoder	Output	Pretraining source
respiratory/	RESP + Pleth (2 ch)	8 s @ 125 Hz	RespJEPA (d=256)	(B, 40, 256)	MC-MED + MIMIC-III (resp + pleth)
ppg/	Pleth (1 ch)	8 s @ 125 Hz	PPGJEPA (ViT-B, d=768, 12 layers)	(B, 40, 768)	MC-MED + MIMIC-III-Ext-PPG
abp/	ABP (1 ch)	8 s @ 125 Hz	PPGJEPA (ViT-B, d=768, 12 layers)	(B, 40, 768)	MIMIC-III ABP
ecg/	12-lead ECG (lead-invariant)	10 s @ 250 Hz	ECG-JEPA (d=768, 12 layers)	(B, 768) + per-lead	MIMIC-IV-ECG (continued from upstream ECG-JEPA)
eeg/	up to 32 ch (channel-tokenised)	30 s @ 255 Hz	EEG-JEPA Small (d=384, 8 layers, factorised channel×time attn)	(B, 32, 150, 384)	P2018 + NCH + HBN + I-CARE + TDBRAIN (3.5M windows)
accelerator/	3-axis accelerometer	3 s @ 50 Hz	ResNet1D morphology encoder (d=256)	(B, 256)	Raw/spectrogram contrastive, multi-source accelerometry
imu6/	6-axis IMU (accel + gyro)	3 s @ 50 Hz	ResNet1D morphology encoder (d=256)	(B, 256)	RelCon/SimCLR contrastive, 18 IMU datasets (2.9M windows)

The respiratory, PPG, and ABP encoders share a common signal contract:

• 1000 samples per window @ 125 Hz (8 s)
• 25-sample patches (200 ms) → 40 patches per window
• Per-patch embeddings (B, 40, D); pool (mean / CLS / attention) for downstream heads.

Repo layout

respiratory/encoder_best.pt   # context encoder + patch embed (downstream-ready)
respiratory/best.pt           # full training state (context + target + predictor + optim)
ppg/encoder_best.pt
ppg/best.pt
abp/encoder_best.pt
abp/best.pt

ecg/best_model.pth            # continued-pretrained ECG-JEPA on MIMIC-IV-ECG (lead-invariant)
ecg/ecg_jepa.py               # upstream ECG-JEPA model (MIT, Sehun Kim 2024)
ecg/pos_encoding.py           # 2D sincos positional encoding
ecg/lead_regimes.py           # 1/3/6/12-lead view definitions + helpers
ecg/extract_embeddings.py     # batch extraction: HDF5 dataset → HDF5 embeddings
ecg/LICENSE_upstream_MIT      # MIT license for the vendored ECG-JEPA files

eeg/final.pt                  # EEG-JEPA Small full training checkpoint
eeg/latest.pt                 # EEG-JEPA Small latest training checkpoint
eeg/modeling_eeg.py           # EEG-JEPA context encoder architecture + loaders
eeg/montage.py                # 217-entry channel-name vocabulary (10-20 / EGI / EOG)
eeg/train.log                 # pretraining log
load_eeg_encoder.py           # repo-root loader wrapper for the EEG encoder

accelerator/encoder.pt        # 3-axis accelerometer morphology encoder
accelerator/modeling_accelerator.py
accelerator/config.json
accelerator/preprocessor_config.json
accelerator/example_encode.py

imu6/encoder_best.pt          # 6-axis IMU compact/default encoder
imu6/encoder_wide.pt          # 6-axis IMU wide/transient alternate encoder
imu6/modeling_imu6.py         # IMU6 ResNet1D architecture + loaders
example_imu6_encode.py        # minimal IMU6 inference example

models/resp_jepa_model.py     # RespJEPA architecture
models/ppg_jepa_model.py      # PPGJEPA architecture (used by both PPG and ABP)

load_encoders.py              # loaders for resp / ppg / abp
load_ecg_encoder.py           # loader + 1/3/6/12-lead helpers for ECG
requirements.txt

For most downstream uses (linear probe, fine-tune, embedding extraction), use the encoder_best.pt files — they contain only the context encoder and patch embedder. The best.pt files are full training-state checkpoints intended for resuming pretraining.

Installation

pip install -r requirements.txt

huggingface-cli download Siddharth63/medical_JEPA --local-dir ./medical_JEPA
cd medical_JEPA

# smoke-test the resp / ppg / abp encoders
python load_encoders.py --repo_dir . --device cpu

Waveform encoders (respiratory / PPG / ABP)

import torch
from load_encoders import load_resp_encoder, load_ppg_encoder, load_abp_encoder

device = "cuda" if torch.cuda.is_available() else "cpu"

resp = load_resp_encoder("respiratory/encoder_best.pt", device)
ppg  = load_ppg_encoder("ppg/encoder_best.pt", device)
abp  = load_abp_encoder("abp/encoder_best.pt", device)

# 8 seconds of signal at 125 Hz
x_resp = torch.randn(4, 2, 1000, device=device)   # (B, [RESP, Pleth], 1000)
x_ppg  = torch.randn(4, 1, 1000, device=device)   # (B, Pleth, 1000)
x_abp  = torch.randn(4, 1, 1000, device=device)   # (B, ABP_normalized, 1000)

with torch.no_grad():
    e_resp = resp.encode(x_resp)   # (4, 40, 256)
    e_ppg  = ppg.encode(x_ppg)     # (4, 40, 768)
    e_abp  = abp.encode(x_abp)     # (4, 40, 768)

    # Attention maps per encoder layer (interpretability)
    attn = ppg.get_attention_maps(x_ppg)   # list of (4, n_heads, 40, 40)

Pooling for downstream tasks

emb_window = e_ppg.mean(dim=1)   # (B, D) — simple mean-pool
# or add a learned attention-pool / CLS token on top

Input preprocessing

Match the preprocessing used during pretraining. All three expect float32 input on the same device as the model.

• Respiratory: bandpass RESP (0.05–1 Hz) and Pleth (0.5–8 Hz); z-score per window, per channel.
• PPG: bandpass 0.5–8 Hz; z-score per window.
• ABP: clip to a plausible range (40–200 mmHg); z-score per window. Do not subtract a global mean — the encoder was trained on per-window standardised input.

ECG-JEPA (continued pretraining)

ecg/best_model.pth is a continued-pretrained checkpoint of ECG-JEPA (Kim et al., 2024) on MIMIC-IV-ECG, trained in mixed input mode so a single encoder serves both standard 12-lead hospital ECGs and reduced-lead (Lead-I, 3-lead, 6-lead) wearable devices. The mixed regime makes the encoder lead-count invariant — the same checkpoint runs in any of the four views below.

Signal contract

• Sampling rate: 250 Hz
• Window: 2500 samples (10 s) per lead, 12 leads
• Input shape: (B, 12, 2500) — leads outside the chosen view are zeroed
• Output: global (B, 768) and per_lead (B, 12, 768)

import torch
from load_ecg_encoder import load_ecg_encoder, apply_input_mode, normalize_per_lead

device = "cuda" if torch.cuda.is_available() else "cpu"
model  = load_ecg_encoder("ecg/best_model.pth", device)

# raw 12-lead ECG window: (B, 12, 2500) float32
sig = torch.randn(2, 12, 2500, device=device)
sig = normalize_per_lead(sig)                  # always per-lead z-score first

# Pick a view: "1lead", "3lead", "6lead", or "12lead"
sig_1lead = apply_input_mode(sig, "1lead")     # zeros all leads except I
sig_12    = apply_input_mode(sig, "12lead")    # no-op

with torch.no_grad():
    g_1lead, pl_1lead = model(sig_1lead)       # smartwatch / single-lead device
    g_12,    pl_12    = model(sig_12)          # standard 12-lead hospital ECG

Use global for whole-window classification heads and per_lead for lead-aware tasks. For inference at scale over a large HDF5 dataset, use ecg/extract_embeddings.py — it streams (B, 12, 2500) records, writes global and per_lead arrays to HDF5, and supports the same four views via --input_mode.

ecg/ecg_jepa.py and ecg/pos_encoding.py are vendored verbatim from the upstream MIT-licensed ECG-JEPA repo (see ecg/LICENSE_upstream_MIT). They depend on timm for DropPath and trunc_normal_.

EEG-JEPA (Small)

EEG-JEPA Small uses factorised channel × time attention (d=384, 8 layers, 6 heads). The downstream-ready context encoder is ~19M params; the full training model (context encoder + EMA target + predictor) is ~40M. It was trained from scratch on 3.5M 30-s windows across five datasets: adult sleep PSG (P2018), pediatric sleep PSG (NCH), pediatric task EEG (HBN, EGI 129-ch downsampled), ICU continuous EEG (I-CARE, post-cardiac-arrest), and adult resting-state psychiatric EEG (TDBRAIN). Trained on an NVIDIA DGX Spark (Blackwell GB10).

The encoder is I-JEPA-style: each channel is patch-embedded by a shared Conv1d(1, 384, kernel=51, stride=51), a learned per-channel embedding (217-entry montage vocab) is added as the only spatial position, and 8 factorised blocks alternate channel attention → patch attention → MLP. Temporal position lives only in the predictor during pretraining, so the context encoder carries no temporal positional embedding.

Signal contract

• Sampling rate: ~255 Hz (256 Hz sources are accepted; the loader fits length)
• Window: 7650 samples (30 s) per channel — 150 patches × 51-sample patches
• Channels: up to 32, channel-tokenised. Fewer channels → zero-pad the channel dim and pass ch_pad=True for padded slots.
• Input shape: (B, C, 7650) float32, per-channel z-scored
• Output: (B, C, 150, 384) — per-(channel, patch) tokens

import torch, numpy as np
from load_eeg_encoder import load_eeg_encoder, prepare_eeg_window

device = "cuda" if torch.cuda.is_available() else "cpu"
eeg    = load_eeg_encoder("eeg", device)          # loads eeg/final.pt (or eeg/latest.pt)

# Raw 30-s EEG window @ ~255 Hz, e.g. 8 channels
sig   = np.random.randn(8, 7650).astype(np.float32)
names = ["Fp1", "Fp2", "F3", "F4", "C3", "C4", "O1", "O2"]

# standardises channel names, robust z-scores, fits length, pads to 32 channels
x, ch_idx, ch_pad = prepare_eeg_window(sig, names)
x, ch_idx, ch_pad = x.to(device), ch_idx.to(device), ch_pad.to(device)

with torch.no_grad():
    tokens = eeg.encode(x, ch_idx, ch_pad)          # (1, 32, 150, 384)
    pooled = eeg.encode_pooled(x, ch_idx, ch_pad)   # (1, 384) — masked mean over real ch + patches

Pooling for downstream heads:

tokens      = eeg.encode(x, ch_idx, ch_pad)   # (B, C, 150, 384)
window_emb  = eeg.encode_pooled(x, ch_idx, ch_pad)   # (B, 384) whole-window
per_channel = tokens.mean(dim=2)                     # (B, C, 384) per-channel

Input preprocessing

prepare_eeg_window() handles channel-name standardisation, robust z-scoring, length-fitting, and channel padding for you. To match pretraining, feed it windows that are already:

• Bandpass 0.5–45 Hz (the full EDF→parquet pipeline in the training repo uses 0.3–45 Hz with mains notch — see eeg/modeling_eeg.py docstring for provenance)
• Resampled to ~255–256 Hz; the loader center-crops/pads the time axis to 7650
• Named per any dataset convention — eeg/montage.py maps names to the 10-20 vocab; unknown names fall back to UNK_IDX (still encoded, without channel-specific info)

Robust per-channel z-score (median + std, clipped to ±5 std) is applied by prepare_eeg_window unless you pass do_zscore=False.

Accelerometer morphology encoder (3-axis)

accelerator/encoder.pt maps a 3-second, 50 Hz, 3-axis accelerometry window into a 256-dimensional movement-morphology embedding. Treat it as a frozen short-window encoder.

input:  [B, 3, 150]   # axis order: x, y, z
output: [B, 256]      # L2-normalized

Quick start

import torch
from accelerator.modeling_accelerator import load_accelerator_encoder, encode_windows

model = load_accelerator_encoder("accelerator")

x = torch.randn(4, 3, 150)          # 4 windows, 3 axes, 3 s @ 50 Hz
z = encode_windows(model, x)

print(z.shape)        # torch.Size([4, 256])
print(z[0].norm())    # ~1.0

Input format & architecture

• 3-axis acceleration, 50 Hz, 3-second windows → 150 samples, shape [B, 3, 150], axis order x, y, z. Resample to 50 Hz if needed; keep units consistent across a dataset/deployment.
• Architecture: 1D ResNet, in-channels 3, widths [48, 96, 192, 384], layers [3, 4, 6, 3], kernel size 7, embedding dim 256, checkpoint key encoder.

Validation summary

Validated as a frozen feature extractor across fall, freezing-of-gait (FoG), Parkinsonian task, rehabilitation, gait, and long-stream fall-history datasets. In several cases the strongest model was a hybrid of frozen encoder embedding + DSP features.

Dataset / task	Main result
SisFall fall-vs-ADL, window level	encoder AUROC 0.940, balanced acc 0.861
SisFall fall-vs-ADL, hybrid encoder+DSP	AUROC 0.952, balanced acc 0.874
SisFall file/event level	hybrid file AUROC 0.997
Daphnet FoG, trunk/lower-back	encoder AUROC 0.847 vs random encoder 0.717
Daphnet FoG, ankle event level	hybrid AUROC 0.868, event recall 0.966
Kaggle Parkinson FoG, 200-file tdcsfog run	DSP AUROC 0.907, encoder 0.898, hybrid 0.905, random 0.836
Kaggle Parkinson FoG, ~0.5 FA/min	DSP recall 0.594, encoder 0.455, hybrid 0.549
LTMM full fall-history association, 71 subjects	hybrid AUROC 0.703, balanced acc 0.653
PD-BioStampRC 7-way task/body morphology	encoder balanced acc 0.481 vs DSP 0.435 vs random 0.325
PD-BioStampRC upper-limb bradykinesia vs tremor	hybrid balanced acc 0.627 vs DSP 0.593 vs random 0.575
UCI Physical Therapy, 8-way exercise recognition	hybrid balanced acc 0.563, encoder 0.506, random 0.456
UCI Physical Therapy, execution quality	hybrid balanced acc 0.436, encoder 0.397, random 0.360, chance 0.333
Clinical gait abnormal-vs-healthy, lower-back	encoder AUROC 0.913, balanced acc 0.830

Takeaways

• FoG is DSP-competitive/DSP-led (strong spectral structure); prefer DSP or hybrid features over encoder-only features for FoG-style alerting. The frozen encoder still clearly beats random controls.
• LTMM shows modest fall-history-associated morphology (hybrid best). This is association with prior fall history, not detection of timestamped falls. LTMM has no timestamped event labels, so it should not be used as event-recall validation, and a naive per-person Mahalanobis baseline produced too many false alerts (2.7–10.4 block alerts/day) to be product-ready.

Intended use & limitations

Good uses (research): fall/instability detection, FoG-like movement-breakdown detection, task/body-morphology and exercise-type morphology, template-vs-current movement comparison. Most useful combined with simple DSP features.

Not validated as: a standalone disease-diagnosis model; a Parkinson's/dystonia/healthy-control diagnostic classifier; a medication ON/OFF detector from arbitrary windows; a clinical-deterioration detector; a production-ready alerting model; or a personal baseline-shift detector. The last is an important future use case but requires continuous data with real timestamped event/shift labels and fixed false-alarm-rate reporting.

This encoder is accelerometer-only (3 input channels). Some validation datasets include gyroscope channels, but they were not used here — for accel+gyro use the separate imu6/ encoder below. Do not treat the two checkpoints as interchangeable.

Recommended downstream protocol

Use frozen encoder embeddings + DSP features, with subject- or file-held-out evaluation, a random-encoder control, and a permutation-label control. For alerting tasks, report event recall, latency, false alarms/min (or /hour), and recall at fixed false-alarm rates — do not rely on AUROC alone.

6-axis IMU morphology encoder (accel + gyro)

The imu6/ encoders map a 3-second, 50 Hz, 6-axis IMU window into a 256-dimensional movement-morphology embedding.

input:  [B, 6, 150]   # channel order: accel_x, accel_y, accel_z, gyro_x, gyro_y, gyro_z
output: [B, 256]      # L2-normalized

Two variants:

Checkpoint	Widths	Layers	Recommended for
imu6/encoder_best.pt	[48, 96, 192, 384]	[3,4,6,3]	Default — general movement morphology
imu6/encoder_wide.pt	[64, 128, 256, 512]	[3,4,6,3]	Transient/fall-sensitive alternate

Quick start

import torch
from imu6.modeling_imu6 import load_imu6_encoder, encode_imu6_windows

device = "cuda" if torch.cuda.is_available() else "cpu"

model = load_imu6_encoder("imu6", device=device)                 # compact/default
wide  = load_imu6_encoder("imu6", device=device, variant="wide") # wide alternate

x = torch.randn(4, 6, 150)   # [B, accel_xyz + gyro_xyz, 3 s @ 50 Hz]

with torch.no_grad():
    z = encode_imu6_windows(model, x)

print(z.shape)        # torch.Size([4, 256])
print(z[0].norm())    # ~1.0

Input preprocessing

Resample to 50 Hz, use 3-second windows shaped [B, 6, 150], channel order accel_x/y/z then gyro_x/y/z, consistent accel/gyro units within a dataset/deployment, float32.

Pretraining

RelCon / SimCLR-style contrastive learning on 2,901,383 clean 3-second windows from 18 datasets ([6, 150]), with dataset-temperature sampling (α=0.5).

Dataset	Windows	Dataset	Windows
extrasensory_top1m	1,000,000	opportunity	57,842
realworld	461,755	payload_imu	55,385
ltmm	358,587	usc_had	52,734
hhar	246,893	mhealth	48,582
comprehensive_imu	224,872	simul	32,855
pamap2	115,085	phys_ex_2024	32,355
ronin	89,156	motionsense	27,360
forth_trace	76,605	uci_har	10,299
		mobifall	9,555
		wisdm	1,463

Training runs: compact = imu6_relcon_v7_clean_temp05 (temperature 0.2, batch 1024); wide = imu6_relcon_v7_clean_wide_temp01 (temperature 0.1, batch 1024).

Diagnostics & compact-vs-wide comparison

Self-supervised diagnostics (compact / wide): pos_sim_mean 0.9987 / 0.9978, neg_sim_mean 0.2834 / 0.3190, top1_aug_retrieval 0.9925 / 0.9941, emb_dim_std_mean 0.0527 / 0.0513, emb_norm_mean 1.0 / 1.0.

Frozen downstream comparison:

Task	Compact	Wide
MobiFall fall-vs-ADL (subject-grouped)	bal-acc 0.9534, F1 0.9354, AUROC 0.9849	bal-acc 0.9563, F1 0.9427, AUROC 0.9859
RealWorld 8-class activity (proband-grp)	bal-acc 0.7296, F1 0.7286, AUROC 0.9564	bal-acc 0.7129, F1 0.7112, AUROC 0.9534

The compact encoder is the recommended default. The wide encoder is a useful alternate for transient/fall-sensitive applications but does not dominate on broad activity morphology, so it should not automatically replace compact.

Checkpoint contents

encoder_best.pt (resp / ppg / abp):

{
  "context_encoder": OrderedDict[str, Tensor],   # transformer block weights
  "patch_embed":     OrderedDict[str, Tensor],   # patch linear + pos_embed + channel embed
  "config":          dict,                       # dataclass fields used at training time
}

best.pt (resp / ppg / abp):

{
  "model":     state_dict,   # full PPGJEPA / RespJEPA (context + target + predictor)
  "optimizer": ...,
  "scaler":    ...,
  "scheduler": ...,
  "config":    dict,
  "args":      dict,
  "metrics":   dict,
  "epoch":     int,
}

ecg/best_model.pth:

{
  "encoder":        OrderedDict[str, Tensor],   # context encoder weights
  "target_encoder": OrderedDict[str, Tensor],   # EMA target encoder
  "predictor":      OrderedDict[str, Tensor],
  "optimizer":      ...,
  "epoch":          int,
  "args":           dict,
}

load_ecg_encoder consumes only the encoder key (per-lead pos_embed is regenerated deterministically via sincos), so the same checkpoint serves any lead view.

eeg/final.pt / eeg/latest.pt: {"model": OrderedDict, "step": int}, where model holds the full EEG-JEPA state dict (context_encoder.*, target_encoder.* EMA, and predictor.*). load_eeg_encoder reads only the context_encoder.* keys and rebuilds the config from their shapes.

accelerator/encoder.pt and imu6/encoder_best.pt / encoder_wide.pt: {"encoder": OrderedDict[str, Tensor]} ResNet1D weights (see config.json / modeling_*.py for the architecture config).

Training details

Pretraining recipes (data prep, masking, EMA schedule, optimiser) live in the source training repos:

• Resp / PPG / ABP: resp_jepa_dataset_pipeline.py, ppg_jepa_pipeline.py, abp_jepa_dataset.py in the upstream training repo.
• ECG: continued from Kim et al., ECG-JEPA, 2024, using the mixed regime (40% Lead-I, 25% 3/6-lead, 35% 12-lead).

License & data use

Pretraining data are derived from MIMIC-III, MIMIC-IV-ECG, MC-MED, and the public EEG/IMU datasets listed above. Users are responsible for complying with the respective PhysioNet and dataset data-use agreements when fine-tuning or redistributing downstream models. The vendored ECG-JEPA files are MIT-licensed (ecg/LICENSE_upstream_MIT).

Clinical disclaimer

These encoders are provided for research and development use only. They are not medical devices and must not be used for diagnosis, treatment decisions, emergency response, or clinical monitoring without additional validation, regulatory review, and appropriate human oversight.