import random from pathlib import Path import numpy as np from scipy.optimize import linprog import torch from torch import nn SEED = 461 random.seed(SEED) np.random.seed(SEED) torch.manual_seed(SEED) try: if torch.cuda.is_available(): torch.cuda.manual_seed_all(SEED) device = torch.device("cuda" if torch.cuda.is_available() else "cpu") torch.zeros(1, device=device) except Exception: device = torch.device("cpu") def exact_free_norm(points, weights, theta=0): """Exact finite-space dual norm by LP, with Euclidean metric.""" points = np.asarray(points, dtype=float) weights = np.asarray(weights, dtype=float) n = len(points) d = np.linalg.norm(points[:, None] - points[None, :], axis=-1) ids = [i for i in range(n) if i != theta] A, b = [], [] for i in range(n): for j in range(i + 1, n): row = np.zeros(len(ids)) if i != theta: row[ids.index(i)] += 1 if j != theta: row[ids.index(j)] -= 1 A.extend([row, -row]) b.extend([d[i, j], d[i, j]]) c = -weights[ids] rpos = linprog(c, A_ub=np.asarray(A), b_ub=np.asarray(b), bounds=[(None, None)] * len(ids), method="highs") rneg = linprog(-c, A_ub=np.asarray(A), b_ub=np.asarray(b), bounds=[(None, None)] * len(ids), method="highs") return max(-rpos.fun, -rneg.fun), d def verify_math(): rng = np.random.default_rng(SEED) rel_err, valid = [], [] for _ in range(30): p = rng.normal(size=(6, 2)) w = rng.random(6); w /= w.sum(); w[0] = 0 exact, d = exact_free_norm(p, w) # Distance-to-anchor is an exactly 1-Lipschitz probe and vanishes at theta. phi = d[0] estimate = abs(np.dot(w, phi)) valid.append(np.max(np.abs(phi[:, None] - phi[None, :]) - d) <= 1e-9) rel_err.append(estimate / max(exact, 1e-12)) return {"valid_probe_fraction": float(np.mean(valid)), "distance_probe_over_exact_mean": float(np.mean(rel_err)), "distance_probe_over_exact_min": float(np.min(rel_err))} class ProbePool(nn.Module): def __init__(self, k=8): super().__init__() self.probes = nn.ModuleList([ nn.utils.spectral_norm(nn.Linear(2, 1)) for _ in range(k) ]) self.out = nn.Linear(k, 2) def forward(self, x, keep=None): # x: batch x items x 2; base point is zero, so subtract phi(0). vals = [] for p in self.probes: y = torch.tanh(p(x)).squeeze(-1) base = torch.tanh(p(torch.zeros_like(x[..., :2]))).squeeze(-1) vals.append(y - base) v = torch.stack(vals, dim=-1) if keep is not None: v = v * keep[..., None] denom = keep.sum(1, keepdim=True).clamp_min(1) else: denom = v.new_full((x.shape[0], 1), x.shape[1]) z = v.sum(1) / denom return self.out(z) class MeanPool(nn.Module): def __init__(self): super().__init__(); self.out = nn.Sequential(nn.Linear(2, 8), nn.Tanh(), nn.Linear(8, 2)) def forward(self, x, keep=None): if keep is None: z = x.mean(1) else: z = (x * keep[..., None]).sum(1) / keep.sum(1, keepdim=True).clamp_min(1) return self.out(z) def make_data(n, items=12, rng=None, keep_fraction=1.0): rng = np.random.default_rng() if rng is None else rng # Class is angular geometry: both classes have near-zero coordinate mean. y = rng.integers(0, 2, n) a = rng.uniform(0, 2*np.pi, (n, items)) # class 0: circle; class 1: two antipodal arcs, same rough mean if True: radial = np.where(y[:, None] == 0, 1.0, 1.0 + 0.35*np.cos(2*a)) x = np.stack([radial*np.cos(a), radial*np.sin(a)], -1) x += rng.normal(0, 0.035, x.shape) keep = np.ones((n, items), dtype=np.float32) if keep_fraction < 1: for i in range(n): idx = rng.choice(items, int(items*keep_fraction), replace=False) keep[i, idx] = 0 return torch.tensor(x, dtype=torch.float32), torch.tensor(y, dtype=torch.long), torch.tensor(keep) def train(model, x, y, steps=350): model.to(device); x=x.to(device); y=y.to(device) opt=torch.optim.Adam(model.parameters(), lr=3e-3) for t in range(steps): ix=torch.randint(0,len(x),(64,),device=device) loss=nn.functional.cross_entropy(model(x[ix]),y[ix]) opt.zero_grad(); loss.backward(); opt.step() with torch.no_grad(): pred=model(x).argmax(1); return float((pred==y).float().mean().cpu()) def run(): math_result=verify_math() rng=np.random.default_rng(SEED) train_x,train_y,_=make_data(768,rng=rng) test_x,test_y,test_keep=make_data(512,rng=rng) results={"math":math_result,"device":str(device),"runs":{}} for run_seed in [461, 462, 463, 464, 465]: rrng=np.random.default_rng(run_seed) train_x,train_y,_=make_data(768,rng=rrng) test_x,test_y,test_keep=make_data(512,rng=rrng,keep_fraction=0.5) results["runs"][str(run_seed)]={} for name, cls in [("mean",MeanPool),("lipschitz_probe",lambda:ProbePool(8))]: torch.manual_seed(run_seed) m=cls(); train_acc=train(m,train_x,train_y) with torch.no_grad(): clean=float((m(test_x.to(device)).argmax(1)==test_y.to(device)).float().mean().cpu()) partial=float((m(test_x.to(device),test_keep.to(device)).argmax(1)==test_y.to(device)).float().mean().cpu()) results["runs"][str(run_seed)][name]={"train_accuracy":train_acc,"clean_test_accuracy":clean, "50_percent_removed_accuracy":partial,"parameters":sum(p.numel() for p in m.parameters())} for name in ["mean","lipschitz_probe"]: vals=[results["runs"][str(s)][name] for s in [461,462,463,464,465]] results[name]={k:float(np.mean([v[k] for v in vals])) for k in ["train_accuracy","clean_test_accuracy","50_percent_removed_accuracy"]} results[name]["std_clean"]=float(np.std([v["clean_test_accuracy"] for v in vals])) results[name]["std_removed"]=float(np.std([v["50_percent_removed_accuracy"] for v in vals])) results[name]["parameters"]=vals[0]["parameters"] Path("results.txt").write_text(repr(results) + "\n") print(results) if __name__ == "__main__": run()