Jacobian-aligned infill for black-box neural tuning / experiment.py

✓✓ Beats tuned baseline

Raw ⬇ ZIP
  1import json
  2import numpy as np
  3
  4
  5def fit_jacobian(z0, r0, Z, R, eta=1e-5, h=1.0):
  6    X = np.asarray(Z) - z0[None, :]
  7    Y = np.asarray(R) - r0[None, :]
  8    w = np.exp(-np.sum(X * X, axis=1) / max(h * h, 1e-12))
  9    H = X.T @ (w[:, None] * X) + eta * np.eye(X.shape[1])
 10    B = X.T @ (w[:, None] * Y)
 11    return np.linalg.solve(H, B).T, w
 12
 13
 14def geometry_steps(J, r0, lam, rho, rng):
 15    p = J.shape[1]
 16    A = J.T @ J + lam * np.eye(p)
 17    dgn = -np.linalg.solve(A, J.T @ r0)
 18    if np.linalg.norm(dgn) > rho:
 19        dgn *= rho / np.linalg.norm(dgn)
 20    ev, Q = np.linalg.eigh(A)
 21    u = rng.normal(size=p)
 22    de = Q @ ((Q.T @ u) / np.sqrt(np.maximum(ev, 1e-12)))
 23    de *= rho / max(np.linalg.norm(de), 1e-12)
 24    return dgn, de, np.linalg.cond(A)
 25
 26
 27def verification(seed=7):
 28    rng = np.random.default_rng(seed)
 29    p, m, n = 4, 6, 80
 30    Jtrue = rng.normal(size=(m, p))
 31    z0, r0 = rng.normal(size=p), rng.normal(size=m)
 32    X = rng.normal(scale=.18, size=(n, p))
 33    Jhat, _ = fit_jacobian(z0, r0, z0 + X, r0 + X @ Jtrue.T, eta=1e-10, h=10.)
 34    rel_err = np.linalg.norm(Jhat - Jtrue) / np.linalg.norm(Jtrue)
 35    d, _, cond = geometry_steps(Jtrue, r0, .03, 10., rng)
 36    ev = np.linalg.eigvalsh(Jtrue.T @ Jtrue + .03 * np.eye(p))
 37    inv_ev = 1.0 / ev
 38    alignment = bool(np.all(np.diff(ev) >= -1e-10) and np.all(np.diff(inv_ev) <= 1e-10))
 39    before, after = np.linalg.norm(r0), np.linalg.norm(r0 + Jtrue @ d)
 40    return {'jacobian_relative_error': float(rel_err),
 41            'residual_norm_before': float(before), 'residual_norm_after': float(after),
 42            'residual_ratio': float(after / before), 'metric_condition': float(cond),
 43            'metric_alignment_eigenvalue_check': alignment,
 44            'sensitivity_eigenvalues': ev.tolist(), 'inverse_metric_eigenvalues': inv_ev.tolist()}
 45
 46
 47def residual(z):
 48    target = np.array([1.2, -1., .7, -.5, .35, -.25, .15, -.1])
 49    scales = np.array([.12, .15, .25, .35, .5, .7, 1., 1.3])
 50    q = (z - target) / scales
 51    r = q.copy()
 52    r[0] += .22 * z[1] ** 2
 53    r[2] += .18 * np.sin(z[3])
 54    r[5] += .12 * z[0] * z[4]
 55    return r
 56
 57
 58def run(seed, guided, eval_budget=500, p=8, pop=20):
 59    rng = np.random.default_rng(seed)
 60    Z = rng.uniform(-3., 3., (pop, p)); R = np.array([residual(z) for z in Z])
 61    trace_z, trace_r, curve = list(Z), list(R), []
 62    while len(trace_z) < eval_budget:
 63        ib = np.argmin(np.sum(R * R, axis=1)); z0, r0 = Z[ib].copy(), R[ib].copy()
 64        nnew = min(pop, eval_budget - len(trace_z)); candidates = []
 65        for k in range(nnew):
 66            use = guided and k < max(1, nnew // 4) and len(trace_z) >= 2 * p
 67            if use:
 68                idx = np.argsort([np.linalg.norm(x-z0) for x in trace_z])[:min(len(trace_z), 80)]
 69                J, _ = fit_jacobian(z0, r0, np.array([trace_z[i] for i in idx]), np.array([trace_r[i] for i in idx]), eta=2e-3, h=2.)
 70                dgn, de, cond = geometry_steps(J, r0, .08, .8, rng)
 71                d = dgn if k == 0 else de
 72                if not np.isfinite(cond) or cond > 1e10:
 73                    d = rng.normal(size=p); d *= .8 / max(np.linalg.norm(d), 1e-12)
 74            else:
 75                d = rng.normal(size=p); d *= .8 / max(np.linalg.norm(d), 1e-12)
 76            candidates.append(np.clip(z0 + d, -3., 3.))
 77        CR = np.array([residual(z) for z in candidates])
 78        allz, allr = np.vstack([Z, candidates]), np.vstack([R, CR])
 79        keep = np.argsort(np.sum(allr * allr, axis=1))[:pop]
 80        Z, R = allz[keep], allr[keep]; trace_z.extend(candidates); trace_r.extend(CR)
 81        curve.append(float(np.min(np.sum(R * R, axis=1))))
 82    return curve
 83
 84
 85def main():
 86    base, guided = [], []
 87    for seed in [11, 22, 33, 44, 55, 66, 77, 88]:
 88        base.append(run(seed, False)); guided.append(run(seed, True))
 89    checkpoints = [100, 200, 300, 400, 500]
 90    def at(curves, ev):
 91        i = max(0, min(len(curves[0])-1, (ev-20)//20-1))
 92        return float(np.mean([c[i] for c in curves]))
 93    out = {'verification': verification(), 'benchmark': {
 94        'seeds': 8, 'checkpoints': checkpoints,
 95        'baseline_best_loss': [at(base, e) for e in checkpoints],
 96        'guided_best_loss': [at(guided, e) for e in checkpoints],
 97        'final_baseline_mean': float(np.mean([c[-1] for c in base])),
 98        'final_guided_mean': float(np.mean([c[-1] for c in guided])),
 99        'final_baseline_std': float(np.std([c[-1] for c in base])),
100        'final_guided_std': float(np.std([c[-1] for c in guided]))}}
101    print(json.dumps(out, indent=2))
102
103if __name__ == '__main__': main()