[M10] Retrain all four probes on the synchronized manifest

[duckyquang]

Task metadata

Field	Value
Task ID	M10
Category	Model
Priority	High
Complexity	Medium
Est. days	3
Blocks	#64, #65, #66, #67, #68
Blocked by	nothing — see #71
Should inherit	#69 pooling decisions

Features / method

Retrain all four probes on the synchronized EchoJEPA + HuBERT manifest and produce a canonical checkpoint set.

Description

PR #61 landed a validated 1,208-row paired manifest (829/131/248 train/val/test, 702/100/201 subjects, zero subject leakage). No probe has been trained on it yet. Every number currently in the README comes from an earlier cohort build with n=245 test, so nothing reported today is reproducible against the manifest that is actually in the repo.

This does not need a GPU reservation. Extraction is already done and the probes are small heads over cached vectors — see #71. It can start now.

Retrain all four probes from src/primed_ai/probes/ against data/processed/echo_hubert_manifest.parquet via the loader in src/primed_ai/data/echo_hubert_dataset.py. Train the unimodal probes too, not just the fused one — the in-cohort ECG-only and echo-only numbers are what E10 needs to tell a cohort problem apart from a fusion problem.

Steps:

1. Wire the four probe training paths to the manifest loader rather than the old cohort files.
2. Train ECG-only (M06), echo-only (M07), concat-MLP (M08), and cross-attention fused (M09) on the same splits and seed.
3. Record val MAE and EF<=40% AUROC per probe; select the fused checkpoint on val, not test.
4. Persist checkpoints plus a run manifest (utils/run_manifest.py) capturing seed, embedding source, dims, and hyperparameters.
5. Confirm EF<=40% prevalence is reportable on all three splits (scripts/check_ef40_prevalence.py) before locking anything.

Acceptance criteria:

• Four checkpoints trained on the 1,208-row manifest, each with a run manifest recording its inputs.
• Val and test MAE + AUROC logged for all four probes in one table.
• Fused checkpoint selected on val and marked as the canonical one for E07-E09.
• --embed-dim, --echo-dim, --ecg-dim recorded so downstream state_dict loads do not fail on shape.

[duckyquang]

Step 1 done. Steps 2–5 blocked on data.

Landed (uncommitted, in the working tree):

• src/primed_ai/probes/manifest.py — reads the joined manifest, reports dims, flattens ECG vectors into ve* columns for the sklearn probe.
• echo_only / concat_mlp / cross_attn / ecg_only — embedding paths are now optional. Omit them and the first argument is read as a manifest. The old two-table path is unchanged, so nothing regresses.
• scripts/train_probes.py — --probe {ecg,echo,concat,fused,all}. Only ecg had a CLI before; the other three were library-only, so "retrain all four" had no entrypoint. Dims are read off the manifest rather than passed as flags, which removes the state_dict shape-mismatch failure the README warns about.
• tests/test_probe_manifest.py — 11 tests, all four probes train from a manifest. Suite is 65 passing.

Half-specified embedding paths (one of two) now raise instead of failing confusingly inside read_table(None).

Blocked from here: steps 2–5 all need data/processed/echo_hubert_manifest.parquet, which is gitignored and absent. Once it exists:

python scripts/train_probes.py --manifest data/processed/echo_hubert_manifest.parquet

trains all four and writes probes/summary.json, which is what #64 consumes.

Worth reading #69 before running this: echo embeddings currently reach the probes as one mean-pooled vector per study, and attentive pooling over that is provably an identity. The retrain will work, but the echo branch will not be doing what its architecture claims until clip-level tokens are wired through.

[duckyquang]

Steps 1, 4 and 5 are now done in code; 2 and 3 are the only ones that need the manifest.

Re-read the acceptance criteria against what actually exists and found two more gaps that would have bitten during the rerun:

