1import json, math, time
  2import numpy as np
  3
  4SEED = 1414
  5rng = np.random.default_rng(SEED)
  6
  7
  8def clip(v, bound):
  9    n = np.linalg.norm(v)
 10    return v if n <= bound else v * (bound / n)
 11
 12
 13def ar1(n, rho, rng, burn=300):
 14    x = np.zeros(n + burn)
 15    noise = rng.normal(size=n + burn)
 16    for i in range(1, len(x)):
 17        x[i] = rho * x[i-1] + math.sqrt(1-rho*rho) * noise[i]
 18    return x[burn:]
 19
 20
 21def means(z, b):
 22    return np.array([z[i:i+b].mean() for i in range(0, len(z)-b+1, b)])
 23
 24
 25def correction_samples(rho, level, trials=6000, b0=4):
 26    b = b0 * (2**level)
 27    # Vectorized independent trajectories; each row is one coupled block.
 28    e = rng.normal(size=(trials, b))
 29    z = np.zeros_like(e)
 30    z[:, 0] = e[:, 0]
 31    scale = math.sqrt(1-rho*rho)
 32    for i in range(1, b):
 33        z[:, i] = rho*z[:, i-1] + scale*e[:, i]
 34    vals = z.mean(axis=1) - z[:, :b//2].mean(axis=1)
 35    return vals
 36
 37
 38def telescoping_check():
 39    # For p_l proportional to 2^-l, E[g0 + Delta_l/p_l] = E[g_K].
 40    K, b0, rho, trials = 4, 4, .8, 30000
 41    p = np.array([2.0**(-l) for l in range(K+1)])
 42    p /= p.sum()
 43    est, fine = [], []
 44    for _ in range(trials):
 45        z = ar1(b0*2**K, rho, rng)
 46        gs = [z[:b0*2**l].mean() for l in range(K+1)]
 47        l = rng.choice(K+1, p=p)
 48        est.append(gs[0] if l == 0 else gs[0] + (gs[l]-gs[l-1])/p[l])
 49        fine.append(gs[K])
 50    est, fine = np.array(est), np.array(fine)
 51    diff = est.mean() - fine.mean()
 52    se = math.sqrt(est.var()/trials + fine.var()/trials)
 53    return {"observed_difference": float(diff), "two_se": float(2*se),
 54            "passes_2se": bool(abs(diff) <= 2*se)}
 55
 56
 57def math_sweeps():
 58    # Prediction A: correction variance is proportional to AR long-run variance
 59    # (1+rho)/(1-rho), and shrinks with level approximately as 1/b.
 60    rows = []
 61    for rho in [0.0, .5, .8, .95]:
 62        vals = correction_samples(rho, 3)
 63        rows.append({"rho": rho, "var_delta_level3": float(vals.var()),
 64                     "predicted_dependence_multiplier": (1+rho)/(1-rho)})
 65    # Compare observed ratios against the iid baseline.
 66    base = rows[0]["var_delta_level3"]
 67    for r in rows: r["observed_ratio_to_iid"] = r["var_delta_level3"] / base
 68    levels = []
 69    for level in range(5):
 70        v = correction_samples(.8, level, b0=32).var()
 71        levels.append({"level": level, "block": 32*2**level, "variance": float(v),
 72                       "variance_times_block": float(v*(4*2**level))})
 73    # Fixed-large-block prediction: Var(mean) ratio tends to (1+rho)/(1-rho).
 74    dep_fixed = []
 75    bfix, trials_fix = 512, 4000
 76    for rho in [0.0, .5, .8, .95]:
 77        e = rng.normal(size=(trials_fix, bfix)); z = np.zeros_like(e); z[:, 0] = e[:, 0]
 78        q = math.sqrt(1-rho*rho)
 79        for i in range(1, bfix): z[:, i] = rho*z[:, i-1] + q*e[:, i]
 80        dep_fixed.append({"rho": rho, "block": bfix,
 81                          "observed_mean_variance_ratio": float(z.mean(1).var() / dep_fixed[0]["_var"] if dep_fixed else 1.0),
 82                          "predicted_ratio": (1+rho)/(1-rho), "_var": float(z.mean(1).var())})
 83    for r in dep_fixed: r.pop("_var", None)
 84    # Prediction B: clipping enforces the pathwise norm bound exactly.
 85    max_norm = 0.0; clipped = 0
 86    for _ in range(10000):
 87        v = rng.normal(size=7)*10
 88        c = clip(v, 1.0)
 89        max_norm = max(max_norm, np.linalg.norm(c))
 90        clipped += np.linalg.norm(v) > 1
 91    return {"dependence": rows, "fixed_large_block_dependence": dep_fixed, "level_scaling_rho_0.8": levels,
 92            "clipping": {"bound": 1.0, "max_observed_norm": float(max_norm),
 93                          "clipping_fraction": clipped/10000}}
 94
 95
 96def grad(theta, z):
 97    # Two-parameter linear neural model, squared loss gradient, ordered stream.
 98    x, y = z
 99    e = theta[0]*x + theta[1] - y
100    return np.array([e*x, e])
101
102
103def stream_data(n, rho, rng):
104    x = ar1(n, rho, rng)
105    # stationary regression target with independent observation noise
106    y = 1.7*x - .35 + .1*rng.normal(size=n)
107    return np.stack([x, y], axis=1)
108
109
110def train(method, rho, seed=1414, steps=350, b0=4, K=3, lr=.08):
111    rg = np.random.default_rng(seed)
112    # Equal expected sample budget: baseline consumes b0*2^K, MLMC expected cost
113    # is b0 * sum p_l 2^l, and we use a fixed reserved fine block for simplicity.
114    data = stream_data(steps*b0*2**K + 20, 0.0 if method == "iid" else rho, rg)
115    th = np.array([0., 0.]); losses=[]; cursor=0; gradseq=[]; clip_count=0
116    p = np.array([2.**(-l) for l in range(K+1)]); p /= p.sum()
117    for _ in range(steps):
118        if method == 'baseline':
119            block = data[cursor:cursor+b0*2**K]; cursor += b0*2**K
120            gh = np.mean([grad(th,z) for z in block], axis=0)
121        elif method == 'coupled':
122            block = data[cursor:cursor+b0*2**K]; cursor += b0*2**K
123            gs=[]
124            for l in range(K+1):
125                b=b0*2**l; gs.append(np.mean([grad(th,z) for z in block[:b]],axis=0))
126            l=rg.choice(K+1,p=p)
127            gh=gs[0] if l==0 else gs[0]+(gs[l]-gs[l-1])/p[l]
128        elif method == 'iid':
129            block = data[cursor:cursor+b0*2**K].copy(); cursor += b0*2**K
130            rg.shuffle(block)
131            gh=np.mean([grad(th,z) for z in block],axis=0)
132        n=np.linalg.norm(gh)
133        if n>4: gh=gh*4/n; clip_count += 1
134        th -= lr*gh
135        gradseq.append(gh)
136        # validation risk on population proxy
137        losses.append((th[0]-1.7)**2 + (th[1]+.35)**2)
138    a=np.array(gradseq); ac=np.corrcoef(a[:-1,0],a[1:,0])[0,1]
139    return {"final_parameter_error":float(losses[-1]), "loss_at_100":float(losses[99]),
140            "loss_at_350":float(losses[-1]), "gradient_lag1_autocorr":float(ac),
141            "clip_fraction":clip_count/steps, "samples":int(cursor)}
142
143
144def main():
145    out={"seed":SEED, "telescoping":telescoping_check(), "math_sweeps":math_sweeps(),
146         "benchmark":{}}
147    for rho in [0.0,.8,.95]:
148        out["benchmark"][str(rho)]={m:train(m,rho) for m in ["baseline","coupled","iid"]}
149    print(json.dumps(out, indent=2))
150
151if __name__ == '__main__': main()