1import json, math, time, random
  2import numpy as np
  3from scipy.linalg import hadamard
  4
  5SEED = 155
  6np.random.seed(SEED); random.seed(SEED)
  7N = 16
  8M = 16
  9
 10# A has columns a_k; a_k are binary sign vectors. Hadamard rows provide
 11# A A^T / N = I, hence exactly N evaluations suffice here.
 12H = hadamard(N).astype(np.float64)
 13A_exact = H.T  # [N coefficients, N selected evaluations]
 14lambda_exact = np.full(M, 1.0 / N, dtype=np.float64)
 15
 16# Cheap algebraic and numerical verification first.
 17residual = A_exact @ np.diag(lambda_exact) @ A_exact.T - np.eye(N)
 18rng = np.random.default_rng(SEED)
 19C = rng.normal(size=(2000, N))
 20exact_E = ((C @ A_exact) ** 2 * lambda_exact).sum(axis=1)
 21true_E = (C*C).sum(axis=1)
 22relative_error = np.max(np.abs(exact_E - true_E) / np.maximum(true_E, 1e-12))
 23# Compare random uniform sign subsampling, unbiased with the same number of evals.
 24A_random = rng.choice([-1.0, 1.0], size=(N, M))
 25random_trials = np.empty((500, len(C)))
 26for t in range(500):
 27    Ar = rng.choice([-1.0, 1.0], size=(N, M))
 28    random_trials[t] = ((C @ Ar) ** 2).mean(axis=1)
 29rand_rel_rmse = float(np.sqrt(np.mean((random_trials - true_E[None,:])**2)) / np.sqrt(np.mean(true_E**2)))
 30rand_rel_std_mean = float(np.mean(np.std(random_trials, axis=0) / np.maximum(true_E, 1e-12)))
 31# Float16 accumulation drift, versus float32 accumulation.
 32Cf = C.astype(np.float16); Af = A_exact.astype(np.float16)
 33E16 = ((Cf @ Af).astype(np.float32)**2 * (1.0/N)).sum(axis=1)
 34mixed_rel = float(np.max(np.abs(E16-true_E)/np.maximum(true_E,1e-12)))
 35
 36# Torch mini experiment. The only stochasticity in the baseline estimator is
 37# the uniformly random sign design drawn independently at each forward pass.
 38import torch
 39from torch import nn
 40
 41def get_device():
 42    if torch.cuda.is_available():
 43        try:
 44            torch.zeros(1, device='cuda')
 45            return torch.device('cuda')
 46        except Exception:
 47            pass
 48    return torch.device('cpu')
 49
 50device = get_device()
 51torch.manual_seed(SEED)
 52if device.type == 'cuda': torch.cuda.manual_seed_all(SEED)
 53
 54g = torch.Generator(device='cpu').manual_seed(SEED)
 55X = torch.randn(96, 8, generator=g)
 56true_w = torch.randn(8, 1, generator=g)
 57y = torch.tanh(X @ true_w) + 0.05*torch.randn(96,1,generator=g)
 58X, y = X.to(device), y.to(device)
 59Aex_t = torch.tensor(A_exact, dtype=torch.float32, device=device)
 60
 61def make_model():
 62    torch.manual_seed(SEED + 7)
 63    return nn.Sequential(nn.Linear(8, 32), nn.Tanh(), nn.Linear(32, N), nn.Tanh(), nn.Linear(N, 1)).to(device)
 64
 65def run(exact, steps=220):
 66    model = make_model(); opt = torch.optim.Adam(model.parameters(), lr=2e-3)
 67    losses=[]; start=time.perf_counter()
 68    for step in range(steps):
 69        opt.zero_grad(set_to_none=True)
 70        c = model[0:4](X)  # coefficient-producing trunk, excluding output layer
 71        if exact:
 72            z = c @ Aex_t
 73            E = (z*z).mean(dim=1, keepdim=True)
 74        else:
 75            # Same 16 activation evaluations and unbiased normalization statistic.
 76            Ar = torch.randint(0, 2, (N,M), device=device, dtype=torch.float32)*2-1
 77            E = ((c @ Ar)**2).mean(dim=1, keepdim=True)
 78        cn = c / torch.sqrt(E + 1e-5)
 79        pred = model[4](cn)
 80        loss = ((pred-y)**2).mean()
 81        loss.backward(); opt.step(); losses.append(float(loss.detach().cpu()))
 82    if device.type == 'cuda': torch.cuda.synchronize()
 83    return {'final_loss': float(np.mean(losses[-20:])), 'best_loss': float(np.min(losses)),
 84            'first_loss': losses[0], 'time_sec': time.perf_counter()-start}
 85
 86# Gradient variance at one fixed c, directly measuring the statistic-induced
 87# noise that the identity claims to remove.
 88c0 = torch.randn(32, N, device=device, generator=torch.Generator(device=device).manual_seed(SEED+11), requires_grad=True)
 89target = torch.randn_like(c0)
 90def grad_samples(exact, count=80):
 91    vals=[]; grads=[]
 92    for i in range(count):
 93        c = c0.detach().clone().requires_grad_(True)
 94        if exact: E=((c@Aex_t)**2).mean(dim=1,keepdim=True)
 95        else:
 96            Ar=torch.randint(0,2,(N,M),device=device,dtype=torch.float32)*2-1
 97            E=((c@Ar)**2).mean(dim=1,keepdim=True)
 98        loss=((c/torch.sqrt(E+1e-5)-target)**2).mean()
 99        grads.append(torch.autograd.grad(loss,c)[0].detach().cpu().numpy())
100        vals.append(float(loss.detach().cpu()))
101    G=np.stack(grads)
102    return float(np.mean(vals)), float(np.mean(np.var(G,axis=0))), float(np.mean(np.linalg.norm(G-G.mean(0),axis=tuple(range(1,G.ndim)))**2))
103
104grad_exact=grad_samples(True); grad_random=grad_samples(False)
105# Make JSON serializable and include enough precision to audit.
106result = {
107 'seed': SEED, 'N':N, 'evaluations':M, 'device':str(device),
108 'math': {'matrix_frobenius_residual':float(np.linalg.norm(residual)),
109          'max_relative_energy_error':float(relative_error),
110          'random_relative_rmse':rand_rel_rmse,
111          'random_mean_relative_std':rand_rel_std_mean,
112          'float16_input_float32_accum_max_relative_error':mixed_rel},
113 'gradient_variance': {'exact_loss_mean':grad_exact[0], 'random_loss_mean':grad_random[0],
114                       'exact_mean_element_variance':grad_exact[1], 'random_mean_element_variance':grad_random[1]},
115 'training': {'exact':run(True), 'random':run(False)}
116}
117print(json.dumps(result, indent=2))