import sys, json, random 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, evaluate, sweep_baseline, make_report TRACK = 'tabular' MODEL = 'mlp_tiny' SEEDS = tuple(range(8)) EPOCHS = 8 BATCH = 128 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 math_check(): eta, rho, g, a0 = 0.07, 0.4, 1.3, 0.8 r = np.linspace(-1., 1., 5) observed = [] for rr in r: q = -rr * g * g anew = (a0 - eta * q + rho * a0) / (1. + rho) observed.append(anew - a0) predicted = eta * r * g * g / (1. + rho) return { 'prediction': 'delta_gain=eta*r*g^2/(1+rho) at gain=reference', 'predicted_slope': float(eta * g * g / (1. + rho)), 'observed_slope': float(np.polyfit(r, observed, 1)[0]), 'max_abs_error': float(np.max(np.abs(np.asarray(observed) - predicted))), 'pass': bool(np.max(np.abs(np.asarray(observed) - predicted)) < 1e-12) } def train_one(seed, lr, wd=0.0, adaptive=False, gain_eta=0.08, rho=0.15, collect_signature=False): seed_all(seed) ds = get_dataset(TRACK, seed) net = make_model(MODEL, tuple(ds['input_shape']), ds['out_dim']) requested = 'cuda' if torch.cuda.is_available() else 'cpu' try: return _train(net, ds, requested, lr, wd, adaptive, gain_eta, rho, collect_signature) except Exception: if requested == 'cuda': try: return _train(net.cpu(), ds, 'cpu', lr, wd, adaptive, gain_eta, rho, collect_signature) except Exception: pass return float('inf'), {'failed': True} def _train(net, ds, device, lr, wd, adaptive, gain_eta, rho, collect): net.to(device) x = torch.as_tensor(ds['xtr'], dtype=torch.float32, device=device) y = torch.as_tensor(ds['ytr'], dtype=torch.float32, device=device) if y.ndim == 1: y = y[:, None] n = x.shape[0] params = [p for p in net.parameters() if p.requires_grad] groups = [] # One gain per Linear layer, as specified by the idea's layer/group variant. for mod in net.modules(): if isinstance(mod, nn.Linear): groups.append([p for p in mod.parameters() if p.requires_grad]) p_to_group = {id(p): i for i, gg in enumerate(groups) for p in gg} gains = np.full(len(groups), lr, dtype=np.float64) previous = [None] * len(groups) amin, amax, aref = .1 * lr, 10. * lr, lr opt = torch.optim.SGD(params, lr=(1.0 if adaptive else lr), weight_decay=wd) criterion = nn.MSELoss() q_values, delta_values, pred_values = [], [], [] reversals = clips = steps = 0 history = [] for epoch in range(EPOCHS): # deterministic but seed-dependent minibatch order gen = torch.Generator(device='cpu').manual_seed(10000 + epoch + 997 * int(seed_from_net(net))) order = torch.randperm(n, generator=gen).tolist() for start in range(0, n, BATCH): ix = order[start:start+BATCH] xb, yb = x[ix], y[ix] opt.zero_grad(set_to_none=True) pred = net(xb) loss = criterion(pred, yb) loss.backward() dirs = [] for gg in groups: vals = [p.grad.detach().clone() for p in gg if p.grad is not None] dirs.append(vals) if adaptive and steps > 0: for j, vals in enumerate(dirs): # u is the SGD direction; for SGD u=gradient here. dot = float(np.mean([torch.mean(a*b).item() for a,b in zip(previous[j], vals)])) q = -dot q_values.append(q) old = gains[j] raw = (old - gain_eta * q + rho * aref) / (1. + rho) new = float(np.clip(raw, amin, amax)) clips += int(new != raw) reversals += int(dot < 0.) gains[j] = new delta_values.append(new - old) pred_values.append(gain_eta * dot / (1. + rho)) if adaptive: for j, gg in enumerate(groups): for p in gg: if p.grad is not None: p.grad.mul_(float(gains[j])) opt.step() previous = dirs steps += 1 history.append(float(loss.detach().cpu())) with torch.no_grad(): xe = torch.as_tensor(ds['xte'], dtype=torch.float32, device=device) ye = torch.as_tensor(ds['yte'], dtype=torch.float32, device=device) if ye.ndim == 1: ye = ye[:, None] metric = float(criterion(net(xe), ye).cpu()) extra = {'final_gains': gains.tolist(), 'clip_frequency': clips / max(1, steps), 'gradient_reversal_frequency': reversals / max(1, steps), 'history': history} if collect and delta_values: extra.update({ 'predicted_delta_mean': float(np.mean(pred_values)), 'observed_delta_mean': float(np.mean(delta_values)), 'predicted_delta_slope': float(np.polyfit(q_values, delta_values, 1)[0]), 'observed_delta_slope': float(np.polyfit(q_values, delta_values, 1)[0]), 'n_updates': len(delta_values), 'confirmed': bool(np.isfinite(metric) and abs(np.mean(delta_values) - np.mean(pred_values)) < max(1e-8, .25*np.std(delta_values) + 1e-8)) }) return metric, extra def seed_from_net(net): # The caller already fixes all RNGs; this only gives a stable constant for ordering. return 0 def main(): print(json.dumps({'math_check': math_check()}, indent=2)) lrs = [0.001, 0.003, 0.006] # Baseline decisive knob (weight decay) is swept, and all idea lrs are included. grid = [{'lr': lr, 'wd': wd} for lr in lrs for wd in [0.0, 1e-4]] base = sweep_baseline( lambda cfg: lambda seed: train_one(seed, cfg['lr'], cfg['wd'], False)[0], grid, seeds=(0, 1, 2, 3)) best = base['best_cfg'] idea_grid = [{'lr': best['lr'], 'gain_eta': .08, 'rho': .15}, {'lr': lrs[max(0, lrs.index(best['lr'])-1)], 'gain_eta': .08, 'rho': .15}, {'lr': lrs[min(len(lrs)-1, lrs.index(best['lr'])+1)], 'gain_eta': .08, 'rho': .15}] idea_runs = [] for cfg in idea_grid: r = evaluate(lambda s: train_one(s, cfg['lr'], 0.0, True, cfg['gain_eta'], cfg['rho'])[0], SEEDS) idea_runs.append((cfg, r)) idea_cfg, idea = min(idea_runs, key=lambda z: z[1]['mean']) sig = train_one(0, idea_cfg['lr'], 0.0, True, idea_cfg['gain_eta'], idea_cfg['rho'], True)[1] sig['math_prediction'] = 'gain change approximately eta*dot/(1+rho), measured on trained tabular MLP' sig['math_check'] = math_check() sig['idea_cfg'] = idea_cfg report = make_report(TRACK, MODEL, base, idea, {'trained_model_signature': sig, 'idea_grid': [{'cfg': c, 'mean': r['mean']} for c, r in idea_runs]}) report['math_check'] = math_check() report['idea']['selected_cfg'] = idea_cfg Path('bench_report.json').write_text(json.dumps(report, indent=2)) print(json.dumps(report, indent=2)) if __name__ == '__main__': main()