1import json, math, random
  2import numpy as np
  3import torch
  4from torch import nn
  5
  6SEED = 447
  7np.random.seed(SEED); random.seed(SEED); torch.manual_seed(SEED)
  8torch.set_num_threads(min(12, torch.get_num_threads()))
  9
 10
 11def free_moments(tau, K=3):
 12    """m_p = 1/p [w^(p-1)] (1+w)^p exp(-p*tau/(1+w))."""
 13    out = []
 14    for p in range(1, K + 1):
 15        n, a = p - 1, p * tau
 16        # exp(-a/(1+w)) = exp(-a) exp(sum_j>=1 a*(-1)^(j+1) w^j)
 17        q = np.zeros(n + 1); q[0] = 1.0
 18        b = np.zeros(n + 1)
 19        for j in range(1, n + 1):
 20            b[j] = a * ((-1) ** (j + 1))
 21        for k in range(1, n + 1):
 22            q[k] = sum(j * b[j] * q[k-j] for j in range(1, k + 1)) / k
 23        coeff = 0.0
 24        for j in range(n + 1):
 25            coeff += math.comb(p, j) * q[n-j]
 26        out.append(math.exp(-a) * coeff / p)
 27    return np.array(out)
 28
 29
 30def haar_unitary(n):
 31    z = (np.random.randn(n, n) + 1j*np.random.randn(n, n)) / np.sqrt(2)
 32    q, r = np.linalg.qr(z)
 33    d = np.diag(r)
 34    return q * (d / np.abs(d)).conj()
 35
 36
 37def product_check(n=48, taus=(.25, .5, 1., 2.), reps=25):
 38    rows = []
 39    for tau in taus:
 40        L = max(1, round(tau * n))
 41        vals, atom = [], []
 42        P = np.diag([1.0] * (n-1) + [0.0])
 43        for _ in range(reps):
 44            B = np.eye(n, dtype=complex)
 45            for _ in range(L):
 46                B = P @ haar_unitary(n) @ B
 47            eig = np.linalg.eigvalsh(B.conj().T @ B).real
 48            vals.append([np.mean(eig), np.mean(eig**2), np.mean(eig**3)])
 49            atom.append(np.mean(eig > 1 - 1e-9))
 50        emp = np.mean(vals, axis=0); target = free_moments(tau, 3)
 51        rows.append(dict(tau=tau, L=L, target=target.tolist(), empirical=emp.tolist(),
 52                         abs_err=np.abs(emp-target).tolist(), atom_emp=float(np.mean(atom)),
 53                         atom_target=max(0., 1.-tau)))
 54    return rows
 55
 56
 57class MLP(nn.Module):
 58    def __init__(self, d=16, width=32, depth=4, out=4):
 59        super().__init__()
 60        layers = []
 61        for i in range(depth):
 62            layers += [nn.Linear(d if i == 0 else width, width), nn.Tanh()]
 63        layers.append(nn.Linear(width, out))
 64        self.net = nn.Sequential(*layers)
 65    def forward(self, x): return self.net(x)
 66
 67
 68def jacobian_moments(model, x, K=2):
 69    """Exact per-example J^T J moments, with scalar output probes replaced by full J."""
 70    x = x.detach().requires_grad_(True)
 71    y = model(x)
 72    estimates = []
 73    for b in range(x.shape[0]):
 74        Jrows = []
 75        for o in range(y.shape[1]):
 76            g = torch.autograd.grad(y[b, o], x, retain_graph=True, create_graph=True)[0][b]
 77            Jrows.append(g)
 78        J = torch.stack(Jrows)
 79        G = J.T @ J
 80        pows = G
 81        estimates.append([torch.trace(pows) / x.shape[1]])
 82        for p in range(2, K + 1):
 83            pows = pows @ G
 84            estimates[-1].append(torch.trace(pows) / x.shape[1])
 85    return torch.stack([torch.stack(a) for a in estimates]).mean(0)
 86
 87
 88def train_variant(use_spec, steps=120):
 89    torch.manual_seed(SEED + (1 if use_spec else 0))
 90    d, out, width, depth = 16, 4, 32, 4
 91    model = MLP(d, width, depth, out)
 92    opt = torch.optim.Adam(model.parameters(), lr=3e-3)
 93    x = torch.randn(64, d)
 94    true_w = torch.randn(d, out)
 95    labels = x @ true_w + .15 * torch.randn(64, out)
 96    tau = depth / width
 97    target = torch.tensor(free_moments(tau, 2), dtype=x.dtype)
 98    history = []
 99    for step in range(steps):
100        pred = model(x)
101        task = ((pred-labels)**2).mean()
102        spec = torch.tensor(0.)
103        if use_spec and step < steps // 3 and step % 8 == 0:
104            moms = jacobian_moments(model, x[:8], 2)
105            spec = ((torch.log(moms + 1e-5) - torch.log(target + 1e-5))**2).sum()
106        loss = task + (0.03 * spec.clamp(max=10.0) if use_spec else 0.)
107        opt.zero_grad(); loss.backward(); opt.step()
108        if step in (0, 19, 39, 79, steps-1):
109            hist_m = jacobian_moments(model, x[:8], 2).detach().numpy()
110            history.append(dict(step=step, task=float(task), spec=float(spec), moments=hist_m.tolist()))
111    final_m = jacobian_moments(model, x[:8], 2).detach().numpy()
112    return dict(final_task=float(((model(x)-labels)**2).mean()), target=target.numpy().tolist(),
113                final_moments=final_m.tolist(), history=history)
114
115
116def main():
117    result = dict(seed=SEED, product_check=product_check(),
118                  baseline=train_variant(False), spectral=train_variant(True))
119    with open('results.json', 'w') as f: json.dump(result, f, indent=2)
120    print(json.dumps(result, indent=2))
121
122if __name__ == '__main__':
123    main()