import sys, os, 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, train_model, evaluate, sweep_baseline, make_report SEEDS = tuple(range(8)) SWEEP_SEEDS = (0, 1, 2, 3) EPOCHS = 20 WARM_STEPS = 3 BATCH = 128 # Union of all learning rates considered by either method. LR_GRID = [1e-3, 3e-3, 6e-3] KAPPA_GRID = [0.01, 0.03, 0.06] 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 baseline_one(cfg, seed, capture=False): seed_all(seed) ds = get_dataset('tabular', seed, n_train=400, n_test=400) model = make_model('mlp_tiny', ds['input_shape'], ds['out_dim']) net, metric, hist = train_model(model, ds, epochs=EPOCHS, lr=cfg['lr'], batch=BATCH, weight_decay=cfg.get('weight_decay', 0.0), log=lambda *_: None) return float(metric) def flat_params(model): return torch.cat([p.detach().flatten() for p in model.parameters()]) def set_flat(model, vec): off = 0 with torch.no_grad(): for p in model.parameters(): n = p.numel(); p.copy_(vec[off:off+n].view_as(p)); off += n def idea_one(cfg, seed, capture=False): seed_all(seed) ds = get_dataset('tabular', seed, n_train=400, n_test=400) # The local loop is necessary because the intervention changes the optimizer. device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') try: target = make_model('mlp_tiny', ds['input_shape'], ds['out_dim']).to(device) packet = make_model('mlp_tiny', ds['input_shape'], ds['out_dim']).to(device) except Exception: device = torch.device('cpu') target = make_model('mlp_tiny', ds['input_shape'], ds['out_dim']).to(device) packet = make_model('mlp_tiny', ds['input_shape'], ds['out_dim']).to(device) # Same initialization, then prepare phi by cancelling the empirically slow # displacement direction v (the displacement accumulated in a short warmup). init = flat_params(target).clone() packet.load_state_dict(target.state_dict()) xtr, ytr = ds['xtr'].to(device), ds['ytr'].to(device) lossf = nn.MSELoss() warm_opt = torch.optim.Adam(target.parameters(), lr=cfg['lr']) target.train() warm_steps = WARM_STEPS for j in range(warm_steps): for a in range(0, len(xtr), BATCH): warm_opt.zero_grad(set_to_none=True) lossf(target(xtr[a:a+BATCH]), ytr[a:a+BATCH]).backward(); warm_opt.step() warm_delta = flat_params(target) - init v = warm_delta / (warm_delta.norm() + 1e-12) # restore target to identical post-warmup state and reflect only v component # in packet: v^T(phi-init) = -v^T(theta-init). theta0 = flat_params(target).clone() packet.load_state_dict(target.state_dict()) packet_vec = theta0 - 2.0 * torch.dot(theta0 - init, v) * v set_flat(packet, packet_vec) residual_before = abs(torch.dot((theta0-init + packet_vec-init)/2, v)).item() / (abs(torch.dot(theta0-init, v)).item()+1e-12) opt_t = torch.optim.Adam(target.parameters(), lr=cfg['lr']) opt_p = torch.optim.Adam(packet.parameters(), lr=cfg['lr']) # Continue with identical minibatch order and symmetric parameter coupling. coupling_start = None; coupling_end = None for ep in range(max(1, EPOCHS - WARM_STEPS)): target.train(); packet.train() perm = torch.randperm(len(xtr), device=device) for ii in range(0, len(xtr), BATCH): ix = perm[ii:ii+BATCH] opt_t.zero_grad(set_to_none=True); opt_p.zero_grad(set_to_none=True) lt = lossf(target(xtr[ix]), ytr[ix]); lp = lossf(packet(xtr[ix]), ytr[ix]) lt.backward(); lp.backward() with torch.no_grad(): tv, pv = flat_params(target), flat_params(packet) diff = cfg['kappa'] * (tv-pv) # Adam applies its update after gradients; coupling is added to # gradients parameter-wise, preserving the stated Euler form. off=0 for p in target.parameters(): n=p.numel(); p.grad.add_(diff[off:off+n].view_as(p)); off += n off=0 for p in packet.parameters(): n=p.numel(); p.grad.add_((-diff[off:off+n]).view_as(p)); off += n opt_t.step(); opt_p.step() if ep == 0: curv = (flat_params(target)+flat_params(packet))/2 - init coupling_start = abs(torch.dot(curv, v)).item() / (abs(torch.dot(flat_params(target)-init, v)).item()+1e-12) with torch.no_grad(): pred = target(ds['xte'].to(device)); metric = lossf(pred, ds['yte'].to(device)).item() final_c = (flat_params(target)+flat_params(packet))/2 - init final_target = flat_params(target)-init coupling_end = abs(torch.dot(final_c, v)).item()/(abs(torch.dot(final_target,v)).item()+1e-12) if capture: return float(metric), {'initial_cancellation_ratio': residual_before, 'observed_c_ratio_epoch1': coupling_start, 'observed_c_ratio_final': coupling_end, 'v_norm': float(v.norm().item())} return float(metric) def make_base(cfg): return lambda seed: baseline_one(cfg, seed) def make_idea(cfg): return lambda seed: idea_one(cfg, seed) def main(): # Baseline includes every lr used by the idea; method-central Adam decay is # explicitly swept as a parity knob (zero is the standard setting). base_grid = [{'lr': lr, 'weight_decay': wd} for lr in LR_GRID for wd in [0.0]] base = sweep_baseline(make_base, base_grid, seeds=SWEEP_SEEDS) idea_grid = [{'lr': lr, 'kappa': k} for lr in LR_GRID for k in KAPPA_GRID] # Nearby lr settings are also evaluated on baseline through base_grid. idea_trials=[] for cfg in idea_grid: r=evaluate(make_idea(cfg), seeds=SWEEP_SEEDS) idea_trials.append({'cfg':cfg, 'mean':r['mean'], 'result':r}) best=min(idea_trials, key=lambda z:z['mean']) idea_full=evaluate(make_idea(best['cfg']), seeds=SEEDS) rep=make_report('tabular','mlp_tiny', {'best_cfg':base['best_cfg'],'sweep':base['sweep'],'full':base['full']}, idea_full, {'predicted': 'packet preparation makes common slow-direction projection near zero', 'observed': idea_one(best['cfg'], 0, capture=True)[1], 'confirmed': idea_one(best['cfg'], 0, capture=True)[1]['initial_cancellation_ratio'] < 0.1, 'idea_sweep': idea_trials}) Path('bench_report.json').write_text(json.dumps(rep, indent=2)) print(json.dumps(rep, indent=2)) if __name__ == '__main__': main()