1import json, math, random
  2from pathlib import Path
  3import numpy as np
  4import torch
  5
  6SEED = 3043
  7
  8def seed_all(seed):
  9    random.seed(seed); np.random.seed(seed); torch.manual_seed(seed)
 10
 11
 12def kendall_tau_a(a, b):
 13    # Efficient Kendall tau over corresponding flattened distance values.
 14    from scipy.stats import kendalltau
 15    return float(kendalltau(np.asarray(a), np.asarray(b), variant="b").statistic)
 16
 17
 18def gaussian_verification():
 19    # P is the first m rows of a random orthogonal matrix. Repeated points
 20    # approximate the paper's independent Gaussian pair distribution.
 21    rng = np.random.default_rng(SEED)
 22    d, n = 64, 500
 23    rows = []
 24    for m in [2, 4, 8, 16, 32]:
 25        vals = []
 26        for rep in range(8):
 27            x = rng.normal(size=(n, d))
 28            q, _ = np.linalg.qr(rng.normal(size=(d, d)))
 29            z = x @ q[:, :m]
 30            D = ((x[:, None, :] - x[None, :, :]) ** 2).sum(-1)
 31            Dm = (d / m) * ((z[:, None, :] - z[None, :, :]) ** 2).sum(-1)
 32            iu = np.triu_indices(n, 1)
 33            vals.append(kendall_tau_a(D[iu], Dm[iu]))
 34        observed = float(np.mean(vals)); sd = float(np.std(vals))
 35        predicted = (2.0 / math.pi) * math.sqrt(m / d)
 36        rows.append({"m": m, "observed_tau": observed, "sd": sd, "prediction": predicted})
 37    # Also verify centering removes the large common distance baseline but
 38    # leaves exactly the same centered ordering / correlation.
 39    x = rng.normal(size=(n, d)); q, _ = np.linalg.qr(rng.normal(size=(d, d))); z = x @ q[:, :8]
 40    D = ((x[:, None, :] - x[None, :, :]) ** 2).sum(-1)
 41    Dm = (d / 8) * ((z[:, None, :] - z[None, :, :]) ** 2).sum(-1)
 42    iu = np.triu_indices(n, 1)
 43    raw_mean_gap = abs(float(D[iu].mean() - Dm[iu].mean()))
 44    centered_mean = float((D[iu] - D[iu].mean()).mean())
 45    centered_projected_mean = float((Dm[iu] - Dm[iu].mean()).mean())
 46    return {"ranking_scaling": rows, "mean_gap_before_centering": raw_mean_gap,
 47            "centered_means": [centered_mean, centered_projected_mean]}
 48
 49
 50def pair_dist(x):
 51    return ((x[:, None, :] - x[None, :, :]) ** 2).sum(-1)
 52
 53
 54def run_one(seed, centered, steps=500):
 55    seed_all(seed)
 56    device = "cuda" if torch.cuda.is_available() else "cpu"
 57    try:
 58        # A compact embedding problem: nuisance dimensions have much larger
 59        # variance than a low-dimensional class/task direction.
 60        g = torch.Generator().manual_seed(seed)
 61        n, d, m = 96, 32, 8
 62        labels = torch.arange(n) % 8
 63        centers = torch.randn(8, 3, generator=g) * 1.8
 64        x = torch.randn(n, d, generator=g) * 2.0
 65        x[:, :3] += centers[labels]
 66        x = x.to(device)
 67        W = torch.randn(m, d, generator=g, device=device) / math.sqrt(d)
 68        W = torch.nn.Parameter(W)
 69        opt = torch.optim.Adam([W], lr=0.035)
 70        # Fixed pair target and fixed normalization make comparisons matched.
 71        with torch.no_grad():
 72            target = pair_dist(x)
 73            target_mean = target.mean()
 74            target_std = target.std() + 1e-8
 75        for step in range(steps):
 76            opt.zero_grad()
 77            z = x @ W.t()
 78            pred = (d / m) * pair_dist(z)
 79            if centered:
 80                # stop-gradient EMA-like batch baseline; using the current
 81                # batch mean is the exact minibatch version in the proposal.
 82                a = (target - target.mean()) / (target.std() + 1e-8)
 83                b = (pred - pred.mean()) / (pred.std() + 1e-8)
 84                loss = ((a - b) ** 2).mean()
 85            else:
 86                loss = ((pred - target) / target_std).pow(2).mean()
 87            loss.backward(); opt.step()
 88        with torch.no_grad():
 89            pred = (d / m) * pair_dist(x @ W.t())
 90            iu = torch.triu_indices(n, n, offset=1, device=device)
 91            a = target[iu[0], iu[1]].detach().cpu().numpy()
 92            b = pred[iu[0], iu[1]].detach().cpu().numpy()
 93            tau = kendall_tau_a(a, b)
 94            centered_mse = float(np.mean(((a-a.mean())/a.std() - (b-b.mean())/b.std())**2))
 95            raw_mse = float(np.mean(((a-b)/a.std())**2))
 96            # Simple downstream retrieval: nearest projected neighbor shares label.
 97            zz = (x @ W.t()).detach().cpu().numpy()
 98            dd = ((zz[:,None,:]-zz[None,:,:])**2).sum(-1); np.fill_diagonal(dd, np.inf)
 99            recall = float(np.mean(labels.cpu().numpy()[dd.argmin(1)] == labels.cpu().numpy()))
100        return {"tau": tau, "centered_mse": centered_mse, "raw_mse": raw_mse, "recall1": recall}
101    except Exception as e:
102        if device == "cuda":
103            torch.cuda.empty_cache()
104            # A CPU retry with the same code path is safer on shared GPUs.
105            torch.cuda.is_available = lambda: False
106            return run_one(seed, centered, steps)
107        raise e
108
109
110def main():
111    verification = gaussian_verification()
112    results = {}
113    for mode, flag in [("raw", False), ("centered", True)]:
114        vals = [run_one(s, flag) for s in [11, 22, 33, 44, 55]]
115        results[mode] = {k: {"mean": float(np.mean([v[k] for v in vals])),
116                             "sd": float(np.std([v[k] for v in vals]))} for k in vals[0]}
117        results[mode]["runs"] = vals
118    out = {"verification": verification, "mini_experiment": results}
119    Path("results.json").write_text(json.dumps(out, indent=2))
120    print(json.dumps(out, indent=2))
121
122if __name__ == "__main__":
123    main()