import os, sys, json, math 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, evaluate, sweep_baseline, make_report SEED0 = 1616 BATCH = 64 EPOCHS = 18 TAUS = [0.0, 0.05, 0.10] LRS = [1e-3, 3e-3, 6e-3] BETAS = [0.8, 0.9] def polar(x): # Batched-free, exact polar factor for a 2-D weight/gradient matrix. u, _, vh = torch.linalg.svd(x, full_matrices=False) return u @ vh def matrix_sign(x): return torch.where(x >= 0, torch.ones_like(x), -torch.ones_like(x)) def direction(m, residual, tau, heldout, gated): post = matrix_sign(polar(m)) pre = polar(matrix_sign(m + residual)) rho = torch.sum(heldout * post) / (torch.linalg.vector_norm(heldout) * torch.linalg.vector_norm(post) + 1e-8) if gated and float(rho.detach()) < tau: return pre, post, pre, float(rho.detach()), False return post, post, pre, float(rho.detach()), True def train_one(seed, lr, beta, method, tau=0.05, collect=False): # Explicit CPU fallback also avoids inheriting a failed CUDA context. torch.manual_seed(seed); np.random.seed(seed) ds = get_dataset('tabular', seed, n_train=400, n_test=200) device = 'cuda' if torch.cuda.is_available() else 'cpu' try: net = make_model('mlp_tiny', ds['input_shape'], ds['out_dim']).to(device) except Exception: device = 'cpu' net = make_model('mlp_tiny', ds['input_shape'], ds['out_dim']).to(device) x, y = ds['xtr'].to(device), ds['ytr'].to(device) xt, yt = ds['xte'].to(device), ds['yte'].to(device) lossf = nn.MSELoss() # SignMuon-like state only for 2-D weight matrices; vectors use Adam-style SGD. mom = {p: torch.zeros_like(p) for p in net.parameters() if p.ndim == 2} residual = {p: torch.zeros_like(p) for p in net.parameters() if p.ndim == 2} rng = np.random.default_rng(seed + 991) logs = {'post_fraction': [], 'rho': [], 'post_negative': [], 'chosen_negative': [], 'predicted_fraction': []} net.train() n = x.shape[0] for epoch in range(EPOCHS): order = rng.permutation(n) # The next cyclic minibatch is an independent routing estimate. for bi in range(0, n, BATCH): ids = order[bi:bi+BATCH] hid = order[(bi+BATCH) % n:(bi+2*BATCH) % n] if bi+BATCH < n else order[:min(BATCH,n)] # avoid accidental empty/wrapped slices if len(hid) == 0: hid = order[:BATCH] net.zero_grad(set_to_none=True) pred = net(x[ids].view(len(ids), -1)) lossf(pred, y[ids]).backward() # Save current gradients, then independently estimate heldout gradients. grads = {p: (p.grad.detach().clone() if p.grad is not None else torch.zeros_like(p)) for p in net.parameters()} net.zero_grad(set_to_none=True) lossf(net(x[hid].view(len(hid), -1)), y[hid]).backward() held = {p: (p.grad.detach().clone() if p.grad is not None else torch.zeros_like(p)) for p in net.parameters()} with torch.no_grad(): for p in net.parameters(): g = grads[p] if p.ndim == 2: mom[p].mul_(beta).add_(g, alpha=1-beta) r = residual[p] d, post, pre, rho, use_post = direction(mom[p], r, tau, held[p], method == 'switch') p.add_(d, alpha=-lr / math.sqrt(max(1, p.shape[1]))) if method == 'switch' and not use_post: residual[p].copy_(mom[p] + r - matrix_sign(mom[p] + r)) elif method == 'pre': residual[p].copy_(mom[p] + r - matrix_sign(mom[p] + r)) if collect: logs['post_fraction'].append(float(use_post)) logs['rho'].append(rho) logs['post_negative'].append(float(torch.sum(held[p]*post).item() < 0)) logs['chosen_negative'].append(float(torch.sum(held[p]*d).item() < 0)) logs['predicted_fraction'].append(float(rho >= tau)) else: # Same simple gradient-side update for all non-matrix parameters. p.add_(g, alpha=-lr) net.eval() with torch.no_grad(): metric = float(lossf(net(xt.view(len(xt), -1)), yt).cpu()) if collect: logs = {k: float(np.mean(v)) if v else 0.0 for k,v in logs.items()} logs['n_observations'] = int(len(order) * EPOCHS / BATCH * max(1, len(mom))) return metric, logs def fn(cfg, method): return lambda seed: train_one(seed, cfg['lr'], cfg['beta'], method, cfg.get('tau', 0.05))[0] def main(): # Union parity: every idea lr/beta is included in baseline's grid. grid = [{'lr': lr, 'beta': beta} for lr in LRS for beta in BETAS] base = sweep_baseline(lambda c: fn(c, 'post'), grid) idea_grid = [dict(c, tau=t) for c in grid for t in TAUS] # Evaluate the idea's 3 tau choices at the selected baseline settings and two nearby lr settings. bcfg = base['best_cfg']; nearby = sorted(set([bcfg['lr']] + [v for v in LRS if v != bcfg['lr']]))[:3] candidates = [{'lr': lr, 'beta': bcfg['beta'], 'tau': tau} for lr in nearby for tau in TAUS] tried = [] for c in candidates: r = evaluate(fn(c, 'switch')) tried.append((c, r)) best_cfg, idea = min(tried, key=lambda z: z[1]['mean']) sigs = [train_one(s, best_cfg['lr'], best_cfg['beta'], 'switch', best_cfg['tau'], True)[1] for s in range(8)] sig = {k: float(np.mean([z[k] for z in sigs])) for k in sigs[0] if k != 'n_observations'} sig['confirmed'] = bool(abs(sig['post_fraction'] - sig['predicted_fraction']) < 1e-9 and sig['chosen_negative'] <= sig['post_negative']) base['sweep'] = base['sweep'] report = make_report('tabular', 'mlp_tiny', base, idea, { 'prediction': 'heldout alignment gating routes post when rho>=tau and reduces negative heldout alignment', 'observed': sig, 'tau': best_cfg['tau'], 'trained_models': True }) report['idea_sweep'] = [{'cfg': c, 'mean': r['mean'], 'per_seed': r['per_seed']} for c,r in tried] report['custom_track'] = None with open('bench_report.json','w') as f: json.dump(report, f, indent=2) print(json.dumps(report, indent=2)) if __name__ == '__main__': main()