[E11] Pooling ablation — HuBERT-ECG and EchoJEPA clip pooling

[duckyquang]

Task metadata

Field	Value
Task ID	E11
Category	Evaluation
Priority	Medium
Complexity	Low
Est. days	2
Replaces	#12 (closed — written against ECG-FM, which never produced a reported number)
Feeds	#63 (M10)

Features / method

Resolve the pooling decision on both modalities before the canonical retrain locks it in.

Description

The closed version of this task asked about ECG-FM pooling. That framing is stale twice over: ECG-FM never produced a reported number, and the ablation was scoped to feed M08/M09, which #63 now supersedes.

Two pooling decisions are actually open, and the echo one is probably the more consequential:

ECG side. TECHNICAL.md 4.2 leaves mean vs attentive pooling unresolved for the 768-d HuBERT-ECG embeddings.

Echo side. docs/echo_hubert_loader.md records that clip-level EchoJEPA embeddings were mean-pooled to one vector per echo study (525K clips down to 7,251 study rows). TECHNICAL.md 4.1 says EchoJEPA embeddings preserve spatial structure across frames and that a plain linear probe is therefore insufficient — an attentive probe is required. Mean-pooling clips into a single vector before the probe ever sees them may be discarding exactly the structure that motivates the attentive probe. If so, it caps echo-branch performance regardless of probe architecture, which is directly relevant to the baseline gap in #67.

Steps:

1. ECG: train the ECG-only probe with mean and with attentive pooling; compare on val.
2. Echo: compare study-level mean-pooled embeddings against retaining clip-level tokens into the attentive probe. n_echo_clips is already in the manifest metadata, so the clip counts per study are known.
3. Pick winners on val only. Do not touch test.
4. Log both decisions to logs/pooling_ablation.json with the val numbers behind them.

Acceptance criteria:

• Val comparison for both pooling axes, with numbers.
• Both decisions recorded in config so [M10] Retrain all four probes on the synchronized manifest #63 inherits them rather than re-deciding.
• TECHNICAL.md 11 pooling checkbox resolved.
• If clip-level echo tokens beat mean-pooling materially, flag it on [E10] Diagnose fused vs single-modality baseline gap #67 — it is a candidate explanation for the baseline gap.

[duckyquang]

The echo-side half of this is answerable from code alone, and the answer is that the choice is currently being made by accident.

_embedding_to_tokens tiles a study-level pooled vector into n_tokens identical copies. AttentivePool scores each token, softmaxes over them, and takes a weighted sum — but with identical tokens the scores are identical, so softmax is exactly uniform and the weighted sum returns the input vector unchanged.

Verified on a 1024-d vector:

attention weights      : [0.125 0.125 0.125 0.125 0.125 0.125 0.125 0.125]
max|attnpool(v) - v|   : 2.4e-07
max|attnpool(v) - mean|: 0.0

So on the current manifest the attentive echo probe is mathematically an MLP on the mean-pooled vector. The AttentivePool.query parameter is learnable but has zero effect on the output. TECHNICAL.md 4.1 names exactly this as the insufficient case.

Two consequences:

• The echo pooling axis in this issue is not really open — mean-pooling is what runs today regardless of which probe is selected. It only becomes a live comparison once clip-level tokens reach the probe.
• echo_only.run reports attentive_beats_linear_val_mae, and that flag compares two things that pool identically. It is not measuring what its name says.

Pinned as a regression test (test_attentive_pooling_is_inert_on_study_level_embeddings) so this fails once clip-level tokens land, which is the signal that the collapse is gone.

The ECG-side pooling question still needs embeddings and is untouched by this.

[duckyquang]

The ECG half of this issue is malformed as written. It cannot be answered by running an experiment, and it is not blocked on data.

_ensure_ecg_tokens (concat_mlp.py:47) takes the pooled 768-d HuBERT-ECG vector and tiles it into n_tokens=4 identical tokens. Measured:

ecg_tokens shape:                    (4, 768)
all 4 ECG tokens identical:          True
attentive(ECG) == mean(ECG):         True   max|diff| = 0.0
AttentivePool.query grad:            0.0

