1import sys, json, random
  2import numpy as np
  3import torch
  4import torch.nn as nn
  5import torch.nn.functional as F
  6sys.path.insert(0, '/home/maxwelhelp/all/math2nn')
  7from bench import get_dataset, train_model, evaluate, sweep_baseline, make_report
  8
  9EPOCHS = 25
 10BATCH = 128
 11GRID = [{'lr': 1e-3}, {'lr': 3e-3}, {'lr': 1e-2}]
 12
 13
 14def seed_all(seed):
 15    np.random.seed(seed)
 16    random.seed(seed)
 17    torch.manual_seed(seed)
 18    if torch.cuda.is_available():
 19        try:
 20            torch.cuda.manual_seed_all(seed)
 21        except Exception:
 22            pass
 23
 24
 25class MatchedReLU(nn.Module):
 26    def __init__(self, d, width=32, depth=2):
 27        super().__init__()
 28        self.layers = nn.ModuleList()
 29        q = d
 30        for _ in range(depth):
 31            self.layers.append(nn.Linear(q, width))
 32            q = width
 33        self.out = nn.Linear(q, 1)
 34
 35    def forward(self, x):
 36        z = x
 37        for layer in self.layers:
 38            z = F.relu(layer(z))
 39        return self.out(z)
 40
 41
 42class ICNN(nn.Module):
 43    """Input-convex MLP with the same width/depth and direct input skips."""
 44    def __init__(self, d, width=32, depth=2):
 45        super().__init__()
 46        self.d, self.width, self.depth = d, width, depth
 47        self.A = nn.ParameterList()
 48        self.U = nn.ParameterList()
 49        self.b = nn.ParameterList()
 50        for k in range(depth):
 51            # z_0 is represented by a width-dimensional first affine embedding;
 52            # subsequent hidden-to-hidden matrices are nonnegative.
 53            self.A.append(nn.Parameter(torch.randn(width, width) * 0.12))
 54            self.U.append(nn.Parameter(torch.randn(width, d) * 0.12))
 55            self.b.append(nn.Parameter(torch.zeros(width)))
 56        self.aw = nn.Parameter(torch.zeros(width))
 57        self.u = nn.Parameter(torch.zeros(d))
 58        self.c = nn.Parameter(torch.zeros(1))
 59
 60    def forward(self, x):
 61        z = F.relu(x @ self.U[0].T + self.b[0])
 62        for k in range(1, self.depth):
 63            W = F.softplus(self.A[k]) + 1e-6
 64            z = F.relu(z @ W.T + x @ self.U[k].T + self.b[k])
 65        return (z @ F.softplus(self.aw) + x @ self.u + self.c).unsqueeze(-1)
 66
 67
 68def train_one(kind, cfg, seed, return_model=False):
 69    seed_all(seed)
 70    ds = get_dataset('tabular', seed=seed)
 71    d = int(ds['xtr'].shape[1])
 72    model = MatchedReLU(d) if kind == 'baseline' else ICNN(d)
 73    net, metric, hist = train_model(model, ds, epochs=EPOCHS,
 74                                    lr=cfg['lr'], batch=BATCH,
 75                                    weight_decay=0.0, log=lambda *_: None)
 76    if net is None:
 77        raise RuntimeError('training failed')
 78    if return_model:
 79        return float(metric), net, ds
 80    return float(metric)
 81
 82
 83def make_fn(kind, cfg):
 84    return lambda seed: train_one(kind, cfg, seed)
 85
 86
 87def mechanism_signature():
 88    # Re-test convexity on predictions from trained benchmark models, not toy weights.
 89    rows = []
 90    for seed in range(8):
 91        _, net, ds = train_one('idea', {'lr': 3e-3}, seed, return_model=True)
 92        device = next(net.parameters()).device
 93        x = ds['xte'][:96].to(device)
 94        rng = torch.Generator().manual_seed(9000 + seed)
 95        perm = torch.randperm(len(x), generator=rng)
 96        a, b = x, x[perm]
 97        t = 0.5
 98        with torch.no_grad():
 99            lhs = net(t*a + (1-t)*b).reshape(-1)
100            rhs = (t*net(a) + (1-t)*net(b)).reshape(-1)
101        gap = lhs - rhs
102        rows.append((float(gap.max()), float((gap > 1e-5).float().mean())))
103    max_gap = max(r[0] for r in rows)
104    rate = float(np.mean([r[1] for r in rows]))
105    return {
106        'claim': 'trained ICNN predictions satisfy midpoint Jensen convexity',
107        'predicted_max_violation': 0.0,
108        'observed_max_violation': max_gap,
109        'observed_violation_rate': rate,
110        'confirmed': bool(max_gap <= 1e-5 and rate == 0.0),
111        'n_models': 8
112    }
113
114
115def main():
116    base = sweep_baseline(lambda cfg: make_fn('baseline', cfg), GRID)
117    # Same three settings are run for the idea; best is selected on the same sweep seeds.
118    idea_sweep = []
119    for cfg in GRID:
120        r = evaluate(make_fn('idea', cfg), seeds=(0,1,2,3))
121        idea_sweep.append({'cfg': cfg, 'mean': r['mean']})
122    best_cfg = min(idea_sweep, key=lambda x: x['mean'])['cfg']
123    idea_full = evaluate(make_fn('idea', best_cfg))
124    idea_res = {'best_cfg': best_cfg, 'sweep': idea_sweep,
125                'full': idea_full, **idea_full}
126    report = make_report('tabular', 'mlp_med', base, idea_full,
127                         extra=mechanism_signature())
128    report['idea'] = idea_res
129    report['protocol_notes'] = {
130        'structural_match': 'tabular is the built-in track for architecture/regularization interventions',
131        'paired_seeds': list(range(8)),
132        'epochs': EPOCHS, 'batch': BATCH,
133        'baseline_and_idea_share_grid': True,
134        'baseline_architecture': '2-layer width-32 ReLU MLP',
135        'idea_architecture': '2-layer width-32 ICNN with softplus W>=0 and nonnegative output weights'
136    }
137    with open('bench_report.json', 'w') as f:
138        json.dump(report, f, indent=2)
139    print(json.dumps(report, indent=2))
140
141
142if __name__ == '__main__':
143    main()