Raniahossam33/knowledge-drift-experiments
022
1#!/usr/bin/env python32"""3cross_model.py — Cross-Model Drift Analysis4=============================================5Runs AFTER analyze_single.py on 2+ models.6Uses probe bundles + caches to compare drift representations across architectures.7 86 Experiments:9 [CM-1] Full-layer CKA matrix (L_A × L_B per pair, not just best layer)10 [CM-2] Drift score correlation (probe A scores vs probe B scores on shared queries)11 [CM-3] Differential facts (queries drifted for A but stable for B)12 [CM-4] Layer correspondence (best layer as % of depth — universal localization?)13 [CM-5] Neuron overlap (same-dim models only: which neuron indices carry drift?)14 [CM-6] Universality score (aggregate metric for paper abstract)15 16Outputs:17 cross_model_results.json Complete results18 figures/fig_cm1_cka.png Layer-wise CKA heatmaps19 figures/fig_cm2_corr.png Score correlation matrix20 figures/fig_cm3_diff.png Differential facts scatter21 figures/fig_cm4_layers.png Layer correspondence bar22 figures/fig_cm5_neurons.png Neuron overlap (same-dim pairs)23 figures/fig_cm6_summary.png Universality summary24 25Usage:26 # Compare two models27 python cross_model.py --models qwen25 llama3128 29 # All available models30 python cross_model.py --all31 32 # Quick mode (skip full-layer CKA, just best-layer)33 python cross_model.py --all --quick34"""35 36import argparse37import json38import logging39import time40import warnings41from pathlib import Path42 43import numpy as np44import yaml45 46warnings.filterwarnings("ignore")47logging.basicConfig(48 level=logging.INFO,49 format="%(asctime)s [%(levelname)s] %(message)s",50 handlers=[logging.StreamHandler()])51logger = logging.getLogger(__name__)52 53 54# ─────────────────────────────────────────────────────────────────────────────55# CONFIG + DATA LOADING56# ─────────────────────────────────────────────────────────────────────────────57 58def load_config(path="models.yaml"):59 with open(path) as f:60 return yaml.safe_load(f)61 62 63def load_cache(model_dir, model_key):64 path = Path(model_dir) / model_key / f"cached_{model_key}.npz"65 if not path.exists():66 logger.error(f"Cache not found: {path}")67 return None68 results = np.load(str(path), allow_pickle=True)["results"].tolist()69 logger.info(f" [{model_key}] Loaded {len(results)} samples")70 return results71 72 73def load_probe_bundle(model_dir, model_key):74 path = Path(model_dir) / model_key / f"probe_bundle_{model_key}.npz"75 if not path.exists():76 logger.warning(f"Probe bundle not found: {path}")77 return None78 d = np.load(str(path), allow_pickle=True)79 bundle = {k: d[k] for k in d.files}80 # Convert scalar items81 for k in ["best_layer", "hidden_dim", "n_samples"]:82 if k in bundle:83 bundle[k] = int(bundle[k])84 for k in ["drift_auroc", "cos_du", "cos_dc"]:85 if k in bundle:86 bundle[k] = float(bundle[k])87 logger.info(f" [{model_key}] Bundle: layer={bundle.get('best_layer')}, "88 f"dim={bundle.get('hidden_dim')}, "89 f"AUROC={bundle.get('drift_auroc', 0):.4f}")90 return bundle91 92 93def load_final_results(model_dir, model_key):94 path = Path(model_dir) / model_key / "final_results.json"95 if not path.exists():96 return None97 with open(path) as f:98 return json.load(f)99 100 101# ─────────────────────────────────────────────────────────────────────────────102# PROBE FITTING (lightweight — for scoring shared queries)103# ─────────────────────────────────────────────────────────────────────────────104 105def soft_threshold(w, lam):106 import torch107 return torch.sign(w) * torch.clamp(torch.abs(w) - lam, min=0.0)108 109 110def fit_quick_probe(X_np, y_np, device="cuda:0", lam=1e-3, max_iter=500):111 """Fast probe fit for cross-model scoring."""112 import torch113 X = np.nan_to_num(X_np.astype(np.float32), nan=0., posinf=1e4, neginf=-1e4)114 X = np.clip(X, -1e4, 1e4)115 m = X.mean(0, keepdims=True)116 s = X.std(0, keepdims=True) + 1e-8117 Xt = torch.tensor((X - m) / s, dtype=torch.float32, device=device)118 yt = torch.tensor(y_np.astype(np.float32), device=device)119 120 w = torch.zeros(Xt.shape[1], device=device)121 b = torch.zeros(1, device=device)122 lr = 1.0123 for _ in range(max_iter):124 z = torch.clamp(Xt @ w + b, -30, 30)125 p = torch.sigmoid(z)126 L = -((yt * torch.log(p + 1e-12)) +127 (1 - yt) * torch.log(1 - p + 1e-12)).mean()128 e = p - yt129 gw = (Xt.T @ e) / len(yt)130 gb = e.mean(keepdim=True)131 wt = soft_threshold(w - lr * gw, lr * lam)132 bt = b - lr * gb133 Lt = -((yt * torch.log(torch.sigmoid(torch.clamp(Xt @ wt + bt, -30, 30)) + 1e-12)) +134 (1 - yt) * torch.log(1 - torch.sigmoid(torch.clamp(Xt @ wt + bt, -30, 30)) + 1e-12)).mean()135 if Lt > L + 1e-4:136 lr *= 0.5137 else:138 lr = min(lr * 1.05, 10.0)139 if (wt - w).abs().max().item() < 1e-6:140 w, b = wt, bt141 break142 w, b = wt, bt143 144 def score(X_new):145 Xn = np.nan_to_num(X_new.astype(np.float32), nan=0., posinf=1e4, neginf=-1e4)146 Xn = np.clip(Xn, -1e4, 1e4)147 Xn = torch.tensor((Xn - m) / s, dtype=torch.float32, device=device)148 with torch.no_grad():149 return torch.sigmoid(torch.clamp(Xn @ w + b, -30, 30)).cpu().numpy()150 151 return score, w.cpu().numpy()152 153 154# ─────────────────────────────────────────────────────────────────────────────155# [CM-1] CKA ANALYSIS156# ─────────────────────────────────────────────────────────────────────────────157 158def linear_cka(Xa, Xb):159 """Centered Kernel Alignment between two representation matrices."""160 def _center(K):161 n = K.shape[0]162 H = np.eye(n) - 1.0 / n163 return H @ K @ H164 Ka = _center(Xa @ Xa.T)165 Kb = _center(Xb @ Xb.T)166 num = np.linalg.norm(Ka.T @ Kb, "fro")167 den = np.linalg.norm(Ka, "fro") * np.linalg.norm(Kb, "fro")168 return float(num / (den + 1e-12))169 170 171def cka_analysis(res_a, res_b, key_a, key_b, quick=False):172 """173 [CM-1] CKA between two models.174 If quick=False: full L_A × L_B heatmap.175 If quick=True: just best-layer CKA.176 """177 logger.info(f"[CM-1] CKA: {key_a} vs {key_b}")178 179 # Build shared query lookup180 qa = {r["query"]: r for r in res_a}181 qb = {r["query"]: r for r in res_b}182 shared = sorted(set(qa) & set(qb))183 logger.info(f" Shared queries: {len(shared)}")184 185 if len(shared) < 50:186 logger.warning(" Too few shared queries for CKA")187 return None188 189 # Subsample for speed (CKA is O(n²))190 if len(shared) > 2000:191 np.random.seed(42)192 shared = list(np.random.choice(shared, 2000, replace=False))193 194 layers_a = sorted(res_a[0]["hidden_states"].keys())195 layers_b = sorted(res_b[0]["hidden_states"].keys())196 197 if quick:198 # Just best layers199 best_a = layers_a[-5:] # top 5 layers200 best_b = layers_b[-5:]201 else:202 # Sample layers evenly (max 10 per model for tractability)203 step_a = max(1, len(layers_a) // 10)204 step_b = max(1, len(layers_b) // 10)205 best_a = layers_a[::step_a]206 best_b = layers_b[::step_b]207 208 cka_mat = np.zeros((len(best_a), len(best_b)))209 for i, la in enumerate(best_a):210 Xa = np.array([qa[q]["hidden_states"][la] for q in shared])211 for j, lb in enumerate(best_b):212 Xb = np.array([qb[q]["hidden_states"][lb] for q in shared])213 cka_mat[i, j] = linear_cka(Xa, Xb)214 if (i + 1) % 3 == 0:215 logger.info(f" CKA row {i+1}/{len(best_a)}")216 217 best_cka = float(cka_mat.max())218 logger.info(f" Best CKA: {best_cka:.4f}")219 220 return {221 "layers_a": best_a, "layers_b": best_b,222 "cka_matrix": cka_mat.tolist(),223 "best_cka": best_cka,224 "n_shared": len(shared),225 }226 227 228# ─────────────────────────────────────────────────────────────────────────────229# [CM-2] SCORE CORRELATION230# ─────────────────────────────────────────────────────────────────────────────231 232def score_correlation(res_a, res_b, key_a, key_b, bundle_a, bundle_b, device):233 """234 [CM-2] Train probe on each model, score shared queries, correlate.235 """236 from sklearn.metrics import roc_auc_score237 logger.info(f"[CM-2] Score correlation: {key_a} vs {key_b}")238 239 qa = {r["query"]: r for r in res_a}240 qb = {r["query"]: r for r in res_b}241 shared = sorted(set(qa) & set(qb))242 logger.info(f" Shared: {len(shared)}")243 244 if len(shared) < 50:245 return None246 247 bl_a = int(bundle_a["best_layer"])248 bl_b = int(bundle_b["best_layer"])249 250 # Train probes on full data251 X_a = np.array([r["hidden_states"][bl_a] for r in res_a])252 y_a = np.array([int(r["is_drifted"]) for r in res_a])253 X_b = np.array([r["hidden_states"][bl_b] for r in res_b])254 y_b = np.array([int(r["is_drifted"]) for r in res_b])255 256 score_a, _ = fit_quick_probe(X_a, y_a, device)257 score_b, _ = fit_quick_probe(X_b, y_b, device)258 259 # Score shared queries260 Xa_shared = np.array([qa[q]["hidden_states"][bl_a] for q in shared])261 Xb_shared = np.array([qb[q]["hidden_states"][bl_b] for q in shared])262 sa = score_a(Xa_shared)263 sb = score_b(Xb_shared)264 265 # Labels for shared266 ya_shared = np.array([int(qa[q]["is_drifted"]) for q in shared])267 yb_shared = np.array([int(qb[q]["is_drifted"]) for q in shared])268 269 corr = float(np.corrcoef(sa, sb)[0, 1])270 try:271 auroc_a = roc_auc_score(ya_shared, sa)272 auroc_b = roc_auc_score(yb_shared, sb)273 except Exception:274 auroc_a = auroc_b = 0.5275 276 logger.info(f" Score corr: {corr:.4f} "277 f"AUROC_a={auroc_a:.4f} AUROC_b={auroc_b:.4f}")278 279 return {280 "correlation": corr,281 "auroc_a_on_shared": auroc_a,282 "auroc_b_on_shared": auroc_b,283 "n_shared": len(shared),284 "scores_a": sa.tolist(),285 "scores_b": sb.tolist(),286 }287 288 289# ─────────────────────────────────────────────────────────────────────────────290# [CM-3] DIFFERENTIAL FACTS291# ─────────────────────────────────────────────────────────────────────────────292 293def differential_facts(res_a, res_b, key_a, key_b, bundle_a, bundle_b, device):294 """295 [CM-3] Queries where is_drifted differs between models.296 Each probe should detect its own model's drift correctly.297 """298 from sklearn.metrics import roc_auc_score299 logger.info(f"[CM-3] Differential facts: {key_a} vs {key_b}")300 301 qa = {r["query"]: r for r in res_a}302 qb = {r["query"]: r for r in res_b}303 shared = sorted(set(qa) & set(qb))304 305 # Find differential: drifted for A but not B, or vice versa306 diff_queries = [q for q in shared307 if qa[q]["is_drifted"] != qb[q]["is_drifted"]]308 logger.info(f" Shared={len(shared)}, Differential={len(diff_queries)}")309 310 if len(diff_queries) < 20:311 logger.warning(" Too few differential facts")312 return None313 314 bl_a = int(bundle_a["best_layer"])315 bl_b = int(bundle_b["best_layer"])316 317 # Train probes318 X_a = np.array([r["hidden_states"][bl_a] for r in res_a])319 y_a = np.array([int(r["is_drifted"]) for r in res_a])320 X_b = np.array([r["hidden_states"][bl_b] for r in res_b])321 y_b = np.array([int(r["is_drifted"]) for r in res_b])322 score_a, _ = fit_quick_probe(X_a, y_a, device)323 score_b, _ = fit_quick_probe(X_b, y_b, device)324 325 # Score differential queries326 Xa_d = np.array([qa[q]["hidden_states"][bl_a] for q in diff_queries])327 Xb_d = np.array([qb[q]["hidden_states"][bl_b] for q in diff_queries])328 sa = score_a(Xa_d)329 sb = score_b(Xb_d)330 la = np.array([int(qa[q]["is_drifted"]) for q in diff_queries])331 lb = np.array([int(qb[q]["is_drifted"]) for q in diff_queries])332 333 try:334 auroc_a = roc_auc_score(la, sa)335 except Exception:336 auroc_a = 0.5337 try:338 auroc_b = roc_auc_score(lb, sb)339 except Exception:340 auroc_b = 0.5341 342 # Anti-correlation: when A says drifted and B says stable,343 # score_a should be high and score_b should be low344 score_corr = float(np.corrcoef(sa, sb)[0, 1])345 346 # Count categories347 a_only = sum(1 for q in diff_queries348 if qa[q]["is_drifted"] and not qb[q]["is_drifted"])349 b_only = sum(1 for q in diff_queries350 if not qa[q]["is_drifted"] and qb[q]["is_drifted"])351 352 logger.info(f" AUROC_a={auroc_a:.4f} AUROC_b={auroc_b:.4f} "353 f"score_corr={score_corr:.4f}")354 logger.info(f" A-only drifted: {a_only} B-only drifted: {b_only}")355 356 return {357 "n_differential": len(diff_queries),358 "n_shared": len(shared),359 "a_only_drifted": a_only,360 "b_only_drifted": b_only,361 "auroc_a": auroc_a,362 "auroc_b": auroc_b,363 "score_correlation": score_corr,364 "scores_a": sa.tolist(),365 "scores_b": sb.tolist(),366 "labels_a": la.tolist(),367 "labels_b": lb.tolist(),368 }369 370 371# ─────────────────────────────────────────────────────────────────────────────372# [CM-4] LAYER CORRESPONDENCE373# ─────────────────────────────────────────────────────────────────────────────374 375def layer_correspondence(all_bundles, all_final):376 """377 [CM-4] Best drift layer as fraction of total depth.378 If all models peak at ~80%, drift localization is universal.379 """380 logger.info("[CM-4] Layer correspondence")381 data = {}382 for key in all_bundles:383 bl = int(all_bundles[key]["best_layer"])384 total = int(all_bundles[key].get("hidden_dim", 0))385 # Get total layers from final results386 fr = all_final.get(key, {})387 n_layers = fr.get("best_layer_results", {}).get("layer", bl) + 1388 # Better: look at probe stability layers389 stab = fr.get("probe_stability", {})390 if "layers" in stab and len(stab["layers"]) > 0:391 n_layers = max(stab["layers"]) + 1392 393 frac = bl / max(n_layers, 1)394 auroc = float(all_bundles[key].get("drift_auroc", 0))395 data[key] = {396 "best_layer": bl,397 "n_layers": n_layers,398 "fraction": frac,399 "auroc": auroc,400 }401 logger.info(f" {key}: L{bl}/{n_layers} = {frac:.1%} "402 f"AUROC={auroc:.4f}")403 404 fracs = [v["fraction"] for v in data.values()]405 mean_frac = float(np.mean(fracs))406 std_frac = float(np.std(fracs))407 logger.info(f" Mean fraction: {mean_frac:.1%} +/- {std_frac:.1%}")408 409 return {410 "per_model": data,411 "mean_fraction": mean_frac,412 "std_fraction": std_frac,413 }414 415 416# ─────────────────────────────────────────────────────────────────────────────417# [CM-5] NEURON OVERLAP (same-dim models only)418# ─────────────────────────────────────────────────────────────────────────────419 420def neuron_overlap(bundle_a, bundle_b, key_a, key_b):421 """422 [CM-5] For same-dimension models: do the same neuron indices carry drift?423 """424 dim_a = int(bundle_a["hidden_dim"])425 dim_b = int(bundle_b["hidden_dim"])426 427 if dim_a != dim_b:428 logger.info(f"[CM-5] {key_a}({dim_a}) vs {key_b}({dim_b}): "429 f"dim mismatch, skipping")430 return None431 432 logger.info(f"[CM-5] Neuron overlap: {key_a} vs {key_b} (dim={dim_a})")433 434 w_a = bundle_a["w_drift"]435 w_b = bundle_b["w_drift"]436 437 active_a = set(np.where(w_a != 0)[0])438 active_b = set(np.where(w_b != 0)[0])439 440 inter = len(active_a & active_b)441 union = len(active_a | active_b)442 jacc = inter / union if union > 0 else 0.0443 444 # Cosine of weight vectors (even though from different models)445 cos = float(np.dot(w_a, w_b) / (np.linalg.norm(w_a) * np.linalg.norm(w_b) + 1e-12))446 447 # Top-k overlap448 top100_a = set(np.argsort(np.abs(w_a))[-100:])449 top100_b = set(np.argsort(np.abs(w_b))[-100:])450 top100_overlap = len(top100_a & top100_b) / 100.0451 452 logger.info(f" Active: A={len(active_a)}, B={len(active_b)}")453 logger.info(f" Jaccard: {jacc:.4f} Cosine: {cos:.4f} "454 f"Top-100 overlap: {top100_overlap:.2%}")455 456 return {457 "dim": dim_a,458 "n_active_a": len(active_a),459 "n_active_b": len(active_b),460 "intersection": inter,461 "union": union,462 "jaccard": jacc,463 "cosine": cos,464 "top100_overlap": top100_overlap,465 }466 467 468# ─────────────────────────────────────────────────────────────────────────────469# [CM-6] UNIVERSALITY SCORE470# ─────────────────────────────────────────────────────────────────────────────471 472def universality_score(all_cka, all_corr, all_diff, all_layer_corr,473 n_bootstrap=1000):474 """475 [CM-6] Aggregate metric: geometric mean of CKA, score correlation,476 differential AUROC, and layer consistency.477 """478 logger.info("[CM-6] Universality score")479 480 components = {}481 482 # Mean best CKA across pairs483 cka_vals = [v["best_cka"] for v in all_cka.values() if v]484 if cka_vals:485 components["mean_cka"] = float(np.mean(cka_vals))486 487 # Mean score correlation488 corr_vals = [v["correlation"] for v in all_corr.values() if v]489 if corr_vals:490 components["mean_score_corr"] = float(np.mean(corr_vals))491 492 # Mean differential AUROC493 diff_aurocs = []494 for v in all_diff.values():495 if v:496 diff_aurocs.extend([v["auroc_a"], v["auroc_b"]])497 if diff_aurocs:498 components["mean_diff_auroc"] = float(np.mean(diff_aurocs))499 500 # Layer consistency (1 - std of fractions)501 if all_layer_corr:502 components["layer_consistency"] = float(503 1.0 - all_layer_corr.get("std_fraction", 0.5))504 505 if not components:506 return None507 508 vals = list(components.values())509 # Geometric mean510 geo_mean = float(np.exp(np.mean(np.log(np.clip(vals, 1e-6, None)))))511 512 # Bootstrap CI513 boot = []514 for _ in range(n_bootstrap):515 idx = np.random.choice(len(vals), len(vals), replace=True)516 boot.append(np.exp(np.mean(np.log(np.clip(np.array(vals)[idx], 1e-6, None)))))517 ci_lo = float(np.percentile(boot, 2.5))518 ci_hi = float(np.percentile(boot, 97.5))519 520 logger.info(f" Components: {components}")521 logger.info(f" Universality: {geo_mean:.4f} [{ci_lo:.4f}, {ci_hi:.4f}]")522 523 return {524 "components": components,525 "universality_score": geo_mean,526 "ci_95": [ci_lo, ci_hi],527 }528 529 530# ─────────────────────────────────────────────────────────────────────────────531# FIGURES532# ─────────────────────────────────────────────────────────────────────────────533 534def save_cross_figures(out_dir, keys, all_cka, all_corr, all_diff,535 layer_data, neuron_data, univ_data):536 import matplotlib537 matplotlib.use("Agg")538 import matplotlib.pyplot as plt539 540 fig_dir = Path(out_dir) / "figures"541 fig_dir.mkdir(parents=True, exist_ok=True)542 543 P = {"drift": "#e74c3c", "unc": "#3498db", "corr": "#2ecc71",544 "null": "#9b59b6", "neu": "#e67e22"}545 546 # ── CM-1: CKA heatmaps ───────────────────────────────────────────────547 cka_pairs = [(k, v) for k, v in all_cka.items() if v]548 if cka_pairs:549 n_pairs = len(cka_pairs)550 fig, axes = plt.subplots(1, n_pairs, figsize=(8 * n_pairs, 7))551 if n_pairs == 1:552 axes = [axes]553 fig.suptitle("[CM-1] Cross-Model CKA", fontsize=16, fontweight="bold")554 for ax, (pair_key, data) in zip(axes, cka_pairs):555 mat = np.array(data["cka_matrix"])556 im = ax.imshow(mat, cmap="viridis", vmin=0, vmax=1, aspect="auto")557 la = data["layers_a"]558 lb = data["layers_b"]559 step_a = max(1, len(la) // 6)560 step_b = max(1, len(lb) // 6)561 ax.set_xticks(range(0, len(lb), step_b))562 ax.set_yticks(range(0, len(la), step_a))563 ax.set_xticklabels([lb[i] for i in range(0, len(lb), step_b)])564 ax.set_yticklabels([la[i] for i in range(0, len(la), step_a)])565 parts = pair_key.split("_vs_")566 ax.set(xlabel=f"{parts[1]} layer", ylabel=f"{parts[0]} layer",567 title=f"{pair_key}\nbest={data['best_cka']:.3f}")568 plt.colorbar(im, ax=ax, shrink=0.8)569 plt.tight_layout()570 plt.savefig(fig_dir / "fig_cm1_cka.png", dpi=300, bbox_inches="tight")571 plt.close()572 logger.info(" fig_cm1 saved")573 574 # ── CM-2: Score correlation matrix ────────────────────────────────────575 if len(keys) >= 2 and all_corr:576 n = len(keys)577 mat = np.eye(n)578 for pair_key, data in all_corr.items():579 if data is None:580 continue581 parts = pair_key.split("_vs_")582 if len(parts) == 2:583 i = keys.index(parts[0]) if parts[0] in keys else -1584 j = keys.index(parts[1]) if parts[1] in keys else -1585 if i >= 0 and j >= 0:586 mat[i, j] = mat[j, i] = data["correlation"]587 588 fig, ax = plt.subplots(figsize=(8, 7))589 im = ax.imshow(mat, cmap="RdBu_r", vmin=-1, vmax=1)590 ax.set_xticks(range(n))591 ax.set_yticks(range(n))592 ax.set_xticklabels(keys, fontsize=12, rotation=20)593 ax.set_yticklabels(keys, fontsize=12)594 for i in range(n):595 for j in range(n):596 c = "white" if abs(mat[i, j]) > 0.5 else "black"597 ax.text(j, i, f"{mat[i,j]:.3f}", ha="center", va="center",598 fontsize=13, fontweight="bold", color=c)599 ax.set_title("[CM-2] Drift Score Correlation Matrix", fontsize=14)600 plt.colorbar(im, ax=ax, shrink=0.8)601 plt.tight_layout()602 plt.savefig(fig_dir / "fig_cm2_corr.png", dpi=300, bbox_inches="tight")603 plt.close()604 logger.info(" fig_cm2 saved")605 606 # ── CM-3: Differential facts ──────────────────────────────────────────607 diff_pairs = [(k, v) for k, v in all_diff.items() if v]608 if diff_pairs:609 n_pairs = min(len(diff_pairs), 4)610 fig, axes = plt.subplots(1, n_pairs, figsize=(7 * n_pairs, 6))611 if n_pairs == 1:612 axes = [axes]613 fig.suptitle("[CM-3] Differential Facts", fontsize=16, fontweight="bold")614 for ax, (pair_key, data) in zip(axes, diff_pairs[:n_pairs]):615 sa = np.array(data["scores_a"])616 sb = np.array(data["scores_b"])617 la = np.array(data["labels_a"])618 lb = np.array(data["labels_b"])619 # Color by which model says drifted620 a_drifted = la.astype(bool) & ~lb.astype(bool)621 b_drifted = ~la.astype(bool) & lb.astype(bool)622 ax.scatter(sa[a_drifted], sb[a_drifted], c=P["drift"], alpha=0.5,623 s=30, label="A=drifted, B=stable")624 ax.scatter(sa[b_drifted], sb[b_drifted], c=P["unc"], alpha=0.5,625 s=30, label="A=stable, B=drifted")626 ax.plot([0, 1], [0, 1], "k--", alpha=0.3)627 ax.axhline(0.5, color="gray", ls=":", alpha=0.3)628 ax.axvline(0.5, color="gray", ls=":", alpha=0.3)629 parts = pair_key.split("_vs_")630 ax.set(xlabel=f"{parts[0]} score", ylabel=f"{parts[1]} score",631 title=f"{pair_key}\nr={data['score_correlation']:.3f}")632 ax.legend(fontsize=8)633 ax.grid(alpha=0.2)634 plt.tight_layout()635 plt.savefig(fig_dir / "fig_cm3_diff.png", dpi=300, bbox_inches="tight")636 plt.close()637 logger.info(" fig_cm3 saved")638 639 # ── CM-4: Layer correspondence ────────────────────────────────────────640 if layer_data and "per_model" in layer_data:641 pm = layer_data["per_model"]642 models = sorted(pm.keys())643 fig, axes = plt.subplots(1, 2, figsize=(14, 6))644 fig.suptitle("[CM-4] Layer Correspondence", fontsize=14,645 fontweight="bold")646 647 # Absolute layers648 x = np.arange(len(models))649 bls = [pm[m]["best_layer"] for m in models]650 nls = [pm[m]["n_layers"] for m in models]651 ax = axes[0]652 ax.bar(x, bls, color=P["drift"], edgecolor="black", lw=0.5,653 label="Best layer")654 ax.bar(x, [n - b for b, n in zip(bls, nls)], bottom=bls,655 color="#ecf0f1", edgecolor="black", lw=0.5, label="Remaining")656 ax.set_xticks(x)657 ax.set_xticklabels(models, fontsize=11)658 ax.set(ylabel="Layer", title="Best Drift Layer (absolute)")659 ax.legend()660 ax.grid(alpha=0.3, axis="y")661 662 # Fraction663 ax = axes[1]664 fracs = [pm[m]["fraction"] for m in models]665 bars = ax.bar(x, fracs, color=P["neu"], edgecolor="black", lw=0.5)666 ax.axhline(layer_data["mean_fraction"], color="red", ls="--", lw=2,667 label=f"Mean: {layer_data['mean_fraction']:.1%}")668 ax.fill_between(669 [-0.5, len(models) - 0.5],670 layer_data["mean_fraction"] - layer_data["std_fraction"],671 layer_data["mean_fraction"] + layer_data["std_fraction"],672 alpha=0.2, color="red")673 ax.set_xticks(x)674 ax.set_xticklabels(models, fontsize=11)675 ax.set(ylabel="Fraction of depth", title="Best Layer as % of Depth",676 ylim=(0, 1))677 ax.legend()678 ax.grid(alpha=0.3, axis="y")679 plt.tight_layout()680 plt.savefig(fig_dir / "fig_cm4_layers.png",681 dpi=300, bbox_inches="tight")682 plt.close()683 logger.info(" fig_cm4 saved")684 685 # ── CM-6: Summary ─────────────────────────────────────────────────────686 if univ_data:687 fig, ax = plt.subplots(figsize=(10, 6))688 comp = univ_data["components"]689 names = list(comp.keys())690 vals = list(comp.values())691 x = np.arange(len(names))692 colors = [P["drift"], P["unc"], P["corr"], P["neu"]][:len(names)]693 ax.bar(x, vals, color=colors, edgecolor="black", lw=0.5, alpha=0.8)694 ax.axhline(univ_data["universality_score"], color="red", ls="--",695 lw=2.5,696 label=f"Geo mean: {univ_data['universality_score']:.3f} "697 f"[{univ_data['ci_95'][0]:.3f}, "698 f"{univ_data['ci_95'][1]:.3f}]")699 ax.set_xticks(x)700 ax.set_xticklabels([n.replace("_", "\n") for n in names], fontsize=10)701 ax.set(ylabel="Score", title="[CM-6] Universality Score Components",702 ylim=(0, 1.1))703 ax.legend(fontsize=11)704 ax.grid(alpha=0.3, axis="y")705 plt.tight_layout()706 plt.savefig(fig_dir / "fig_cm6_summary.png",707 dpi=300, bbox_inches="tight")708 plt.close()709 logger.info(" fig_cm6 saved")710 711 logger.info(f"All cross-model figures -> {fig_dir}")712 713 714# ─────────────────────────────────────────────────────────────────────────────715# MAIN716# ─────────────────────────────────────────────────────────────────────────────717 718def main():719 p = argparse.ArgumentParser(720 description="Cross-model drift analysis",721 formatter_class=argparse.ArgumentDefaultsHelpFormatter)722 p.add_argument("--models", nargs="+", default=None,723 help="Model keys to compare")724 p.add_argument("--all", action="store_true",725 help="Use all models with available caches")726 p.add_argument("--config", default="models.yaml")727 p.add_argument("--output_dir", default=None)728 p.add_argument("--device", default="cuda:0")729 p.add_argument("--quick", action="store_true",730 help="Skip full-layer CKA, just best-layer")731 args = p.parse_args()732 733 cfg = load_config(args.config)734 defaults = cfg.get("defaults", {})735 output_dir = args.output_dir or defaults.get("output_dir",736 "data/experiments/v4")737 738 # Determine which models to use739 if args.all:740 model_keys = list(cfg["models"].keys())741 elif args.models:742 model_keys = args.models743 else:744 logger.error("Specify --models or --all")745 return746 747 # Load caches and bundles748 all_results = {}749 all_bundles = {}750 all_final = {}751 for key in model_keys:752 res = load_cache(output_dir, key)753 bundle = load_probe_bundle(output_dir, key)754 final = load_final_results(output_dir, key)755 if res and bundle:756 all_results[key] = res757 all_bundles[key] = bundle758 if final:759 all_final[key] = final760 761 keys = sorted(all_results.keys())762 logger.info(f"\nModels available: {keys}")763 764 if len(keys) < 2:765 logger.error("Need at least 2 models with caches + bundles")766 return767 768 cross_dir = Path(output_dir) / "cross_model"769 cross_dir.mkdir(parents=True, exist_ok=True)770 771 # Run all 6 experiments772 all_cka = {}773 all_corr = {}774 all_diff = {}775 all_neuron = {}776 777 for i, ka in enumerate(keys):778 for j, kb in enumerate(keys):779 if i >= j:780 continue781 pair = f"{ka}_vs_{kb}"782 logger.info(f"\n{'─'*50}")783 logger.info(f" {pair}")784 logger.info(f"{'─'*50}")785 786 # [CM-1] CKA787 all_cka[pair] = cka_analysis(788 all_results[ka], all_results[kb], ka, kb, quick=args.quick)789 790 # [CM-2] Score correlation791 all_corr[pair] = score_correlation(792 all_results[ka], all_results[kb], ka, kb,793 all_bundles[ka], all_bundles[kb], args.device)794 795 # [CM-3] Differential facts796 all_diff[pair] = differential_facts(797 all_results[ka], all_results[kb], ka, kb,798 all_bundles[ka], all_bundles[kb], args.device)799 800 # [CM-5] Neuron overlap801 all_neuron[pair] = neuron_overlap(802 all_bundles[ka], all_bundles[kb], ka, kb)803 804 # [CM-4] Layer correspondence805 layer_data = layer_correspondence(all_bundles, all_final)806 807 # [CM-6] Universality score808 univ_data = universality_score(all_cka, all_corr, all_diff, layer_data)809 810 # Save results811 results = {812 "models": keys,813 "cka": {k: v for k, v in all_cka.items()},814 "score_correlation": {k: v for k, v in all_corr.items()},815 "differential_facts": {k: v for k, v in all_diff.items()},816 "neuron_overlap": {k: v for k, v in all_neuron.items() if v},817 "layer_correspondence": layer_data,818 "universality": univ_data,819 "timestamp": datetime.now().isoformat(),820 }821 822 from datetime import datetime823 out_path = cross_dir / "cross_model_results.json"824 with open(out_path, "w") as f:825 json.dump(results, f, indent=2, default=str)826 logger.info(f"\nResults saved: {out_path}")827 828 # Figures829 save_cross_figures(str(cross_dir), keys, all_cka, all_corr, all_diff,830 layer_data, all_neuron, univ_data)831 832 # Print summary833 print(f"\n{'='*70}")834 print(f" CROSS-MODEL SUMMARY")835 print(f"{'='*70}")836 for pair, data in all_corr.items():837 if data:838 print(f" {pair}: score_corr={data['correlation']:.4f}")839 for pair, data in all_diff.items():840 if data:841 print(f" {pair}: diff_AUROC_a={data['auroc_a']:.4f} "842 f"diff_AUROC_b={data['auroc_b']:.4f} "843 f"n_diff={data['n_differential']}")844 if layer_data:845 print(f"\n Layer correspondence: "846 f"{layer_data['mean_fraction']:.1%} +/- "847 f"{layer_data['std_fraction']:.1%}")848 if univ_data:849 print(f"\n UNIVERSALITY SCORE: "850 f"{univ_data['universality_score']:.4f} "851 f"[{univ_data['ci_95'][0]:.4f}, {univ_data['ci_95'][1]:.4f}]")852 print(f"{'='*70}")853 854 855if __name__ == "__main__":856 main()