1import json
  2import math
  3import random
  4import time
  5import numpy as np
  6
  7SEED = 2261
  8np.random.seed(SEED)
  9random.seed(SEED)
 10
 11
 12def projection(d, r, rng):
 13    return rng.normal(0.0, 1.0 / math.sqrt(r), size=(r, d)).astype(np.float64)
 14
 15
 16def pairwise_ratios(X, P, max_pairs=200000, rng=None):
 17    n = len(X)
 18    rng = np.random.default_rng(SEED) if rng is None else rng
 19    total = n * (n - 1) // 2
 20    if total <= max_pairs:
 21        ii, jj = np.triu_indices(n, 1)
 22    else:
 23        ii = rng.integers(0, n, size=max_pairs)
 24        jj = rng.integers(0, n, size=max_pairs)
 25        keep = ii != jj
 26        ii, jj = ii[keep], jj[keep]
 27    dx = X[ii] - X[jj]
 28    dz = (X[ii] - X[jj]) @ P.T
 29    return np.linalg.norm(dz, axis=1) / np.maximum(np.linalg.norm(dx, axis=1), 1e-12)
 30
 31
 32def width_rule(d, n, eps, C=1.0):
 33    raw = min(d, n - 1, math.log(2.0 + eps * eps * n) / (eps * eps))
 34    return max(1, min(d, int(math.ceil(C * raw)))), raw
 35
 36
 37def toy_math_verification():
 38    rng = np.random.default_rng(SEED + 1)
 39    d, n = 128, 180
 40    X = rng.normal(size=(n, d))
 41    # Prediction 1: Gaussian scaling is unbiased: E[||Pv||^2/||v||^2] = 1.
 42    v = rng.normal(size=d)
 43    norm_trials = []
 44    for _ in range(300):
 45        P = projection(d, 32, rng)
 46        norm_trials.append(np.sum((P @ v) ** 2) / np.sum(v ** 2))
 47    norm_mean = float(np.mean(norm_trials))
 48    # Prediction 2: for a fixed pair, ratio has SD approximately 1/sqrt(2r).
 49    fixed_v = X[0] - X[1]
 50    fixed_ratios = []
 51    for _ in range(500):
 52        P = projection(d, 32, rng)
 53        fixed_ratios.append(np.linalg.norm(P @ fixed_v) / np.linalg.norm(fixed_v))
 54    fixed_sd = float(np.std(fixed_ratios, ddof=1))
 55    fixed_sd_pred = 1.0 / math.sqrt(2 * 32)
 56    # Prediction 3: maximum of many pair errors scales as sqrt(log(number pairs)/r).
 57    max_rows = []
 58    for r in [8, 16, 32, 64, 128]:
 59        P = projection(d, r, rng)
 60        ratios = pairwise_ratios(X, P, rng=rng)
 61        max_rows.append({
 62            "r": r,
 63            "max_abs_error": float(np.max(np.abs(ratios - 1))),
 64            "p95_abs_error": float(np.percentile(np.abs(ratios - 1), 95)),
 65            "predicted_max_scale": math.sqrt(2.0 * math.log(max(2, n * (n - 1) / 2)) / r),
 66            "predicted_fixed_sd": 1.0 / math.sqrt(2 * r),
 67        })
 68    # Prediction 4: sharp schedule transitions with n and saturates at d or n-1.
 69    schedule = []
 70    for eps in [0.1, 0.2, 0.3]:
 71        for nn in [8, 32, 128, 512]:
 72            rr, raw = width_rule(128, nn, eps, C=1)
 73            schedule.append({"epsilon": eps, "n": nn, "r": rr, "raw": raw})
 74    return {
 75        "unbiased_norm": {"observed_mean": norm_mean, "predicted": 1.0, "relative_error": abs(norm_mean - 1.0)},
 76        "fixed_pair_sd": {"observed": fixed_sd, "predicted": fixed_sd_pred, "ratio_observed_to_predicted": fixed_sd / fixed_sd_pred},
 77        "max_error_scaling": max_rows,
 78        "sharp_width_schedule": schedule,
 79    }
 80
 81
 82def softmax_loss(W, X, y):
 83    logits = X @ W
 84    logits -= logits.max(axis=1, keepdims=True)
 85    exp = np.exp(logits)
 86    probs = exp / exp.sum(axis=1, keepdims=True)
 87    loss = -np.log(np.maximum(probs[np.arange(len(y)), y], 1e-12)).mean()
 88    grad = X.T @ (probs - np.eye(W.shape[1])[y]) / len(y)
 89    return float(loss), grad
 90
 91
 92def train_linear(Xtr, ytr, Xva, yva, steps=250, lr=0.4):
 93    W = np.zeros((Xtr.shape[1], 2), dtype=np.float64)
 94    for _ in range(steps):
 95        _, g = softmax_loss(W, Xtr, ytr)
 96        W -= lr * g
 97    loss, _ = softmax_loss(W, Xva, yva)
 98    acc = float(np.mean(np.argmax(Xva @ W, axis=1) == yva))
 99    return loss, acc
100
101
102def downstream_experiment():
103    rng = np.random.default_rng(SEED + 2)
104    ntr, nva, d, n_tokens = 2400, 800, 96, 24
105    # Two classes differ in a low-dimensional token mean; unrelated token noise
106    # makes the full-width representation deliberately width-sensitive.
107    ytr = rng.integers(0, 2, ntr)
108    yva = rng.integers(0, 2, nva)
109    def make(y):
110        x = rng.normal(size=(len(y), n_tokens, d))
111        x[:, :, :8] += (2 * y[:, None, None] - 1) * 0.9
112        return x
113    Xtr0, Xva0 = make(ytr), make(yva)
114    full_tr, full_va = Xtr0.mean(axis=1), Xva0.mean(axis=1)
115    full_loss, full_acc = train_linear(full_tr, ytr, full_va, yva)
116    rows = []
117    for eps in [0.1, 0.2, 0.3]:
118        for C in [1, 2, 4]:
119            r, raw = width_rule(d, n_tokens, eps, C)
120            P = projection(d, r, rng)
121            ztr, zva = Xtr0 @ P.T, Xva0 @ P.T
122            loss, acc = train_linear(ztr.mean(axis=1), ytr, zva.mean(axis=1), yva)
123            ratios = pairwise_ratios(Xva0.reshape(-1, d)[:160], P, rng=rng)
124            rows.append({"epsilon": eps, "C": C, "r": r, "raw": raw, "val_loss": loss, "val_accuracy": acc,
125                         "relative_accuracy_loss": (full_acc - acc) / max(full_acc, 1e-12),
126                         "p95_abs_distortion": float(np.percentile(np.abs(ratios - 1), 95)),
127                         "activation_width_ratio": r / d})
128    return {"full_width": {"d": d, "val_loss": full_loss, "val_accuracy": full_acc, "activation_width": d}, "compressed": rows}
129
130
131def main():
132    t0 = time.time()
133    out = {"seed": SEED, "math": toy_math_verification(), "downstream": downstream_experiment()}
134    out["runtime_sec"] = time.time() - t0
135    with open("results.json", "w") as f:
136        json.dump(out, f, indent=2)
137    print(json.dumps(out, indent=2))
138
139
140if __name__ == "__main__":
141    main()