1import json
  2import math
  3from pathlib import Path
  4import numpy as np
  5
  6
  7def tail_coefficients(eta):
  8    eta = np.asarray(eta, dtype=float)
  9    tau = len(eta)
 10    c = np.empty(tau)
 11    for t in range(tau - 1):
 12        c[t] = eta[t] ** 2 / (2.0 * np.sum(eta[t + 1:]))
 13    c[-1] = eta[-1] / 2.0
 14    return c
 15
 16
 17def continuous_batches(c, s, budget):
 18    r = np.sqrt(np.maximum(c * s, 1e-30))
 19    return budget * r / r.sum()
 20
 21
 22def integer_batches(weights, budget, bmin=1, bmax=None):
 23    """Allocate integer examples, preserving budget and bounds as closely as possible."""
 24    n = len(weights)
 25    if bmax is None:
 26        bmax = budget
 27    if budget < n * bmin or budget > n * bmax:
 28        raise ValueError("infeasible bounds")
 29    x = budget * np.asarray(weights, dtype=float) / np.sum(weights)
 30    x = np.clip(x, bmin, bmax)
 31    # Iterative largest-remainder allocation with bounds.
 32    b = np.floor(x).astype(int)
 33    b = np.maximum(b, bmin)
 34    b = np.minimum(b, bmax)
 35    while b.sum() < budget:
 36        candidates = np.where(b < bmax)[0]
 37        score = x[candidates] - b[candidates]
 38        b[candidates[np.argmax(score)]] += 1
 39    while b.sum() > budget:
 40        candidates = np.where(b > bmin)[0]
 41        score = b[candidates] - x[candidates]
 42        b[candidates[np.argmax(score)]] -= 1
 43    return b
 44
 45
 46def objective(b, c, s):
 47    return float(np.sum(np.asarray(c) * np.asarray(s) / np.asarray(b)))
 48
 49
 50def cosine_lr(tau, eta_max=0.12, eta_min=0.012):
 51    x = np.arange(tau) / max(1, tau - 1)
 52    return eta_min + 0.5 * (eta_max - eta_min) * (1 + np.cos(np.pi * x))
 53
 54
 55def verify_math(seed=7):
 56    rng = np.random.default_rng(seed)
 57    tau = 32
 58    eta = cosine_lr(tau)
 59    c = tail_coefficients(eta)
 60    s = np.exp(rng.normal(0, 0.8, tau))
 61    C = 32 * 16
 62    b_star = continuous_batches(c, s, C)
 63    b_int = integer_batches(np.sqrt(c * s), C, 1, 64)
 64    uniform = np.full(tau, C / tau)
 65    # Predictions: (P1) ratios match sqrt(c*s); (P2) continuous optimum has the
 66    # Cauchy-Schwarz value (sum sqrt(cs))^2/C; (P3) multiplying all noise by k
 67    # scales batches by sqrt(k), while relative allocations do not change.
 68    ratio_err = np.max(np.abs((b_star / b_star[0]) / (np.sqrt(c*s) / np.sqrt(c[0]*s[0])) - 1))
 69    predicted = np.sum(np.sqrt(c * s)) ** 2 / C
 70    opt_gap = objective(b_star, c, s) / predicted - 1
 71    scales = np.array([0.25, 1.0, 4.0, 16.0])
 72    scale_rows = []
 73    for k in scales:
 74        bk = continuous_batches(c, k*s, C)
 75        # Relative allocation is predicted invariant; because C is fixed, absolute
 76        # batches are also invariant. The meaningful scaling prediction is objective.
 77        scale_rows.append({"noise_multiplier": float(k),
 78                           "objective_ratio_observed": objective(bk, c, k*s) / objective(b_star, c, s),
 79                           "objective_ratio_predicted": float(k),
 80                           "allocation_relative_max_error": float(np.max(np.abs(bk/bk.sum()-b_star/b_star.sum())))})
 81    # Tail sweep: for constant noise, compare first/last coefficient under cosine
 82    # horizons. The reported quantity is directly predicted by c_t formula.
 83    tail_rows = []
 84    for T in [8, 16, 32, 64, 128]:
 85        et = cosine_lr(T)
 86        ct = tail_coefficients(et)
 87        bt = continuous_batches(ct, np.ones(T), T * 16)
 88        tail_rows.append({"horizon": T,
 89                          "c_first_over_c_last": float(ct[0]/ct[-1]),
 90                          "batch_first_over_last": float(bt[0]/bt[-1]),
 91                          "prediction_error": float(abs(bt[0]/bt[-1] - math.sqrt(ct[0]/ct[-1])))})
 92    return {"P1_ratio_max_abs_error": float(ratio_err),
 93            "P2_optimality_relative_error": float(opt_gap),
 94            "P3_noise_scale_sweep": scale_rows,
 95            "tail_weight_sweep": tail_rows,
 96            "integer_objective_over_continuous": objective(b_int,c,s)/objective(b_star,c,s),
 97            "eta": eta.tolist(), "c": c.tolist(), "s": s.tolist(),
 98            "continuous_batches": b_star.tolist(), "integer_batches": b_int.tolist()}
 99
