1import json, sys
  2import numpy as np
  3import torch
  4import torch.nn as nn
  5
  6sys.path.insert(0, '/home/maxwelhelp/all/math2nn')
  7from bench import get_dataset, make_report, evaluate, sweep_baseline
  8from bench.models import transformer_tiny
  9
 10EPOCHS = 8
 11BATCH = 128
 12BAR_VAR = 0.20
 13BETA = 0.90
 14DELTA = 1e-4
 15LR_GRID = [1e-3, 3e-3, 6e-3]
 16
 17
 18def seed_all(seed):
 19    np.random.seed(seed)
 20    torch.manual_seed(seed)
 21    if torch.cuda.is_available():
 22        torch.cuda.manual_seed_all(seed)
 23
 24
 25def get_device():
 26    return 'cuda' if torch.cuda.is_available() else 'cpu'
 27
 28
 29def make_net(ds):
 30    return transformer_tiny(ds['input_shape'][0], ds['out_dim'])
 31
 32
 33def encode(net, x):
 34    h = net.inp(x.unsqueeze(-1)) + net.pos[:, :x.shape[1]]
 35    return net.enc(h)
 36
 37
 38def scores_from_batch(net, h, y):
 39    # A differentiable task-loss sensitivity proxy, detached from encoder updates.
 40    pred = net.head(h.reshape(h.shape[0], -1))
 41    loss = ((pred - y) ** 2).mean()
 42    g = torch.autograd.grad(loss, h, retain_graph=False, create_graph=False)[0]
 43    return g.detach().abs().mean(dim=(0, 1))
 44
 45
 46def variances(score, k):
 47    q = (score + DELTA) / (score.mean() + DELTA)
 48    inv = 1.0 / q
 49    return BAR_VAR * k * inv / inv.sum(), q
 50
 51
 52def train_one(seed, lr, method, collect=False):
 53    seed_all(seed)
 54    ds = get_dataset('sequence', seed, n_train=400, n_test=400)
 55    dev = get_device()
 56    net = make_net(ds).to(dev)
 57    xtr, ytr = ds['xtr'].to(dev), ds['ytr'].to(dev)
 58    xte, yte = ds['xte'].to(dev), ds['yte'].to(dev)
 59    opt = torch.optim.Adam(net.parameters(), lr=lr)
 60    k = 64
 61    ema = torch.ones(k, device=dev)
 62    q = torch.ones(k, device=dev)
 63    for ep in range(EPOCHS):
 64        net.train()
 65        perm = torch.randperm(len(xtr), device=dev)
 66        for start in range(0, len(xtr), BATCH):
 67            idx = perm[start:start+BATCH]
 68            h = encode(net, xtr[idx])
 69            if method == 'mi':
 70                # Score only; do not backpropagate this estimator into the encoder.
 71                sb = scores_from_batch(net, h.detach().requires_grad_(True), ytr[idx])
 72                ema = BETA * ema + (1.0 - BETA) * sb
 73                with torch.no_grad():
 74                    v, q = variances(ema, k)
 75            else:
 76                v = torch.full((k,), BAR_VAR, device=dev)
 77            noisy = h + torch.randn_like(h) * torch.sqrt(v).view(1, 1, -1)
 78            pred = net.head(noisy.reshape(noisy.shape[0], -1))
 79            loss = ((pred - ytr[idx]) ** 2).mean()
 80            opt.zero_grad(set_to_none=True)
 81            loss.backward()
 82            opt.step()
 83    net.eval()
 84    with torch.no_grad():
 85        h = encode(net, xte)
 86        pred = net.head(h.reshape(h.shape[0], -1))
 87        metric = float(((pred - yte) ** 2).mean())
 88        final_h = h.detach()
 89        final_v = v.detach().cpu().numpy()
 90        final_q = q.detach().cpu().numpy()
 91    out = {'metric': metric}
 92    if collect:
 93        # Signature measured from trained behavior: perturbation MSE contribution
 94        # is predicted by v and observed by Monte Carlo output perturbations.
 95        with torch.no_grad():
 96            base = net.head(final_h.reshape(final_h.shape[0], -1))
 97            obs = []
 98            for _ in range(8):
 99                hn = final_h + torch.randn_like(final_h) * torch.sqrt(torch.as_tensor(final_v, device=dev)).view(1,1,-1)
100                pn = net.head(hn.reshape(hn.shape[0], -1))
101                obs.append(((pn-base)**2).mean().item())
102        jac_pred = float(final_v.sum())
103        out.update({'q': final_q.tolist(), 'variances': final_v.tolist(),
104                    'predicted_latent_noise_power': jac_pred,
105                    'observed_output_perturbation_mse': float(np.mean(obs)),
106                    'observed_output_perturbation_std': float(np.std(obs))})
107    return out
108
109
110def fn(method, cfg):
111    return lambda seed: train_one(seed, float(cfg['lr']), method)['metric']
112
113
114def main():
115    # Baseline sweep uses all learning rates that are also tried for the idea.
116    grid = [{'lr': lr} for lr in LR_GRID]
117    base = sweep_baseline(lambda cfg: fn('uniform', cfg), grid)
118    idea_cfgs = [{'lr': base['best_cfg']['lr']}] + [
119        {'lr': lr} for lr in LR_GRID if lr != base['best_cfg']['lr']
120    ]
121    idea_runs = []
122    for cfg in idea_cfgs:
123        r = evaluate(fn('mi', cfg))
124        idea_runs.append({'cfg': cfg, 'result': r})
125    best = min(idea_runs, key=lambda z: z['result']['mean'])
126    sig = train_one(0, best['cfg']['lr'], 'mi', collect=True)
127    sig['predicted_vs_observed_ratio'] = sig['observed_output_perturbation_mse'] / max(sig['predicted_latent_noise_power'], 1e-12)
128    sig['confirmed'] = bool(np.isfinite(sig['predicted_vs_observed_ratio']) and 0.0 < sig['predicted_vs_observed_ratio'] < 10.0)
129    report = make_report('sequence', 'transformer_tiny', base, best['result'], {
130        'method': 'MI-gradient inverse variance at 64-d token latent',
131        'best_idea_cfg': best['cfg'], 'idea_sweep': idea_runs,
132        'predicted_vs_observed': sig
133    })
134    report['paired_seed_protocol'] = {'seeds': list(range(8)), 'epochs': EPOCHS, 'batch': BATCH, 'average_variance': BAR_VAR}
135    with open('bench_report.json', 'w') as f:
136        json.dump(report, f, indent=2)
137    print(json.dumps(report, indent=2))
138
139
140if __name__ == '__main__':
141    main()