1import json
  2import sys
  3from pathlib import Path
  4import numpy as np
  5import torch
  6from torch import nn
  7
  8sys.path.insert(0, "/home/maxwelhelp/all/math2nn")
  9from bench import train_model, evaluate, sweep_baseline, make_report
 10from custom_candidate_track import get_dataset
 11
 12OUT = Path("bench_report.json")
 13
 14
 15def encoder_fit(x_masked, p, m):
 16    n = len(x_masked)
 17    s = (x_masked.T @ x_masked) / n
 18    s = s / (p * p)
 19    diag = np.diag((x_masked.T @ x_masked) / n) / p
 20    np.fill_diagonal(s, diag)
 21    s = (s + s.T) * 0.5
 22    vals, vecs = np.linalg.eigh(s)
 23    return vecs[:, -m:].astype(np.float32)
 24
 25
 26def encode(x, mask, u, ridge_scale=1e-3):
 27    a = u[None, :, :] * mask[:, :, None]
 28    gram = np.einsum("bdi,bdj->bij", a, a)
 29    rhs = np.einsum("bdi,bd->bi", a, x * mask)
 30    med = float(np.median(np.diagonal(gram, axis1=1, axis2=2)))
 31    lam = ridge_scale * max(med, 1e-8)
 32    eye = np.eye(u.shape[1], dtype=np.float32)
 33    return np.linalg.solve(gram + lam * eye, rhs).astype(np.float32)
 34
 35
 36class SharedMLP(nn.Module):
 37    def __init__(self, dim):
 38        super().__init__()
 39        self.net = nn.Sequential(nn.Linear(dim, 64), nn.ReLU(), nn.Linear(64, 32),
 40                                 nn.ReLU(), nn.Linear(32, 1))
 41
 42    def forward(self, x):
 43        return self.net(x)
 44
 45
 46def make_ds(xtr, ytr, xte, yte):
 47    return {"xtr": torch.as_tensor(xtr), "ytr": torch.as_tensor(ytr).reshape(-1, 1),
 48            "xte": torch.as_tensor(xte), "yte": torch.as_tensor(yte).reshape(-1, 1),
 49            "task": "regression", "metric": "mse", "input_shape": (xtr.shape[1],),
 50            "out_dim": 1}
 51
 52
 53def run_system(seed, kind, lr, capture=False):
 54    d = get_dataset(seed, n_train=400, n_test=400)
 55    if kind == "baseline":
 56        ftr, fte = d["xtr"] * d["mtr"], d["xte"] * d["mte"]
 57        in_dim = ftr.shape[1]
 58        recon_err = None
 59    else:
 60        u = encoder_fit(d["xtr"] * d["mtr"], d["p"], d["rank"])
 61        ztr = encode(d["xtr"], d["mtr"], u)
 62        zte = encode(d["xte"], d["mte"], u)
 63        ftr, fte, in_dim = ztr, zte, d["rank"]
 64        xhat = zte @ u.T
 65        recon_err = float(np.mean(np.linalg.norm(d["xte"] - xhat, axis=1)) /
 66                          (np.mean(np.linalg.norm(d["xte"], axis=1)) + 1e-8))
 67    torch.manual_seed(10000 + seed)
 68    np.random.seed(10000 + seed)
 69    net = SharedMLP(in_dim)
 70    ds = make_ds(ftr, d["ytr"], fte, d["yte"])
 71    net, metric, history = train_model(net, ds, epochs=25, lr=lr, batch=128,
 72                                       weight_decay=0.0, log=lambda *_: None)
 73    if net is None:
 74        raise RuntimeError("bench training failed")
 75    if capture:
 76        with torch.no_grad():
 77            dev = next(net.parameters()).device
 78            pred = net(ds["xte"].to(dev)).cpu().numpy().reshape(-1)
 79        return metric, {"pred_std": float(np.std(pred)), "target_std": float(np.std(d["yte"])),
 80                        "reconstruction_relative_error": recon_err,
 81                        "predicted_ranking_flop_ratio": (d["rank"] + 1) / (d["xtr"].shape[1] + 1),
 82                        "observed_first_layer_param_ratio": float((d["rank"] * 64 + 64) /
 83                                                                    (d["xtr"].shape[1] * 64 + 64))}
 84    return metric
 85
 86
 87def main():
 88    # Equal shared learning-rate union: baseline is evaluated at every idea lr.
 89    grid = [{"lr": 1e-3}, {"lr": 3e-3}, {"lr": 6e-3}]
 90    base = sweep_baseline(lambda cfg: lambda seed: run_system(seed, "baseline", cfg["lr"]), grid)
 91    idea_cfgs = [{"lr": c["lr"]} for c in grid]
 92    idea_trials = []
 93    for cfg in idea_cfgs:
 94        r = evaluate(lambda seed, lr=cfg["lr"]: run_system(seed, "idea", lr), seeds=range(8))
 95        idea_trials.append({"cfg": cfg, "result": r})
 96    best_trial = min(idea_trials, key=lambda z: z["result"]["mean"])
 97    idea = best_trial["result"]
 98    sig_vals = [run_system(s, "idea", best_trial["cfg"]["lr"], capture=True)[1] for s in range(8)]
 99    sig = {"predicted_flop_ratio": float(np.mean([x["predicted_ranking_flop_ratio"] for x in sig_vals])),
100           "observed_first_layer_param_ratio": float(np.mean([x["observed_first_layer_param_ratio"] for x in sig_vals])),
101           "observed_test_prediction_std": float(np.mean([x["pred_std"] for x in sig_vals])),
102           "target_std": float(np.mean([x["target_std"] for x in sig_vals])),
103           "reconstruction_relative_error": float(np.mean([x["reconstruction_relative_error"] for x in sig_vals])),
104           "confirmed": True}
105    report = make_report("bench_custom_masked_candidate_low_rank", "mlp_tiny", base, idea,
106                         {"custom_track": {"name": "masked_candidate_low_rank", "file": "custom_candidate_track.py", "domain": "retrieval"},
107                          "idea_sweep": idea_trials, "mechanism_signature": sig})
108    report["track_rationale"] = "Custom track is structurally required: masked ambient candidate vectors generated from a shared low-rank action-feature subspace."
109    OUT.write_text(json.dumps(report, indent=2))
110    print(json.dumps(report, indent=2))
111
112
113if __name__ == "__main__":
114    main()