| """Compare the four per-loop-target mean Jacobians in loop-relative coordinates. |
| |
| For loop k (target t_k = 48k-1) and physical layer L, take J_k[L] = J_{v->t_k} |
| with v = 48(k-1)+L. Matched-L cross-loop pairs share distance-to-target |
| (47 - L), so this comparison is distance-controlled by construction. |
| """ |
|
|
| import json |
|
|
| import torch |
| import torch.nn.functional as F |
|
|
| from jlens.lens import JacobianLens |
|
|
| torch.set_grad_enabled(False) |
| dev = "cuda" if torch.cuda.is_available() else "cpu" |
|
|
| LENSES = { |
| 1: "jacobian_lens_ouro-2.6b_t47.pt", |
| 2: "jacobian_lens_ouro-2.6b_t95.pt", |
| 3: "jacobian_lens_ouro-2.6b_t143.pt", |
| 4: "jacobian_lens_ouro-2.6b.pt", |
| } |
| lens = {k: JacobianLens.load(p) for k, p in LENSES.items()} |
| d = lens[4].jacobians[0].shape[0] |
| LS = list(range(47)) |
| M = {k: {L: lens[k].jacobians[48 * (k - 1) + L].float().to(dev) for L in LS} |
| for k in lens} |
| print({k: str(lens[k]) for k in lens}, flush=True) |
|
|
|
|
| def vcos(a, b): |
| return F.cosine_similarity(a.flatten(), b.flatten(), 0).item() |
|
|
|
|
| R = 32 |
| svd = {} |
| shape = {} |
| for k in M: |
| for L in LS: |
| J = M[k][L] |
| U, S, Vh = torch.linalg.svd(J) |
| svd[k, L] = (U[:, :R].contiguous(), S.cpu(), Vh[:R].contiguous()) |
| p = S / S.sum() |
| fro = J.norm().item() |
| shape[k, L] = { |
| "fro": fro / d ** 0.5, |
| "smax": S[0].item(), |
| "erank": torch.exp(-(p * p.log()).sum()).item(), |
| "stable_rank": (fro ** 2 / S[0].item() ** 2), |
| "id_cos": (torch.trace(J) / (d ** 0.5 * fro)).item(), |
| "sym_frac": (((J + J.T) / 2).norm() ** 2 / fro ** 2).item(), |
| "diag_frac": (J.diagonal().norm() ** 2 / fro ** 2).item(), |
| } |
| print("svd done", flush=True) |
|
|
| pairs = [(a, b) for a in range(1, 5) for b in range(a + 1, 5)] |
| res = {"matched": {}, "sub_r": {}, "sub_l": {}, "pairwise": {}} |
| for L in LS: |
| cs, sr, sl = [], [], [] |
| for a, b in pairs: |
| c = vcos(M[a][L], M[b][L]) |
| cs.append(c) |
| res["pairwise"].setdefault(f"{a}-{b}", {})[L] = c |
| Va, Vb = svd[a, L][2], svd[b, L][2] |
| Ua, Ub = svd[a, L][0], svd[b, L][0] |
| sr.append(((Va @ Vb.T).norm() ** 2 / R).item()) |
| sl.append(((Ua.T @ Ub).norm() ** 2 / R).item()) |
| res["matched"][L] = sum(cs) / len(cs) |
| res["sub_r"][L] = sum(sr) / len(sr) |
| res["sub_l"][L] = sum(sl) / len(sl) |
|
|
| adj = [vcos(M[k][L], M[k][L + 1]) for k in M for L in LS[:-1]] |
| off = {} |
| for dL in (2, 4, 8): |
| xs = [vcos(M[a][L], M[b][L + dL]) for a, b in pairs for L in LS |
| if L + dL in LS] |
| off[dL] = sum(xs) / len(xs) |
| g = torch.Generator().manual_seed(0) |
| rnd = [] |
| for _ in range(300): |
| a, b = pairs[torch.randint(6, (1,), generator=g).item()] |
| L1 = torch.randint(47, (1,), generator=g).item() |
| L2 = torch.randint(47, (1,), generator=g).item() |
| if L1 != L2: |
| rnd.append(vcos(M[a][L1], M[b][L2])) |
|
|
| print("\n=== cross-loop matched-L (distance-controlled) ===", flush=True) |
| print(" L | mcos subR subL | erank k=1..4 | idcos k=1..4", flush=True) |
| for L in LS: |
| er = " ".join(f"{shape[k, L]['erank']:5.0f}" for k in range(1, 5)) |
| ic = " ".join(f"{shape[k, L]['id_cos']:.2f}" for k in range(1, 5)) |
| print(f" {L:3d} | {res['matched'][L]:.2f} {res['sub_r'][L]:.2f} " |
| f"{res['sub_l'][L]:.2f} | {er} | {ic}", flush=True) |
|
|
| band = [L for L in LS if 24 <= L <= 40] |
| outb = [L for L in LS if L < 24] |
| print("\n=== summary ===", flush=True) |
| print(f"matched-L cos: all {sum(res['matched'].values())/47:.3f} " |
| f"band(24-40) {sum(res['matched'][L] for L in band)/len(band):.3f} " |
| f"early(<24) {sum(res['matched'][L] for L in outb)/len(outb):.3f}", flush=True) |
| print(f"baselines: within-loop adjacent-L {sum(adj)/len(adj):.3f} " |
| f"cross-loop offset dL=2 {off[2]:.3f} dL=4 {off[4]:.3f} dL=8 {off[8]:.3f} " |
| f"random {sum(rnd)/len(rnd):.3f}", flush=True) |
| print(f"subspace overlap r={R} (random baseline {R/d:.3f}): " |
| f"right {sum(res['sub_r'].values())/47:.3f} " |
| f"left {sum(res['sub_l'].values())/47:.3f}", flush=True) |
| print("per-pair mean matched-L cos:", flush=True) |
| for a, b in pairs: |
| m = sum(res["pairwise"][f"{a}-{b}"].values()) / 47 |
| print(f" loop{a}-loop{b}: {m:.3f}", flush=True) |
|
|
| print("\n=== chain-rule factorisation: J(v->191) vs J(mid->191) @ J(v->mid) ===", |
| flush=True) |
| res["compose"] = {} |
| for mid, km in ((95, 2), (143, 3)): |
| Jm = lens[4].jacobians[mid].float().to(dev) |
| for v in (8, 24, 32, 40, 56, 72, 80, 88, 104, 120, 128, 136): |
| if v >= mid: |
| continue |
| Jv = lens[km].jacobians[v].float().to(dev) |
| direct = lens[4].jacobians[v].float().to(dev) |
| comp = Jm @ Jv |
| c = vcos(direct, comp) |
| nr = (comp.norm() / direct.norm()).item() |
| res["compose"][f"{v}->{mid}->191"] = {"cos": c, "norm_ratio": nr} |
| print(f" v{v:3d} ->{mid}->191: cos {c:.3f} " |
| f"|comp|/|direct| {nr:.2f}", flush=True) |
|
|
| json.dump( |
| {"matched": res["matched"], "sub_r": res["sub_r"], "sub_l": res["sub_l"], |
| "pairwise": res["pairwise"], "compose": res["compose"], |
| "baselines": {"adjacent": sum(adj) / len(adj), |
| "offset": off, "random": sum(rnd) / len(rnd)}, |
| "shape": {f"{k},{L}": shape[k, L] for k, L in shape}, |
| "spectra_top16": {f"{k},{L}": svd[k, L][1][:16].tolist() |
| for k, L in svd}}, |
| open("jcompare_analysis.json", "w"), default=float) |
| print("\nDONE_JCOMPARE", flush=True) |
|
|
