Binary-form symmetric-power equivariant layer / symmetric_power_experiment.py
Mechanism confirmed, baseline not beaten
1import json
2import math
3from pathlib import Path
4
5import numpy as np
6from sklearn.neural_network import MLPClassifier
7from sklearn.metrics import accuracy_score
8
9
10def symmetric_power_matrix(A, n=4):
11 """R_n(A), coefficient order X^n, X^(n-1)Y, ..., Y^n."""
12 a, b, c, d = np.asarray(A, dtype=float).reshape(2, 2).ravel()
13 R = np.zeros((n + 1, n + 1), dtype=float)
14 for k in range(n + 1):
15 for u in range(n - k + 1):
16 left = math.comb(n - k, u) * a ** (n - k - u) * c ** u
17 for v in range(k + 1):
18 R[u + v, k] += left * math.comb(k, v) * b ** (k - v) * d ** v
19 return R
20
21
22def polynomial_value(p, x, y):
23 n = len(p) - 1
24 return sum(p[k] * x ** (n - k) * y ** k for k in range(n + 1))
25
26
27def random_sl2(rng):
28 theta = rng.uniform(-math.pi, math.pi)
29 rot = np.array([[math.cos(theta), -math.sin(theta)],
30 [math.sin(theta), math.cos(theta)]])
31 t = rng.uniform(-1.1, 1.1)
32 diag = np.diag([math.exp(t), math.exp(-t)])
33 s = rng.uniform(-0.8, 0.8)
34 return rot @ diag @ np.array([[1., s], [0., 1.]])
35
36
37def monomials(points, n=4):
38 x, y = points[:, 0], points[:, 1]
39 return np.stack([x ** (n-k) * y ** k for k in range(n+1)], axis=1)
40
41
42def rotation_points(rng, radii, angles):
43 return np.stack([radii * np.cos(angles), radii * np.sin(angles)], axis=1)
44
45
46def mini_classification(rng, n=4):
47 # Binary radius task. The proposed representation has an exact SO(2)-invariant
48 # quadratic norm in the binomially normalized coefficient basis.
49 n_train = 500
50 radii = np.where(rng.integers(0, 2, n_train) == 0, .8, 1.2)
51 labels = (radii > 1).astype(int)
52 angles = rng.uniform(-math.pi, math.pi, n_train)
53 train_xy = rotation_points(rng, radii, angles)
54 train_p = monomials(train_xy, n)
55 weights = np.array([math.comb(n, k) for k in range(n+1)], dtype=float)
56 train_inv = np.sum(weights[None, :] * train_p**2, axis=1)[:, None]
57
58 # Same small MLP, trained only on the available rotated examples.
59 baseline = MLPClassifier(hidden_layer_sizes=(16, 16), max_iter=500,
60 random_state=123, solver='lbfgs')
61 baseline.fit(train_xy, labels)
62
63 # Calibrate the equivariant scalar threshold on training data; no learned
64 # transformation-specific parameters are used by this feature.
65 threshold = (train_inv[labels == 0].mean() + train_inv[labels == 1].mean()) / 2
66
67 # Held-out transformations are rotations in a disjoint angular sector.
68 test_r = np.repeat([.8, 1.2], 400)
69 test_y = np.repeat([0, 1], 400)
70 test_angles = np.concatenate([rng.uniform(0.7, math.pi, 400),
71 rng.uniform(0.7, math.pi, 400)])
72 test_xy = rotation_points(rng, test_r, test_angles)
73 test_p = monomials(test_xy, n)
74 test_inv = np.sum(weights[None, :] * test_p**2, axis=1)
75 equiv_pred = (test_inv > threshold).astype(int)
76 base_pred = baseline.predict(test_xy)
77
78 # Directly verify invariance of the proposed scalar on random rotations.
79 invariant_err = []
80 for _ in range(100):
81 pxy = rng.normal(size=(1, 2))
82 phi = rng.uniform(-math.pi, math.pi)
83 A = np.array([[math.cos(phi), -math.sin(phi)],
84 [math.sin(phi), math.cos(phi)]])
85 p = monomials(pxy, n)[0]
86 pa = monomials(pxy @ A, n)[0]
87 invariant_err.append(abs(np.dot(weights, p*p)-np.dot(weights, pa*pa)) /
88 (1 + abs(np.dot(weights, p*p))))
89 return {
90 'baseline_mlp_test_accuracy': float(accuracy_score(test_y, base_pred)),
91 'symmetric_power_invariant_test_accuracy': float(accuracy_score(test_y, equiv_pred)),
92 'invariant_scalar_max_relative_error': float(max(invariant_err)),
93 'train_examples': n_train,
94 }
95
96
97def run(seed=17, n=4, copies=3, trials=200):
98 rng = np.random.default_rng(seed)
99 poly_err, comp_err = [], []
100 for _ in range(trials):
101 A, B = random_sl2(rng), random_sl2(rng)
102 p = rng.normal(size=n + 1)
103 q = symmetric_power_matrix(A, n) @ p
104 for _ in range(3):
105 x, y = rng.normal(size=2)
106 lhs = polynomial_value(q, x, y)
107 z = np.array([x, y]) @ A
108 rhs = polynomial_value(p, z[0], z[1])
109 poly_err.append(abs(lhs-rhs) / (1 + abs(rhs)))
110 comp_err.append(np.linalg.norm(symmetric_power_matrix(A @ B, n) -
111 symmetric_power_matrix(A, n) @ symmetric_power_matrix(B, n)))
112
113 Wcopy = rng.normal(size=(copies, copies))
114 Wgood = np.kron(Wcopy, np.eye(n + 1))
115 Rmany = np.kron(np.eye(copies), symmetric_power_matrix(random_sl2(rng), n))
116 p = rng.normal(size=copies * (n + 1))
117 good_err = np.linalg.norm(Wgood @ Rmany @ p - Rmany @ Wgood @ p) / (1e-12 + np.linalg.norm(Rmany @ Wgood @ p))
118 Wbad = rng.normal(size=Wgood.shape)
119 bad_err = np.linalg.norm(Wbad @ Rmany @ p - Rmany @ Wbad @ p) / (1e-12 + np.linalg.norm(Rmany @ Wbad @ p))
120 good, bad = [], []
121 for _ in range(trials):
122 Rm = np.kron(np.eye(copies), symmetric_power_matrix(random_sl2(rng), n))
123 x = rng.normal(size=copies*(n+1))
124 good.append(np.linalg.norm(Wgood @ Rm @ x - Rm @ Wgood @ x) / (1e-12 + np.linalg.norm(Rm @ Wgood @ x)))
125 bad.append(np.linalg.norm(Wbad @ Rm @ x - Rm @ Wbad @ x) / (1e-12 + np.linalg.norm(Rm @ Wbad @ x)))
126 result = {
127 'seed': seed, 'degree': n, 'copies': copies, 'trials': trials,
128 'max_polynomial_relative_error': float(max(poly_err)),
129 'max_composition_absolute_error': float(max(comp_err)),
130 'single_good_layer_relative_error': float(good_err),
131 'single_bad_layer_relative_error': float(bad_err),
132 'median_good_layer_error': float(np.median(good)),
133 'median_bad_layer_error': float(np.median(bad)),
134 'mean_good_layer_error': float(np.mean(good)),
135 'mean_bad_layer_error': float(np.mean(bad)),
136 'mini_experiment': mini_classification(rng, n),
137 }
138 return result
139
140
141if __name__ == '__main__':
142 out = run()
143 print(json.dumps(out, indent=2))
144 Path('results.json').write_text(json.dumps(out, indent=2) + '\n')