| """Joint training of the multi-encoder (late-fusion) four-modality model: does the J8 two-stage gain of 42_ stack with |
| the encoder fusion of 47_? Real pairs only; virtual / synthetic data only at test time (same splits / test items). |
| |
| Stage 1 (separate, frozen UNI2 anchor; = 47_): one head per encoder - VirTues mIF -> MLP; every RNA / text encoder -> |
| linear (FM inputs z-scored with train stats, hvg / BiomedCLIP raw). |
| Stage 2 (joint, = J8 of 42_): all stage-1 heads + a shared zero-initialised residual adapter on UNI2 (H&E = UNI2 + |
| adapter(UNI2)) fine-tuned together, sum of the three pair losses (H&E-mIF, H&E-RNA, H&E-text), similarity |
| distillation lambda 1, AdamW lr 1e-4, 1000 steps, joint early stopping on the geometric mean of the three val R@1. |
| Configurations (per fold, all from the SAME stage-1 heads): |
| S4 stage 1 only, RNA hvg, text BiomedCLIP (single encoders, separate) |
| J8 S4 -> stage 2 (single encoders, joint) |
| E stage 1 only, ens-small (RNA hvg+Tahoe-x1+SE-600M+AIDO.Cell, text KEEP+BiomedCLIP+Qwen3-0.6B) |
| EJ-fused E -> stage 2, loss on the FUSED modality embedding normalise(sum_e normalise(head_e(x_e))) |
| EJ-each E -> stage 2, loss = mean over encoders of each head's own loss (fusion only at test time) |
| EJ-frz EJ-each with the RNA / text heads frozen (only the H&E adapter, the mIF head and the temperatures move) |
| Output results/joint_ens_long.csv, results/joint_ens_summary.csv, results/joint_ens_tests.csv |
| """ |
| import copy |
| 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 import stats |
| from scipy.optimize import linear_sum_assignment |
|
|
| from common import CACHE, ROOT, SPLITS |
|
|
| dev = "cuda" |
| H, RF = os.path.join(CACHE, "hest"), os.path.join(CACHE, "rnafm") |
| G100, DRAWS, BS, LR, EVAL_EVERY = 100, 50, 256, 3e-4, 50 |
| f32 = lambda a: torch.from_numpy(np.asarray(a, dtype=np.float32)) |
| nz = lambda x: F.normalize(x, dim=-1) |
| RNA_E, TXT_E = ["hvg", "tahoex1", "se600m", "aidocell"], ["keep", "biomedclip", "qwen3-0.6b"] |
| CONF = {"S4": (["hvg"], ["biomedclip"], None), "J8": (["hvg"], ["biomedclip"], "fused"), |
| "E": (RNA_E, TXT_E, None), "EJ-fused": (RNA_E, TXT_E, "fused"), "EJ-each": (RNA_E, TXT_E, "each"), |
| "EJ-frz": (RNA_E, TXT_E, "frz")} |
|
|
| |
| 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) |
| HE_T = f32(np.load(os.path.join(CACHE, "allimg_he_uni2.npy"))).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] |
| hvg = pd.read_csv(os.path.join(H, "hvg_genes.csv")).gene.values |
|
|
|
|
| def hvg_inputs(): |
| def virt(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 c[:, pd.Series(range(len(g)), index=g).loc[hvg].values] |
| return {"hest": np.load(os.path.join(H, "hvg.npy")), "orion_v": virt("orion"), "tcga_v": virt("tcga_text")} |
|
|
|
|
| RNA_IN = {e: ({s: f32(x).to(dev) for s, x in hvg_inputs().items()} if e == "hvg" else |
| {s: f32(np.load(os.path.join(RF, f"{e}_{s}.npy"))).to(dev) for s in ["hest", "orion_v", "tcga_v"]}) |
| for e in RNA_E} |
| TXT_IN = {e: {s: f32(np.load(os.path.join(CACHE, f"{s}_text_{e}.npy"))[:, 0]).to(dev) |
| for s in ["alltxt", "orion_gentext"]} for e in TXT_E} |
|
|
|
|
| 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) |
|
|
|
|
| def clip_loss(lg, pos=None): |
| if pos is None: |
| t = torch.arange(len(lg), device=dev) |
| return (F.cross_entropy(lg, t) + F.cross_entropy(lg.T, t)) / 2 |
| tgt = pos / pos.sum(1, keepdim=True) |
| return (torch.sum(-tgt * F.log_softmax(lg, 1), 1).mean() + torch.sum(-tgt * F.log_softmax(lg.T, 1), 1).mean()) / 2 |
|
|
|
|
| def zs(X, ref_set, idx, raw): |
| if raw: |
| return dict(X) |
| mu, sd = X[ref_set][idx].mean(0), X[ref_set][idx].std(0) + 1e-6 |
| return {s: (v - mu) / sd for s, v in X.items()} |
|
|
|
|
| class Joint(nn.Module): |
| """heads: 'mif' + 'rna:<enc>' + 'text:<enc>'; logits per head + fused 'rna' / 'text'; residual adapter on UNI2""" |
|
|
| def __init__(self, heads, logits, adapter): |
| super().__init__() |
| self.heads = nn.ModuleDict({k.replace(".", "_"): copy.deepcopy(h) for k, h in heads.items()}) |
| self.keys = {k.replace(".", "_"): k for k in heads} |
| self.logit = nn.ParameterDict({k.replace(":", "__").replace(".", "_"): nn.Parameter(v.detach().clone()) |
| for k, v in logits.items()}) |
| self.ah = None |
| if adapter: |
| self.ah = mlp(1536) |
| nn.init.zeros_(self.ah[-1].weight); nn.init.zeros_(self.ah[-1].bias) |
|
|
| def h(self, k): |
| return self.heads[k.replace(".", "_")] |
|
|
| def lg(self, k): |
| return self.logit[k.replace(":", "__").replace(".", "_")].exp().clamp(max=100) |
|
|
| def he(self, x): |
| return x if self.ah is None else x + self.ah(x) |
|
|
| def fuse(self, mod, encs, X): |
| """X: enc -> input rows; returns normalised fused embedding and the per-encoder embeddings""" |
| z = {e: nz(self.h(f"{mod}:{e}")(X[e])) for e in encs} |
| return nz(sum(z.values())), z |
|
|
|
|
| f_out = os.path.join(ROOT, "results", "joint_ens_long.csv") |
| rows = pd.read_csv(f_out).to_dict("records") if os.path.exists(f_out) else [] |
| done = {(r["config"], r["fold"]) for r in rows} |
|
|
| for k in range(5): |
| if all((c, k) in done for c in CONF): |
| continue |
| torch.manual_seed(k) |
| s = pd.read_csv(os.path.join(SPLITS, "patient_5fold", f"fold{k}.csv")) |
| o = {x: s.tile_id[s.split == x].values for x in ["train", "val", "test"]} |
| mu_if, sd_if = IF_O[o["train"]].mean((0, 1)), IF_O[o["train"]].std((0, 1)) + 1e-3 |
| nif = lambda x: (x - mu_if) / sd_if |
| hh = pd.read_csv(os.path.join(SPLITS, "final", f"hest_fold{k}.csv")) |
| hs = {x: hh.spot_idx[hh.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 |
| ts = {x: T.index[tsp == x].values for x in ["train", "val", "test"]} |
| vi, vt = (torch.as_tensor(x, device=dev) for x in one_text_per_image(ts["val"], np.random.default_rng(0))) |
| ii, tt = (torch.as_tensor(x, device=dev) for x in one_text_per_image(ts["test"], np.random.default_rng(k))) |
| trt = torch.as_tensor(ts["train"], device=dev) |
| Xr = {e: zs(RNA_IN[e], "hest", hs["train"], e == "hvg") for e in RNA_E} |
| Xt = {e: zs(TXT_IN[e], "alltxt", trt, e == "biomedclip") for e in TXT_E} |
| o_tr = torch.as_tensor(o["train"], device=dev) |
| h_tr = torch.as_tensor(hs["train"], device=dev) |
|
|
| def batch(task, rng): |
| """one training batch: (H&E anchor rows, modality inputs (enc -> rows), multi-positive matrix or None)""" |
| if task == "mif": |
| b = o_tr[torch.as_tensor(rng.choice(len(o_tr), BS, replace=False), device=dev)] |
| d = torch.randint(0, 8, (BS,), device=dev) |
| return HE_O[b, d], {"mif": nif(IF_O[b, d])}, None |
| if task == "rna": |
| b = h_tr[torch.as_tensor(rng.choice(len(h_tr), BS, replace=False), device=dev)] |
| return HE_H[b], {e: Xr[e]["hest"][b] for e in RNA_E}, None |
| b = trt[torch.as_tensor(rng.choice(len(trt), BS, replace=False), device=dev)] |
| im = T_IMG[b] |
| return (HE_T[im, torch.randint(0, 8, (BS,), device=dev)], {e: Xt[e]["alltxt"][b] for e in TXT_E}, |
| (im[:, None] == im[None, :]).float()) |
|
|
| VAL = {"mif": (HE_O[o["val"], 0], {"mif": nif(IF_O[o["val"], 0])}), |
| "rna": (HE_H[hs["val"]], {e: Xr[e]["hest"][hs["val"]] for e in RNA_E}), |
| "text": (HE_T[vi, 0], {e: Xt[e]["alltxt"][vt] for e in TXT_E})} |
|
|
| |
| def train_head(task, enc, kind, steps): |
| torch.manual_seed(k); rng = np.random.default_rng(k) |
| din = VAL[task][1][enc].shape[1] |
| h = (mlp(din) if kind == "mlp" else nn.Linear(din, 1536)).to(dev) |
| logit = nn.Parameter(torch.tensor(np.log(1 / 0.07), dtype=torch.float32, device=dev)) |
| opt = torch.optim.AdamW(list(h.parameters()) + [logit], lr=LR, weight_decay=0.05) |
| sch = torch.optim.lr_scheduler.OneCycleLR(opt, max_lr=LR, total_steps=steps, pct_start=0.05) |
| za_va = nz(VAL[task][0]) |
| best, state = -1, None |
| for step in range(1, steps + 1): |
| h.train() |
| a, X, pos = batch(task, rng) |
| loss = clip_loss(logit.exp().clamp(max=100) * nz(a) @ nz(h(X[enc])).T, pos) |
| opt.zero_grad(set_to_none=True); loss.backward(); opt.step(); sch.step() |
| if step % EVAL_EVERY == 0 or step == steps: |
| h.eval() |
| with torch.no_grad(): |
| v = r1(za_va, nz(h(VAL[task][1][enc]))) |
| if v > best: |
| best, state = v, ({k_: x.detach().clone() for k_, x in h.state_dict().items()}, logit.detach().clone()) |
| h.load_state_dict(state[0]); h.eval() |
| return h, state[1] |
|
|
| HEADS, LOGITS = {}, {} |
| HEADS["mif"], LOGITS["mif"] = train_head("mif", "mif", "mlp", 400) |
| for e in RNA_E: |
| HEADS[f"rna:{e}"], LOGITS[f"rna:{e}"] = train_head("rna", e, "linear", 2700) |
| for e in TXT_E: |
| HEADS[f"text:{e}"], LOGITS[f"text:{e}"] = train_head("text", e, "linear", 1000) |
| print(f"fold {k} stage 1 done", flush=True) |
|
|
| |
| def val_score(m, re_, te_): |
| m.eval() |
| with torch.no_grad(): |
| v = {"mif": r1(nz(m.he(VAL["mif"][0])), nz(m.h("mif")(VAL["mif"][1]["mif"]))), |
| "rna": r1(nz(m.he(VAL["rna"][0])), m.fuse("rna", re_, VAL["rna"][1])[0]), |
| "text": r1(nz(m.he(VAL["text"][0])), m.fuse("text", te_, VAL["text"][1])[0])} |
| m.train() |
| return float(np.exp(np.mean(np.log([max(x, 1e-4) for x in v.values()])))), v |
|
|
| def joint(re_, te_, mode, steps=1000, lr=1e-4, distill=1.0): |
| keys = ["mif"] + [f"rna:{e}" for e in re_] + [f"text:{e}" for e in te_] |
| lg = {kk: LOGITS[kk] for kk in keys} |
| lg["rna"] = torch.stack([LOGITS[f"rna:{e}"] for e in re_]).mean() |
| lg["text"] = torch.stack([LOGITS[f"text:{e}"] for e in te_]).mean() |
| m = Joint({kk: HEADS[kk] for kk in keys}, lg, adapter=True).to(dev) |
| torch.manual_seed(k); rng = np.random.default_rng(k) |
| if mode == "frz": |
| for kk in keys[1:]: |
| m.h(kk).requires_grad_(False) |
| opt = torch.optim.AdamW([p for p in m.parameters() if p.requires_grad], lr=lr, weight_decay=0.05) |
| sch = torch.optim.lr_scheduler.OneCycleLR(opt, max_lr=lr, total_steps=steps, pct_start=0.05) |
| best, state, bv = -1, None, None |
| m.train() |
| for step in range(1, steps + 1): |
| loss = 0 |
| for task in ("mif", "rna", "text"): |
| a, X, pos = batch(task, rng) |
| za = nz(m.he(a)) |
| if task == "mif": |
| loss = loss + clip_loss(m.lg("mif") * za @ nz(m.h("mif")(X["mif"])).T, pos) |
| else: |
| encs = re_ if task == "rna" else te_ |
| X = {e: X[e] for e in encs} |
| zf, ze = m.fuse(task, encs, X) |
| if mode == "fused": |
| loss = loss + clip_loss(m.lg(task) * za @ zf.T, pos) |
| else: |
| loss = loss + sum(clip_loss(m.lg(f"{task}:{e}") * za @ ze[e].T, pos) for e in encs) / len(encs) |
| u = nz(a) |
| loss = loss + distill * F.mse_loss(za @ za.T, u @ u.T) |
| opt.zero_grad(set_to_none=True); loss.backward(); opt.step(); sch.step() |
| if step % EVAL_EVERY == 0 or step == steps: |
| sc, v = val_score(m, re_, te_) |
| if sc > best: |
| best, bv, state = sc, 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(m, re_, te_): |
| R = lambda st, idx: m.fuse("rna", re_, {e: Xr[e][st][idx] for e in re_})[0] |
| Tx = lambda st, idx: m.fuse("text", te_, {e: Xt[e][st][idx] for e in te_})[0] |
| zm = nz(m.h("mif")(nif(IF_O[o["test"], 0]))) |
| zt = Tx("alltxt", tt) |
| return {("H&E<->mIF (ORION)", "trained"): g100(nz(m.he(HE_O[o["test"], 0])), zm), |
| ("H&E<->RNA (HEST)", "trained"): g100(nz(m.he(HE_H[hs["test"]])), R("hest", hs["test"])), |
| ("H&E<->text (TCGA)", "trained"): g100(nz(m.he(HE_T[ii, 0])), zt), |
| ("mIF<->RNA(virtual) (ORION)", "emergent"): g100(zm, R("orion_v", o["test"])), |
| ("mIF<->text(synthetic) (ORION)", "emergent"): g100(zm, Tx("orion_gentext", o["test"])), |
| ("text<->RNA(virtual) (TCGA)", "emergent"): g100(zt, R("tcga_v", ii)), |
| ("text<->mIF(MIPHEI) (TCGA)", "emergent"): g100(zt, nz(m.h("mif")(nif(IF_V_T[ii]))))} |
|
|
| for cname, (re_, te_, mode) in CONF.items(): |
| if (cname, k) in done: |
| continue |
| if mode is None: |
| keys = ["mif"] + [f"rna:{e}" for e in re_] + [f"text:{e}" for e in te_] |
| m = Joint({kk: HEADS[kk] for kk in keys}, {kk: LOGITS[kk] for kk in keys}, adapter=False).to(dev).eval() |
| val = {} |
| else: |
| m, val = joint(re_, te_, mode) |
| for (pair, kind), r in evaluate(m, re_, te_).items(): |
| rows.append({"config": cname, "fold": k, "pair": pair, "kind": kind, **r, |
| **{f"val_{t}": v for t, v in val.items()}}) |
| pd.DataFrame(rows).to_csv(f_out, index=False) |
| print(f"fold {k} {cname:<9} " + " ".join(f"{r['pair'].split(' (')[0]} {r['R@1']:.3f}" |
| for r in rows if r["config"] == cname and r["fold"] == k), flush=True) |
|
|
| R = pd.DataFrame(rows) |
| S = R.groupby(["config", "kind", "pair"], sort=False)[["R@1", "R@10", "pair_acc"]].agg(["mean", "std"]).round(4) |
| S.to_csv(os.path.join(ROOT, "results", "joint_ens_summary.csv")) |
| tests = [] |
| for c, ref in [("J8", "S4"), ("E", "S4"), ("EJ-fused", "S4"), ("EJ-each", "S4"), ("EJ-fused", "E"), ("EJ-each", "E"), |
| ("EJ-fused", "J8"), ("EJ-each", "EJ-fused"), ("EJ-frz", "S4"), |
| ("EJ-frz", "E"), ("EJ-frz", "EJ-each")]: |
| for p in R.pair.unique(): |
| a = R[(R.config == c) & (R.pair == p)].sort_values("fold")["R@1"].values |
| b = R[(R.config == ref) & (R.pair == p)].sort_values("fold")["R@1"].values |
| tests.append({"config": c, "vs": ref, "pair": p, "ref": b.mean(), "new": a.mean(), "delta": (a - b).mean(), |
| "folds_won": int((a > b).sum()), "p_paired_t": stats.ttest_rel(a, b).pvalue}) |
| pd.DataFrame(tests).round(4).to_csv(os.path.join(ROOT, "results", "joint_ens_tests.csv"), index=False) |
| pd.set_option("display.width", 250) |
| print(R.pivot_table(index="config", columns="pair", values="R@1", sort=False).round(3).T.to_string()) |
| print(R.pivot_table(index="config", columns="pair", values="pair_acc", sort=False).round(3).T.to_string()) |
| print(pd.DataFrame(tests).round(4).to_string()) |
|
|
