import argparse import json import os import time import numpy as np import torch from torch import nn from torchvision import datasets, transforms def operator_check(d=24, gamma=0.35, lam=0.4, trials=300000, seed=17): rng = np.random.default_rng(seed) eig = np.linspace(0.03, 1.0, d) / d idx = rng.integers(0, d, trials) scale2 = d * eig[idx] p = np.full((trials, d), 1.0 - gamma * lam) p[np.arange(trials), idx] -= gamma * scale2 ep2 = np.mean(p * p, axis=0) A = 1.0 - gamma * (eig + lam) rhs1 = (1.0 - gamma * lam) * A rhs2 = (1.0 - gamma * lam) ** 2 gap1 = float(np.max(ep2 - rhs1)) gap2 = float(np.max(rhs1 - rhs2)) return {"max_E_P2_minus_(1-gamma-lambda)A": gap1, "max_(1-gamma-lambda)A_minus_scalar_bound": gap2, "gamma_lambda": gamma * lam, "A_min": float(A.min()), "A_max": float(A.max()), "passed": bool(gap1 <= 0.01 and gap2 <= 1e-12 and 0 <= gamma*lam <= 1)} def load_mnist(root): tf = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,)), transforms.Lambda(lambda x: x.view(-1))]) tr = datasets.MNIST(root, train=True, download=True, transform=tf) te = datasets.MNIST(root, train=False, download=True, transform=tf) # Materialize once: direct tensor indexing is substantially faster than repeated PIL transforms. Xtr = torch.stack([tr[i][0] for i in range(len(tr))]) ytr = torch.tensor(tr.targets, dtype=torch.long) Xte = torch.stack([te[i][0] for i in range(len(te))]) yte = torch.tensor(te.targets, dtype=torch.long) return Xtr, ytr, Xte, yte def accuracy(model, x, y, device): model.eval(); correct = 0 with torch.inference_mode(): for i in range(0, len(x), 512): correct += int((model(x[i:i+512].to(device)).argmax(1).cpu() == y[i:i+512]).sum()) model.train() return correct / len(x) def train_one(data, seed, mode, device, T=3000, width=128, batch=128, wd=.05, p=.5, q=.2): Xtr, ytr, Xte, yte = data torch.manual_seed(seed); np.random.seed(seed) model = nn.Sequential(nn.Linear(784, width), nn.ReLU(), nn.Linear(width, width), nn.ReLU(), nn.Linear(width, 10)).to(device) opt = torch.optim.AdamW(model.parameters(), lr=1e-3, betas=(.9, .999), weight_decay=wd) sched = torch.optim.lr_scheduler.CosineAnnealingLR(opt, T_max=T) tail_start, tail_end = int((1-q)*T), T avg = [torch.zeros_like(x, dtype=torch.float32, device='cpu') for x in model.parameters()] avg_n = 0; best = -1.; best_step = 0; grad_mid = 0.; grad_n = 0 rng = torch.Generator().manual_seed(seed + 10000) t0 = time.perf_counter() for step in range(T): ix = torch.randint(len(Xtr), (batch,), generator=rng) xb, yb = Xtr[ix].to(device), ytr[ix].to(device) opt.param_groups[0]['weight_decay'] = wd if (mode in ('baseline','wd_tail') or step < int(p*T)) else 0. opt.zero_grad(set_to_none=True); loss = nn.functional.cross_entropy(model(xb), yb); loss.backward() gn = torch.sqrt(sum((x.grad.detach()**2).sum() for x in model.parameters() if x.grad is not None)).item() if 2000 <= step < 3000: grad_mid += gn; grad_n += 1 opt.step(); sched.step() if mode.endswith('tail') and tail_start <= step < tail_end: with torch.no_grad(): for a, x in zip(avg, model.parameters()): a.add_(x.detach().cpu()) avg_n += 1 # Test every 300 steps plus final: adequate for a small deterministic super-check. if (step + 1) % 300 == 0 or step == T-1: acc = accuracy(model, Xte, yte, device) if acc > best: best, best_step = acc, step + 1 last_acc = acc; tail_acc = None if mode.endswith('tail') and avg_n: saved = [x.detach().clone() for x in model.parameters()] with torch.no_grad(): for x, a in zip(model.parameters(), avg): x.copy_((a / avg_n).to(device)) tail_acc = accuracy(model, Xte, yte, device) with torch.no_grad(): for x, s in zip(model.parameters(), saved): x.copy_(s) norm = torch.sqrt(sum((x.detach()**2).sum() for x in model.parameters())).item() return {'test_accuracy': float(tail_acc if tail_acc is not None else last_acc), 'last_nonaveraged_accuracy': float(last_acc), 'best_test_accuracy': float(best), 'steps_to_best': int(best_step), 'final_weight_norm': float(norm), 'mid_gradient_norm': float(grad_mid/max(1,grad_n)), 'runtime_sec': time.perf_counter()-t0} def main(): ap = argparse.ArgumentParser(); ap.add_argument('--runs', type=int, default=5) ap.add_argument('--T', type=int, default=3000); ap.add_argument('--width', type=int, default=128) ap.add_argument('--out', default='super_results.json'); args = ap.parse_args() try: device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') except Exception: device = torch.device('cpu') root = os.path.join(os.path.dirname(__file__), 'data'); data = load_mnist(root) check = operator_check(); modes = ['baseline','initwd','initwd_tail','wd_tail']; raw={m:[] for m in modes}; start=time.perf_counter() for seed in range(args.runs): for mode in modes: try: row=train_one(data, seed, mode, device, T=args.T, width=args.width) except RuntimeError as e: if device.type=='cuda': torch.cuda.empty_cache(); device=torch.device('cpu'); row=train_one(data, seed, mode, device, T=args.T, width=args.width) else: raise raw[mode].append(row); print(mode, seed, row, flush=True) summary={} for mode, rows in raw.items(): summary[mode]={} for key in ['test_accuracy','best_test_accuracy','steps_to_best','final_weight_norm','mid_gradient_norm','runtime_sec']: v=np.array([r[key] for r in rows],float); summary[mode][key+'_mean']=float(v.mean()); summary[mode][key+'_std']=float(v.std(ddof=1)) out={'config':{'dataset':'MNIST train=60000 test=10000','T':args.T,'width':args.width,'batch':128,'lr':1e-3,'betas':[.9,.999],'weight_decay':.05,'p':.5,'q':.2,'seeds':args.runs,'device':str(device),'total_runtime_sec':time.perf_counter()-start},'operator_check':check,'summary':summary,'raw':raw} with open(args.out,'w') as f: json.dump(out,f,indent=2) print(json.dumps({'config':out['config'],'operator_check':check,'summary':summary},indent=2)) if __name__ == '__main__': main()