1import json, time
  2import numpy as np
  3
  4SEED = 2052
  5rng = np.random.default_rng(SEED)
  6
  7
  8def softmax(a):
  9    a = a - np.max(a, axis=-1, keepdims=True)
 10    z = np.exp(a)
 11    return z / z.sum(axis=-1, keepdims=True)
 12
 13
 14def make_frozen(d, dh, m, seed):
 15    r = np.random.default_rng(seed)
 16    # Variance-normalized frozen Gaussian maps.
 17    Q = r.normal(0, 1 / np.sqrt(dh), size=(dh, dh))
 18    K = r.normal(0, 1 / np.sqrt(dh), size=(dh, dh))
 19    V = r.normal(0, 1 / np.sqrt(dh), size=(m, dh))
 20    A = (Q.T @ K)[: d + 1] / np.sqrt(dh)
 21    return Q, K, V, A
 22
 23
 24def targets(Xs, sigma):
 25    # b_i(x), modulo the query-only term -||x||^2/(2 sigma^2).
 26    return np.column_stack((Xs / sigma**2, -np.sum(Xs * Xs, axis=1) / (2 * sigma**2)))
 27
 28
 29def construct(Xs, ys, dh, sigma, ridge=0.0, seed=0):
 30    d = Xs.shape[1]
 31    m = ys.shape[1]
 32    _, _, V, A = make_frozen(d, dh, m, seed)
 33    C = targets(Xs, sigma)
 34    B = np.vstack((A, V))
 35    rhs = np.hstack((C, ys)).T  # (d+1+m, N)
 36    if ridge == 0:
 37        P = np.linalg.lstsq(B, rhs, rcond=None)[0]
 38    else:
 39        P = np.linalg.solve(B.T @ B + ridge * np.eye(dh), B.T @ rhs)
 40    return P.T, A, V, C
 41
 42
 43def predict(Xq, P, A, V, d):
 44    E = np.column_stack((Xq, np.ones(len(Xq)), np.zeros((len(Xq), A.shape[1] - d - 1))))
 45    logits = E[:, : d + 1] @ (A @ P.T)
 46    W = softmax(logits)
 47    vals = (V @ P.T).T
 48    return W @ vals, W, logits
 49
 50
 51def nw(Xq, Xs, ys, sigma):
 52    b = -((Xq[:, None, :] - Xs[None, :, :]) ** 2).sum(axis=2) / (2 * sigma**2)
 53    return softmax(b) @ ys, softmax(b)
 54
 55
 56def knn(Xq, Xs, ys, k=3):
 57    dist = ((Xq[:, None, :] - Xs[None, :, :]) ** 2).sum(axis=2)
 58    ix = np.argpartition(dist, k - 1, axis=1)[:, :k]
 59    return ys[ix].mean(axis=1)
 60
 61
 62def holder(x):
 63    return (np.sin(2.3 * x[:, :1]) + 0.6 * np.cos(3.1 * x[:, 1:2]) +
 64            0.25 * np.sin(5 * (x[:, :1] + x[:, 1:2])))
 65
 66
 67def math_verification():
 68    d, m, N, sigma = 2, 1, 24, 0.35
 69    Xs = rng.uniform(-1, 1, (N, d)); ys = holder(Xs)
 70    rows = []
 71    # Prediction 1: generic full row rank appears at dh >= d+1+m = 4.
 72    for dh in [3, 4, 5, 8, 16]:
 73        P, A, V, C = construct(Xs, ys, dh, sigma, 0, 11)
 74        B = np.vstack((A, V)); rhs = np.vstack((C.T, ys.T))
 75        residual = np.linalg.norm((B @ P.T - rhs).T) / np.sqrt(N)
 76        rank = int(np.linalg.matrix_rank(B, tol=1e-10))
 77        rows.append({'dh': dh, 'joint_rank': rank, 'joint_rows': d+1+m,
 78                     'rms_joint_residual': float(residual),
 79                     'mean_prompt_norm': float(np.linalg.norm(P, axis=1).mean())})
 80    # Prediction 2: ridge regularization trades residual for prompt norm monotonically.
 81    ridge_rows = []
 82    for lam in [0, 1e-5, 1e-3, 1e-1, 1.0]:
 83        P, A, V, C = construct(Xs, ys, 16, sigma, lam, 11)
 84        B = np.vstack((A, V)); rhs = np.vstack((C.T, ys.T))
 85        ridge_rows.append({'lambda': lam,
 86            'rms_residual': float(np.linalg.norm((B @ P.T-rhs).T)/np.sqrt(N)),
 87            'mean_prompt_norm': float(np.linalg.norm(P,axis=1).mean())})
 88    # Prediction 3: target coefficient and prompt norm scale as sigma^-2, while
 89    # normalized kernel agreement remains exact in the full-rank regime.
 90    sigma_rows = []
 91    Xq = rng.uniform(-1, 1, (100, d))
 92    for s in [0.25, 0.35, 0.5, 0.7]:
 93        P, A, V, C = construct(Xs, ys, 16, s, 0, 11)
 94        pred, W, logits = predict(Xq, P, A, V, d)
 95        ref, Wref = nw(Xq, Xs, ys, s)
 96        sigma_rows.append({'sigma': s, 'mean_prompt_norm': float(np.linalg.norm(P,axis=1).mean()),
 97                           'weight_max_abs_error': float(np.max(np.abs(W-Wref))),
 98                           'prediction_mse_to_nw': float(np.mean((pred-ref)**2))})
 99    return {'rank_prediction': rows, 'ridge_prediction': ridge_rows,
100            'bandwidth_prediction': sigma_rows}
101
102
103def mini_experiment():
104    d, m, dh, sigma = 2, 1, 16, 0.38
105    r = np.random.default_rng(77)
106    Xtest = r.uniform(-1, 1, (1000, d)); ytest = holder(Xtest)
107    result = []
108    for N in [32, 64, 128, 256]:
109        Xs = r.uniform(-1, 1, (N, d)); ys = holder(Xs)
110        t0 = time.perf_counter()
111        P, A, V, C = construct(Xs, ys, dh, sigma, 0, 19)
112        adaptation_ms = (time.perf_counter()-t0)*1000
113        pred, W, logits = predict(Xtest, P, A, V, d)
114        ref, Wref = nw(Xtest, Xs, ys, sigma)
115        result.append({'N': N, 'analytic_mse': float(np.mean((pred[:,0]-ytest[:,0])**2)),
116          'nw_mse': float(np.mean((ref[:,0]-ytest[:,0])**2)),
117          'knn3_mse': float(np.mean((knn(Xtest,Xs,ys)[:,0]-ytest[:,0])**2)),
118          'attention_weight_error': float(np.max(np.abs(W-Wref))),
119          'logit_rms_residual_mod_common': float(np.sqrt(np.mean((logits - (Xtest @ Xs.T / sigma**2 - (Xs**2).sum(axis=1)[None,:] / (2*sigma**2)))**2))),
120          'prompt_memory_floats': int(P.size), 'adaptation_ms': adaptation_ms})
121    return result
122
123
124def main():
125    out = {'seed': SEED, 'math_verification': math_verification(),
126           'mini_experiment': mini_experiment()}
127    with open('results.json', 'w') as f: json.dump(out, f, indent=2)
128    print(json.dumps(out, indent=2))
129
130if __name__ == '__main__':
131    main()