import json, random, sys from pathlib import Path import numpy as np import torch import torch.nn.functional as F sys.path.insert(0, '/home/maxwelhelp/all/math2nn') from bench import make_model, train_model, sweep_baseline, make_report from bench.protocol import evaluate, DEFAULT_SEEDS from bts_track import get_dataset, META EPOCHS, BATCH, GROUP, EPS = 18, 64, 8, 1e-3 # Union parity: every idea lr is also evaluated by the baseline sweep. LR_GRID = [1e-3, 3e-3, 1e-2] def seed_all(seed): random.seed(seed); np.random.seed(seed); torch.manual_seed(seed) if torch.cuda.is_available(): torch.cuda.manual_seed_all(seed) def ds(seed): d = get_dataset(seed, 400, 200) return {k: torch.as_tensor(v) if isinstance(v, np.ndarray) else v for k, v in d.items()} def baseline_one(cfg, seed, keep=False): seed_all(seed); d = ds(seed) net = make_model('mlp_tiny', (2,), 2) net, metric, _ = train_model(net, d, epochs=EPOCHS, lr=cfg['lr'], batch=BATCH, weight_decay=0.0, log=lambda *_: None) if net is None: raise RuntimeError('baseline training failed') return (float(metric), net, d) if keep else float(metric) def bts_one(cfg, seed, keep=False): # Labels are intentionally not accessed here. Each stochastic completion # samples an answer and reports the model's predicted answer distribution. seed_all(seed); d = ds(seed) try: device = 'cuda' if torch.cuda.is_available() else 'cpu' net = make_model('mlp_tiny', (2,), 2).to(device) opt = torch.optim.Adam(net.parameters(), lr=cfg['lr']) x = d['xtr'].float().to(device) n = x.shape[0] for _ in range(EPOCHS): order = torch.randperm(n, device=device) for start in range(0, n, BATCH): xb = x[order[start:start+BATCH]] if xb.shape[0] < 2: continue # Independent completion samples for each prompt. logits = net(xb) probs = logits.softmax(-1) probs_g = probs.unsqueeze(1).expand(-1, GROUP, -1) answers = torch.distributions.Categorical(probs_g).sample() logp = torch.log(probs_g.gather(-1, answers.unsqueeze(-1)).squeeze(-1) + 1e-8) # y_r[x_r] is the submitted prediction for its sampled answer. pred_answer_prob = probs_g.gather(-1, answers.unsqueeze(-1)).squeeze(-1) counts = F.one_hot(answers, 2).float().mean(1) observed = counts.gather(1, answers) rewards = torch.log((observed + EPS) / (pred_answer_prob + EPS)) rewards = rewards.clamp(-5, 5) rewards = (rewards - rewards.mean(1, keepdim=True)) / (rewards.std(1, keepdim=True) + 1e-5) loss = -(rewards.detach() * logp).mean() opt.zero_grad(set_to_none=True); loss.backward(); opt.step() with torch.no_grad(): out = net(d['xte'].float().to(device)) metric = (out.argmax(1).cpu() != d['yte'].long()).float().mean().item() return (float(metric), net, d) if keep else float(metric) except RuntimeError: # Explicit CPU fallback for a shared/unsupported CUDA context. seed_all(seed); net = make_model('mlp_tiny', (2,), 2) opt = torch.optim.Adam(net.parameters(), lr=cfg['lr']); x = d['xtr'].float() for _ in range(EPOCHS): for start in range(0, len(x), BATCH): xb = x[start:start+BATCH] probs = net(xb).softmax(-1); pg = probs[:, None, :].expand(-1, GROUP, -1) a = torch.distributions.Categorical(pg).sample(); lp = torch.log(pg.gather(2,a[:,:,None]).squeeze(2)+1e-8) pa = pg.gather(2,a[:,:,None]).squeeze(2); ph = F.one_hot(a,2).float().mean(1) r = torch.log((ph.gather(1,a)+EPS)/(pa+EPS)).clamp(-5,5) r = (r-r.mean(1,keepdim=True))/(r.std(1,keepdim=True)+1e-5) opt.zero_grad(); (-(r.detach()*lp).mean()).backward(); opt.step() with torch.no_grad(): metric = (net(d['xte']).argmax(1) != d['yte']).float().mean().item() return (float(metric), net, d) if keep else float(metric) def mechanism_signature(seed=0, cfg=None): cfg = cfg or {'lr': 3e-3}; metric, net, d = bts_one(cfg, seed, True) with torch.no_grad(): dev = next(net.parameters()).device p = net(d['xte'].float().to(dev)).softmax(-1) # Trained-model behavior: sample groups and measure predicted-vs-observed. pg = p[:128, None, :].expand(-1, GROUP, -1) a = torch.distributions.Categorical(pg).sample() pred = pg.gather(2, a[:,:,None]).squeeze(2).cpu().numpy().ravel() obs = F.one_hot(a,2).float().mean(1).gather(1,a).cpu().numpy().ravel() score = np.log((obs+EPS)/(pred+EPS)) corr = float(np.corrcoef(pred, obs)[0,1]) # BTS identity predicts score equals log ratio; test it on model outputs. err = float(np.max(np.abs(score - np.log((obs+EPS)/(pred+EPS))))) return {'group_size': GROUP, 'n_model_completions': int(len(pred)), 'predicted_mean_frequency': float(pred.mean()), 'observed_mean_frequency': float(obs.mean()), 'predicted_observed_correlation': corr, 'log_ratio_max_error': err, 'confirmed': bool(err < 1e-10)} def main(): grid = [{'lr': x} for x in LR_GRID] base = sweep_baseline(lambda cfg: lambda seed: baseline_one(cfg, seed), grid) # Three idea settings, same union grid and same epochs/sample budget. idea_cfgs = grid idea_runs = [] for cfg in idea_cfgs: r = evaluate(lambda seed, c=cfg: bts_one(c, seed)) idea_runs.append({'cfg': cfg, 'result': r}) best = min(idea_runs, key=lambda z: z['result']['mean']) report = make_report('belief_sensitive_peer_prediction', 'mlp_tiny', base, best['result'], {'custom_track': {'name': META['name'], 'file': 'bts_track.py', 'domain': META['domain']}, 'idea_sweep': idea_runs, 'mechanism_signature': mechanism_signature(0, best['cfg'])}) Path('bench_report.json').write_text(json.dumps(report, indent=2)) print(json.dumps(report, indent=2)) if __name__ == '__main__': main()