| """ |
| subspace_angle.py -- principal angles between BASE and OURS truncation subspaces. |
| |
| Efficiency / subspace-drift evidence: BASE and OURS run the identical whitening |
| math at the same cost; the ONLY thing that changes is which subspace the |
| calibration selects. This quantifies that difference per layer. |
| |
| For each matching layer we take the retained input subspace = row space of the |
| B factor (k x in), orthonormalize both, and compute principal angles via the SVD |
| of Q_base^T Q_ours (its singular values are cos of the principal angles). |
| Reports mean/max angle (degrees) per layer + an overall summary to |
| results/efficiency/. |
| """ |
| import os |
| import sys |
| import glob |
| import argparse |
| import math |
| import torch |
|
|
| sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) |
| import common as C |
|
|
|
|
| def ortho_rowspace(B): |
| """Return an orthonormal basis (in x k) of the row space of B (k x in).""" |
| |
| Q, _ = torch.linalg.qr(B.t().to(torch.float64), mode="reduced") |
| return Q |
|
|
|
|
| def principal_angles_deg(B1, B2): |
| Q1 = ortho_rowspace(B1) |
| Q2 = ortho_rowspace(B2) |
| M = Q1.t() @ Q2 |
| s = torch.linalg.svdvals(M).clamp(-1.0, 1.0) |
| angles = torch.arccos(s) * (180.0 / math.pi) |
| return angles |
|
|
|
|
| def main(): |
| ap = argparse.ArgumentParser() |
| ap.add_argument("--base_weights", required=True) |
| ap.add_argument("--ours_weights", required=True) |
| ap.add_argument("--out", default=None) |
| args = ap.parse_args() |
|
|
| base_Bs = sorted(glob.glob(os.path.join(args.base_weights, "*_B.pt"))) |
| per_layer = {} |
| all_means = [] |
| for bp in base_Bs: |
| fn = os.path.basename(bp) |
| op = os.path.join(args.ours_weights, fn) |
| if not os.path.exists(op): |
| continue |
| B1 = torch.load(bp, map_location="cpu").float() |
| B2 = torch.load(op, map_location="cpu").float() |
| if B1.shape != B2.shape: |
| per_layer[fn] = {"skipped": f"shape {tuple(B1.shape)} vs {tuple(B2.shape)}"} |
| continue |
| ang = principal_angles_deg(B1, B2) |
| per_layer[fn] = {"k": B1.shape[0], "mean_deg": float(ang.mean()), |
| "max_deg": float(ang.max()), "median_deg": float(ang.median())} |
| all_means.append(float(ang.mean())) |
|
|
| summary = { |
| "base_weights": args.base_weights, |
| "ours_weights": args.ours_weights, |
| "n_layers": len(all_means), |
| "overall_mean_deg": sum(all_means) / len(all_means) if all_means else None, |
| "overall_min_layer_mean_deg": min(all_means) if all_means else None, |
| "overall_max_layer_mean_deg": max(all_means) if all_means else None, |
| "git_hash": C.git_hash(), |
| "per_layer": per_layer, |
| } |
| root = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) |
| out = args.out or os.path.join(root, "results", "efficiency", |
| f"{C.git_hash()}_subspace_angle.json") |
| C.dump_json(summary, out) |
| print(f"[subspace] {len(all_means)} layers " |
| f"overall_mean={summary['overall_mean_deg']}deg -> {out}") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
