import json, random from pathlib import Path import numpy as np import torch import sys sys.path.insert(0, '/home/maxwelhelp/all/math2nn') from bench import get_dataset, make_model, evaluate, sweep_baseline, make_report SEEDS = tuple(range(8)) SWEEP_SEEDS = (0, 1, 2, 3) EPOCHS = 12 BATCH = 128 MASK_RATE = 0.70 def math_check(): r = np.random.default_rng(123) lam, beta = 1.0, 0.9 ratios = [] for _ in range(2000): v, g = r.exponential(size=32), r.normal(size=32) ratios.append(np.linalg.norm(g/(lam+v))/(np.linalg.norm(g)/lam)) q = np.array([.05, .2, 1., 4.]) v = np.zeros(4) for _ in range(500): v = beta*v + (1-beta)*q pred = np.abs(1-.8*q/(lam+q)) obs = np.abs(1-.8*q/(lam+v)) return {'bound_max_ratio':float(max(ratios)), 'bound':1.0, 'contraction_predicted':pred.tolist(), 'contraction_measured':obs.tolist(), 'max_contraction_error':float(np.max(np.abs(pred-obs)))} def feature_mask(x, seed, epoch, start): g = torch.Generator().manual_seed(int(seed*1000003 + epoch*9176 + start)) n,d = x.shape width = max(1, int(round(d*(1-MASK_RATE)))) starts = torch.randint(0,d,(n,),generator=g) cols = torch.arange(d).view(1,-1) return (((cols-starts[:,None]) % d) < width).to(x.dtype) def run(seed, lr, weight_decay=0.0, method='adam', lam=1.0, beta=.95, return_sig=False): torch.manual_seed(seed); np.random.seed(seed); random.seed(seed) d = get_dataset('tabular', seed, n_train=400, n_test=100) model = make_model('mlp_tiny', d['input_shape'], d['out_dim']) device = 'cpu' model.to(device) x,y = d['xtr'].float(), d['ytr'].float() params = list(model.parameters()) state = [torch.zeros_like(p) for p in params] m = [torch.zeros_like(p) for p in params] vv = [torch.zeros_like(p) for p in params] step = 0 logged_pred, logged_obs, logged_sens = [], [], [] model.train() for ep in range(EPOCHS): gen = torch.Generator().manual_seed(seed*1009+ep) order = torch.randperm(len(x), generator=gen) for st in range(0,len(x),BATCH): ix=order[st:st+BATCH]; xb=x[ix]; yb=y[ix] mask=feature_mask(xb,seed,ep,st); xb=xb*mask model.zero_grad(set_to_none=True) pred=model(xb); loss=((pred-yb)**2).mean(); loss.backward() grads=[p.grad.detach().clone() for p in params] if method=='mop': # Diagonal masked-Jacobian Gramian. For scalar output, per-sample # parameter sensitivities are obtained by one backward pass/sample. sens=[torch.zeros_like(p) for p in params] bsz=len(xb) for k in range(bsz): model.zero_grad(set_to_none=True) model(xb[k:k+1]).sum().backward() for j,p in enumerate(params): sens[j] += p.grad.detach()**2 / bsz for j,p in enumerate(params): state[j].mul_(beta).add_(sens[j],alpha=1-beta) before = p.data.clone() p.data.add_(grads[j]/(lam+state[j]), alpha=-lr) if ep == EPOCHS-1: good = grads[j].abs() > 1e-12 if bool(good.any()): logged_pred.append(float(torch.mean((1/(lam+state[j]))[good]))) logged_obs.append(float(torch.mean(((before-p.data).abs()/(lr*grads[j].abs()))[good]))) logged_sens.append(float(torch.mean(state[j]))) else: step += 1 for j,p in enumerate(params): m[j].mul_(.9).add_(grads[j],alpha=.1) vv[j].mul_(.999).addcmul_(grads[j],grads[j],value=.001) upd=m[j]/(1-.9**step)/(torch.sqrt(vv[j]/(1-.999**step))+1e-8) p.data.add_(upd + weight_decay*p.data, alpha=-lr) model.eval() with torch.no_grad(): metric=float(((model(d['xte'].float())-d['yte'].float())**2).mean()) if not return_sig: return metric return metric, {'mean_observed_diag_gramian':float(np.mean(logged_sens)), 'mean_update_factor_predicted':float(np.mean(logged_pred)), 'mean_update_factor_observed':float(np.mean(logged_obs)), 'relative_factor_error':float(abs(np.mean(logged_pred)-np.mean(logged_obs))/max(np.mean(logged_pred),1e-12)), 'confirmed':bool(abs(np.mean(logged_pred)-np.mean(logged_obs))/max(np.mean(logged_pred),1e-12) < 0.05)} def main(): # Baseline decisive knobs: Adam learning rate and weight decay. The idea uses # the identical lr union and the same weight-decay choices for parity. grid=[{'lr':lr,'weight_decay':wd} for lr in [0.001,0.002,0.003,0.006] for wd in [0.0,1e-4]] def base_fn(cfg): return lambda s: run(s,cfg['lr'],cfg['weight_decay'],'adam') base=sweep_baseline(base_fn,grid,seeds=SWEEP_SEEDS) best=base['best_cfg'] # Evaluate the full union on baseline and idea; baseline selection is retained, # while all idea candidate lrs were also run by the baseline sweep. idea_cfgs=[best, {'lr':0.002,'weight_decay':best['weight_decay']}, {'lr':0.006,'weight_decay':best['weight_decay']}] idea_runs=[] for cfg in idea_cfgs: vals=evaluate(lambda s:run(s,cfg['lr'],cfg['weight_decay'],'mop',lam=1.0,beta=.95),seeds=SEEDS) idea_runs.append({'cfg':cfg,'result':vals}) chosen=min(idea_runs,key=lambda z:z['result']['mean']) base_full=evaluate(base_fn(best),seeds=SEEDS) sig_metric,sig=run(0,best['lr'],best['weight_decay'],'mop',return_sig=True) report=make_report('tabular','mlp_tiny',{'best_cfg':best,'sweep':base['sweep'],'full':base_full},chosen['result'],{'math_check':math_check(),'trained_model_signature':sig,'idea_configs':idea_runs}) report['idea']['selected_cfg']=chosen['cfg'] report['protocol_notes']='Optimizer modification on structurally matched Friedman tabular regression; paired seeds 0-7, batch 128, 12 epochs, same architecture and lr/weight-decay union.' Path('bench_report.json').write_text(json.dumps(report,indent=2)) print(json.dumps(report,indent=2)) if __name__=='__main__': main()