""" Token probe — the GO / NO-GO experiment before touching the LLM ============================================================== Question: do the 3x3x3 = 27 spatial positions sitting before the encoder's AdaptiveAvgPool3d(1) carry information worth handing to an LLM as "visual tokens"? Why this comes first: if the 27 tokens say nothing beyond the pooled 512-d vector, wiring the encoder into the LLM as a token sequence is pointless — you would just be moving the pipeline inside a transformer. Learning that in an hour beats learning it after a week of projector training. Method: 1. Load the checkpoint and FREEZE the entire encoder. 2. For every scan, compute the ASPP output (B, 512, 3, 3, 3) once and cache it as a (B, 27, 512) token sequence. The encoder is frozen, so it cannot change between epochs and the 3D CNN never has to run twice. 3. Train two small probes on those tokens: - mean-pool probe : discards spatial information (what the model does now) - attention probe : weighted pooling with a learned query over 27 tokens 4. Compare on the test split, and measure whether the attention is degenerate (uniform vs selective) through its entropy. Reading the result: attention probe clearly ahead → the token grid is meaningful, wire it in. probes equal and entropy ~max → no spatial information; fix the encoder's strides / ASPP dilations first. NOTE: the normalizer is loaded from the checkpoint and never re-fitted on the evaluation data — re-fitting would leak. This script uses the model modules (config, model, dataset); the path bootstrap is below. YOU MUST SET YOUR OWN PATH via VBAI_MODEL_DIR if they are elsewhere. Run: python token_probe.py --ckpt Vbai-2.6AD.pt python token_probe.py --ckpt ... --epochs 40 --batch-size 8 """ from __future__ import annotations import argparse import os import sys # ---------------------------------------------------------------------- # The modality must be chosen BEFORE config is imported: config decides at # import time which visit manifest to read (USE_TBM). # # CRITICAL: the checkpoint and the input modality must match. If they do not, # the encoder drops to chance and the probe result is meaningless. # ---------------------------------------------------------------------- _ap = argparse.ArgumentParser(add_help=False) _ap.add_argument("--tbm", action="store_true") _ap.add_argument("--t1", action="store_true") _known, _ = _ap.parse_known_args() if _known.tbm == _known.t1: sys.exit("ERROR: pass exactly one of --tbm / --t1 " "(match whichever modality the checkpoint was trained on).") os.environ["VBAI_USE_TBM"] = "1" if _known.tbm else "0" MODALITY = "TBM" if _known.tbm else "raw T1" # ---------------------------------------------------------------------- # Make the model modules importable. # ---------------------------------------------------------------------- def _bootstrap_model_path() -> str: env = os.environ.get("VBAI_MODEL_DIR") candidates = [] if env: candidates.append(env) here = os.path.dirname(os.path.abspath(__file__)) candidates.append(here) # next to this file candidates.append(os.path.join(here, "Vbai-2.6AD")) # ./Vbai-2.6AD for c in candidates: if os.path.isfile(os.path.join(c, "config.py")): if c not in sys.path: sys.path.insert(0, c) return c raise ImportError( "Could not locate the model modules (config.py).\n" f"Tried: {candidates}\n" "YOU MUST SET YOUR OWN PATH: point VBAI_MODEL_DIR at the directory " "holding config.py / model.py / dataset.py." ) MODEL_DIR = _bootstrap_model_path() import numpy as np import torch import torch.nn as nn import torch.nn.functional as F from sklearn.metrics import accuracy_score, f1_score import config as C from model import Vbai26ADModel from dataset import (PairedVisitDataset, TabularNormalizer, collate_pad, subject_split, load_paired) # ---------------------------------------------------------------------- # Token extraction — the encoder's forward pass without the pooling step # ---------------------------------------------------------------------- @torch.no_grad() def encode_tokens(encoder, mri: torch.Tensor): """(B, 1, 96, 96, 96) → (B, 27, 512) sequence of spatial tokens.""" x = encoder.stem(mri) x = encoder.stage1(x) x = encoder.stage2(x) x = encoder.stage3(x) x = encoder.stage4(x) x = encoder.aspp(x) # (B, 512, 3, 3, 3) ch = x.shape[1] return x.flatten(2).transpose(1, 2).contiguous(), (ch, tuple(x.shape[2:])) @torch.no_grad() def build_token_cache(model, loader, device): """The encoder is frozen, so compute the tokens once and keep them in RAM.""" toks, labels, shape_info = [], [], None for batch in loader: if "mri" not in batch: continue mri = batch["mri"].to(device, non_blocking=True) t, shape_info = encode_tokens(model.mri_encoder, mri) toks.append(t.float().cpu()) labels.append(batch["label"].clone()) if not toks: raise RuntimeError("No MRI batch could be loaded — check your data paths.") return torch.cat(toks), torch.cat(labels), shape_info # ---------------------------------------------------------------------- # Probes # ---------------------------------------------------------------------- class MeanPoolProbe(nn.Module): """Discards spatial information — equivalent to the current AdaptiveAvgPool3d(1).""" def __init__(self, dim, n_cls=3): super().__init__() self.head = nn.Sequential(nn.LayerNorm(dim), nn.Linear(dim, 128), nn.GELU(), nn.Dropout(0.2), nn.Linear(128, n_cls)) def forward(self, tok): # (B, N, D) return self.head(tok.mean(dim=1)), None class AttnPoolProbe(nn.Module): """Weighted pooling over the 27 tokens with a single learned query.""" def __init__(self, dim, n_cls=3): super().__init__() self.q = nn.Parameter(torch.randn(1, 1, dim) * 0.02) self.norm = nn.LayerNorm(dim) self.attn = nn.MultiheadAttention(dim, num_heads=8, batch_first=True) self.head = nn.Sequential(nn.LayerNorm(dim), nn.Linear(dim, 128), nn.GELU(), nn.Dropout(0.2), nn.Linear(128, n_cls)) def forward(self, tok): # (B, N, D) x = self.norm(tok) q = self.q.expand(x.size(0), -1, -1) pooled, w = self.attn(q, x, x, need_weights=True, average_attn_weights=True) return self.head(pooled.squeeze(1)), w.squeeze(1) # w: (B, N) def train_probe(probe, tr_tok, tr_y, te_tok, te_y, device, epochs, bs, lr=3e-4): probe = probe.to(device) opt = torch.optim.AdamW(probe.parameters(), lr=lr, weight_decay=1e-4) sched = torch.optim.lr_scheduler.CosineAnnealingLR(opt, T_max=epochs) # Compensate for class imbalance — MCI is already the weakest class counts = torch.bincount(tr_y, minlength=3).float().clamp(min=1) w = (counts.sum() / (3 * counts)).to(device) n = tr_tok.size(0) for _ in range(epochs): probe.train() perm = torch.randperm(n) for i in range(0, n, bs): idx = perm[i:i + bs] xb = tr_tok[idx].to(device) yb = tr_y[idx].to(device) logits, _ = probe(xb) loss = F.cross_entropy(logits, yb, weight=w) opt.zero_grad(set_to_none=True) loss.backward() opt.step() sched.step() probe.eval() preds, attns = [], [] with torch.no_grad(): for i in range(0, te_tok.size(0), bs): logits, a = probe(te_tok[i:i + bs].to(device)) preds.append(logits.argmax(-1).cpu()) if a is not None: attns.append(a.cpu()) preds = torch.cat(preds).numpy() y = te_y.numpy() attn = torch.cat(attns) if attns else None return { "acc": accuracy_score(y, preds), "f1_macro": f1_score(y, preds, average="macro"), "f1_per": f1_score(y, preds, average=None, labels=[0, 1, 2]), "attn": attn, } def remap_nifti_paths(df): """ The visit manifest stores absolute volume paths from the machine that built it, which will not resolve anywhere else. The tail of each path (after the dataset root) is re-attached to the roots configured here. Harmless when the original paths already resolve — it returns them unchanged. """ def _fix(p): p0 = str(p) if os.path.exists(p0): return p0 q = p0.replace("\\", "/") # Dataset root i = q.find("/Datasets/") if i >= 0: cand = os.path.join(C.DATASET_ROOT, q[i + len("/Datasets/"):]) if os.path.exists(cand): return cand # The stored tail may or may not include the top volume folder, # so both spellings are tried. j = q.find("/volumes/") if j >= 0: rest = q[j + len("/volumes/"):] for cand in (os.path.join(C.TBM_ROOT, rest), os.path.join(C.TBM_ROOT, "volumes", rest)): if os.path.exists(cand): return cand return p0 df = df.copy() df["nifti_path"] = df["nifti_path"].map(_fix) ok = int(sum(os.path.exists(str(p)) for p in df["nifti_path"])) print(f"[path] reachable images ({MODALITY}): {ok}/{len(df)}") if ok == 0: raise FileNotFoundError( f"No image is reachable ({MODALITY}).\n" f" DATASET_ROOT = {C.DATASET_ROOT}\n" f" VOLUME_ROOT = {C.TBM_ROOT}\n" "YOU MUST SET YOUR OWN PATHS: see VBAI_DATASET_ROOT / " "VBAI_VOLUME_ROOT in config.py." ) return df[df["nifti_path"].map(lambda p: os.path.exists(str(p)))].reset_index(drop=True) @torch.no_grad() def baseline_from_checkpoint(model, loader, device): """The trained mri_classifier's own score — the reference baseline.""" preds, ys = [], [] for batch in loader: if "mri" not in batch: continue out = model(mri=batch["mri"].to(device)) preds.append(out["mri_logits"].argmax(-1).cpu()) ys.append(batch["label"]) preds, ys = torch.cat(preds).numpy(), torch.cat(ys).numpy() return {"acc": accuracy_score(ys, preds), "f1_macro": f1_score(ys, preds, average="macro")} def main(): ap = argparse.ArgumentParser() ap.add_argument("--ckpt", default="Vbai-2.6AD.pt", help="Vbai-2.6AD checkpoint") ap.add_argument("--tbm", action="store_true", help="TBM input") ap.add_argument("--t1", action="store_true", help="raw T1 input") ap.add_argument("--epochs", type=int, default=30) ap.add_argument("--batch-size", type=int, default=8) ap.add_argument("--workers", type=int, default=2) args = ap.parse_args() device = torch.device("cuda" if torch.cuda.is_available() else "cpu") print(f"[device] {device}") print(f"[model dir] {MODEL_DIR}") # --- model + normalizer (from the checkpoint; NEVER re-fitted) --- sd = torch.load(args.ckpt, map_location=device, weights_only=False) mcfg = C.ModelConfig() for k, v in sd.get("model_cfg", {}).items(): if hasattr(mcfg, k): setattr(mcfg, k, v) model = Vbai26ADModel(mcfg).to(device) # Any key mismatch is fatal on purpose: strict=False silently swallows a # checkpoint from a different architecture, leaving the model on random # weights and producing below-chance results. Older checkpoints from a # different architecture must fail loudly here. res = model.load_state_dict(sd["model"], strict=False) if res.missing_keys or res.unexpected_keys: raise RuntimeError( f"Checkpoint does NOT match this architecture: " f"{len(res.missing_keys)} missing / {len(res.unexpected_keys)} " f"unexpected keys.\n" f" first missing : {res.missing_keys[:3]}\n" f" first unexpected : {res.unexpected_keys[:3]}\n" "Checkpoints from a different architecture are not compatible.\n" "Expected file: Vbai-2.6AD.pt" ) print(f"[ckpt] {len(sd['model'])} keys matched exactly") if "extra" in sd and sd["extra"].get("metrics"): print(f"[ckpt] stored metrics: {sd['extra']['metrics']}") model.eval() for p in model.parameters(): p.requires_grad_(False) norm = TabularNormalizer() norm.load_state_dict(sd["norm"]) # --- subject-level split (no leakage between train and test) --- df = load_paired() df = remap_nifti_paths(df) train_ids, val_ids, test_ids = subject_split(df) tr_df = df[df["ptid"].isin(train_ids | val_ids)] # train+val to fit the probes te_df = df[df["ptid"].isin(test_ids)] print(f"[data] probe-train {len(tr_df)} scans / test {len(te_df)} scans") def make_loader(d, shuffle=False): ds = PairedVisitDataset(d, norm, mode="mri", augment=False, mcfg=mcfg) return torch.utils.data.DataLoader(ds, batch_size=args.batch_size, shuffle=shuffle, collate_fn=collate_pad, num_workers=args.workers) tr_loader, te_loader = make_loader(tr_df), make_loader(te_df) # --- baseline: the checkpoint's own MRI head --- print("\n[1/3] Baseline (trained mri_classifier, pooled)...") base = baseline_from_checkpoint(model, te_loader, device) print(f" acc {base['acc']:.4f} | macro-F1 {base['f1_macro']:.4f}") # --- token cache --- print("\n[2/3] Building the token cache (encoder frozen, single pass)...") tr_tok, tr_y, shape_info = build_token_cache(model, tr_loader, device) te_tok, te_y, _ = build_token_cache(model, te_loader, device) ch, grid = shape_info n_tok, dim = tr_tok.shape[1], tr_tok.shape[2] print(f" ASPP output: {ch} channels @ {grid} → {n_tok} tokens x {dim}-d") print(f" cache: train {tuple(tr_tok.shape)} / test {tuple(te_tok.shape)}") # --- problar --- print("\n[3/3] Training the probes (encoder frozen)...") r_mean = train_probe(MeanPoolProbe(dim), tr_tok, tr_y, te_tok, te_y, device, args.epochs, args.batch_size) r_attn = train_probe(AttnPoolProbe(dim), tr_tok, tr_y, te_tok, te_y, device, args.epochs, args.batch_size) # --- is the attention degenerate? --- a = r_attn["attn"].clamp_min(1e-9) ent = float((-(a * a.log()).sum(dim=1)).mean()) max_ent = float(np.log(n_tok)) print("\n" + "=" * 66) print(" TOKEN PROBE RESULT") print("=" * 66) print(f" {'':22s} {'acc':>8s} {'macro-F1':>10s} F1 (CN/MCI/AD)") print(f" {'checkpoint (pooled)':22s} {base['acc']:8.4f} {base['f1_macro']:10.4f}") print(f" {'probe: mean-pool':22s} {r_mean['acc']:8.4f} {r_mean['f1_macro']:10.4f}" f" {'/'.join(f'{v:.3f}' for v in r_mean['f1_per'])}") print(f" {'probe: attention':22s} {r_attn['acc']:8.4f} {r_attn['f1_macro']:10.4f}" f" {'/'.join(f'{v:.3f}' for v in r_attn['f1_per'])}") print(f"\n attention entropy: {ent:.3f} / {max_ent:.3f} (max = fully uniform)") delta = r_attn["f1_macro"] - r_mean["f1_macro"] print(f" attention - mean difference (macro-F1): {delta:+.4f}") print("-" * 66) if delta > 0.02 and ent < 0.95 * max_ent: print(" → GO. The spatial tokens carry extra information; wiring them") print(" into the LLM as a sequence is justified.") elif ent >= 0.95 * max_ent: print(" → STOP. The attention is nearly uniform: the 27 positions do") print(" not separate. Lower the ASPP dilations to (1,2,3), or drop the") print(" stage-4 stride to reach a 6^3 grid, then measure again.") else: print(" → WEAK. The tokens add nothing meaningful over the pooled") print(" vector. Wiring them in without raising the encoder resolution") print(" will not help.") print("=" * 66) if __name__ == "__main__": main()