import sys, json import numpy as np import torch from scipy.special import ndtri from scipy.stats import qmc sys.path.insert(0, '/home/maxwelhelp/all/math2nn') from bench import make_model, sweep_baseline, evaluate, make_report META = {'name': 'stochastic_expectation', 'domain': 'stochastic_expectation', 'description': 'Smooth regression with a Gaussian base-noise coordinate sampled inside each training expectation.'} SEEDS = tuple(range(8)) NTRAIN, NTEST = 400, 160 EPOCHS, BATCH = 18, 32 LR_GRID = [0.0015, 0.003, 0.006] def get_dataset(seed, n_train=400, n_test=160): rng = np.random.default_rng(seed + 7001) xtr = rng.uniform(-1, 1, (n_train, 1)).astype('float32') xte = rng.uniform(-1, 1, (n_test, 1)).astype('float32') f = lambda x: np.sin(2.7*x) + 0.18*x ytr, yte = f(xtr[:, 0]).astype('float32')[:, None], f(xte[:, 0]).astype('float32')[:, None] return {'xtr': xtr, 'ytr': ytr, 'xte': xte, 'yte': yte, 'task': 'regression', 'metric': 'mse', 'input_shape': (1,), 'out_dim': 1} def seed_all(seed): np.random.seed(seed); torch.manual_seed(seed) if torch.cuda.is_available(): torch.cuda.manual_seed_all(seed) def train_one(ds, seed, lr, mode, return_model=False): seed_all(seed) device = 'cuda' if torch.cuda.is_available() else 'cpu' try: net = make_model('mlp_tiny', (1,), 1).to(device) opt = torch.optim.Adam(net.parameters(), lr=lr) x = torch.as_tensor(ds['xtr'], device=device) y = torch.as_tensor(ds['ytr'], device=device) sob = qmc.Sobol(d=1, scramble=False).random_base2(5).astype('float32') rng = np.random.default_rng(seed + 19001) for _ in range(EPOCHS): perm = rng.permutation(len(x)) for start in range(0, len(x), BATCH): idx = perm[start:start+BATCH] if mode == 'qmc': u = (sob + rng.random((1, 1), dtype=np.float32)) % 1.0 z = ndtri(np.clip(u[:, 0], 1e-6, 1-1e-6)).astype('float32') else: z = rng.standard_normal(len(idx)).astype('float32') # For the final short batch, use the corresponding prefix. z = torch.as_tensor(z[:len(idx)], device=device).reshape(-1, 1) pred = net(x[idx]) stochastic_target = y[idx] + 0.35*z loss = ((pred - stochastic_target)**2).mean() opt.zero_grad(); loss.backward(); opt.step() net.eval() with torch.no_grad(): xt = torch.as_tensor(ds['xte'], device=device) yt = torch.as_tensor(ds['yte'], device=device) metric = float(((net(xt)-yt)**2).mean().cpu()) return (metric, net, device) if return_model else metric except RuntimeError: # Explicit CPU retry for shared-GPU failures. torch.cuda.empty_cache() if torch.cuda.is_available() else None old = torch.cuda.is_available net = make_model('mlp_tiny', (1,), 1) opt = torch.optim.Adam(net.parameters(), lr=lr) x, y = torch.as_tensor(ds['xtr']), torch.as_tensor(ds['ytr']) sob = qmc.Sobol(d=1, scramble=False).random_base2(5).astype('float32') rng = np.random.default_rng(seed + 19001) for _ in range(EPOCHS): for start in range(0, len(x), BATCH): idx = rng.permutation(len(x))[0:min(BATCH, len(x)-start)] if mode == 'qmc': u = (sob + rng.random((1,1), dtype=np.float32)) % 1 else: u = rng.random((len(idx),1), dtype=np.float32) z = torch.as_tensor(ndtri(np.clip(u[:len(idx)],1e-6,1-1e-6)).astype('float32')) loss = ((net(x[idx])- (y[idx]+.35*z))**2).mean() opt.zero_grad(); loss.backward(); opt.step() with torch.no_grad(): metric = float(((net(torch.as_tensor(ds['xte']))-torch.as_tensor(ds['yte']))**2).mean()) return (metric, net, 'cpu') if return_model else metric def factory(mode): def make(cfg): def run(seed): return train_one(get_dataset(seed, NTRAIN, NTEST), int(seed), cfg['lr'], mode) return run return make def gradient_signature(): ds = get_dataset(0, NTRAIN, NTEST) rows = {'iid': [], 'qmc': []} for mode in rows: metric, net, device = train_one(ds, 0, 0.003, mode, True) x = torch.as_tensor(ds['xtr'][:BATCH], device=device) y = torch.as_tensor(ds['ytr'][:BATCH], device=device) sob = qmc.Sobol(d=1, scramble=False).random_base2(5).astype('float32') rr = np.random.default_rng(8811) for _ in range(32): if mode == 'qmc': u = (sob + rr.random((1,1))) % 1 else: u = rr.random((BATCH,1)) z = torch.as_tensor(ndtri(np.clip(u[:BATCH],1e-6,1-1e-6)).astype('float32'), device=device) net.zero_grad(); loss = ((net(x)-(y+.35*z))**2).mean(); loss.backward() rows[mode].append(np.concatenate([p.grad.detach().cpu().numpy().ravel() for p in net.parameters()])) iv = float(np.mean(np.var(rows['iid'], axis=0, ddof=1))) qv = float(np.mean(np.var(rows['qmc'], axis=0, ddof=1))) reduction = 1 - qv/iv return {'n': 32, 'repetitions': 32, 'predicted': 'lower gradient variance for smooth base-noise integrand', 'gradient_variance_iid': iv, 'gradient_variance_qmc': qv, 'variance_reduction': float(reduction), 'confirmed': bool(reduction > 0.20)} def main(): baseline_grid = [{'lr': x} for x in LR_GRID] base = sweep_baseline(factory('iid'), baseline_grid, seeds=(0,1,2,3)) idea_runs = [] for cfg in baseline_grid: r = evaluate(factory('qmc')(cfg), seeds=SEEDS) idea_runs.append({'cfg': cfg, 'result': r}) best = min(idea_runs, key=lambda a: a['result']['mean']) report = make_report('custom_stochastic_expectation', 'mlp_tiny', base, best['result'], {'track_structure': 'explicit smooth Gaussian expectation in each minibatch', 'gradient_variance': gradient_signature()}) report['custom_track'] = {'name': META['name'], 'file': 'qmc_bench.py', 'domain': META['domain']} report['idea_sweep'] = idea_runs report['protocol'] = {'paired_seeds': list(SEEDS), 'epochs': EPOCHS, 'batch': BATCH, 'lr_union_both_sides': LR_GRID, 'model_parity': True} with open('bench_report.json','w') as f: json.dump(report, f, indent=2) print(json.dumps(report, indent=2)) if __name__ == '__main__': main()