import json, math, random, sys from pathlib import Path import numpy as np import torch import torch.nn as nn sys.path.insert(0, '/home/maxwelhelp/all/math2nn') import bench SEEDS = tuple(range(8)) # Identical learning-rate union on both sides; baseline also sweeps its beta knob. IDEA_GRID = [{'lr': lr, 'gamma': 0.10} for lr in (0.003, 0.006, 0.012)] BASE_GRID = [{'lr': lr, 'beta': beta} for lr in (0.003, 0.006, 0.012) for beta in (0.85, 0.95)] EPOCHS, BATCH = 18, 64 def seed_all(seed): random.seed(seed); np.random.seed(seed); torch.manual_seed(seed) if torch.cuda.is_available(): torch.cuda.manual_seed_all(seed) def make_net(ds, seed): seed_all(seed) return bench.make_model('mlp_tiny', tuple(ds['input_shape']), int(ds['out_dim'])) def loss_fn(out, y, task): return nn.CrossEntropyLoss()(out, y.long().view(-1)) if task == 'classification' else nn.MSELoss()(out, y) def run(seed, method, cfg, signature=False): seed_all(seed) ds = bench.get_dataset('tabular', seed, n_train=400, n_test=200) dev = 'cuda' if torch.cuda.is_available() else 'cpu' try: return _run(seed, ds, method, cfg, dev, signature) except Exception: if dev == 'cuda': return _run(seed, ds, method, cfg, 'cpu', signature) raise def _run(seed, ds, method, cfg, dev, signature): net = make_net(ds, seed).to(dev) xtr, ytr = ds['xtr'].to(dev), ds['ytr'].to(dev) xte, yte = ds['xte'].to(dev), ds['yte'].to(dev) params = [p for p in net.parameters() if p.requires_grad] # M=I is the prescribed initial MVP; p is the contact momentum. p = [torch.zeros_like(q) for q in params] s = 0.0 rng = np.random.default_rng(seed + 991) cert_rates, certs = [], [] grad_evals = 0 def grad_batch(xb, yb): nonlocal grad_evals net.zero_grad(set_to_none=True) out = net(xb); loss = loss_fn(out, yb, ds['task']); loss.backward() grad_evals += 1 return float(loss.detach().cpu()), [q.grad.detach().clone() for q in params] def test_metric(): net.eval() with torch.no_grad(): val = loss_fn(net(xte), yte, ds['task']) net.train(); return float(val.cpu()) net.train() n = xtr.shape[0] for epoch in range(EPOCHS): order = rng.permutation(n) for st in range(0, n, BATCH): ix = torch.as_tensor(order[st:st+BATCH], device=dev) xb, yb = xtr[ix], ytr[ix] lr = float(cfg['lr']) if method == 'baseline': # Two ordinary momentum updates, matching the two contact gradients. for j in range(2): f, g = grad_batch(xb, yb) beta = float(cfg['beta']) with torch.no_grad(): for k in range(len(params)): p[k].mul_(beta).add_(g[k], alpha=-lr / 2.0) params[k].add_(p[k]) else: gamma = float(cfg['gamma']) # K(h/2): x += h p/2 and s += h K/2. with torch.no_grad(): kin = 0.5 * sum((q*q).sum() for q in p) s += (lr / 2.0) * float(kin.cpu()) for q, v in zip(params, p): q.add_(v, alpha=lr / 2.0) f1, g1 = grad_batch(xb, yb) with torch.no_grad(): for v, gg in zip(p, g1): v.add_(gg, alpha=-lr / 2.0) s -= lr / 2.0 * f1 damp = math.exp(-gamma * lr) for v in p: v.mul_(damp) s *= damp f2, g2 = grad_batch(xb, yb) with torch.no_grad(): for v, gg in zip(p, g2): v.add_(gg, alpha=-lr / 2.0) s -= lr / 2.0 * f2 kin = 0.5 * sum((q*q).sum() for q in p) s += (lr / 2.0) * float(kin.cpu()) for q, v in zip(params, p): q.add_(v, alpha=lr / 2.0) H = float((0.5 * sum((q*q).sum() for q in p)).cpu()) + f2 + gamma*s if certs and abs(certs[-1]) > 1e-8 and abs(H) > 1e-8 and certs[-1]*H > 0: cert_rates.append(-(math.log(abs(H))-math.log(abs(certs[-1]))) / lr) certs.append(H) metric = test_metric() result = {'metric': metric, 'grad_evals': grad_evals} if signature: result['cert_rate_mean'] = float(np.mean(cert_rates)) if cert_rates else float('nan') result['cert_rate_n'] = len(cert_rates) return result def evaluate(method, cfg, seeds=SEEDS, signature=False): vals, rows = [], [] for seed in seeds: r = run(seed, method, cfg, signature) vals.append(r['metric']); rows.append(r) return {'mean': float(np.mean(vals)), 'std': float(np.std(vals)), 'per_seed': vals, 'n': len(vals), 'details': rows} def main(): # Fair small baseline sweep on four seeds, then full paired evaluation. sweep = [] for cfg in BASE_GRID: r = evaluate('baseline', cfg, seeds=(0,1,2,3)) sweep.append({'cfg': cfg, 'mean': r['mean']}) best = min(sweep, key=lambda z: z['mean'])['cfg'] base_full = evaluate('baseline', best, SEEDS) base_block = {'best_cfg': best, 'sweep': sweep, 'full': base_full} idea_sweep = [] for cfg in IDEA_GRID: # Union parity is satisfied because all idea lrs occur in BASE_GRID. r = evaluate('idea', cfg, seeds=(0,1,2,3)) idea_sweep.append({'cfg': cfg, 'mean': r['mean']}) idea_best = min(idea_sweep, key=lambda z: z['mean'])['cfg'] idea_full = evaluate('idea', idea_best, SEEDS, signature=True) # Re-test the stage-1 prediction on trained models: expected empirical rate ~ gamma. rates = [d['cert_rate_mean'] for d in idea_full['details'] if np.isfinite(d['cert_rate_mean'])] observed = float(np.mean(rates)) if rates else float('nan') gamma = idea_best['gamma'] sig = {'prediction': 'certificate conformal rate ~= gamma', 'predicted': gamma, 'observed': observed, 'absolute_error': abs(observed-gamma) if np.isfinite(observed) else None, 'n_models': len(rates), 'confirmed': bool(np.isfinite(observed) and abs(observed-gamma) <= 0.08)} report = bench.make_report('tabular', 'mlp_tiny', base_block, idea_full, {'idea_sweep': idea_sweep, 'mechanism_signature': sig, 'budget': {'epochs': EPOCHS, 'batch': BATCH, 'baseline_gradient_evals': base_full['details'][0]['grad_evals'], 'idea_gradient_evals': idea_full['details'][0]['grad_evals']}}) Path('bench_report.json').write_text(json.dumps(report, indent=2, allow_nan=False)) print(json.dumps(report, indent=2, allow_nan=False)) if __name__ == '__main__': main()