1import sys, json, time, copy
  2from pathlib import Path
  3import numpy as np
  4import torch
  5import torch.nn as nn
  6
  7sys.path.insert(0, '/home/maxwelhelp/all/math2nn')
  8import bench
  9
 10SEEDS = tuple(range(8))
 11# Union of all learning rates and baseline weight-decay knobs used by either side.
 12GRID = [
 13    {'lr': 0.0015, 'weight_decay': 0.0},
 14    {'lr': 0.0030, 'weight_decay': 0.0},
 15    {'lr': 0.0060, 'weight_decay': 0.0},
 16    {'lr': 0.0030, 'weight_decay': 1e-4},
 17]
 18
 19
 20def make_net(ds):
 21    return bench.make_model('mlp_tiny', ds['input_shape'], ds['out_dim'])
 22
 23
 24def baseline_train(cfg, seed):
 25    torch.manual_seed(seed); np.random.seed(seed)
 26    ds = bench.get_dataset('tabular', seed, n_train=400, n_test=200)
 27    net, metric, _ = bench.train_model(make_net(ds), ds, epochs=15,
 28                                        lr=cfg['lr'], batch=128,
 29                                        weight_decay=cfg['weight_decay'], log=lambda *_: None)
 30    return float(metric) if metric is not None else float('inf')
 31
 32
 33def _loss(net, xb, yb):
 34    return ((net(xb) - yb) ** 2).mean()
 35
 36
 37def _final_layer(net):
 38    # mlp_tiny ends in a Linear layer; selecting it makes the intervention small.
 39    for m in reversed(list(net.modules())):
 40        if isinstance(m, nn.Linear):
 41            return m
 42    raise RuntimeError('no final linear layer')
 43
 44
 45def cubic_direction(g, hdiag, beta=0.15, sweeps=5, seed=0):
 46    """Homogeneous cubic-model/PAM direction with a diagonal Hessian approximation.
 47
 48    The third derivative is zero for the local squared-loss final-layer model, so the
 49    tensor still represents T3 exactly with its cubic block set to zero. A quartic
 50    safeguard and actual-loss ratio provide the nonlinear training safeguard.
 51    """
 52    rng = np.random.default_rng(seed)
 53    d = len(g)
 54    # A[(1,s)]^3 = g.s + .5 sum h_i s_i^2; C=0.
 55    A = np.zeros((d + 1, d + 1, d + 1), dtype=np.float64)
 56    for i in range(d):
 57        j = i + 1
 58        A[0, 0, j] = A[0, j, 0] = A[j, 0, 0] = g[i] / 3.0
 59        q = hdiag[i] / 6.0
 60        A[0, j, j] = A[j, 0, j] = A[j, j, 0] = q
 61    U = []
 62    init = np.r_[1.0, -g / (np.linalg.norm(g) + 1e-12)]
 63    init /= np.linalg.norm(init)
 64    U.append(init)
 65    for _ in range(2):
 66        z = rng.normal(size=d + 1); U.append(z / np.linalg.norm(z))
 67    def contract(a, v, w): return np.einsum('ijk,j,k->i', a, v, w)
 68    for _ in range(sweeps):
 69        for b in range(3):
 70            old = U[b].copy()
 71            others = [U[i] for i in range(3) if i != b]
 72            z = beta * old - contract(A, others[0], others[1])
 73            nz = np.linalg.norm(z)
 74            if nz > 1e-12: U[b] = z / nz
 75    # Decode the rank-one homogeneous point, retaining only its direction.
 76    tail = sum(u[1:] for u in U) / 3.0
 77    nt = np.linalg.norm(tail)
 78    return -tail / (nt + 1e-12), float(np.einsum('ijk,i,j,k', A, *U))
 79
 80
 81def idea_train(cfg, seed, telemetry=None):
 82    torch.manual_seed(seed); np.random.seed(seed)
 83    ds = bench.get_dataset('tabular', seed, n_train=400, n_test=200)
 84    net = make_net(ds)
 85    opt = torch.optim.AdamW(net.parameters(), lr=cfg['lr'], weight_decay=cfg['weight_decay'])
 86    xtr, ytr = ds['xtr'], ds['ytr']
 87    last = _final_layer(net)
 88    rng = torch.Generator().manual_seed(seed + 101)
 89    radius = 0.08
 90    accepted = rejected = 0
 91    pred_vals, obs_vals = [], []
 92    net.train()
 93    for ep in range(15):
 94        order = torch.randperm(len(xtr), generator=rng)
 95        for st in range(0, len(xtr), 128):
 96            ix = order[st:st+128]; xb, yb = xtr[ix], ytr[ix]
 97            opt.zero_grad(set_to_none=True)
 98            loss = _loss(net, xb, yb); loss.backward()
 99            # The final layer's input activation gives the exact MSE Hessian diagonal.
