1import json, math, os, random, time
  2import numpy as np
  3import torch
  4import torch.nn as nn
  5
  6SEED = 442
  7random.seed(SEED); np.random.seed(SEED); torch.manual_seed(SEED)
  8torch.set_num_threads(min(8, os.cpu_count() or 1))
  9try:
 10    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 11except Exception:
 12    device = torch.device('cpu')
 13
 14
 15def normalized_silu(t, r):
 16    s = torch.nn.functional.softplus(r) + 1e-3
 17    return torch.nn.functional.silu(s * t) / s, s
 18
 19
 20class VarMLP(nn.Module):
 21    def __init__(self, width=32, depth=2):
 22        super().__init__()
 23        self.width, self.depth = width, depth
 24        self.weights = nn.ParameterList()
 25        self.biases = nn.ParameterList()
 26        self.scales = nn.ParameterList()
 27        in_dim = 1
 28        for _ in range(depth):
 29            self.weights.append(nn.Parameter(torch.randn(width, in_dim) * math.sqrt(2 / in_dim)))
 30            self.biases.append(nn.Parameter(torch.zeros(width)))
 31            self.scales.append(nn.Parameter(torch.zeros(width)))
 32            in_dim = width
 33        self.out_w = nn.Parameter(torch.randn(1, width) * math.sqrt(2 / width))
 34        self.out_b = nn.Parameter(torch.zeros(1))
 35
 36    def forward(self, x, return_v=False):
 37        h = x
 38        q = torch.zeros(self.width, device=x.device)
 39        for W, b, r in zip(self.weights, self.biases, self.scales):
 40            h = h @ W.t() + b
 41            s = torch.nn.functional.softplus(r) + 1e-3
 42            h = torch.nn.functional.silu(s * h) / s
 43            q = q + torch.sqrt(W.square() + 1e-8).sum(1) + torch.sqrt(b.square() + 1e-8)
 44        y = h @ self.out_w.t() + self.out_b
 45        if return_v:
 46            V = (torch.sqrt(self.out_w.square() + 1e-8) * (1 + q[None, :])).sum() + torch.sqrt(self.out_b.square() + 1e-8).sum()
 47            return y, V, q
 48        return y
 49
 50    def l2(self):
 51        return sum(p.square().sum() for p in self.parameters())
 52
 53
 54@torch.no_grad()
 55def math_checks():
 56    t = torch.linspace(-3, 3, 1001)
 57    r = torch.tensor(0.37)
 58    s = torch.nn.functional.softplus(r) + 1e-3
 59    direct = torch.nn.functional.silu(s * t) / s
 60    impl, _ = normalized_silu(t, r)
 61    formula_err = float((direct - impl).abs().max())
 62    m = VarMLP(width=1, depth=2)
 63    expected = sum(torch.sqrt(W.square() + 1e-8).sum() + torch.sqrt(b.square() + 1e-8).sum() for W, b in zip(m.weights, m.biases))
 64    _, _, q = m(torch.zeros(1, 1), True)
 65    q_err = float((q[0] - expected).abs())
 66    tt = torch.tensor([-2., -0.5, 0.5, 2.])
 67    vals = []
 68    for ss in [0.25, 1., 4.]:
 69        rr = torch.log(torch.expm1(torch.tensor(ss - 1e-3)))
 70        vals.append((torch.nn.functional.silu(ss * tt) / ss).numpy())
 71    shape_delta = float(np.max(np.abs(vals[0] - vals[2])))
 72    return {'normalized_formula_max_error': formula_err, 'recursive_q_max_error': q_err, 'silu_shape_delta_s025_vs_s4': shape_delta}
 73
 74
 75def train(kind, lam, steps=700):
 76    torch.manual_seed(SEED + (0 if kind == 'l2' else 100) + int(lam * 1e6))
 77    ntr, nva = 128, 256
 78    xtr = torch.linspace(-1, 1, ntr, device=device).unsqueeze(1)
 79    ytr = torch.sin(2 * math.pi * 3 * xtr)
 80    g = torch.Generator(device='cpu').manual_seed(SEED)
 81    xv = torch.rand(nva, 1, generator=g).to(device) * 2 - 1
 82    yv = torch.sin(2 * math.pi * 3 * xv)
 83    model = VarMLP(width=32, depth=2).to(device)
 84    if kind == 'l2':
 85        opt = torch.optim.AdamW(model.parameters(), lr=3e-3, weight_decay=lam)
 86    else:
 87        opt = torch.optim.Adam(model.parameters(), lr=3e-3)
 88    t0 = time.time()
 89    for _ in range(steps):
 90        opt.zero_grad(set_to_none=True)
 91        pred, V, _ = model(xtr, True)
 92        mse = ((pred - ytr) ** 2).mean()
 93        loss = mse + (lam * V if kind == 'var' else 0)
 94        loss.backward()
 95        opt.step()
 96    with torch.no_grad():
 97        pv, V, _ = model(xv, True)
 98        val = float(((pv - yv) ** 2).mean())
 99        train_mse = float(((model(xtr) - ytr) ** 2).mean())
100        l2 = float(model.l2())
101        maxpred = float(pv.abs().max())
102    return {'kind': kind, 'lambda': lam, 'train_mse': train_mse, 'val_mse': val, 'V': float(V), 'param_l2': l2, 'max_prediction': maxpred, 'seconds': time.time() - t0, 'steps': steps}
103
104
105def main():
106    global device
107    checks = math_checks()
108    results = []
109    for lam in [1e-4, 1e-3, 1e-2]:
110        for kind in ['l2', 'var']:
111            try:
112                results.append(train(kind, lam))
113            except Exception:
114                if device.type == 'cuda':
115                    device = torch.device('cpu')
116                    results.append(train(kind, lam))
117                else:
118                    raise
119    best_l2 = min((r for r in results if r['kind'] == 'l2'), key=lambda r: r['val_mse'])
120    best_var = min((r for r in results if r['kind'] == 'var'), key=lambda r: r['val_mse'])
121    out = {'device': str(device), 'checks': checks, 'results': results, 'best_l2': best_l2, 'best_var': best_var}
122    with open('results.json', 'w') as f:
123        json.dump(out, f, indent=2)
124    print(json.dumps(out, indent=2))
125
126
127if __name__ == '__main__':
128    main()