""" New mapping methods on top of the existing 25-anchor experiment data: (f) procrustes Per-tensor orthogonal Procrustes alignment between X-anchor and Y-anchor matrices in tensor-flat space. Respects the LoRA gauge ambiguity better than ridge. (g) topk_global_ridge Pick top-K anchors by X-side cosine similarity to X_target, then run global anchor-basis ridge on that subset. (h) topk_pertensor_ridge Same, but per-tensor ridge. (i) topk_pertensor_mlp Top-K + per-tensor PCA-MLP hypernet (re-trained per task). Reuses all training artifacts from /app/scaled/X and /app/scaled/Y. """ import os, json, gc, shutil, re, collections, math from pathlib import Path import numpy as np import torch import torch.nn.functional as F from datasets import load_dataset from transformers import AutoTokenizer, AutoModelForCausalLM, set_seed from peft import PeftModel from safetensors.torch import load_file, save_file # Reuse main pipeline's task definitions import sys; sys.path.insert(0, "/app") import scaled_pipeline as sp set_seed(42) OUT = sp.OUT # /app/scaled ANCHOR_NAMES = sp.ANCHOR_NAMES HELDOUT_NAMES = sp.HELDOUT_NAMES MODEL_Y = sp.MODEL_Y EVAL_PER_TASK = sp.EVAL_PER_TASK def load_sd(p): return {k: v.float().cpu() for k,v in load_file(str(p/"adapter_model.safetensors")).items()} def parse_key(k): m = re.search(r'layers\.(\d+)\.self_attn\.(\w+)_proj\.lora_(A|B)', k) return int(m.group(1)), m.group(2), m.group(3) def group_by_layer(sd): g = collections.defaultdict(dict) for k, v in sd.items(): L, mod, AB = parse_key(k) g[(mod, AB)][L] = (k, v) return g def align_X_to_Y(X_anchors_list, X_target, Y_template_sd): Y_g = group_by_layer(Y_template_sd) X_gs = [group_by_layer(s) for s in X_anchors_list] X_t = group_by_layer(X_target) nLY = max(Y_g[("q","A")].keys()) + 1 nLX = max(X_gs[0][("q","A")].keys()) + 1 out = {} for (mod, AB), layers in Y_g.items(): for L_y, (yk, _) in layers.items(): L_x = round(L_y * (nLX-1) / max(1, nLY-1)) x_anchors = [X_gs[i][(mod,AB)][L_x][1] for i in range(len(X_gs))] x_target = X_t[(mod,AB)][L_x][1] out[yk] = (x_anchors, x_target) return out # ----------------- (f) Procrustes ----------------- def procrustes_pred(X_anchors, Y_anchors, X_target): """Per-tensor orthogonal Procrustes (rectangular), avoiding dx*dy materialization. Cross-covariance M = Xc^T Yc has rank ≤ N. Compute its SVD efficiently: Xc = U_x S_x V_x^T (compact, V_x^T is (N, dx), V_x is (dx, N)) Yc = U_y S_y V_y^T Then SVD(M) <- via N-by-N inner SVD; never materialize dx-by-dy. """ align = align_X_to_Y(X_anchors, X_target, Y_anchors[0]) out = {} for yk, (x_anchors, x_target) in align.items(): Y_mat = torch.stack([s[yk].reshape(-1) for s in Y_anchors]) # (N, dy) X_mat = torch.stack([t.reshape(-1) for t in x_anchors]) # (N, dx) Ym, Xm = Y_mat.mean(0), X_mat.mean(0) Yc, Xc = Y_mat - Ym, X_mat - Xm Ux_, Sx_, Vxh = torch.linalg.svd(Xc, full_matrices=False) Uy_, Sy_, Vyh = torch.linalg.svd(Yc, full_matrices=False) C = (Sx_.unsqueeze(1) * Ux_.T) @ (Uy_ * Sy_.unsqueeze(0)) # (N, N) Uz, Sz, Vzh = torch.linalg.svd(C, full_matrices=False) scale = (Sz.sum() / (Sx_.pow(2).sum() + 1e-8)).item() xt_c = x_target.reshape(-1) - Xm a = Vxh @ xt_c # (N,) b = (Uz @ Vzh) @ a # (N,) delta = b @ Vyh # (dy,) out[yk] = (Ym + scale * delta).reshape(Y_anchors[0][yk].shape) return out # ----------------- (g/h/i) top-K nearest anchor selection ----------------- def topk_anchor_idx(X_anchors, X_target, k=8): """Return indices of top-K anchors by cosine similarity to X_target.""" keys = sorted(X_anchors[0].keys()) Xa = torch.stack([torch.cat([s[kk].flatten() for kk in keys]) for s in X_anchors]) # (N, dx) xt = torch.cat([X_target[kk].flatten() for kk in keys]) sims = F.cosine_similarity(Xa, xt.unsqueeze(0), dim=1) idx = torch.topk(sims, k=min(k, len(X_anchors))).indices.tolist() return idx, sims.tolist() def topk_global_ridge(X_anchors, Y_anchors, X_target, k=8, ridge=1e-3): idx, sims = topk_anchor_idx(X_anchors, X_target, k=k) Xa = [X_anchors[i] for i in idx] Ya = [Y_anchors[i] for i in idx] pred, alpha = sp.global_ridge_pred(Xa, Ya, X_target, ridge=ridge) return pred, idx, sims def topk_pertensor_ridge(X_anchors, Y_anchors, X_target, k=8, ridge=1e-3): idx, sims = topk_anchor_idx(X_anchors, X_target, k=k) Xa = [X_anchors[i] for i in idx] Ya = [Y_anchors[i] for i in idx] pred = sp.pertensor_ridge_pred(Xa, Ya, X_target, ridge=ridge) return pred, idx, sims def topk_pertensor_mlp(X_anchors, Y_anchors, X_target, k=12, k_lat=6, epochs=300, lr=1e-3): idx, sims = topk_anchor_idx(X_anchors, X_target, k=k) Xa = [X_anchors[i] for i in idx] Ya = [Y_anchors[i] for i in idx] pred, _ = sp.pertensor_pca_mlp_pred(Xa, Ya, X_target, k_lat=k_lat, epochs=epochs, lr=lr) return pred, idx, sims # ----------------- save / eval ----------------- def save_adapter(sd, dirname, src_template): d = OUT/"Y_pred"/dirname d.mkdir(parents=True, exist_ok=True) shutil.copy(src_template/"adapter_config.json", d/"adapter_config.json") save_file({k: v.to(torch.bfloat16) for k,v in sd.items()}, str(d/"adapter_model.safetensors")) return d @torch.no_grad() def eval_adapter(adapter_dir, eval_ds, labels, tok, max_n=300): base = AutoModelForCausalLM.from_pretrained(MODEL_Y, torch_dtype=torch.bfloat16, attn_implementation="eager").cuda() if adapter_dir is None: m = base else: m = PeftModel.from_pretrained(base, str(adapter_dir)) acc = sp.eval_classification(m, tok, eval_ds, labels, max_n=max_n) del m, base if adapter_dir is None: pass gc.collect(); torch.cuda.empty_cache() return acc def cos_sd(sd1, sd2): keys = sorted(sd1.keys()) a = torch.cat([sd1[k].float().flatten() for k in keys]) b = torch.cat([sd2[k].float().flatten() for k in keys]) return F.cosine_similarity(a.unsqueeze(0), b.unsqueeze(0)).item() # ----------------- main ----------------- def main(): print("Loading anchors...") X_anchors = [load_sd(OUT/"X"/n) for n in ANCHOR_NAMES] Y_anchors = [load_sd(OUT/"Y"/n) for n in ANCHOR_NAMES] print(f" N={len(X_anchors)} anchors") tokY = AutoTokenizer.from_pretrained(MODEL_Y) if tokY.pad_token is None: tokY.pad_token = tokY.eos_token tokY.padding_side = "left" out_results = {} template = OUT/"Y"/ANCHOR_NAMES[0] for t_name in HELDOUT_NAMES: print(f"\n=== Held-out task: {t_name} ===") X_target = load_sd(OUT/"X"/t_name) Y_oracle = load_sd(OUT/"Y"/t_name) results = {"cosines": {}, "acc": {}, "selected_anchors": {}} # build predictions preds = {} print(" procrustes ...") preds["procrustes"] = procrustes_pred(X_anchors, Y_anchors, X_target) for k_top in [5, 8, 12]: print(f" topk_global_ridge k={k_top} ...") p, idx, sims = topk_global_ridge(X_anchors, Y_anchors, X_target, k=k_top) preds[f"topk{k_top}_global_ridge"] = p results["selected_anchors"][f"topk{k_top}"] = [ANCHOR_NAMES[i] for i in idx] print(f" selected: {[ANCHOR_NAMES[i] for i in idx]}") print(f" topk_pertensor_ridge k={k_top} ...") p, _, _ = topk_pertensor_ridge(X_anchors, Y_anchors, X_target, k=k_top) preds[f"topk{k_top}_pertensor_ridge"] = p print(" topk12_pertensor_mlp ...") p, _, _ = topk_pertensor_mlp(X_anchors, Y_anchors, X_target, k=12, k_lat=6, epochs=300) preds["topk12_pertensor_mlp"] = p # save and eval _, eval_ds, labels = sp.build_task(t_name, n_train=10, n_eval=EVAL_PER_TASK) # cosines for n, sd in preds.items(): results["cosines"][n] = cos_sd(sd, Y_oracle) # accuracies results["acc"]["base_Y"] = eval_adapter(None, eval_ds, labels, tokY) print(f" base_Y={results['acc']['base_Y']:.4f}") for n, sd in preds.items(): d = save_adapter(sd, f"{t_name}_{n}", template) acc = eval_adapter(d, eval_ds, labels, tokY) results["acc"][n] = acc print(f" {n:32s} cos={results['cosines'][n]:.4f} acc={acc:.4f}") results["acc"]["oracle_Y"] = eval_adapter(OUT/"Y"/t_name, eval_ds, labels, tokY) print(f" oracle_Y={results['acc']['oracle_Y']:.4f}") out_results[t_name] = results # aggregate methods = ["base_Y","procrustes","topk5_global_ridge","topk5_pertensor_ridge", "topk8_global_ridge","topk8_pertensor_ridge", "topk12_global_ridge","topk12_pertensor_ridge","topk12_pertensor_mlp","oracle_Y"] avg = {m: float(np.mean([out_results[t]["acc"][m] for t in HELDOUT_NAMES])) for m in methods} print("\n=== AVG (new methods) ===") for m in methods: print(f" {m:32s} {avg[m]:.4f}") out_results["avg"] = avg (OUT/"results_phaseA.json").write_text(json.dumps(out_results, indent=2, default=float)) if __name__ == "__main__": main()