[I08] Right-size compute for probe training and evaluation

[duckyquang]

Task metadata

Field	Value
Task ID	I08
Category	Infrastructure
Priority	Low
Complexity	Low
Est. days	1
Replaces	#15 (closed — assumed a full-corpus extraction that is no longer needed)

Features / method

Right-size compute. The GPU reservation the closed issue asked for is no longer on the critical path.

Description

The closed version reserved an H200 through "Day 6" to run embedding extraction over the paired cohort. That premise no longer holds, and it was quietly blocking the wrong things.

Extraction is done. EchoJEPA embeddings were pulled pre-extracted from Hugging Face and HuBERT-ECG embeddings came across as Parquet — see docs/echo_hubert_loader.md. Nothing in #62-#68 re-runs a foundation model.

What remains is probe training and evaluation on a 1,208-row manifest of cached vectors. CrossAttnFusedProbe is two linear projections, one cross-attention block, and a 256-hidden MLP head; scripts/train_ecg_probe.py does not even expose a device flag. TECHNICAL.md 9 already says probe training after caching is CPU/GPU-light and completes in minutes. #63 does not need a GPU reservation and should not wait on one.

So this issue shrinks to confirming the two cases where compute actually matters:

Steps:

1. Confirm probe training and evaluation for [M10] Retrain all four probes on the synchronized manifest #63-[W07] Restate reported numbers from the canonical run #68 run acceptably on local or modest shared CPU. Record actual wall-clock so the assumption is measured rather than asserted.
2. Storage: size the cached embedding + checkpoint footprint and confirm a writable location. This is the part of the closed issue that is still real.
3. Identify the one case that does need real GPU — if [D07] Paired cohort recovery — 4,103 shared subjects to 1,208 rows #62 concludes the cohort can be materially expanded, more echo studies must be pooled from the EchoJEPA shards, and that is a GPU job. Scope it only if [D07] Paired cohort recovery — 4,103 shared subjects to 1,208 rows #62 says yes.
4. If step 3 triggers, reserve then, sized to the actual expansion rather than to the full 525K corpus.

Acceptance criteria:

• Measured wall-clock for a full probe train + eval cycle, on stated hardware.
• Writable storage path for checkpoints and cache documented.
• Explicit statement that [M10] Retrain all four probes on the synchronized manifest #63-[W07] Restate reported numbers from the canonical run #68 are unblocked on compute, or evidence to the contrary.
• GPU reserved only if [D07] Paired cohort recovery — 4,103 shared subjects to 1,208 rows #62 requires re-pooling, and sized to that.

[duckyquang]

Measured. #63–#68 are unblocked on compute — no GPU reservation needed.

Provenance, stated plainly: this is a hardware timing benchmark, not a training run. I built a parquet with the real manifest geometry (n=1208, echo 1024-d, ECG 768-d, splits 829/131/248) filled with random tensors, and timed a full train + eval cycle for each probe at its default epoch count. Wall-clock depends on tensor shapes and model size, not on what the numbers mean, so the timings transfer. No MAE or AUROC from this run is meaningful and none is reported — those still require the real manifest via #63.

Hardware: Apple silicon (arm), CPU only, torch 2.12.1, cuda=False, 8 threads.

ecg    (sklearn, 2000 bootstrap)     3.2s
echo   (30 epochs)                   4.3s
concat (40 epochs)                   4.2s
fused  (50 epochs)                  44.3s
-----------------------------------------
ALL FOUR (train + eval, CPU)        56.0s

The entire M10 retrain is under a minute on a laptop CPU. The original framing of the closed #15 reserved an H200 through "Day 6" for this. That reservation was for embedding extraction, which is already done — nothing in #62–#68 runs a foundation model.

Storage, measured at the same geometry:

manifest parquet          9.5 MB   (raw embedding bytes: 8.7 MB)
cross_attn_fused.pt      37   MB   <- dominated by 1024-d projections
concat_mlp.pt             1.8 MB
echo_only.pt            520   KB
ecg_only.joblib         396   KB
--------------------------------
all four checkpoints     40   MB

Total footprint for a full retrain is well under 100 MB. Any writable scratch works; this does not need a storage allocation.

Remaining acceptance item is step 3 — GPU is only in play if #62 concludes the cohort can be materially expanded, since re-pooling more echo studies from the EchoJEPA shards is a real GPU job. Scope that only if #62 says yes.

One caveat for whoever runs the real thing: the fused probe's 44s is with embed_dim=1024, which makes the projections identity and the attention wide. If #69 lands clip-level echo tokens, the token dimension grows and this number will rise — still nowhere near needing a reservation.

[duckyquang]

Rescoping. Two updates, pulling in opposite directions.

Probe training genuinely needs no GPU — now confirmed on real data. @kevzho ran all four probes against the real 1,208-row manifest: 49s / 81s on CPU for train+eval. That matches the shape-matched benchmark I ran and settles the original question. The closed #15 was reserving an H200 for an extraction that had already happened.

But the GPU ask returns for #72. Kevin also confirmed there are no pre-pooling HuBERT-ECG hidden states on disk, so getting token-level ECG — the thing that would make the fused probe's echo_to_ecg direction work at all — means re-running HuBERT-ECG over the MIMIC-IV-ECG waveforms. That is a GPU job, and its output is much larger than one pooled vector per record.