So "mean vs attentive pooling on the ECG side" is a comparison between two names for the same tensor. Attentive pooling over identical tokens is a bit-exact identity with exactly zero gradient on the query — the same degeneracy documented on #67 for echo, for the same reason. Running the ablation on the current embeddings would produce two identical numbers and an ablation table implying a choice was made.

This is only answerable once ECG has genuine tokens, which is #72's re-extraction — there are no pre-pooling HuBERT-ECG hidden states on disk (@kevzho confirmed on #73). So the ECG half depends on #72, not on the manifest.

Status of the two halves:

half	state
Echo	Mechanism answered and fixed. Mean-pooling clips to one vector per study made attentive pooling an identity; --max-clips retains clip vectors and attention goes live (ecg_to_echo weight spread 0.000 → 0.037). The empirical val comparison still needs the manifest.
ECG	Not runnable. Blocked on #72, not on data.

Revised acceptance criteria:

• Echo: val comparison of study-level mean-pooled vs clip-level, on the real manifest. Unblocked the moment the manifest is reachable.
• ECG: deferred to [M11] Fused probe has no cross-modal interaction — one direction still dead #72. Reinstate once token-level ECG exists; until then record in logs/pooling_ablation.json that the ECG axis was not evaluated and why, rather than leaving it looking unexamined.
• TECHNICAL.md 11's pooling checkbox should be split — the echo axis can be closed independently of the ECG one.

Worth flagging for #68: the concat probe's ecg_tokens.mean(dim=1) is not a modelling choice anyone made, it is the only thing a tiled vector can do. Nothing in the docs should imply the ECG pooling strategy was selected on evidence.

[duckyquang]

Ran the pooled anchor now that the real manifest is available.

The manifest Kevin sent is pooled, not clip-level: echo_tokens_are_clip_level=false, echo embeddings are 1-D vectors, and there is no --max-clips retained-token matrix in this file. So this gives the baseline half of the echo comparison, but not the clip-level half.

Pooled run, local Mac CPU, seed 42, fusion_dim 256:

probe	val MAE	val EF<=40 AUROC	test MAE	test EF<=40 AUROC
Echo-only	12.1181	0.8245	11.1033	0.7683
Fused	10.6228	0.8898	10.4218	0.7708

What remains:

• Echo axis: still needs a clip-level manifest built with --max-clips from the raw EchoJEPA shards, or the equivalent file from someone who can build it. The current manifest cannot answer whether retained clip tokens help.
• ECG axis: still deferred to [M11] Fused probe has no cross-modal interaction — one direction still dead #72. The HuBERT-ECG parquet has only pooled 768-d vectors, so attentive pooling over tiled ECG tokens is the same tensor as mean pooling.

I updated TECHNICAL.md in the PR branch to say exactly that: ECG pooling is deferred until token-level ECG exists, while the missing-modality rerun uses inference-time branch masks.

[duckyquang]

Tried to move this past the pooled baseline locally.

What exists here:

• the 1,208-row manifest from PR Retain clip-level echo tokens so attentive pooling actually pools #73, but it is pooled (echo_embedding: list<double>, one 1024-d vector per study)
• no raw 1024-d EchoJEPA clip shard under data/, Downloads, the repo tree, or the searched local criticaldata tree
• no huggingface_hub package in the venv and no cached MITCriticalData/mimic-iv-echo-jepa-embeddings dataset under ~/.cache/huggingface

A large-Parquet search under the local tree did not find candidate raw shards either. So the clip-level half cannot be run on this machine from the files currently present. The command we need once the raw shards are available is still:

python scripts/build_echo_hubert_manifest.py build-echo \
  --input <raw EchoJEPA clip parquet dir> \
  --output data/interim/echo_study_embeddings_vjepa2.1-vitl-mimic-pt-100-clips16.parquet \
  --max-clips 16

then join/train with that clip-level echo parquet. Until that file exists, #69 has a real pooled anchor but not the empirical retained-clip comparison.