100
101def run_sgd(seed, schedule, eta, s_profile, dim=8):
102    rng = np.random.default_rng(seed)
103    # Strongly convex quadratic; gradient noise has E||noise||^2 approximately s_t/B.
104    lam = 0.5
105    w = rng.normal(0, 1, dim)
106    target = np.zeros(dim)
107    losses = []
108    for t, b in enumerate(schedule):
109        true_g = lam * (w - target)
110        noise = rng.normal(0, math.sqrt(s_profile[t] / b / dim), dim)
111        w = w - eta[t] * (true_g + noise)
112        losses.append(0.5 * lam * float(np.dot(w, w)))
113    return np.asarray(losses)
114
115
116def mini_experiment(seed=19):
117    tau, per_step = 64, 16
118    C = tau * per_step
119    eta = cosine_lr(tau, 0.16, 0.016)
120    c = tail_coefficients(eta)
121    # Noise falls during optimization but remains deliberately heterogeneous.
122    s = 1.0 + 8.0 * np.exp(-np.arange(tau) / 18.0)
123    idea = integer_batches(np.sqrt(c*s), C, bmin=4, bmax=48)
124    static = np.full(tau, per_step, dtype=int)
125    # Common hand-designed comparator: linear growth, same total examples.
126    linear_weights = np.linspace(0.5, 1.5, tau)
127    linear = integer_batches(linear_weights, C, bmin=4, bmax=48)
128    all_results = {}
129    for name, sched in [("static", static), ("linear_growth", linear), ("tail_weighted", idea)]:
130        curves = np.array([run_sgd(seed+i, sched, eta, s) for i in range(40)])
131        all_results[name] = {"final_loss_mean": float(curves[:, -1].mean()),
132                             "final_loss_se": float(curves[:, -1].std(ddof=1)/math.sqrt(len(curves))),
133                             "loss_at_step_32": float(curves[:,31].mean()),
134                             "batch_min": int(sched.min()), "batch_max": int(sched.max()),
135                             "batch_first": int(sched[0]), "batch_last": int(sched[-1])}
136    all_results["theoretical_batch_objectives"] = {
137        "static": objective(static,c,s), "linear_growth": objective(linear,c,s), "tail_weighted": objective(idea,c,s)}
138    all_results["eta"] = eta.tolist()
139    all_results["c"] = c.tolist()
140    all_results["s_profile"] = s.tolist()
141    return all_results
142
143
144if __name__ == "__main__":
145    out = {"verification": verify_math(), "mini_experiment": mini_experiment()}
146    Path("results.json").write_text(json.dumps(out, indent=2))
147    print(json.dumps({"verification": out["verification"], "mini_experiment": out["mini_experiment"]}, indent=2))