Uniform-Certificate Bayesian Feature Head / experiment.py

✓✓ Beats tuned baseline

Raw ⬇ ZIP
  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()