| """Training-regime comparison: SEPARATE vs JOINT training of the four-modality model (real pairs only, final protocol). |
| |
| Real paired data (fold k; train / val (early stopping only) / test, patients / study groups never cross splits): |
| ORION H&E <-> mIF splits/patient_5fold/fold<k>.csv UNI2 (8 views) <-> VirTues 16-ch MIPHEI panel (8 aligned views) |
| HEST H&E <-> RNA splits/final/hest_fold<k>.csv UNI2 <-> log1p CP10k of the 2000 HVGs |
| TCGA H&E <-> text splits/final/tcga_fold<k>.csv UNI2 (8 views) <-> BiomedCLIP, CLEAN real texts (40_) |
| Regimes (same MLP heads, InfoNCE, AdamW 3e-4 / wd 0.05 / OneCycle, batch 256, one temperature per pair): |
| S1 separate, frozen anchor modality head -> frozen UNI2 space; three independent trainings (= current method) |
| S2 separate, two-tower per pair a trainable H&E head + modality head (CLIP-style); no shared space |
| J1 joint, frozen anchor S1 heads trained in one loop (sum of the 3 losses), joint early stopping |
| J2 joint, shared H&E head ONE trainable H&E head shared by all pairs + 3 modality heads (ImageBind with |
| a learned anchor projection) |
| J3 J2, size-proportional one task per step, sampled proportional to its training-set size |
| J4 J2 + anchor distillation + lambda * MSE between in-batch cosine-similarity matrices of UNI2 and of the projected |
| H&E (keeps the shared head close to the UNI2 geometry) |
| J5 joint, residual adapter H&E = UNI2 + shared zero-initialised MLP adapter (frozen anchor + small shared change) |
| S3 separate, frozen, linear S1 with a single linear layer per modality (head capacity) |
| J4b J4 with lambda = 10 stronger distillation |
| J6 J2 + uncertainty weighting learnable per-task log-variance s_t: sum_t exp(-s_t) L_t + s_t (Kendall et al. 2018) |
| S4 separate, hybrid heads S1 with MLP head for mIF, linear heads for RNA and text |
| J8 two-stage from S4 J7 starting from the S4 heads |
| J7 two-stage S1 heads, then joint fine-tuning of S1 heads + residual adapter (lr 1e-4, 1000 steps, |
| distillation lambda 1) |
| Early stopping: separate = own val R@1; joint = geometric mean of the three val R@1 (every 50 steps). |
| Test (fold k): trained pairs H&E<->mIF / RNA / text, and emergent pairs with test-time-only virtual data: |
| ORION mIF<->RNA(STPath virtual), mIF<->text(GPT synthetic); TCGA text<->RNA(STPath virtual), text<->mIF(MIPHEI virtual). |
| Usage: python 42_train_regimes.py [--regimes S1 S2 J1 J2 J3 J4 J5] [--folds 0 1 2 3 4] |
| Output results/regimes_long.csv, results/regimes_summary.csv |
| """ |
| import argparse |
| import os |
|
|
| import numpy as np |
| import pandas as pd |
| import torch |
| import torch.nn as nn |
| import torch.nn.functional as F |
| from scipy.optimize import linear_sum_assignment |
|
|
| from common import CACHE, ROOT, SPLITS |
|
|
| ap = argparse.ArgumentParser() |
| ap.add_argument("--regimes", nargs="+", default=["S1", "S2", "S3", "S4", "J1", "J2", "J3", "J4", "J4b", "J5", "J6", "J7", "J8"]) |
| ap.add_argument("--folds", nargs="+", type=int, default=[0, 1, 2, 3, 4]) |
| args = ap.parse_args() |
| dev = "cuda" |
| H = os.path.join(CACHE, "hest") |
| G100, DRAWS, BS, LR, EVAL_EVERY = 100, 50, 256, 3e-4, 50 |
| SEP_STEPS = {"mif": 400, "rna": 2700, "text": 1000} |
| JOINT_STEPS = {"equal": 2000, "prop": 4000} |
| f32 = lambda a: torch.from_numpy(np.asarray(a, dtype=np.float32)) |
|
|
| |
| HE_O = f32(np.load(os.path.join(CACHE, "he_uni2.npy"))).to(dev) |
| IF_O = f32(np.load(os.path.join(CACHE, "if_virtues_m16.npy"))).to(dev) |
| IF_V_T = f32(np.load(os.path.join(CACHE, "if_virtues_m16_virtual_miphei_tcga_text.npy"))[:, 0]).to(dev) |
| HE_H = f32(np.load(os.path.join(H, "he_uni2.npy"))[:, 0]).to(dev) |
| RNA_H = f32(np.load(os.path.join(H, "hvg.npy"))).to(dev) |
| HE_T = f32(np.load(os.path.join(CACHE, "allimg_he_uni2.npy"))).to(dev) |
| TXT = f32(np.load(os.path.join(CACHE, "alltxt_text_biomedclip.npy"))[:, 0]).to(dev) |
| T = pd.read_csv(os.path.join(CACHE, "alltxt_meta.csv")) |
| T_IMG = torch.as_tensor(T.img_idx.values, device=dev) |
| T = T[~pd.read_csv(os.path.join(CACHE, "alltxt_junk.csv")).junk.values] |
| GT_O = f32(np.load(os.path.join(CACHE, "orion_gentext_text_biomedclip.npy"))[:, 0]).to(dev) |
| hvg = pd.read_csv(os.path.join(H, "hvg_genes.csv")).gene.values |
|
|
|
|
| def virtual_rna(st): |
| g = pd.read_csv(os.path.join(CACHE, "stpath", f"{st}_genes.csv")).gene.values |
| c = np.expm1(np.load(os.path.join(CACHE, "stpath", f"{st}_pred.npy")).astype(np.float32)).clip(0) |
| c = np.log1p(c / c.sum(1, keepdims=True) * 1e4) |
| return f32(c[:, pd.Series(range(len(g)), index=g).loc[hvg].values]).to(dev) |
|
|
|
|
| RNA_V_O, RNA_V_T = virtual_rna("orion"), virtual_rna("tcga_text") |
|
|
|
|
| def mlp(din, dout=1536, hid=512): |
| return nn.Sequential(nn.Linear(din, hid), nn.GELU(), nn.Dropout(0.1), nn.Linear(hid, hid), nn.GELU(), |
| nn.LayerNorm(hid), nn.Linear(hid, dout)) |
|
|
|
|
| def one_text_per_image(t_ids, rng): |
| pick = T.loc[t_ids].groupby("img_idx").apply(lambda g: g.index[rng.integers(len(g))], include_groups=False) |
| return pick.index.values, pick.values |
|
|
|
|
| def g100(a, b, seed=0): |
| a, b = a.cpu(), b.cpu() |
| rng = np.random.default_rng(seed) |
| out = {m: [] for m in ["R@1", "R@5", "R@10", "MRR", "pair_acc"]} |
| for _ in range(DRAWS): |
| s = torch.from_numpy(rng.choice(len(a), min(G100, len(a)), replace=False)) |
| sim = (a[s] @ b[s].T).numpy(); d = np.diag(sim) |
| for S in (sim, sim.T): |
| rank = (S > d[:, None]).sum(1) |
| for k_ in (1, 5, 10): |
| out[f"R@{k_}"].append((rank < k_).mean() / 2) |
| out["MRR"].append((1 / (rank + 1)).mean() / 2) |
| ri, ci = linear_sum_assignment(-sim) |
| out["pair_acc"].append((ci[np.argsort(ri)] == np.arange(len(s))).mean()) |
| return {m: float(np.sum(v) / DRAWS) for m, v in out.items()} |
|
|
|
|
| def r1(a, b): |
| s = a @ b.T |
| return float((((s > s.diag()[:, None]).sum(1) == 0).float().mean() + ((s.T > s.diag()[:, None]).sum(1) == 0).float().mean()) / 2) |
|
|
|
|
| class Model(nn.Module): |
| def __init__(self, anchor, tasks, head="mlp", uw=False): |
| super().__init__() |
| self.anchor = anchor |
| kind = head if isinstance(head, dict) else {t: head for t in ("mif", "rna", "text")} |
| mk = lambda t, d: mlp(d) if kind[t] == "mlp" else nn.Linear(d, 1536) |
| self.heads = nn.ModuleDict({"mif": mk("mif", 512), "rna": mk("rna", 2000), "text": mk("text", 512)}) |
| self.uw = nn.ParameterDict({t: nn.Parameter(torch.zeros(())) for t in tasks}) if uw else None |
| self.logit = nn.ParameterDict({t: nn.Parameter(torch.tensor(np.log(1 / 0.07), dtype=torch.float32)) for t in tasks}) |
| if anchor == "shared": |
| self.ah = mlp(1536) |
| elif anchor == "pair": |
| self.ah = nn.ModuleDict({t: mlp(1536) for t in tasks}) |
| elif anchor == "residual": |
| self.ah = mlp(1536) |
| nn.init.zeros_(self.ah[-1].weight); nn.init.zeros_(self.ah[-1].bias) |
|
|
| def he(self, x, task): |
| if self.anchor == "frozen": |
| return x |
| if self.anchor == "shared": |
| return self.ah(x) |
| if self.anchor == "pair": |
| return self.ah[task](x) |
| return x + self.ah(x) |
|
|
|
|
| class Fold: |
| def __init__(self, k): |
| s = pd.read_csv(os.path.join(SPLITS, "patient_5fold", f"fold{k}.csv")) |
| self.o = {x: s.tile_id[s.split == x].values for x in ["train", "val", "test"]} |
| self.if_mu = IF_O[self.o["train"]].mean((0, 1)); self.if_sd = IF_O[self.o["train"]].std((0, 1)) + 1e-3 |
| h = pd.read_csv(os.path.join(SPLITS, "final", f"hest_fold{k}.csv")) |
| self.h = {x: h.spot_idx[h.split == x].values for x in ["train", "val", "test"]} |
| sp = pd.read_csv(os.path.join(SPLITS, "final", f"tcga_fold{k}.csv")).set_index("img_idx").split |
| tsp = sp.loc[T.img_idx].values |
| self.t = {x: T.index[tsp == x].values for x in ["train", "val", "test"]} |
| self.t_va = one_text_per_image(self.t["val"], np.random.default_rng(0)) |
| self.t_te = one_text_per_image(self.t["test"], np.random.default_rng(k)) |
| self.n = {"mif": len(self.o["train"]), "rna": len(self.h["train"]), "text": len(self.t["train"])} |
|
|
| def nif(self, x): |
| return (x - self.if_mu) / self.if_sd |
|
|
|
|
| def batch_loss(m, F_, task, rng): |
| if task == "mif": |
| b = torch.as_tensor(rng.choice(F_.o["train"], BS, replace=False), device=dev) |
| d = torch.randint(0, 8, (BS,), device=dev) |
| x_he, x_m = HE_O[b, d], F_.nif(IF_O[b, d]); pos = None |
| elif task == "rna": |
| b = torch.as_tensor(rng.choice(F_.h["train"], BS, replace=False), device=dev) |
| x_he, x_m = HE_H[b], RNA_H[b]; pos = None |
| else: |
| b = torch.as_tensor(rng.choice(F_.t["train"], BS, replace=False), device=dev) |
| ii = T_IMG[b] |
| x_he, x_m = HE_T[ii, torch.randint(0, 8, (BS,), device=dev)], TXT[b] |
| pos = (ii[:, None] == ii[None, :]).float() |
| za = F.normalize(m.he(x_he, task), dim=-1) |
| zm = F.normalize(m.heads[task](x_m), dim=-1) |
| lg = m.logit[task].exp().clamp(max=100) * za @ zm.T |
| if pos is None: |
| t = torch.arange(BS, device=dev) |
| loss = (F.cross_entropy(lg, t) + F.cross_entropy(lg.T, t)) / 2 |
| else: |
| tgt = pos / pos.sum(1, keepdim=True) |
| loss = (torch.sum(-tgt * F.log_softmax(lg, 1), 1).mean() + torch.sum(-tgt * F.log_softmax(lg.T, 1), 1).mean()) / 2 |
| if m.uw is not None: |
| loss = torch.exp(-m.uw[task]) * loss + m.uw[task] |
| if m.distill and m.anchor != "frozen": |
| u = F.normalize(x_he, dim=-1) |
| loss = loss + m.distill * F.mse_loss(za @ za.T, u @ u.T) |
| return loss |
|
|
|
|
| @torch.no_grad() |
| def embed(m, F_, task, split): |
| if task == "mif": |
| idx = F_.o[split] |
| return F.normalize(m.he(HE_O[idx, 0], task), dim=-1), F.normalize(m.heads["mif"](F_.nif(IF_O[idx, 0])), dim=-1) |
| if task == "rna": |
| idx = F_.h[split] |
| return F.normalize(m.he(HE_H[idx], task), dim=-1), F.normalize(m.heads["rna"](RNA_H[idx]), dim=-1) |
| ii, tt = F_.t_va if split == "val" else F_.t_te |
| ii, tt = torch.as_tensor(ii, device=dev), torch.as_tensor(tt, device=dev) |
| return F.normalize(m.he(HE_T[ii, 0], task), dim=-1), F.normalize(m.heads["text"](TXT[tt]), dim=-1) |
|
|
|
|
| def val_scores(m, F_, tasks): |
| m.eval() |
| v = {t: r1(*embed(m, F_, t, "val")) for t in tasks} |
| m.train() |
| return v |
|
|
|
|
| def train(anchor, tasks, F_, steps, sampling="equal", distill=0.0, seed=0, head="mlp", uw=False, init=None, lr=LR): |
| torch.manual_seed(seed); rng = np.random.default_rng(seed) |
| m = Model(anchor, tasks, head, uw).to(dev); m.distill = distill |
| if init is not None: |
| for t, src in init.items(): |
| m.heads[t].load_state_dict(src.heads[t].state_dict()); m.logit[t].data.copy_(src.logit[t].data) |
| opt = torch.optim.AdamW(m.parameters(), lr=lr, weight_decay=0.05) |
| sched = torch.optim.lr_scheduler.OneCycleLR(opt, max_lr=lr, total_steps=steps, pct_start=0.05) |
| p = np.array([F_.n[t] for t in tasks], float); p /= p.sum() |
| best, state, bv = -1, None, None |
| m.train() |
| for step in range(1, steps + 1): |
| todo = tasks if sampling == "equal" else [tasks[rng.choice(len(tasks), p=p)]] |
| loss = sum(batch_loss(m, F_, t, rng) for t in todo) |
| opt.zero_grad(set_to_none=True); loss.backward(); opt.step(); sched.step() |
| if step % EVAL_EVERY == 0 or step == steps: |
| v = val_scores(m, F_, tasks) |
| score = float(np.exp(np.mean(np.log([max(x, 1e-4) for x in v.values()])))) |
| if score > best: |
| best, bv, state = score, v, {k_: x.detach().clone() for k_, x in m.state_dict().items()} |
| m.load_state_dict(state); m.eval() |
| return m, bv |
|
|
|
|
| @torch.no_grad() |
| def evaluate(models, F_): |
| """models: task -> model holding that task's head (same object for joint regimes)""" |
| out = {} |
| for t, lab in [("mif", "H&E<->mIF (ORION)"), ("rna", "H&E<->RNA (HEST)"), ("text", "H&E<->text (TCGA)")]: |
| out[(lab, "trained")] = g100(*embed(models[t], F_, t, "test")) |
| te = F_.o["test"] |
| zm = F.normalize(models["mif"].heads["mif"](F_.nif(IF_O[te, 0])), dim=-1) |
| zr = F.normalize(models["rna"].heads["rna"](RNA_V_O[te]), dim=-1) |
| zt = F.normalize(models["text"].heads["text"](GT_O[te]), dim=-1) |
| out[("mIF<->RNA(virtual) (ORION)", "emergent")] = g100(zm, zr) |
| out[("mIF<->text(synthetic) (ORION)", "emergent")] = g100(zm, zt) |
| ii, tt = F_.t_te |
| ii, tt = torch.as_tensor(ii, device=dev), torch.as_tensor(tt, device=dev) |
| zt = F.normalize(models["text"].heads["text"](TXT[tt]), dim=-1) |
| zr = F.normalize(models["rna"].heads["rna"](RNA_V_T[ii]), dim=-1) |
| zv = F.normalize(models["mif"].heads["mif"](F_.nif(IF_V_T[ii])), dim=-1) |
| out[("text<->RNA(virtual) (TCGA)", "emergent")] = g100(zt, zr) |
| out[("text<->mIF(MIPHEI) (TCGA)", "emergent")] = g100(zt, zv) |
| return out |
|
|
|
|
| HYB = {"mif": "mlp", "rna": "linear", "text": "linear"} |
| REG = {"S1": dict(anchor="frozen", mode="separate"), "S2": dict(anchor="pair", mode="separate"), |
| "S3": dict(anchor="frozen", mode="separate", head="linear"), |
| "J1": dict(anchor="frozen", mode="joint"), "J2": dict(anchor="shared", mode="joint"), |
| "J3": dict(anchor="shared", mode="joint", sampling="prop"), "J4": dict(anchor="shared", mode="joint", distill=1.0), |
| "J4b": dict(anchor="shared", mode="joint", distill=10.0), "J5": dict(anchor="residual", mode="joint"), |
| "J6": dict(anchor="shared", mode="joint", uw=True), "J7": dict(anchor="residual", mode="two-stage", distill=1.0), |
| "S4": dict(anchor="frozen", mode="separate", head=HYB), "J8": dict(anchor="residual", mode="two-stage", distill=1.0, head=HYB)} |
| NAME = {"S1": "separate, frozen anchor", "S2": "separate, two-tower", "J1": "joint, frozen anchor", |
| "J2": "joint, shared H&E head", "J3": "joint, shared H&E head, size-prop.", "J4": "joint, shared H&E head + distill", |
| "J5": "joint, residual adapter", "S3": "separate, frozen anchor, linear head", |
| "J4b": "joint, shared H&E head + distill x10", "J6": "joint, shared H&E head, uncertainty wt.", |
| "J7": "two-stage: separate -> joint residual", "S4": "separate, frozen, MLP mIF + linear RNA/text", |
| "J8": "two-stage from S4 heads -> joint residual"} |
| TASKS = ["mif", "rna", "text"] |
| f_out = os.path.join(ROOT, "results", "regimes_long.csv") |
| rows = pd.read_csv(f_out).to_dict("records") if os.path.exists(f_out) else [] |
| done = {(r["regime"], r["fold"]) for r in rows} |
| for k in args.folds: |
| F_ = Fold(k) |
| for rg in args.regimes: |
| if (rg, k) in done: |
| continue |
| c = REG[rg] |
| if c["mode"] in ("separate", "two-stage"): |
| models, val = {}, {} |
| for t in TASKS: |
| models[t], v = train("frozen" if c["mode"] == "two-stage" else c["anchor"], [t], F_, SEP_STEPS[t], |
| head=c.get("head", "mlp")) |
| val.update(v) |
| if c["mode"] == "joint": |
| sm = c.get("sampling", "equal") |
| m, val = train(c["anchor"], TASKS, F_, JOINT_STEPS[sm], sm, c.get("distill", 0.0), uw=c.get("uw", False)) |
| models = {t: m for t in TASKS} |
| elif c["mode"] == "two-stage": |
| m, val = train(c["anchor"], TASKS, F_, 1000, "equal", c["distill"], init=models, lr=1e-4, head=c.get("head", "mlp")) |
| models = {t: m for t in TASKS} |
| for (pair, kind), r in evaluate(models, F_).items(): |
| rows.append({"regime": rg, "name": NAME[rg], "fold": k, "pair": pair, "kind": kind, **r, |
| "val_mif": val["mif"], "val_rna": val["rna"], "val_text": val["text"]}) |
| pd.DataFrame(rows).to_csv(f_out, index=False) |
| last = [r for r in rows if r["regime"] == rg and r["fold"] == k] |
| print(f"fold {k} {rg:>2} {NAME[rg]:<36} " + " ".join(f"{r['pair'].split(' (')[0]} {r['R@1']:.3f}" for r in last), flush=True) |
|
|
| R = pd.DataFrame(rows) |
| S = R.groupby(["regime", "name", "kind", "pair"], sort=False)[["R@1", "R@10", "pair_acc"]].agg(["mean", "std"]).round(4) |
| S.to_csv(os.path.join(ROOT, "results", "regimes_summary.csv")) |
| pd.set_option("display.width", 250) |
| print(R.pivot_table(index=["regime", "name"], columns="pair", values="R@1", sort=False).round(3).to_string()) |
|
|