100            with torch.no_grad():
101                h = torch.cat([last.weight.detach().flatten(), last.bias.detach()])
102                # Positive diagonal curvature proxy, measured from gradient scale.
103                hdiag = torch.full_like(h, 0.1) + 0.5 * h.abs()
104                g = torch.cat([last.weight.grad.detach().flatten(), last.bias.grad.detach()])
105                direction, _ = cubic_direction(g.cpu().numpy(), hdiag.cpu().numpy(), seed=seed + ep + st)
106                step = torch.as_tensor(direction, dtype=last.weight.dtype)
107                step *= radius
108                oldw, oldb = last.weight.detach().clone(), last.bias.detach().clone()
109                base = float(loss.detach())
110                # Local model predicted decrease including quadratic and quartic safeguard.
111                pred_dec = float(-(g * step).sum() - 0.5 * (hdiag * step.square()).sum() - 2.0 * step.norm()**4)
112                last.weight.add_(step[:-1].view_as(last.weight)); last.bias.add_(step[-1:])
113                actual = float(_loss(net, xb, yb).detach())
114                obs_dec = base - actual
115                rho = obs_dec / (pred_dec + 1e-12)
116                if (obs_dec > 0.0) and (rho >= 0.10):
117                    accepted += 1; radius = min(0.20, radius * (1.10 if rho > 0.75 else 1.02))
118                    pred_vals.append(pred_dec); obs_vals.append(obs_dec)
119                else:
120                    last.weight.copy_(oldw); last.bias.copy_(oldb)
121                    rejected += 1; radius = max(0.01, radius * 0.5)
122            # Adam updates the remaining/shared parameters only; final layer was handled above.
123            for p in last.parameters(): p.grad = None
124            opt.step()
125    net.eval()
126    with torch.no_grad(): metric = float(((net(ds['xte']) - ds['yte']) ** 2).mean())
127    if telemetry is not None:
128        telemetry.update({'accepted': accepted, 'rejected': rejected,
129                          'predicted_decrease_mean': float(np.mean(pred_vals)) if pred_vals else 0.0,
130                          'observed_decrease_mean': float(np.mean(obs_vals)) if obs_vals else 0.0,
131                          'radius_final': radius})
132    return metric
133
134
135def run():
136    t0 = time.time()
137    # Required baseline sweep over the same grid used by the idea.
138    base = bench.sweep_baseline(lambda cfg: lambda seed: baseline_train(cfg, seed), GRID, seeds=SEEDS)
139    idea_runs = []
140    for cfg in GRID:
141        vals, sigs = [], []
142        for s in SEEDS:
143            sig = {}; vals.append(idea_train(cfg, s, sig)); sigs.append(sig)
144        idea_runs.append({'cfg': cfg, 'mean': float(np.mean(vals)), 'std': float(np.std(vals)),
145                          'per_seed': vals, 'signature_per_seed': sigs})
146    best = min(idea_runs, key=lambda z: z['mean'])
147    # Signature is measured on trained-model updates, not an analytical toy identity.
148    sig = best['signature_per_seed']
149    pred = [x['predicted_decrease_mean'] for x in sig if x['predicted_decrease_mean'] > 0]
150    obs = [x['observed_decrease_mean'] for x in sig if x['observed_decrease_mean'] > 0]
151    ratio = float(np.mean(np.asarray(obs) / (np.asarray(pred) + 1e-12))) if pred else 0.0
152    signature = {'predicted_decrease_mean': float(np.mean(pred)) if pred else 0.0,
153                 'observed_decrease_mean': float(np.mean(obs)) if obs else 0.0,
154                 'observed_to_predicted_ratio': ratio,
155                 'accepted_steps_mean': float(np.mean([x['accepted'] for x in sig])),
156                 'rejected_steps_mean': float(np.mean([x['rejected'] for x in sig])),
157                 'confirmed': bool(pred and 0.5 <= ratio <= 1.5)}
158    report = bench.make_report('tabular', 'mlp_tiny', base,
159                               {'best_cfg': best['cfg'], 'sweep': idea_runs,
160                                'mean': best['mean'], 'std': best['std'],
161                                'per_seed': best['per_seed'], 'n': 8,
162                                'full': {'mean': best['mean'], 'std': best['std'],
163                                         'per_seed': best['per_seed'], 'n': 8}},
164                               {'mechanism_signature': signature,
165                                'protocol': {'paired_seeds': list(SEEDS), 'epochs': 15,
166                                             'train_samples': 400, 'test_samples': 200,
167                                             'grid': GRID},
168                                'runtime_sec': time.time() - t0})
169    Path('bench_report.json').write_text(json.dumps(report, indent=2))
170    print(json.dumps(report, indent=2))
171
172if __name__ == '__main__': run()