Step 4 — run manifest (was missing). utils/run_manifest.py existed but nothing called it, so checkpoints would have been written with no record of what produced them. train_probes.py now writes run_metadata.json next to summary.json (82adbf9): manifest path, seed, epochs, fusion width, encoder dims, per-probe n_dropped_nonfinite, git SHA.

The field that matters is echo_tokens_are_clip_level. A pooled run and a --max-clips run produce checkpoints that are indistinguishable on disk, and which regime was used is the entire provenance question for the pooled-vs-clip comparison in #69. Without it, two months from now nobody can tell which numbers came from the degenerate path.

Step 5 — prevalence check (already fine). scripts/check_ef40_prevalence.py --input <manifest>.parquet works unchanged; it only requires split and ef_le_40, both present in the manifest. Verified against a fixture. No code needed.

Step 3 — val-based selection (already correct). cross_attn.run keeps best_state on val MAE and never looks at test during selection. Confirmed by reading, not assumed.

Also fixed downstream: evaluate_missing_modality couldn't read a manifest at all, so #64 could not have scored whatever this issue produces. Fixed in 30ab118 — details on #64.

Remaining, and genuinely blocked on the manifest:

• Step 2: train all four on the same splits and seed
• Step 3: the actual val/test table and picking the canonical fused checkpoint

Everything those need is in place: python scripts/train_probes.py --manifest <path> runs all four, writes checkpoints, results, summary and run metadata in one go.

One decision to make when running it. Do a pooled run and a --max-clips 16 run rather than only the latter. The pooled numbers are the baseline that makes #69's echo comparison meaningful, and @kevzho's existing figures (MAE 10.90 / AUROC 0.742 at fusion-dim 1024, 10.98 / 0.734 at 256) are almost certainly pooled — so a matching pooled run under the new default gives a clean anchor before the clip-level comparison.

[duckyquang]

Ran M10 on the real 1,208-row pooled EchoJEPA+HuBERT manifest from PR #73.

Command shape:

python scripts/train_probes.py --manifest <1,208-row manifest> --out-dir <ignored local out> --fusion-dim 256

Device/config: local Mac CPU, seed 42, 50 epochs for the torch probes, fusion_dim 256, echo_dim 1024, ecg_dim 768. Wall-clock for all four train+eval: 44.84s.

The finite filter from #73 dropped the 12 HuBERT-ECG rows Kevin found: 8 train, 1 val, 3 test. So ECG/concat/fused trained and evaluated on 821/130/245, while echo-only stays at the full 829/131/248.

probe	val MAE	val EF<=40 AUROC	test MAE	test EF<=40 AUROC
ECG-only	11.8704	0.7785	11.5965	0.6713
Echo-only	12.1181	0.8245	11.1033	0.7683
Concat-MLP	11.2445	0.8328	10.9285	0.7164
Fused	10.6228	0.8898	10.4218	0.7708

Fused is the val-MAE winner, so that is the canonical pooled checkpoint for #64.

Prevalence check: raw split positives are 175/829 train, 26/131 val, 50/248 test. After the finite filter, fused sees 175/821 train, 26/130 val, 49/245 test. Train/test are reportable; val AUROC is thin and should carry the caveat that it has only 26 positives.

Artifacts are local and ignored. Nothing from the manifest or checkpoints is being committed.

[duckyquang]

Promoted the M10 pooled run artifacts into the default ignored local probes/ layout too, not just the PR-specific scratch path:

probes/ecg/results.json
probes/ecg/ecg_only.joblib
probes/echo/results.json
probes/echo/echo_only.pt
probes/concat/results.json
probes/concat/concat_mlp.pt
probes/fused/results.json
probes/fused/cross_attn_fused.pt
probes/summary.json
probes/run_metadata.json

probes/run_metadata.json records seed 42, fusion_dim 256, echo_dim 1024, ecg_dim 768, echo_tokens_are_clip_level=false, and n_dropped_nonfinite = 12 for every ECG-consuming probe.

These are local artifacts and remain gitignored.