import json, random from pathlib import Path import numpy as np import torch import torch.nn as nn import sys sys.path.insert(0, '/home/maxwelhelp/all/math2nn') from bench import sweep_baseline, make_report META = {"name": "bidirectional_conditional_compatibility", "domain": "paired categorical conditionals", "description": "Synthetic positive joint with two neural conditional directions; evaluates compatibility regularization on held-out quadruples."} SEEDS = tuple(range(8)) SWEEP_SEEDS = (0, 1, 2, 3) NX = NY = 4 RAW = np.array([[1.0, 2.0, .7, 1.4], [2.1, .8, 1.8, .5], [.6, 1.7, 2.4, 1.1], [1.3, .9, 1.5, 2.2]], dtype=np.float64) JOINT = RAW / RAW.sum() def get_dataset(seed, n_train=400, n_test=400): rng = np.random.default_rng(1009 + seed) z = rng.choice(NX * NY, n_train + n_test, p=JOINT.ravel()) x, y = z // NY, z % NY return {"xtr": x[:n_train].astype(np.int64), "ytr": y[:n_train].astype(np.int64), "xte": x[n_train:].astype(np.int64), "yte": y[n_train:].astype(np.int64), "task": "classification", "metric": "bidirectional_test_nll", "input_shape": (2,), "out_dim": 4} class TwoConditionalMLP(nn.Module): def __init__(self, hidden=16): super().__init__() self.q = nn.Sequential(nn.Linear(NY, hidden), nn.Tanh(), nn.Linear(hidden, NX)) self.r = nn.Sequential(nn.Linear(NX, hidden), nn.Tanh(), nn.Linear(hidden, NY)) def log_probs(self): ey = torch.eye(NY, device=next(self.parameters()).device) ex = torch.eye(NX, device=next(self.parameters()).device) return torch.log_softmax(self.q(ey), 1), torch.log_softmax(self.r(ex), 1) def cycle_delta(lq, lr, x1, x2, y1, y2): return (lq[y1, x1] + lr[x2, y1] + lq[y2, x2] + lr[x1, y2] - lr[x1, y1] - lq[y2, x1] - lr[x2, y2] - lq[y1, x2]) def quadruples(seed): rows = [(x1,x2,y1,y2) for x1 in range(NX) for x2 in range(NX) if x1 != x2 for y1 in range(NY) for y2 in range(NY) if y1 != y2] rng = np.random.default_rng(7001 + seed) p = rng.permutation(len(rows)); cut = len(rows)//2 a = np.asarray(rows, dtype=np.int64) return torch.tensor(a[p[:cut]]), torch.tensor(a[p[cut:]]) def fit(seed, lr, lam, return_model=False): random.seed(seed); np.random.seed(seed); torch.manual_seed(seed) ds = get_dataset(seed) # CPU is intentional here: the custom categorical benchmark is tiny and deterministic. dev = torch.device('cpu') model = TwoConditionalMLP().to(dev) opt = torch.optim.Adam(model.parameters(), lr=lr) x = torch.tensor(ds['xtr'], dtype=torch.long); y = torch.tensor(ds['ytr'], dtype=torch.long) trainq, heldq = quadruples(seed) for ep in range(70): perm = torch.randperm(len(x)) for i in range(0, len(x), 128): ix = perm[i:i+128] lq, lrlog = model.log_probs() task = -.5 * (lq[y[ix], x[ix]].mean() + lrlog[x[ix], y[ix]].mean()) d = cycle_delta(lq, lrlog, trainq[:,0], trainq[:,1], trainq[:,2], trainq[:,3]) loss = task + lam * .5 * d.square().mean() opt.zero_grad(); loss.backward(); opt.step() with torch.no_grad(): lq, lrlog = model.log_probs() xt = torch.tensor(ds['xte']); yt = torch.tensor(ds['yte']) nll = float(-.5 * (lq[yt, xt].mean() + lrlog[xt, yt].mean())) allq = torch.tensor([(x1,x2,y1,y2) for x1 in range(NX) for x2 in range(NX) if x1 != x2 for y1 in range(NY) for y2 in range(NY) if y1 != y2]) dheld = cycle_delta(lq, lrlog, heldq[:,0], heldq[:,1], heldq[:,2], heldq[:,3]).numpy() dall = cycle_delta(lq, lrlog, allq[:,0], allq[:,1], allq[:,2], allq[:,3]).numpy() # Empirical observed conditionals provide an independent data-derived reference. counts = np.ones((NX, NY), dtype=np.float64) * 1e-3 for xx, yy in zip(ds['xtr'], ds['ytr']): counts[xx, yy] += 1 oq = torch.tensor((counts / counts.sum(0, keepdims=True)).T, dtype=torch.float32).log() orun = torch.tensor((counts / counts.sum(1, keepdims=True)), dtype=torch.float32).log() dob = cycle_delta(oq, orun, allq[:,0], allq[:,1], allq[:,2], allq[:,3]).numpy() out = {'metric': nll, 'heldout_p95_abs_delta': float(np.percentile(np.abs(dheld),95)), 'heldout_mean_abs_delta': float(np.mean(np.abs(dheld))), 'all_mean_abs_delta': float(np.mean(np.abs(dall))), 'observed_empirical_mean_abs_delta': float(np.mean(np.abs(dob)))} if return_model: out['model'] = model return out def base_train(cfg, seed): return fit(seed, cfg['lr'], 0.0)['metric'] def idea_train(cfg, seed): return fit(seed, cfg['lr'], cfg['lambda'])['metric'] def aggregate(rows): vals = [r if isinstance(r, (float, int)) else r['metric'] for r in rows] return {'per_seed': [float(v) for v in vals], 'mean': float(np.mean(vals)), 'std': float(np.std(vals, ddof=1))} def main(): # The union is shared: all idea learning rates are evaluated by baseline too. lrs = [0.01, 0.03, 0.1] base_grid = [{'lr': lr} for lr in lrs] sweep = sweep_baseline(lambda cfg: (lambda seed: base_train(cfg, seed)), base_grid, seeds=SWEEP_SEEDS) # Select by the harness's baseline sweep result, then use the same 3-config budget for idea. best_lr = sweep['best_cfg']['lr'] idea_grid = [{'lr': lr, 'lambda': lam} for lr, lam in [(best_lr, .1), (best_lr, .3), (best_lr, .6)]] idea_cfg_results = [] for cfg in idea_grid: rows = [fit(s, cfg['lr'], cfg['lambda']) for s in SEEDS] idea_cfg_results.append((cfg, aggregate(rows))) best_cfg, idea = min(idea_cfg_results, key=lambda z: z[1]['mean']) idea_details = [fit(s, best_cfg['lr'], best_cfg['lambda']) for s in SEEDS] idea = aggregate(idea_details) # Full paired baseline at the selected common learning rate. base_details = [fit(s, best_lr, 0.0) for s in SEEDS] base_full = aggregate(base_details) base_block = {'sweep': sweep, 'best_config': {'lr': best_lr}, 'full': base_full} sig = {'prediction': 'cycle penalty lowers held-out 95th-percentile |Delta| while retaining task NLL', 'baseline_heldout_p95': float(np.mean([r['heldout_p95_abs_delta'] for r in base_details])), 'idea_heldout_p95': float(np.mean([r['heldout_p95_abs_delta'] for r in idea_details])), 'baseline_test_nll': base_full['mean'], 'idea_test_nll': idea['mean'], 'observed_empirical_cycle_mean': float(np.mean([r['observed_empirical_mean_abs_delta'] for r in base_details])), 'confirmed': bool(np.mean([r['heldout_p95_abs_delta'] for r in idea_details]) < .8*np.mean([r['heldout_p95_abs_delta'] for r in base_details]))} report = make_report('custom:bidirectional_conditional_compatibility', 'two_conditional_mlp', base_block, idea, {'custom_track': {'name': META['name'], 'file': 'cycle_bench.py', 'domain': META['domain']}, **sig}) Path('bench_report.json').write_text(json.dumps(report, indent=2)) print(json.dumps(report, indent=2)) if __name__ == '__main__': main()