Lipschitz-Free Metric Pooling / experiment.py
Failed on benchmark
1import random
2from pathlib import Path
3import numpy as np
4from scipy.optimize import linprog
5import torch
6from torch import nn
7
8SEED = 461
9random.seed(SEED)
10np.random.seed(SEED)
11torch.manual_seed(SEED)
12try:
13 if torch.cuda.is_available():
14 torch.cuda.manual_seed_all(SEED)
15 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
16 torch.zeros(1, device=device)
17except Exception:
18 device = torch.device("cpu")
19
20
21def exact_free_norm(points, weights, theta=0):
22 """Exact finite-space dual norm by LP, with Euclidean metric."""
23 points = np.asarray(points, dtype=float)
24 weights = np.asarray(weights, dtype=float)
25 n = len(points)
26 d = np.linalg.norm(points[:, None] - points[None, :], axis=-1)
27 ids = [i for i in range(n) if i != theta]
28 A, b = [], []
29 for i in range(n):
30 for j in range(i + 1, n):
31 row = np.zeros(len(ids))
32 if i != theta:
33 row[ids.index(i)] += 1
34 if j != theta:
35 row[ids.index(j)] -= 1
36 A.extend([row, -row])
37 b.extend([d[i, j], d[i, j]])
38 c = -weights[ids]
39 rpos = linprog(c, A_ub=np.asarray(A), b_ub=np.asarray(b),
40 bounds=[(None, None)] * len(ids), method="highs")
41 rneg = linprog(-c, A_ub=np.asarray(A), b_ub=np.asarray(b),
42 bounds=[(None, None)] * len(ids), method="highs")
43 return max(-rpos.fun, -rneg.fun), d
44
45
46def verify_math():
47 rng = np.random.default_rng(SEED)
48 rel_err, valid = [], []
49 for _ in range(30):
50 p = rng.normal(size=(6, 2))
51 w = rng.random(6); w /= w.sum(); w[0] = 0
52 exact, d = exact_free_norm(p, w)
53 # Distance-to-anchor is an exactly 1-Lipschitz probe and vanishes at theta.
54 phi = d[0]
55 estimate = abs(np.dot(w, phi))
56 valid.append(np.max(np.abs(phi[:, None] - phi[None, :]) - d) <= 1e-9)
57 rel_err.append(estimate / max(exact, 1e-12))
58 return {"valid_probe_fraction": float(np.mean(valid)),
59 "distance_probe_over_exact_mean": float(np.mean(rel_err)),
60 "distance_probe_over_exact_min": float(np.min(rel_err))}
61
62
63class ProbePool(nn.Module):
64 def __init__(self, k=8):
65 super().__init__()
66 self.probes = nn.ModuleList([
67 nn.utils.spectral_norm(nn.Linear(2, 1)) for _ in range(k)
68 ])
69 self.out = nn.Linear(k, 2)
70
71 def forward(self, x, keep=None):
72 # x: batch x items x 2; base point is zero, so subtract phi(0).
73 vals = []
74 for p in self.probes:
75 y = torch.tanh(p(x)).squeeze(-1)
76 base = torch.tanh(p(torch.zeros_like(x[..., :2]))).squeeze(-1)
77 vals.append(y - base)
78 v = torch.stack(vals, dim=-1)
79 if keep is not None:
80 v = v * keep[..., None]
81 denom = keep.sum(1, keepdim=True).clamp_min(1)
82 else:
83 denom = v.new_full((x.shape[0], 1), x.shape[1])
84 z = v.sum(1) / denom
85 return self.out(z)
86
87
88class MeanPool(nn.Module):
89 def __init__(self):
90 super().__init__(); self.out = nn.Sequential(nn.Linear(2, 8), nn.Tanh(), nn.Linear(8, 2))
91 def forward(self, x, keep=None):
92 if keep is None: z = x.mean(1)
93 else: z = (x * keep[..., None]).sum(1) / keep.sum(1, keepdim=True).clamp_min(1)
94 return self.out(z)
95
96
97def make_data(n, items=12, rng=None, keep_fraction=1.0):
98 rng = np.random.default_rng() if rng is None else rng
99 # Class is angular geometry: both classes have near-zero coordinate mean.
100 y = rng.integers(0, 2, n)
101 a = rng.uniform(0, 2*np.pi, (n, items))
102 # class 0: circle; class 1: two antipodal arcs, same rough mean
103 if True:
104 radial = np.where(y[:, None] == 0, 1.0, 1.0 + 0.35*np.cos(2*a))
105 x = np.stack([radial*np.cos(a), radial*np.sin(a)], -1)
106 x += rng.normal(0, 0.035, x.shape)
107 keep = np.ones((n, items), dtype=np.float32)
108 if keep_fraction < 1:
109 for i in range(n):
110 idx = rng.choice(items, int(items*keep_fraction), replace=False)
111 keep[i, idx] = 0
112 return torch.tensor(x, dtype=torch.float32), torch.tensor(y, dtype=torch.long), torch.tensor(keep)
113
114
115def train(model, x, y, steps=350):
116 model.to(device); x=x.to(device); y=y.to(device)
117 opt=torch.optim.Adam(model.parameters(), lr=3e-3)
118 for t in range(steps):
119 ix=torch.randint(0,len(x),(64,),device=device)
120 loss=nn.functional.cross_entropy(model(x[ix]),y[ix])
121 opt.zero_grad(); loss.backward(); opt.step()
122 with torch.no_grad():
123 pred=model(x).argmax(1); return float((pred==y).float().mean().cpu())
124
125
126def run():
127 math_result=verify_math()
128 rng=np.random.default_rng(SEED)
129 train_x,train_y,_=make_data(768,rng=rng)
130 test_x,test_y,test_keep=make_data(512,rng=rng)
131 results={"math":math_result,"device":str(device),"runs":{}}
132 for run_seed in [461, 462, 463, 464, 465]:
133 rrng=np.random.default_rng(run_seed)
134 train_x,train_y,_=make_data(768,rng=rrng)
135 test_x,test_y,test_keep=make_data(512,rng=rrng,keep_fraction=0.5)
136 results["runs"][str(run_seed)]={}
137 for name, cls in [("mean",MeanPool),("lipschitz_probe",lambda:ProbePool(8))]:
138 torch.manual_seed(run_seed)
139 m=cls(); train_acc=train(m,train_x,train_y)
140 with torch.no_grad():
141 clean=float((m(test_x.to(device)).argmax(1)==test_y.to(device)).float().mean().cpu())
142 partial=float((m(test_x.to(device),test_keep.to(device)).argmax(1)==test_y.to(device)).float().mean().cpu())
143 results["runs"][str(run_seed)][name]={"train_accuracy":train_acc,"clean_test_accuracy":clean,
144 "50_percent_removed_accuracy":partial,"parameters":sum(p.numel() for p in m.parameters())}
145 for name in ["mean","lipschitz_probe"]:
146 vals=[results["runs"][str(s)][name] for s in [461,462,463,464,465]]
147 results[name]={k:float(np.mean([v[k] for v in vals])) for k in ["train_accuracy","clean_test_accuracy","50_percent_removed_accuracy"]}
148 results[name]["std_clean"]=float(np.std([v["clean_test_accuracy"] for v in vals]))
149 results[name]["std_removed"]=float(np.std([v["50_percent_removed_accuracy"] for v in vals]))
150 results[name]["parameters"]=vals[0]["parameters"]
151 Path("results.txt").write_text(repr(results) + "\n")
152 print(results)
153
154if __name__ == "__main__":
155 run()