Uniform-Certificate Bayesian Feature Head / experiment.py
Beats tuned baseline
1import json, time
2from pathlib import Path
3import numpy as np
4
5SEED = 2399
6rng = np.random.default_rng(SEED)
7
8
9def phi(x, K=8):
10 x = np.asarray(x).reshape(-1)
11 # Frequencies are deterministic, as required by the proposal.
12 out = [np.ones_like(x)]
13 for k in range(1, K + 1):
14 out += [np.cos(np.pi * k * x), np.sin(np.pi * k * x)]
15 return np.stack(out, axis=1)
16
17
18def true_fn(x):
19 x = np.asarray(x)
20 return 0.4 + 0.8*np.sin(np.pi*x) - 0.35*np.cos(2*np.pi*x) + 0.25*np.sin(4*np.pi*x)
21
22
23def update_stats(X, y, lam=1e-2, K=8):
24 P = phi(X, K)
25 V = lam*np.eye(P.shape[1]) + P.T @ P
26 b = P.T @ y
27 return V, b
28
29
30def predict(V, b, X, sigma, beta=2.5, K=8):
31 P = phi(X, K)
32 theta = np.linalg.solve(V, b)
33 # solve once, then form diagonal quadratic forms
34 VinvP = np.linalg.solve(V, P.T).T
35 q = np.maximum(0.0, np.sum(P * VinvP, axis=1))
36 mu = P @ theta
37 sd = sigma*np.sqrt(q)
38 return mu, sd
39
40
41def online_stats(X, y, lam=1e-2, K=8):
42 m = 2*K+1
43 V = lam*np.eye(m)
44 b = np.zeros(m)
45 for x, yy in zip(X, y):
46 p = phi([x], K)[0]
47 V += np.outer(p, p)
48 b += p*yy
49 return V, b
50
51
52def contraction_sweep():
53 # Repeated identical observations give q_n=q_0/(1+n*q_0/lambda-like scale),
54 # hence posterior sd is asymptotically proportional to n^-1/2.
55 sigma, lam, K = 0.12, 1e-2, 8
56 x0 = 0.15
57 ns = np.array([1, 2, 4, 8, 16, 32, 64, 128, 256, 512, 1024])
58 sds = []
59 for n in ns:
60 X = np.full(n, x0)
61 y = true_fn(X) # y values do not affect posterior variance
62 V, b = update_stats(X, y, lam, K)
63 sds.append(predict(V, b, [x0], sigma, K=K)[1][0])
64 # Fit slope only in the data-dominated regime.
65 slope = np.polyfit(np.log(ns[2:]), np.log(sds[2:]), 1)[0]
66 # Uncovered point after observations at x0.
67 X = np.full(1024, x0); y = true_fn(X)
68 V, b = update_stats(X, y, lam, K)
69 covered = predict(V, b, [x0], sigma, K=K)[1][0]
70 uncovered = predict(V, b, [-0.93], sigma, K=K)[1][0]
71 return {"n": ns.tolist(), "sd": np.array(sds).tolist(),
72 "loglog_slope": float(slope), "predicted_slope": -0.5,
73 "covered_sd_n1024": float(covered), "uncovered_sd_n1024": float(uncovered),
74 "uncovered_to_covered": float(uncovered/covered)}
75
76
77def coverage_sweep():
78 # Repeated trials assess simultaneous grid coverage of beta*s.
79 K, lam, sigma, beta = 8, 1e-2, 0.12, 2.5
80 grid = np.linspace(-1, 1, 201)
81 rows = []
82 for n, region in [(40, "concentrated"), (160, "concentrated"), (40, "broad"), (160, "broad")]:
83 covers = []
84 widths = []
85 rmses = []
86 for t in range(100):
87 rr = np.random.default_rng(SEED + 10000 + t + n + (0 if region == "concentrated" else 1000))
88 X = rr.uniform(-0.25, 0.25, n) if region == "concentrated" else rr.uniform(-1, 1, n)
89 y = true_fn(X) + rr.normal(0, sigma, n)
90 V, b = update_stats(X, y, lam, K)
91 mu, sd = predict(V, b, grid, sigma, beta, K)
92 err = np.abs(true_fn(grid)-mu)
93 covers.append(float(np.all(err <= beta*sd)))
94 widths.append(float(np.mean(2*beta*sd)))
95 rmses.append(float(np.sqrt(np.mean((true_fn(grid)-mu)**2))))
96 rows.append({"n": n, "region": region, "simultaneous_coverage": float(np.mean(covers)),
97 "mean_interval_width": float(np.mean(widths)), "grid_rmse": float(np.mean(rmses)),
98 "target_coverage": 0.95, "beta": beta})
99 return rows
100
101
102def baseline_comparison():
103 # Same feature family and same data: baseline gives only a point prediction;
104 # proposed head additionally supplies a calibrated, location-dependent interval.
105 rr = np.random.default_rng(SEED)
106 n = 160
107 X = rr.uniform(-0.25, 0.25, n)
108 sigma = 0.12
109 y = true_fn(X) + rr.normal(0, sigma, n)
110 grid = np.linspace(-1, 1, 401)
111 V, b = update_stats(X, y, 1e-2, 8)
112 mu, sd = predict(V, b, grid, sigma, beta=2.5, K=8)
113 baseline_rmse = float(np.sqrt(np.mean((true_fn(grid)-mu)**2)))
114 # deterministic ridge has identical mean but no uncertainty/rejection signal
115 feature_dim = 17
116 baseline_interval_width = 0.0
117 proposed_width = float(np.mean(2*2.5*sd))
118 unsafe_fraction = float(np.mean((mu - 2.5*sd > true_fn(grid)) | (mu + 2.5*sd < true_fn(grid))))
119 return {"deterministic_ridge_grid_rmse": baseline_rmse,
120 "certificate_grid_rmse": baseline_rmse,
121 "deterministic_interval_width": baseline_interval_width,
122 "certificate_mean_interval_width": proposed_width,
123 "certificate_grid_miss_fraction": unsafe_fraction,
124 "feature_dim": feature_dim}
125
126
127def main():
128 # Algebraic/numerical identity: rank-one updates and batch sufficient statistics.
129 rr = np.random.default_rng(SEED)
130 X = rr.uniform(-1, 1, 300); y = true_fn(X) + rr.normal(0, .12, len(X))
131 Vb, bb = update_stats(X, y)
132 Vo, bo = online_stats(X, y)
133 identity = {"max_abs_V_difference": float(np.max(np.abs(Vb-Vo))),
134 "max_abs_b_difference": float(np.max(np.abs(bb-bo))),
135 "max_abs_theta_difference": float(np.max(np.abs(np.linalg.solve(Vb,bb)-np.linalg.solve(Vo,bo))))}
136 result = {"seed": SEED, "identity_check": identity,
137 "contraction_sweep": contraction_sweep(),
138 "coverage_sweep": coverage_sweep(),
139 "baseline_comparison": baseline_comparison(),
140 "parameter_scaling_sweeps": parameter_scaling_sweeps()}
141 Path("results.json").write_text(json.dumps(result, indent=2))
142 print(json.dumps(result, indent=2))
143
144def parameter_scaling_sweeps():
145 rr = np.random.default_rng(SEED + 77)
146 X = rr.uniform(-1, 1, 240)
147 y0 = true_fn(X) + rr.normal(0, 1, len(X))
148 # Prediction: posterior s is exactly linear in the supplied noise scale sigma.
149 V, b = update_stats(X, y0)
150 sigmas = np.array([0.03, 0.06, 0.12, 0.24])
151 mean_s = []
152 for sig in sigmas:
153 mean_s.append(float(np.mean(predict(V, b, X, sig, K=8)[1])))
154 sigma_ratio = np.array(mean_s) / mean_s[2]
155 expected_sigma_ratio = sigmas / sigmas[2]
156 # Prediction: increasing lambda increases uncertainty, while reducing data fit strength.
157 lambdas = np.array([1e-3, 1e-2, 1e-1, 1.0])
158 grid = np.linspace(-1, 1, 201)
159 lambda_s = []
160 lambda_rmse = []
161 for lam in lambdas:
162 Vl, bl = update_stats(X, true_fn(X) + rr.normal(0, .12, len(X)), lam)
163 mu, sd = predict(Vl, bl, grid, .12, K=8)
164 lambda_s.append(float(np.mean(sd)))
165 lambda_rmse.append(float(np.sqrt(np.mean((mu-true_fn(grid))**2))))
166 return {"sigma": sigmas.tolist(), "mean_sd": mean_s,
167 "observed_sigma_ratios": sigma_ratio.tolist(),
168 "predicted_sigma_ratios": expected_sigma_ratio.tolist(),
169 "lambda": lambdas.tolist(), "mean_sd_by_lambda": lambda_s,
170 "grid_rmse_by_lambda": lambda_rmse}
171
172
173if __name__ == "__main__":
174 main()