Randomized-QMC gradient batches / qmc_bench.py
Beats tuned baseline
1import sys, json
2import numpy as np
3import torch
4from scipy.special import ndtri
5from scipy.stats import qmc
6
7sys.path.insert(0, '/home/maxwelhelp/all/math2nn')
8from bench import make_model, sweep_baseline, evaluate, make_report
9
10META = {'name': 'stochastic_expectation', 'domain': 'stochastic_expectation',
11 'description': 'Smooth regression with a Gaussian base-noise coordinate sampled inside each training expectation.'}
12
13SEEDS = tuple(range(8))
14NTRAIN, NTEST = 400, 160
15EPOCHS, BATCH = 18, 32
16LR_GRID = [0.0015, 0.003, 0.006]
17
18
19def get_dataset(seed, n_train=400, n_test=160):
20 rng = np.random.default_rng(seed + 7001)
21 xtr = rng.uniform(-1, 1, (n_train, 1)).astype('float32')
22 xte = rng.uniform(-1, 1, (n_test, 1)).astype('float32')
23 f = lambda x: np.sin(2.7*x) + 0.18*x
24 ytr, yte = f(xtr[:, 0]).astype('float32')[:, None], f(xte[:, 0]).astype('float32')[:, None]
25 return {'xtr': xtr, 'ytr': ytr, 'xte': xte, 'yte': yte,
26 'task': 'regression', 'metric': 'mse', 'input_shape': (1,), 'out_dim': 1}
27
28
29def seed_all(seed):
30 np.random.seed(seed); torch.manual_seed(seed)
31 if torch.cuda.is_available(): torch.cuda.manual_seed_all(seed)
32
33
34def train_one(ds, seed, lr, mode, return_model=False):
35 seed_all(seed)
36 device = 'cuda' if torch.cuda.is_available() else 'cpu'
37 try:
38 net = make_model('mlp_tiny', (1,), 1).to(device)
39 opt = torch.optim.Adam(net.parameters(), lr=lr)
40 x = torch.as_tensor(ds['xtr'], device=device)
41 y = torch.as_tensor(ds['ytr'], device=device)
42 sob = qmc.Sobol(d=1, scramble=False).random_base2(5).astype('float32')
43 rng = np.random.default_rng(seed + 19001)
44 for _ in range(EPOCHS):
45 perm = rng.permutation(len(x))
46 for start in range(0, len(x), BATCH):
47 idx = perm[start:start+BATCH]
48 if mode == 'qmc':
49 u = (sob + rng.random((1, 1), dtype=np.float32)) % 1.0
50 z = ndtri(np.clip(u[:, 0], 1e-6, 1-1e-6)).astype('float32')
51 else:
52 z = rng.standard_normal(len(idx)).astype('float32')
53 # For the final short batch, use the corresponding prefix.
54 z = torch.as_tensor(z[:len(idx)], device=device).reshape(-1, 1)
55 pred = net(x[idx])
56 stochastic_target = y[idx] + 0.35*z
57 loss = ((pred - stochastic_target)**2).mean()
58 opt.zero_grad(); loss.backward(); opt.step()
59 net.eval()
60 with torch.no_grad():
61 xt = torch.as_tensor(ds['xte'], device=device)
62 yt = torch.as_tensor(ds['yte'], device=device)
63 metric = float(((net(xt)-yt)**2).mean().cpu())
64 return (metric, net, device) if return_model else metric
65 except RuntimeError:
66 # Explicit CPU retry for shared-GPU failures.
67 torch.cuda.empty_cache() if torch.cuda.is_available() else None
68 old = torch.cuda.is_available
69 net = make_model('mlp_tiny', (1,), 1)
70 opt = torch.optim.Adam(net.parameters(), lr=lr)
71 x, y = torch.as_tensor(ds['xtr']), torch.as_tensor(ds['ytr'])
72 sob = qmc.Sobol(d=1, scramble=False).random_base2(5).astype('float32')
73 rng = np.random.default_rng(seed + 19001)
74 for _ in range(EPOCHS):
75 for start in range(0, len(x), BATCH):
76 idx = rng.permutation(len(x))[0:min(BATCH, len(x)-start)]
77 if mode == 'qmc': u = (sob + rng.random((1,1), dtype=np.float32)) % 1
78 else: u = rng.random((len(idx),1), dtype=np.float32)
79 z = torch.as_tensor(ndtri(np.clip(u[:len(idx)],1e-6,1-1e-6)).astype('float32'))
80 loss = ((net(x[idx])- (y[idx]+.35*z))**2).mean()
81 opt.zero_grad(); loss.backward(); opt.step()
82 with torch.no_grad(): metric = float(((net(torch.as_tensor(ds['xte']))-torch.as_tensor(ds['yte']))**2).mean())
83 return (metric, net, 'cpu') if return_model else metric
84
85
86def factory(mode):
87 def make(cfg):
88 def run(seed): return train_one(get_dataset(seed, NTRAIN, NTEST), int(seed), cfg['lr'], mode)
89 return run
90 return make
91
92
93def gradient_signature():
94 ds = get_dataset(0, NTRAIN, NTEST)
95 rows = {'iid': [], 'qmc': []}
96 for mode in rows:
97 metric, net, device = train_one(ds, 0, 0.003, mode, True)
98 x = torch.as_tensor(ds['xtr'][:BATCH], device=device)
99 y = torch.as_tensor(ds['ytr'][:BATCH], device=device)
100 sob = qmc.Sobol(d=1, scramble=False).random_base2(5).astype('float32')
101 rr = np.random.default_rng(8811)
102 for _ in range(32):
103 if mode == 'qmc': u = (sob + rr.random((1,1))) % 1
104 else: u = rr.random((BATCH,1))
105 z = torch.as_tensor(ndtri(np.clip(u[:BATCH],1e-6,1-1e-6)).astype('float32'), device=device)
106 net.zero_grad(); loss = ((net(x)-(y+.35*z))**2).mean(); loss.backward()
107 rows[mode].append(np.concatenate([p.grad.detach().cpu().numpy().ravel() for p in net.parameters()]))
108 iv = float(np.mean(np.var(rows['iid'], axis=0, ddof=1)))
109 qv = float(np.mean(np.var(rows['qmc'], axis=0, ddof=1)))
110 reduction = 1 - qv/iv
111 return {'n': 32, 'repetitions': 32, 'predicted': 'lower gradient variance for smooth base-noise integrand',
112 'gradient_variance_iid': iv, 'gradient_variance_qmc': qv,
113 'variance_reduction': float(reduction), 'confirmed': bool(reduction > 0.20)}
114
115
116def main():
117 baseline_grid = [{'lr': x} for x in LR_GRID]
118 base = sweep_baseline(factory('iid'), baseline_grid, seeds=(0,1,2,3))
119 idea_runs = []
120 for cfg in baseline_grid:
121 r = evaluate(factory('qmc')(cfg), seeds=SEEDS)
122 idea_runs.append({'cfg': cfg, 'result': r})
123 best = min(idea_runs, key=lambda a: a['result']['mean'])
124 report = make_report('custom_stochastic_expectation', 'mlp_tiny', base, best['result'],
125 {'track_structure': 'explicit smooth Gaussian expectation in each minibatch',
126 'gradient_variance': gradient_signature()})
127 report['custom_track'] = {'name': META['name'], 'file': 'qmc_bench.py', 'domain': META['domain']}
128 report['idea_sweep'] = idea_runs
129 report['protocol'] = {'paired_seeds': list(SEEDS), 'epochs': EPOCHS, 'batch': BATCH,
130 'lr_union_both_sides': LR_GRID, 'model_parity': True}
131 with open('bench_report.json','w') as f: json.dump(report, f, indent=2)
132 print(json.dumps(report, indent=2))
133
134if __name__ == '__main__': main()