#!/usr/bin/env python3 """Certified robustness on a 2-D toy problem: Lipschitz margins and randomized smoothing, checked against a real attack. Two small numpy MLPs f: R^2 -> R classify 'two moons' by sign(f): * wiggly unconstrained, trained to fit noisy labels (high Lipschitz); * Lipschitz every weight matrix is projected to spectral norm <= 1 after each Adam step; with 1-Lipschitz activations (ReLU) the network is L-Lipschitz in l2 with L = prod_i ||W_i||_2 <= 1. Lipschitz certificate. If |f(x) - f(x')| <= L |x - x'|, then sign(f) cannot change within radius |f(x)| / L of x. This is a *guarantee*: no attack can succeed inside it. We verify that with a projected-gradient (PGD) attack. For the wiggly net, the same product-of-norms bound is valid but enormous, so its certificates are useless even where it happens to be robust. Randomized smoothing (Cohen, Rosenfeld & Kolter 2019). For any base classifier h, the smoothed classifier g(x) = argmax_c P(h(x + sigma Z) = c), Z ~ N(0, I), is certifiably robust: if the top class has probability >= pA > 1/2 then g is constant on the l2 ball of radius sigma * Phi^-1(pA). We estimate pA by Monte Carlo and use a one-sided Clopper-Pearson lower bound, so the certificate holds with probability >= 1 - alpha (abstaining when pA_lower <= 1/2). Applied here to the wiggly network. python certified_robustness.py [--quick] Writes out/certified_robustness.png and prints certified accuracies. """ from __future__ import annotations import argparse import os import numpy as np from scipy import stats OUT = os.path.join(os.path.dirname(os.path.abspath(__file__)), "out") # -------------------------------------------------------------------------- # data # -------------------------------------------------------------------------- def make_moons(n, noise, rng): """Two interleaving half circles, labels in {-1, +1}, roughly centred.""" n1 = n // 2 t1 = rng.uniform(0, np.pi, n1) t2 = rng.uniform(0, np.pi, n - n1) a = np.stack([np.cos(t1), np.sin(t1)], 1) b = np.stack([1 - np.cos(t2), 0.5 - np.sin(t2)], 1) X = np.vstack([a, b]) + noise * rng.standard_normal((n, 2)) y = np.r_[np.ones(n1), -np.ones(n - n1)] X = X - np.array([0.5, 0.25]) return X, y # -------------------------------------------------------------------------- # a small MLP in numpy, with input gradients # -------------------------------------------------------------------------- class MLP: def __init__(self, sizes, seed=0, init_scale=1.0): rng = np.random.default_rng(seed) self.W = [init_scale * rng.standard_normal((a, b)) * np.sqrt(2.0 / a) for a, b in zip(sizes, sizes[1:])] self.b = [np.zeros(b) for b in sizes[1:]] self.m = [np.zeros_like(p) for p in self.W + self.b] self.v = [np.zeros_like(p) for p in self.W + self.b] self.t = 0 def forward(self, X): acts, pre = [X], [] for i, (W, b) in enumerate(zip(self.W, self.b)): z = acts[-1] @ W + b pre.append(z) acts.append(z if i == len(self.W) - 1 else np.maximum(z, 0)) return acts, pre def __call__(self, X): out = X for i, (W, b) in enumerate(zip(self.W, self.b)): out = out @ W + b if i < len(self.W) - 1: out = np.maximum(out, 0) return out[:, 0] def backward(self, acts, pre, g_out): """g_out: dLoss/df (n,). Returns param grads and input grad.""" g = g_out[:, None] gW, gb = [None] * len(self.W), [None] * len(self.W) for i in reversed(range(len(self.W))): gW[i] = acts[i].T @ g gb[i] = g.sum(axis=0) g = g @ self.W[i].T if i > 0: g = g * (pre[i - 1] > 0) return gW, gb, g def input_grad(self, X): acts, pre = self.forward(X) return self.backward(acts, pre, np.ones(len(X)))[2] def adam_step(self, gW, gb, lr): self.t += 1 for j, (p, g) in enumerate(zip(self.W + self.b, gW + gb)): self.m[j] = 0.9 * self.m[j] + 0.1 * g self.v[j] = 0.999 * self.v[j] + 0.001 * g**2 p -= lr * (self.m[j] / (1 - 0.9**self.t)) / (np.sqrt(self.v[j] / (1 - 0.999**self.t)) + 1e-8) def project_spectral(self, max_norm=1.0): for W in self.W: s = np.linalg.norm(W, 2) if s > max_norm: W /= s / max_norm def lipschitz_bound(self): """Valid global l2 Lipschitz constant (ReLU is 1-Lipschitz).""" return float(np.prod([np.linalg.norm(W, 2) for W in self.W])) def train(net, X, y, steps, lr, temperature=1.0, lipschitz=False, seed=0, batch=None): """Logistic loss on temperature * y * f(x); optional spectral projection.""" rng = np.random.default_rng(seed) for _ in range(steps): idx = rng.integers(len(X), size=batch) if batch else slice(None) xb, yb = X[idx], y[idx] acts, pre = net.forward(xb) m = temperature * yb * acts[-1][:, 0] # d/df of mean log(1 + exp(-m)) g = -temperature * yb * (1 / (1 + np.exp(np.clip(m, -50, 50)))) / len(xb) gW, gb, _ = net.backward(acts, pre, g) net.adam_step(gW, gb, lr) if lipschitz: net.project_spectral(1.0) return net # -------------------------------------------------------------------------- # certificates and attack # -------------------------------------------------------------------------- def lipschitz_radius(net, X): return np.abs(net(X)) / net.lipschitz_bound() def pgd_attack(net, X, y, radius, steps=60, restarts=3, seed=0): """l2 PGD on the margin y f(x+delta), ||delta|| <= radius (per point). Returns the smallest margin found (negative = attack succeeded) and the adversarial points.""" rng = np.random.default_rng(seed) radius = np.broadcast_to(np.asarray(radius, float), (len(X),)) best_m = y * net(X) best_x = X.copy() for r in range(restarts): if r == 0: delta = np.zeros_like(X) else: u = rng.standard_normal(X.shape) delta = u / np.linalg.norm(u, axis=1, keepdims=True) * radius[:, None] * rng.random((len(X), 1)) step = 2.5 * radius / steps for _ in range(steps): g = -y[:, None] * net.input_grad(X + delta) # ascent dir. of -margin gn = np.linalg.norm(g, axis=1, keepdims=True) delta = delta + step[:, None] * g / np.maximum(gn, 1e-12) nrm = np.linalg.norm(delta, axis=1, keepdims=True) delta = delta * np.minimum(1, radius[:, None] / np.maximum(nrm, 1e-12)) m = y * net(X + delta) better = m < best_m best_m = np.where(better, m, best_m) best_x[better] = (X + delta)[better] return best_m, best_x def smoothed_certify(net, X, sigma, n0=100, n=10000, alpha=0.001, seed=0, chunk=2000): """Cohen et al. CERTIFY for binary h = sign(f). Returns (prediction in {-1, 0 (abstain), +1}, certified radius).""" rng = np.random.default_rng(seed) preds, radii = np.zeros(len(X)), np.zeros(len(X)) for i, x in enumerate(X): sel = np.sign(net(x + sigma * rng.standard_normal((n0, 2)))) cA = 1.0 if (sel > 0).sum() >= (sel < 0).sum() else -1.0 k = 0 for s in range(0, n, chunk): m = min(chunk, n - s) k += int((np.sign(net(x + sigma * rng.standard_normal((m, 2)))) == cA).sum()) pA_low = stats.beta.ppf(alpha, k, n - k + 1) if k > 0 else 0.0 # one-sided CP if pA_low > 0.5: preds[i], radii[i] = cA, sigma * stats.norm.ppf(pA_low) return preds, radii def certified_accuracy(correct, radii, grid): return np.array([np.mean(correct & (radii >= r)) for r in grid]) def empirical_robust_accuracy(net, X, y, grid, seed=0, steps=40): clean = y * net(X) > 0 out = [] for r in grid: if r == 0: out.append(clean.mean()) continue m, _ = pgd_attack(net, X, y, r, steps=steps, restarts=2, seed=seed) out.append(np.mean(clean & (m > 0))) return np.array(out) # -------------------------------------------------------------------------- # demo # -------------------------------------------------------------------------- def build_models(quick=False, seed=0): rng = np.random.default_rng(seed) Xtr, ytr = make_moons(200, 0.15, rng) Xte, yte = make_moons(200, 0.15, rng) wiggly = train(MLP([2, 64, 64, 1], seed=1, init_scale=2.0), Xtr, ytr, steps=1500 if quick else 4000, lr=1e-2) lip = train(MLP([2, 64, 64, 1], seed=2), Xtr, ytr, steps=1500 if quick else 4000, lr=3e-3, temperature=20.0, lipschitz=True) return (Xtr, ytr, Xte, yte), wiggly, lip def main(quick=False): import matplotlib matplotlib.use("Agg") import matplotlib.pyplot as plt os.makedirs(OUT, exist_ok=True) (Xtr, ytr, Xte, yte), wiggly, lip = build_models(quick) sigma = 0.25 grid = np.linspace(0, 0.5, 11) print(f"{'model':<14}{'train acc':>10}{'test acc':>10}{'Lipschitz bound':>17}") for name, net in (("wiggly", wiggly), ("Lipschitz", lip)): print(f"{name:<14}{np.mean(ytr * net(Xtr) > 0):>10.3f}{np.mean(yte * net(Xte) > 0):>10.3f}" f"{net.lipschitz_bound():>17.3g}") res = {} for name, net in (("wiggly", wiggly), ("Lipschitz", lip)): correct = yte * net(Xte) > 0 rad = lipschitz_radius(net, Xte) res[name] = {"cert": certified_accuracy(correct, rad, grid), "pgd": empirical_robust_accuracy(net, Xte, yte, grid, steps=25 if quick else 40), "rad": rad, "correct": correct} # sanity: attack the Lipschitz net right at (just inside) its certificates m, _ = pgd_attack(lip, Xte, yte, 0.999 * res["Lipschitz"]["rad"], steps=60, restarts=3) viol = int(np.sum(res["Lipschitz"]["correct"] & (m <= 0))) print(f"PGD inside certified radii of the Lipschitz net: {viol} violations " f"out of {int(res['Lipschitz']['correct'].sum())} certified points (must be 0)") n_mc = 2000 if quick else 10000 preds, srad = smoothed_certify(wiggly, Xte, sigma, n=n_mc, alpha=0.001) s_correct = preds == yte res["smoothed"] = {"cert": certified_accuracy(s_correct, srad, grid)} print(f"smoothed wiggly net (sigma={sigma}, n={n_mc}, alpha=0.001): " f"accuracy {np.mean(s_correct):.3f}, abstain {np.mean(preds == 0):.3f}") print(f"\n{'radius':>7}{'Lip cert':>10}{'Lip PGD':>9}{'wig cert':>10}{'wig PGD':>9}{'smooth cert':>13}") for i, r in enumerate(grid): print(f"{r:>7.2f}{res['Lipschitz']['cert'][i]:>10.3f}{res['Lipschitz']['pgd'][i]:>9.3f}" f"{res['wiggly']['cert'][i]:>10.3f}{res['wiggly']['pgd'][i]:>9.3f}{res['smoothed']['cert'][i]:>13.3f}") print("(cert = guaranteed lower bound on robust accuracy; PGD = attack-based upper bound)") # ---- figure ------------------------------------------------------------ fig, axes = plt.subplots(1, 4, figsize=(20, 4.8)) xx, yy = np.meshgrid(np.linspace(-2, 2, 160), np.linspace(-1.5, 1.5, 120)) G = np.c_[xx.ravel(), yy.ravel()] rng = np.random.default_rng(3) show = rng.choice(len(Xte), 30, replace=False) def region(ax, Z, title): ax.contourf(xx, yy, Z.reshape(xx.shape), levels=[-1e9, 0, 1e9], colors=["#f4cccc", "#cfe2f3"]) ax.contour(xx, yy, Z.reshape(xx.shape), levels=[0], colors="k", linewidths=1) ax.scatter(Xte[:, 0], Xte[:, 1], c=np.where(yte > 0, "C0", "C3"), s=6) ax.set_title(title) ax.set_aspect("equal") ax.set_xlim(-2, 2) ax.set_ylim(-1.5, 1.5) def circles(ax, centers, radii, color): for c, r in zip(centers, radii): if r > 0: ax.add_patch(plt.Circle(c, r, fill=False, color=color, lw=1)) region(axes[0], wiggly(G), f"Wiggly net (L ≤ {wiggly.lipschitz_bound():.0f})\n" "Lipschitz certificates |f|/L: tiny") circles(axes[0], Xte[show], res["wiggly"]["rad"][show] * res["wiggly"]["correct"][show], "k") region(axes[1], lip(G), f"Spectral-normalised net (L ≤ {lip.lipschitz_bound():.2f})\n" "circles = certified radius |f|/L") circles(axes[1], Xte[show], res["Lipschitz"]["rad"][show] * res["Lipschitz"]["correct"][show], "k") # smoothed classifier on the grid (cheap estimate, for display only) n_grid = 100 votes = np.zeros(len(G)) for _ in range(n_grid): votes += np.sign(wiggly(G + sigma * rng.standard_normal(G.shape))) region(axes[2], votes, f"Randomized smoothing of wiggly net (σ={sigma})\n" "circles = certified radius σΦ⁻¹(p_A,lower)") circles(axes[2], Xte[show], srad[show] * s_correct[show], "k") ax = axes[3] ax.plot(grid, res["Lipschitz"]["cert"], "C2-", lw=2, label="Lipschitz net: certified") ax.plot(grid, res["Lipschitz"]["pgd"], "C2--", label="Lipschitz net: PGD (upper bd)") ax.plot(grid, res["wiggly"]["cert"], "C3-", lw=2, label="wiggly net: certified") ax.plot(grid, res["wiggly"]["pgd"], "C3--", label="wiggly net: PGD (upper bd)") ax.plot(grid, res["smoothed"]["cert"], "C0-", lw=2, label="smoothed wiggly: certified") ax.set_xlabel("l2 radius ε") ax.set_ylabel("accuracy") ax.set_ylim(0, 1) ax.set_title("Certified vs empirical robust accuracy") ax.grid(alpha=0.3) ax.legend(fontsize=8) fig.tight_layout() path = os.path.join(OUT, "certified_robustness.png") fig.savefig(path, dpi=130) plt.close(fig) print(f"saved {path}") return res, viol if __name__ == "__main__": ap = argparse.ArgumentParser() ap.add_argument("--quick", action="store_true") main(ap.parse_args().quick)