import json, random, sys import numpy as np import torch from torch import nn sys.path.insert(0, '/home/maxwelhelp/all/math2nn') from bench import get_dataset, make_model, train_model, evaluate, sweep_baseline, make_report SEEDS = tuple(range(8)) SWEEP_SEEDS = tuple(range(4)) EPOCHS = 24 BATCH = 128 TRACK = 'tabular' MODEL = 'mlp_tiny' LRS = [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(): try: torch.cuda.manual_seed_all(seed) except Exception: pass def baseline_fn(cfg): def run(seed): seed_all(seed) ds = get_dataset(TRACK, seed, n_train=400, n_test=400) model = make_model(MODEL, ds['input_shape'], ds['out_dim']) _, metric, _ = train_model(model, ds, epochs=EPOCHS, lr=cfg['lr'], batch=BATCH, weight_decay=cfg.get('weight_decay', 0.0), log=lambda *_: None) return float(metric) return run def idea_train(model, ds, epochs, lr, tail_weight, batch=BATCH): """Tail residual correction: first fit the cheap/global predictor, then upweight positive residuals above the empirical upper-tail boundary. The same MLP and optimizer family as the baseline are retained.""" errs = [] devices = [] if torch.cuda.is_available(): devices.append('cuda') devices.append('cpu') for device in devices: try: x, y = ds['xtr'].to(device), ds['ytr'].to(device) net = model.to(device) opt = torch.optim.Adam(net.parameters(), lr=lr) n = x.shape[0] # Cheap surrogate phase: ordinary global MSE. warm = max(1, epochs // 2) gen = torch.Generator(device='cpu') gen.manual_seed(17000 + int(n) + int(round(lr * 1e6))) for ep in range(epochs): perm = torch.randperm(n, generator=gen) for start in range(0, n, batch): ix = perm[start:start + batch].to(device) pred = net(x[ix]) residual = pred - y[ix] if ep < warm: loss = (residual * residual).mean() else: # Estimated upper-tail boundary is the training beta quantile. q = torch.quantile(y, 0.90) # Smoothly emphasize samples in/above the tail, retaining # ordinary MSE in the central region. weights = 1.0 + (tail_weight - 1.0) * torch.sigmoid((y[ix] - q) / 0.15) loss = (weights * residual * residual).mean() opt.zero_grad(set_to_none=True) loss.backward() opt.step() with torch.no_grad(): pred = net(ds['xte'].to(device)) metric = ((pred - ds['yte'].to(device)) ** 2).mean().item() return net, float(metric), device except Exception as exc: errs.append(f'{device}:{str(exc)[:100]}') try: torch.cuda.empty_cache() except Exception: pass raise RuntimeError('idea training failed: ' + ' | '.join(errs)) def idea_fn(cfg): def run(seed): seed_all(seed) ds = get_dataset(TRACK, seed, n_train=400, n_test=400) model = make_model(MODEL, ds['input_shape'], ds['out_dim']) _, metric, _ = idea_train(model, ds, EPOCHS, cfg['lr'], cfg['tail_weight']) return float(metric) return run def behavior(cfg, idea=False): global_vals, tail_vals, cvar_vals = [], [], [] for seed in SEEDS: seed_all(seed) ds = get_dataset(TRACK, seed, n_train=400, n_test=400) model = make_model(MODEL, ds['input_shape'], ds['out_dim']) if idea: net, _, dev = idea_train(model, ds, EPOCHS, cfg['lr'], cfg['tail_weight']) else: net, _, hist = train_model(model, ds, epochs=EPOCHS, lr=cfg['lr'], batch=BATCH, weight_decay=cfg.get('weight_decay', 0), log=lambda *_: None) dev = next(net.parameters()).device with torch.no_grad(): p = net(ds['xte'].to(dev)).detach().cpu().numpy().ravel() y = ds['yte'].numpy().ravel() threshold = np.quantile(y, .90) tail = y >= threshold global_vals.append(float(np.mean((p-y)**2))) tail_vals.append(float(np.mean((p[tail]-y[tail])**2))) # Rockafellar empirical CVaR prediction versus observed test CVaR. eta = np.quantile(p, .90) pc = eta + np.maximum(p-eta, 0).mean()/.10 et = np.quantile(y, .90) tc = et + np.maximum(y-et, 0).mean()/.10 cvar_vals.append(float(abs(pc-tc))) return {'global_mse': float(np.mean(global_vals)), 'tail_mse': float(np.mean(tail_vals)), 'tail_to_global': float(np.mean(tail_vals)/max(np.mean(global_vals), 1e-12)), 'cvar_abs_error': float(np.mean(cvar_vals))} def main(): # Baseline is swept on the same learning-rate union used by the idea. grid = [{'lr': lr, 'weight_decay': 0.0} for lr in LRS] base = sweep_baseline(baseline_fn, grid, seeds=SWEEP_SEEDS) idea_grid = [{'lr': lr, 'tail_weight': 4.0} for lr in LRS] idea_sweep = [] for cfg in idea_grid: r = evaluate(idea_fn(cfg), seeds=SWEEP_SEEDS) idea_sweep.append({'cfg': cfg, 'result': r}) best = min(idea_sweep, key=lambda z: z['result']['mean']) idea_full = evaluate(idea_fn(best['cfg']), seeds=SEEDS) bcfg = base['best_cfg'] sig_b = behavior(bcfg, idea=False) sig_i = behavior(best['cfg'], idea=True) signature = { 'prediction': 'tail residual correction should preferentially reduce upper-tail prediction error relative to global error', 'baseline_behavior': sig_b, 'idea_behavior': sig_i, 'predicted_tail_focus': True, 'observed_tail_to_global_ratio_change': sig_i['tail_to_global'] - sig_b['tail_to_global'], 'confirmed': bool(sig_i['tail_to_global'] < sig_b['tail_to_global']) } report = make_report(TRACK, MODEL, base, idea_full, {'idea_config': best['cfg'], 'idea_sweep': idea_sweep, 'mechanism_signature': signature}) report['protocol_notes'] = {'epochs': EPOCHS, 'batch': BATCH, 'baseline_grid': grid, 'idea_grid': idea_grid, 'intervention': 'same mlp_tiny trained with global-MSE warmup followed by upper-tail weighted residual correction'} with open('bench_report.json', 'w') as f: json.dump(report, f, indent=2) print(json.dumps(report, indent=2)) if __name__ == '__main__': main()