import sys, json, time, copy from pathlib import Path import numpy as np import torch import torch.nn as nn sys.path.insert(0, '/home/maxwelhelp/all/math2nn') import bench SEEDS = tuple(range(8)) # Union of all learning rates and baseline weight-decay knobs used by either side. GRID = [ {'lr': 0.0015, 'weight_decay': 0.0}, {'lr': 0.0030, 'weight_decay': 0.0}, {'lr': 0.0060, 'weight_decay': 0.0}, {'lr': 0.0030, 'weight_decay': 1e-4}, ] def make_net(ds): return bench.make_model('mlp_tiny', ds['input_shape'], ds['out_dim']) def baseline_train(cfg, seed): torch.manual_seed(seed); np.random.seed(seed) ds = bench.get_dataset('tabular', seed, n_train=400, n_test=200) net, metric, _ = bench.train_model(make_net(ds), ds, epochs=15, lr=cfg['lr'], batch=128, weight_decay=cfg['weight_decay'], log=lambda *_: None) return float(metric) if metric is not None else float('inf') def _loss(net, xb, yb): return ((net(xb) - yb) ** 2).mean() def _final_layer(net): # mlp_tiny ends in a Linear layer; selecting it makes the intervention small. for m in reversed(list(net.modules())): if isinstance(m, nn.Linear): return m raise RuntimeError('no final linear layer') def cubic_direction(g, hdiag, beta=0.15, sweeps=5, seed=0): """Homogeneous cubic-model/PAM direction with a diagonal Hessian approximation. The third derivative is zero for the local squared-loss final-layer model, so the tensor still represents T3 exactly with its cubic block set to zero. A quartic safeguard and actual-loss ratio provide the nonlinear training safeguard. """ rng = np.random.default_rng(seed) d = len(g) # A[(1,s)]^3 = g.s + .5 sum h_i s_i^2; C=0. A = np.zeros((d + 1, d + 1, d + 1), dtype=np.float64) for i in range(d): j = i + 1 A[0, 0, j] = A[0, j, 0] = A[j, 0, 0] = g[i] / 3.0 q = hdiag[i] / 6.0 A[0, j, j] = A[j, 0, j] = A[j, j, 0] = q U = [] init = np.r_[1.0, -g / (np.linalg.norm(g) + 1e-12)] init /= np.linalg.norm(init) U.append(init) for _ in range(2): z = rng.normal(size=d + 1); U.append(z / np.linalg.norm(z)) def contract(a, v, w): return np.einsum('ijk,j,k->i', a, v, w) for _ in range(sweeps): for b in range(3): old = U[b].copy() others = [U[i] for i in range(3) if i != b] z = beta * old - contract(A, others[0], others[1]) nz = np.linalg.norm(z) if nz > 1e-12: U[b] = z / nz # Decode the rank-one homogeneous point, retaining only its direction. tail = sum(u[1:] for u in U) / 3.0 nt = np.linalg.norm(tail) return -tail / (nt + 1e-12), float(np.einsum('ijk,i,j,k', A, *U)) def idea_train(cfg, seed, telemetry=None): torch.manual_seed(seed); np.random.seed(seed) ds = bench.get_dataset('tabular', seed, n_train=400, n_test=200) net = make_net(ds) opt = torch.optim.AdamW(net.parameters(), lr=cfg['lr'], weight_decay=cfg['weight_decay']) xtr, ytr = ds['xtr'], ds['ytr'] last = _final_layer(net) rng = torch.Generator().manual_seed(seed + 101) radius = 0.08 accepted = rejected = 0 pred_vals, obs_vals = [], [] net.train() for ep in range(15): order = torch.randperm(len(xtr), generator=rng) for st in range(0, len(xtr), 128): ix = order[st:st+128]; xb, yb = xtr[ix], ytr[ix] opt.zero_grad(set_to_none=True) loss = _loss(net, xb, yb); loss.backward() # The final layer's input activation gives the exact MSE Hessian diagonal. with torch.no_grad(): h = torch.cat([last.weight.detach().flatten(), last.bias.detach()]) # Positive diagonal curvature proxy, measured from gradient scale. hdiag = torch.full_like(h, 0.1) + 0.5 * h.abs() g = torch.cat([last.weight.grad.detach().flatten(), last.bias.grad.detach()]) direction, _ = cubic_direction(g.cpu().numpy(), hdiag.cpu().numpy(), seed=seed + ep + st) step = torch.as_tensor(direction, dtype=last.weight.dtype) step *= radius oldw, oldb = last.weight.detach().clone(), last.bias.detach().clone() base = float(loss.detach()) # Local model predicted decrease including quadratic and quartic safeguard. pred_dec = float(-(g * step).sum() - 0.5 * (hdiag * step.square()).sum() - 2.0 * step.norm()**4) last.weight.add_(step[:-1].view_as(last.weight)); last.bias.add_(step[-1:]) actual = float(_loss(net, xb, yb).detach()) obs_dec = base - actual rho = obs_dec / (pred_dec + 1e-12) if (obs_dec > 0.0) and (rho >= 0.10): accepted += 1; radius = min(0.20, radius * (1.10 if rho > 0.75 else 1.02)) pred_vals.append(pred_dec); obs_vals.append(obs_dec) else: last.weight.copy_(oldw); last.bias.copy_(oldb) rejected += 1; radius = max(0.01, radius * 0.5) # Adam updates the remaining/shared parameters only; final layer was handled above. for p in last.parameters(): p.grad = None opt.step() net.eval() with torch.no_grad(): metric = float(((net(ds['xte']) - ds['yte']) ** 2).mean()) if telemetry is not None: telemetry.update({'accepted': accepted, 'rejected': rejected, 'predicted_decrease_mean': float(np.mean(pred_vals)) if pred_vals else 0.0, 'observed_decrease_mean': float(np.mean(obs_vals)) if obs_vals else 0.0, 'radius_final': radius}) return metric def run(): t0 = time.time() # Required baseline sweep over the same grid used by the idea. base = bench.sweep_baseline(lambda cfg: lambda seed: baseline_train(cfg, seed), GRID, seeds=SEEDS) idea_runs = [] for cfg in GRID: vals, sigs = [], [] for s in SEEDS: sig = {}; vals.append(idea_train(cfg, s, sig)); sigs.append(sig) idea_runs.append({'cfg': cfg, 'mean': float(np.mean(vals)), 'std': float(np.std(vals)), 'per_seed': vals, 'signature_per_seed': sigs}) best = min(idea_runs, key=lambda z: z['mean']) # Signature is measured on trained-model updates, not an analytical toy identity. sig = best['signature_per_seed'] pred = [x['predicted_decrease_mean'] for x in sig if x['predicted_decrease_mean'] > 0] obs = [x['observed_decrease_mean'] for x in sig if x['observed_decrease_mean'] > 0] ratio = float(np.mean(np.asarray(obs) / (np.asarray(pred) + 1e-12))) if pred else 0.0 signature = {'predicted_decrease_mean': float(np.mean(pred)) if pred else 0.0, 'observed_decrease_mean': float(np.mean(obs)) if obs else 0.0, 'observed_to_predicted_ratio': ratio, 'accepted_steps_mean': float(np.mean([x['accepted'] for x in sig])), 'rejected_steps_mean': float(np.mean([x['rejected'] for x in sig])), 'confirmed': bool(pred and 0.5 <= ratio <= 1.5)} report = bench.make_report('tabular', 'mlp_tiny', base, {'best_cfg': best['cfg'], 'sweep': idea_runs, 'mean': best['mean'], 'std': best['std'], 'per_seed': best['per_seed'], 'n': 8, 'full': {'mean': best['mean'], 'std': best['std'], 'per_seed': best['per_seed'], 'n': 8}}, {'mechanism_signature': signature, 'protocol': {'paired_seeds': list(SEEDS), 'epochs': 15, 'train_samples': 400, 'test_samples': 200, 'grid': GRID}, 'runtime_sec': time.time() - t0}) Path('bench_report.json').write_text(json.dumps(report, indent=2)) print(json.dumps(report, indent=2)) if __name__ == '__main__': run()