So this issue is no longer "does probe training need compute" (answered: no). It's sizing the re-extraction in #72:

1. Probe training + evaluation for [M10] Retrain all four probes on the synchronized manifest #63–[W07] Restate reported numbers from the canonical run #68: resolved, CPU is sufficient, measured twice.
2. Storage for the current cached embeddings and checkpoints: still worth documenting.
3. New: GPU + storage estimate for token-level ECG re-extraction over the paired cohort — record count, token granularity, resulting footprint. Blocks [M11] Fused probe has no cross-modal interaction — one direction still dead #72.

Revised acceptance criteria:

• Statement that [M10] Retrain all four probes on the synchronized manifest #63–[W07] Restate reported numbers from the canonical run #68 are unblocked on compute — met, with numbers from two independent runs.
• Writable storage path for checkpoints and cache documented.
• A sized GPU + storage ask for [M11] Fused probe has no cross-modal interaction — one direction still dead #72's re-extraction, or a decision not to pursue it.

[duckyquang]

Closing the remaining acceptance criteria: storage footprint, measured.

Checkpoints. Instantiated at the real encoder dims (echo 1024, ECG 768) and serialized:

checkpoint	size	params
echo_only (attentive)	0.53 MB	0.13M
echo_only linear baseline	0.01 MB	0.00M
concat_mlp	1.84 MB	0.46M
cross_attn @ --fusion-dim 256	4.48 MB	1.12M
all four, one seed	6.86 MB
cross_attn @ --fusion-dim 1024	38.87 MB	9.72M

This independently reproduces @kevzho's 1.1M vs 9.7M figure from #73 — 1.12M and 9.72M. Checkpoints are negligible: even ten seeds of all four probes is under 70 MB.

Embeddings. The real cost, and it scales with --max-clips. Payload arithmetic at float32 (echo studies = 7,251, joined cohort = 1,208 rows):

clips/study	echo embeddings parquet	joined manifest
pooled (1)	30 MB	9 MB
8	238 MB	43 MB
16	475 MB	83 MB
32	950 MB	162 MB

Measured peak RSS during a clip-level build at 7,251 studies × 16 clips × 1024 dim was 1,807 MB against a 476 MB payload, after the write-stage fix in #73 (it was ~7 GB before). Rule of thumb documented in the builder docstring: payload is n_studies × max_clips × dim × 4, peak RSS roughly 4×.

Practical read: the pooled manifest is ~10 MB and fits anywhere. --max-clips 16 puts the echo parquet at ~476 MB and the joined manifest at ~79 MB — still laptop-scale. Nothing here needs cluster storage.

Status against the acceptance criteria:

• ✅ Measured wall-clock on stated hardware — 49s/81s CPU on the real 1,208 rows (@kevzho, Retain clip-level echo tokens so attentive pooling actually pools #73), and a shape-matched 56s benchmark independently.
• ✅ Storage footprint documented, above.
• ✅ [M10] Retrain all four probes on the synchronized manifest #63–[W07] Restate reported numbers from the canonical run #68 unblocked on compute. Probe training does not need a GPU.
• ➡️ GPU sizing for token-level ECG re-extraction moved to [M11] Fused probe has no cross-modal interaction — one direction still dead #72, where it belongs.

The only thing left is a writable path decision, and that is a deployment choice rather than a measurement — configs/paths/default.yaml still carries /path/to/mimic placeholders, which #70 covers. I'd close this once someone fills those in.

[duckyquang]

Filled in the remaining compute/storage measurements from the real pooled run.

Real-data commands were the M10 all-probe train/eval and the E07 missing-modality rerun, both on local Mac CPU:

job	wall-clock	notes
all four probes, --fusion-dim 256	44.84s	train + eval on the 1,208-row pooled manifest
missing-modality eval, 1000 bootstraps	5.71s	from the pooled fused checkpoint

Storage measured locally:

artifact	size
pooled manifest	14 MB
all four checkpoints + summaries	7.0 MB
fused checkpoint	4.3 MB
missing-modality JSON + figures	216 KB

So #63-#68 do not need GPU or cluster storage for the cached-vector path. The current manifest and one-seed checkpoint set are laptop-scale.

I also updated TECHNICAL.md on the PR branch so the resource table says the same thing: no GPU for cached-vector probe training; GPU/storage sizing only returns for #72's token-level ECG re-extraction or any new EchoJEPA extraction.

[duckyquang]

Also reflected the real compute/storage conclusion in tracked docs on PR #73.

README now reports the canonical pooled run as a real CPU run with seed/fusion width/bootstrap provenance. TECHNICAL.md now says cached-vector probe training does not need GPU or cluster storage; GPU is only back in scope for #72 token-level ECG re-extraction or new EchoJEPA extraction.

Local artifact sizes after promoting the canonical pooled run:

pooled manifest                         14 MB
all four probe outputs under probes/     7 MB
fused checkpoint                         4.3 MB
missing-modality JSON + figures          < 1 MB

So this issue's original compute concern is handled for #63-#68. The remaining GPU/storage work belongs to #72.