Retain clip-level echo tokens so attentive pooling actually pools - #73
Retain clip-level echo tokens so attentive pooling actually pools#73duckyquang wants to merge 5 commits into
Conversation
The probes predate the manifest: each expected a cohort table plus one or two separate embedding tables to join on study/record id. The manifest already carries both embeddings inline, so make the embedding paths optional and read the first argument as a manifest when they are omitted. The old two-table path is untouched. Only the ECG probe had a CLI, so "retrain all four" had no entrypoint. Adds scripts/train_probes.py covering all four, reading dims off the manifest instead of taking --embed-dim, which removes the state_dict shape mismatch the README warns about. Also pins a regression test for a problem this surfaced: echo embeddings arrive as one mean-pooled vector per study, which _embedding_to_tokens tiles into identical tokens. AttentivePool softmaxes over identical scores, so the weights come out uniform and the layer is an exact identity. The attentive echo probe is currently an MLP on a mean-pooled vector, which TECHNICAL.md 4.1 names as the insufficient case. The test fails once clip-level tokens land, which is the signal the collapse is gone.
Mean-pooling clips to one vector per study made every downstream attention layer degenerate. AttentivePool over identical tokens is a bit-exact identity (uniform softmax) and its query gets exactly zero gradient, so the echo probe was an MLP on a mean vector. In the fused probe the cross-attention degenerates twice over -- identical keys force uniform weights, identical values make the weighted sum equal the value -- leaving echo_ctx a function of ECG alone and ecg_ctx a function of echo alone, with no cross-modal interaction at all. build_echo_study_embeddings gains an opt-in max_clips that keeps an even-stride subsample of clip vectors per study. Default stays mean-pooling. The write stage now fills one contiguous float32 block and hands pyarrow a zero-copy column, draining the accumulator as it goes; the old np.stack().tolist() round-trip boxed every float into a Python object and peaked around 7 GB at full scale against 1.8 GB now. Ragged clip counts mean pad_sequence appends zero rows, which score 0 and take real softmax weight, so echo_mask is threaded through AttentivePool, CrossAttentionFusion, the three torch probes and fairness. Ragged ECG token counts are rejected rather than masked, since nothing can produce them yet. Parsing now runs before the stride decision so null clips cannot consume retention slots or drop a study the pooled path would keep, and n_echo_clips counts parsed clips in both modes. Known limits: the pyarrow list-offset ceiling is guarded with an explicit precondition, not lifted -- pa.ListArray narrows offsets to int32 regardless of the dtype passed in, and build_joined_manifest re-infers plain list<> through pandas, so LargeListArray in the builder alone would buy nothing. The default pooled build now writes list<float> where it wrote list<double>; values are bit-identical. An all-null shard now raises instead of writing an empty parquet.
There was a problem hiding this comment.
Nice work, the findings are solid and properly labeled as synthetic-only.
I have the manifest locally, so I ran all four probes against the real 1,208 rows. Requesting changes to one thing that only appears with the data.
concat_mlp.py crashes, as twelve rows carry non-finite HuBERT-ECG vectors (8 train, 1 val, 3 test). cross_attn filters them and ecg_only filters them, but concat_mlp and echo_only have no finite filter on main or here. echo_only survives by luck (echo has zero non-finite rows); concat_mlp trains on NaN and dies at concat_mlp.py:85 -> common.py:44 with ValueError: Input contains NaN. manifest.load's require_both doesn't catch it — .notna() tests whether the cell is null, not whether the parsed vector contains NaN. Since this PR is what makes "retrain all four" runnable, I think it belongs here. Cleanest fix is one shared finite filter that also logs its drop count, which closes the silent row-loss gap at the same time so #68 doesn't inherit it.
Also, expose the fusion width as a CLI flag. train_probes.py passes embed_dim=echo_dim, pinning the fused probe at 1024 — 9.7M params vs 1.1M at the 256 the current checkpoint used, on 821 training rows. #67 step 3 is "check probe capacity and training budget," and I'd rather an 8.7x capacity jump not enter as an implicit default. Measured: @1024 gets MAE 10.90 / AUROC 0.742 vs @256 at 10.98 /0.734 — basically nothing for 8.7x the parameters, so 256 looks like the right default. Runtime shouldn't be an issue either way (all four train+eval in 49s / 81s on CPU).
Answers to your two questions:
- Yes, and it's
parse_embeddingspecifically that makes that work. Onmainit doesn't, as pyarrow returns alist<list<double>>column as an object-dtype array of per-cliip arrays and.astype(float)raises, soEchoHubertDatasetcan't read a clip-level manifest at all. The version you have parses fine. - I don't think we can get token-level ecg without re-extraction. The interim parquet is the pooled 768-d vector per record; there aren't any pre-pooling hidden states on disk. HuBERT-ECG is wav2vec2-family, so re-running it on the MIMIC-IV-ECG waveforms would give us genuine tokens but that means #72 neds to be modified into a re-extraction task instead of a loader change (and that also means #71 should be rescoped to match asw for that).
I can just push both fixes if that's faster (I can verify on real data in a minute).
concat_mlp crashed on the real 1,208-row cohort: twelve rows carry non-finite HuBERT-ECG vectors and only ecg_only and cross_attn filtered them. echo_only was surviving on luck, since the echo side happens to have none. manifest.load's require_both does not help -- notna() asks whether the cell is null, not what the vector holds. One drop_non_finite in probes/common.py now serves all four, and each reports n_dropped_nonfinite in its results.json so the row loss is visible rather than silent. prepare_fused_probe_data returns the count alongside the splits, which changes its signature for its two callers. train_probes.py passed embed_dim=echo_dim, pinning the fused probe to the encoder width: 9.7M params against 1.1M at 256, on 821 training rows. That is a capacity choice, not a property of the data, so it becomes --fusion-dim defaulting to the 256 the existing checkpoint used.
|
Thanks — both were real and both are fixed in 3973bff. I did them myself rather than take you up on the offer, since you'd have been fixing my bug on your Saturday, but I do need your data for the verification (below). Non-finite rows. You were right about the root cause and I'd got One Regression test builds a manifest with NaN/inf/-inf inside otherwise-present ECG cells and runs all four probes. I checked it actually bites rather than just going green: pull the Fusion width. Agreed, and thanks for measuring it — 8.7x the parameters for 0.08 MAE is not a trade I'd have wanted to make silently. It's I left the encoder dims read off the manifest, since a mismatch there is the One thing I want to check before we read anything into your numbers. Was that run against the pooled manifest? It has to have been unless you rebuilt with Also noting 10.90 sits right next to the old 10.28, so #67's gap against EchoJEPA's 5.97 hasn't moved yet. Not surprising if the echo branch was inert for both runs. On your answer to (2) — that's decisive, thank you. I've rescoped #72 from a loader change to a re-extraction task and #71 to match: probe training genuinely doesn't need a GPU (your 49s/81s confirms what I measured), but re-running HuBERT-ECG over the MIMIC-IV-ECG waveforms does, so the compute ask comes back for that specifically. And the ask that matters most: can you get me the manifest, or point me at where it lives? It's the blocker on #63, #64 and #69 for me, and you having it locally is the only reason any of this has real numbers attached. |
evaluate_missing_modality still required a cohort plus two separate embedding tables, so E07 could not score the checkpoint M10 produces even after the probes learned to read the manifest. Both embedding paths are optional now, matching the probes, and --manifest supersedes the two-table flags on the CLI. checkpoint_path keeps its old positional slot but is now keyword-optional, so it raises a named error instead of failing inside Path(None).
M10 asks for a run manifest per training run, and utils/run_manifest existed but nothing called it. train_probes now writes run_metadata.json alongside summary.json with the manifest path, seed, epochs, fusion width, encoder dims and the per-probe non-finite drop counts. The field that matters is echo_tokens_are_clip_level. Pooled and clip-level runs produce checkpoints that look identical on disk, and telling them apart after the fact is the whole provenance question for the pooled-vs-clip comparison. check_ef40_prevalence already reads a manifest unchanged -- it only needs split and ef_le_40, both present -- so M10 step 5 needs no code.
2 participants
Summary
Two things, and the second one is why this PR is bigger than I planned.
First: the probes couldn't read the manifest from #61. Each one wanted a cohort table plus one or two separate embedding tables to join on study/record id, and the manifest already carries both embeddings inline. Only the ECG probe had a CLI at all, so "retrain all four" had nothing to run. Both fixed — embedding paths are optional now, omit them and the first arg is read as a manifest, and
scripts/train_probes.pydoes all four.Second: while wiring that up I found the attentive pooling has never done anything. We mean-pool clips down to one vector per study, then tile that vector into identical tokens, and attention over identical tokens is a bit-exact identity — the softmax comes out uniform and the query gets exactly zero gradient, so it can't even learn its way out. The echo probe has been an MLP on a mean-pooled vector this whole time, which is the case TECHNICAL.md 4.1 specifically says isn't good enough.
Same thing hits the fused probe harder. Both cross-attention blocks degenerate, and because each one's keys/values come from the other modality,
echo_ctxended up depending only on ECG andecg_ctxonly on echo — the two halves of the concat are swapped and there's no cross-modal interaction at all.So this adds an opt-in
--max-clipsthat keeps clip-level vectors per study instead of averaging them away, plus padding masks, since ragged clip counts meanpad_sequenceappends zero rows that would otherwise get real softmax weight.Result
All measured on this Mac, CPU, on synthetic test fixtures —
data/is empty here, so none of this is a model result and there's no MAE or AUROC I'd stand behind yet.Attention weight spread, D=32, eval mode (0 = uniform = degenerate):
ecg_to_echoecho_to_ecgOnly one direction is fixed — see Notes.
The first cut of the clip writer would have OOM'd.
np.stack(...).tolist()boxed every float32 into a Python object and held all studies at once; measured ~7 GB projected at 7,251 studies. Rewrote it as one contiguous float32 block handed to pyarrow zero-copy:Tests 65 → 90.
Notes
The fused probe is only half fixed.
echo_to_ecgtakes its keys and values from the ECG tokens, and those are still one pooled HuBERT-ECG vector tiled into 4 identical copies, so that direction stays dead no matter what the echo side does. Fixing it needs token-level ECG embeddings, which probably means re-extraction rather than a loader change. Split out as #72.This is mechanism work, not validation. It shows the architecture can now attend. It says nothing about whether attending helps LVEF. That needs the manifest, which isn't on my machine — so I haven't run #63 or #64 against any of this.
Knock-on for #64: the missing-modality numbers were measured on a model with no cross-modal interaction, so they describe degradation of a concat model. Worth a rerun once the above settles.
Also worth knowing:
attentive_beats_linear_val_maeinecho_only's results.json compares an MLP and a linear head on a bit-identical vector on the pooled path. It's a capacity comparison, not an attention one — don't read it as evidence for attentive pooling.Things I didn't fix, deliberately:
pa.ListArraynarrows offsets to int32 whatever dtype you hand it, andbuild_joined_manifestre-infers plainlist<>through pandas, so switching the builder toLargeListArrayalone buys nothing end to end. A build past the limit now fails immediately with a clear message instead of an opaqueArrowInvalid.list<float>where it wrotelist<double>— values bit-identical, type changed. An all-null shard now raises instead of writing an empty parquet.@kevzho you've got the most history in
probes/, so you're the one I'd most want eyes on this. Two specific asks: does the clip-level manifest schema fit how you expect to consume it downstream, and do you know off-hand whether we can get token-level ECG out of HuBERT-ECG without re-extracting? That's the blocker on #72.