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') from bench import get_dataset, make_model, train_model, make_report from bench.protocol import evaluate, sweep_baseline SEEDS = tuple(range(8)) SWEEP_SEEDS = (0, 1, 2, 3) EPOCHS = 8 BATCH = 64 RHO = 0.035 EPS = 0.08 def sinkhorn_ot(a, b, eps=EPS, iters=18): c = ((a[:, None, :] - b[None, :, :]) ** 2).sum(-1) logk = -c / eps la = torch.full((a.shape[0],), -math.log(a.shape[0]), device=a.device, dtype=a.dtype) lb = torch.full((b.shape[0],), -math.log(b.shape[0]), device=b.device, dtype=b.dtype) u = torch.zeros_like(la) v = torch.zeros_like(lb) for _ in range(iters): u = la - torch.logsumexp(logk + v[None, :], dim=1) v = lb - torch.logsumexp(logk + u[:, None], dim=0) return (torch.exp(logk + u[:, None] + v[None, :]) * c).sum() def sinkhorn_div(a, b): return sinkhorn_ot(a, b) - 0.5 * sinkhorn_ot(a, a) - 0.5 * sinkhorn_ot(b, b) def seed_all(seed): random.seed(seed) np.random.seed(seed) torch.manual_seed(seed) def make_basis(win=32): t = torch.linspace(-1, 1, win) b1 = torch.sin(math.pi * (t + 1) / 2) b2 = torch.cos(math.pi * (t + 1) / 2) return torch.stack((b1 / b1.norm(), b2 / b2.norm()), dim=1) def math_check(): z = torch.linspace(-1.0, 1.0, 16)[:, None] vals = [] for delta in np.linspace(0, 0.6, 7): vals.append(float(sinkhorn_div(z + delta, z))) ds = np.linspace(0.1, 0.6, 6) yy = np.asarray(vals[1:]) k = float(np.dot(ds ** 2, yy) / np.dot(ds ** 2, ds ** 2)) rel = float(np.max(np.abs(yy - k * ds ** 2) / (np.abs(yy) + 1e-8))) ident = float(sinkhorn_div(z, z)) return {'prediction': 'translation S approximately k*delta^2', 'fitted_k': k, 'max_relative_error': rel, 'S_values': vals, 'debiased_identity': ident, 'confirmed': bool(rel < 0.03 and abs(ident) < 1e-6)} def baseline_run(cfg, seed, return_info=False): seed_all(seed) ds = get_dataset('sequence', seed, n_train=400, n_test=200) net = make_model('transformer_tiny', ds['input_shape'], ds['out_dim']) net, metric, hist = train_model(net, ds, epochs=cfg['epochs'], lr=cfg['lr'], batch=BATCH, weight_decay=cfg['weight_decay'], log=lambda *_: None) return metric def idea_run(cfg, seed, return_info=False): seed_all(seed) ds = get_dataset('sequence', seed, n_train=400, n_test=200) dev = torch.device('cuda' if torch.cuda.is_available() else 'cpu') try: xtr, ytr = ds['xtr'].to(dev), ds['ytr'].to(dev) xte, yte = ds['xte'].to(dev), ds['yte'].to(dev) net = make_model('transformer_tiny', ds['input_shape'], ds['out_dim']).to(dev) basis = make_basis(xtr.shape[1]).to(dev) coeff = torch.zeros(2, device=dev, requires_grad=True) lam = 0.0 opt = torch.optim.Adam(net.parameters(), lr=cfg['lr'], weight_decay=cfg['weight_decay']) lossf = nn.MSELoss() last_s = 0.0 for ep in range(cfg['epochs']): net.train() perm = torch.randperm(len(xtr), device=dev) for start in range(0, len(xtr), BATCH): idx = perm[start:start+BATCH] xb, yb = xtr[idx], ytr[idx] nominal = xb # Conditional adversarial generator: smooth residual trajectory. adv = nominal + cfg['radius_scale'] * (basis @ coeff).view(1, -1) task = lossf(net(adv), yb) s = sinkhorn_div(adv[:min(32, len(adv))], nominal[:min(32, len(nominal))]) lag = task - lam * torch.relu(s - RHO) g = torch.autograd.grad(lag, coeff, retain_graph=True)[0] with torch.no_grad(): coeff.add_(cfg['adv_lr'] * g).clamp_(-1.5, 1.5) last_s = float(s.detach()) lam = max(0.0, lam + cfg['lambda_lr'] * (last_s - RHO)) adv2 = (nominal + cfg['radius_scale'] * (basis @ coeff).view(1, -1)).detach() loss = lossf(net(adv2), yb) opt.zero_grad(set_to_none=True) loss.backward(); opt.step() net.eval() with torch.no_grad(): metric = float(lossf(net(xte), yte)) if return_info: return metric, {'sinkhorn': last_s, 'rho': RHO, 'lambda': lam, 'coeff_norm': float(coeff.norm())} return metric except RuntimeError: # Explicit CPU fallback, preserving the same algorithm and seed. torch.cuda.empty_cache() if torch.cuda.is_available() else None old = torch.cuda.is_available torch.cuda.is_available = lambda: False try: return idea_run(cfg, seed, return_info) finally: torch.cuda.is_available = old def main(): # Union parity: every lr and method knob used by the idea is in baseline grid. grid = [] for lr in [1e-3, 3e-3, 6e-3]: for wd in [0.0, 1e-4]: grid.append({'lr': lr, 'weight_decay': wd, 'epochs': EPOCHS, 'adv_lr': 0.08, 'lambda_lr': 0.2, 'radius_scale': 0.10}) base = sweep_baseline(lambda cfg: lambda s: baseline_run(cfg, s), grid, SWEEP_SEEDS) best_lr = base['best_cfg']['lr']; best_wd = base['best_cfg']['weight_decay'] idea_grid = [dict(lr=lr, weight_decay=best_wd, epochs=EPOCHS, adv_lr=al, lambda_lr=0.2, radius_scale=rs) for lr, al, rs in [(best_lr, .08, .10), (1e-3 if best_lr != 1e-3 else 3e-3, .08, .10), (best_lr, .04, .07)]] # Select idea hyperparameters on the same sweep seeds, then evaluate full paired seeds. idea_trials = [{'cfg': c, 'mean': evaluate(lambda s, c=c: idea_run(c, s), SWEEP_SEEDS)['mean']} for c in idea_grid] idea_cfg = min(idea_trials, key=lambda x: x['mean'])['cfg'] idea_res = evaluate(lambda s: idea_run(idea_cfg, s), SEEDS) sig = math_check() observed = [] for s in SEEDS: _, inf = idea_run(idea_cfg, s, True) observed.append(inf) sig.update({'nn_observed_mean_sinkhorn': float(np.mean([q['sinkhorn'] for q in observed])), 'nn_observed_mean_lambda': float(np.mean([q['lambda'] for q in observed])), 'nn_predicted_radius': RHO}) report = make_report('sequence', 'transformer_tiny', base, idea_res, {'math_and_nn_signature': sig, 'idea_sweep': idea_trials, 'custom_track': None, 'track_rationale': 'Sequence forecasting has context-conditioned multi-token windows; residual trajectory augmentation acts on the generated sequence manifold.'}) 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()