Conditioned PSD sensing bottleneck / experiment.py

Failed on benchmark

Raw ⬇ ZIP
  1import json
  2import math
  3import random
  4import numpy as np
  5import torch
  6
  7SEED = 58
  8random.seed(SEED); np.random.seed(SEED); torch.manual_seed(SEED)
  9torch.set_num_threads(4)
 10DT = torch.float64
 11EPS = 1e-10
 12
 13
 14def pairs(n, d, seed):
 15    g = torch.Generator().manual_seed(seed)
 16    u = torch.randn(n, d, generator=g, dtype=DT)
 17    v = torch.randn(n, d, generator=g, dtype=DT)
 18    # Rank-one PSD features X=uu^T, Y=vv^T, normalized to remove trivial scale.
 19    u = u / u.norm(dim=1, keepdim=True)
 20    v = v / v.norm(dim=1, keepdim=True)
 21    return u, v
 22
 23
 24def ratios(A, u, v):
 25    # A contains sensing vectors a_i; A_i=a_i a_i^T is PSD.
 26    zu = (u @ A.T).square()
 27    zv = (v @ A.T).square()
 28    numerator = (zu - zv).abs().sum(dim=1)
 29    # ||uu^T-vv^T||_F, computed exactly for unit u,v.
 30    den = torch.sqrt(2.0 - 2.0 * (u * v).sum(dim=1).square())
 31    return numerator / (den + EPS)
 32
 33
 34def stats(q):
 35    q = q.detach().cpu().numpy()
 36    lo, hi = np.quantile(q, [.05, .95])
 37    return {"L_q05": float(lo), "U_q95": float(hi),
 38            "beta": float(hi / max(lo, EPS)), "median": float(np.median(q)),
 39            "mean": float(np.mean(q))}
 40
 41
 42def random_A(m, d, seed, gamma=1.0):
 43    g = torch.Generator().manual_seed(seed)
 44    A = torch.randn(m, d, generator=g, dtype=DT)
 45    # A controlled anisotropy: amplify one coordinate before row normalization.
 46    A[:, 0] *= gamma
 47    return A / A.norm(dim=1, keepdim=True)
 48
 49
 50def optimize_A(m, d, train_u, train_v, test_u, test_v, steps=500):
 51    g = torch.Generator().manual_seed(9000 + m)
 52    W = torch.randn(m, d, generator=g, dtype=DT, requires_grad=True)
 53    opt = torch.optim.Adam([W], lr=.035)
 54    history = []
 55    for step in range(steps):
 56        A = W / (W.norm(dim=1, keepdim=True) + EPS)
 57        q = ratios(A, train_u, train_v)
 58        lo = torch.quantile(q, .05)
 59        hi = torch.quantile(q, .95)
 60        loss = torch.log(hi + EPS) - torch.log(lo + EPS)
 61        opt.zero_grad(); loss.backward(); opt.step()
 62        if step in (0, steps - 1): history.append(float(loss.detach()))
 63    A = (W / (W.norm(dim=1, keepdim=True) + EPS)).detach()
 64    return A, history
 65
 66
 67def main():
 68    d = 8
 69    train_u, train_v = pairs(2400, d, 10)
 70    test_u, test_v = pairs(10000, d, 11)
 71    out = {"seed": SEED, "d": d, "predictions": {}}
 72
 73    # Prediction 1: common positive scaling changes L and U equally, hence beta is invariant.
 74    A = random_A(64, d, 21)
 75    base = stats(ratios(A, test_u, test_v))
 76    scaled = stats(ratios(3.7 * A, test_u, test_v))
 77    out["scale_invariance"] = {"base": base, "scaled_3.7": scaled,
 78        "observed_beta_relative_change": scaled["beta"] / base["beta"] - 1.0,
 79        "observed_L_scale": scaled["L_q05"] / base["L_q05"],
 80        "observed_U_scale": scaled["U_q95"] / base["U_q95"]}
 81    out["predictions"]["scale"] = "beta unchanged; L and U multiply by 3.7^2 because A_i=a_i a_i^T"
 82
 83    # Prediction 2: more independent nonnegative PSD measurements narrow the empirical ratio spread.
 84    widths = [8, 16, 32, 64, 128, 256]
 85    width_rows = []
 86    for m in widths:
 87        s = stats(ratios(random_A(m, d, 100 + m), test_u, test_v))
 88        width_rows.append({"m": m, **s})
 89    out["width_sweep"] = width_rows
 90    out["predictions"]["width"] = "beta should generally fall as independent measurements average fluctuations"
 91
 92    # Prediction 3: directional anisotropy creates collapsed directions and increases beta.
 93    anis_rows = []
 94    for gamma in [1., 2., 4., 8., 16.]:
 95        s = stats(ratios(random_A(128, d, 333, gamma), test_u, test_v))
 96        anis_rows.append({"gamma": gamma, **s})
 97    out["anisotropy_sweep"] = anis_rows
 98    out["predictions"]["anisotropy"] = "increasing directional imbalance raises U/L"
 99
100    # Mini experiment: same width, random unregularized PSD sensing vs conditioned sensing.
101    m = 32
102    baseline_A = random_A(m, d, 700)
103    cond_A, hist = optimize_A(m, d, train_u, train_v, test_u, test_v)
104    out["mini_experiment"] = {
105        "m": m, "baseline_random": stats(ratios(baseline_A, test_u, test_v)),
106        "conditioned_optimized": stats(ratios(cond_A, test_u, test_v)),
107        "train_log_condition_start_end": hist,
108        "mean_row_norm_baseline": float(baseline_A.norm(dim=1).mean()),
109        "mean_row_norm_conditioned": float(cond_A.norm(dim=1).mean())}
110
111    # Simple pass/fail checks stated quantitatively.
112    scale = out["scale_invariance"]
113    scale_ok = abs(scale["observed_beta_relative_change"]) < 1e-8 and abs(scale["observed_L_scale"]-3.7**2) < 1e-8 and abs(scale["observed_U_scale"]-3.7**2) < 1e-8
114    b = [x["beta"] for x in width_rows]
115    width_ok = b[-1] < b[0]
116    a = [x["beta"] for x in anis_rows]
117    anis_ok = a[-1] > a[0]
118    mini_ok = out["mini_experiment"]["conditioned_optimized"]["beta"] < out["mini_experiment"]["baseline_random"]["beta"]
119    out["checks"] = {"scale_exact": scale_ok, "width_endpoint": width_ok,
120                      "anisotropy_endpoint": anis_ok, "optimization_test_beta": mini_ok,
121                      "mechanism_manifestations": int(scale_ok) + int(width_ok) + int(anis_ok)}
122    with open("results.json", "w") as f: json.dump(out, f, indent=2)
123    print(json.dumps(out, indent=2))
124
125if __name__ == "__main__":
126    main()