repro-bundle-v0/experiments/claim5_comm_efficient_floor.py

"""Claim 5 (Theorem 5): the communication-efficient distributed Lion variant
(1-bit unbiased sign compression both directions, Q1=S_G, Q2=sign, Algorithm
3 v1) converges as O(d^{1/2} T^{-1/2} + d n^{-1/2}) -- i.e. it decreases like
a plain rate for a while, then hits a HARD FLOOR of order d/sqrt(n) that more
iterations cannot shrink. This is qualitatively different from Claims 1-4
(which decay to zero as T -> infinity) and is the paper's own motivation for
introducing the further-improved Theorem 6/7 variants (Claim 6).

Two checks:
  (a) T-sweep at fixed (n,d): metric should decrease then plateau near the
      floor -- NOT continue shrinking like Claims 1-4.
  (b) n-sweep at large fixed T (deep in the floor regime): the plateau LEVEL
      should scale as n^{-1/2} (at fixed d).
"""
import sys, os, json, csv, time
sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..', 'src'))

import numpy as np
from objective import make_heterogeneous_nodes, empirical_grad_bound
import lion, rates
from analysis import loglog_slope, mean_sem, plot_scaling, plot_collapse

RESULTS = os.path.join(os.path.dirname(__file__), '..', 'results')
os.makedirs(RESULTS, exist_ok=True)

N_PER_NODE = 200
BATCH = 32
SEEDS = list(range(8))
D_FIXED = 20
N_FIXED = 8

T_GRID = [200, 1000, 5000, 25000, 75000, 150000]

N_GRID = [2, 4, 8, 16, 32, 64]
SEEDS_N = list(range(6))
T_LARGE = 60000  # deep in the floor regime for all n in N_GRID

t0 = time.time()
rows = []


def run_point(n, d, T, seed):
    nodes = make_heterogeneous_nodes(n, N_PER_NODE, d, seed=seed)
    G = empirical_grad_bound(*nodes, x_radius=2 * np.sqrt(d), seed=seed)
    hp = rates.thm5_comm_eff(T, n, d, regime='decay')
    val = lion.run_comm_efficient(nodes, T, hp['eta'], hp['beta1'], hp['beta2'], hp['lam'],
                                   BATCH, 'v1', seed=seed, G=G, q2_op='sign')
    return val


print("(a) T-sweep at n=%d, d=%d -- expect decrease then plateau" % (N_FIXED, D_FIXED))
for T in T_GRID:
    for s in SEEDS:
        val = run_point(N_FIXED, D_FIXED, T, s)
        rows.append(dict(sweep='T', T=T, n=N_FIXED, d=D_FIXED, seed=s, metric=val))
    print(f"  T={T:>7}  done  ({time.time()-t0:.1f}s elapsed)", flush=True)

print("(b) n-sweep at T=%d (floor regime), d=%d" % (T_LARGE, D_FIXED))
for n in N_GRID:
    for s in SEEDS_N:
        val = run_point(n, D_FIXED, T_LARGE, s)
        rows.append(dict(sweep='n', T=T_LARGE, n=n, d=D_FIXED, seed=s, metric=val))
    print(f"  n={n:>3}  done  ({time.time()-t0:.1f}s elapsed)", flush=True)

csv_path = os.path.join(RESULTS, 'claim5_raw.csv')
with open(csv_path, 'w', newline='') as f:
    w = csv.DictWriter(f, fieldnames=['sweep', 'T', 'n', 'd', 'seed', 'metric'])
    w.writeheader()
    w.writerows(rows)

t_runs = {}
for r in rows:
    if r['sweep'] == 'T':
        t_runs.setdefault(r['T'], []).append(r['metric'])
xs, means, sems = mean_sem(t_runs)
# fit slope on ONLY the early/decaying portion vs the whole range, to quantify the plateau
t_slope_all, t_intercept_all, t_r2_all = loglog_slope(xs, means)
half = len(xs) // 2
t_slope_late, _, t_r2_late = loglog_slope(xs[half:], means[half:])
t_slope_early, _, t_r2_early = loglog_slope(xs[:half + 1], means[:half + 1])

plot_scaling(xs, means, sems, t_slope_all, t_intercept_all, 'T (iterations)',
             f'Claim 5: T-scaling at n={N_FIXED}, d={D_FIXED} (floor test)',
             os.path.join(RESULTS, 'claim5_T_floor.png'),
             claimed_slope=-0.5)

n_runs = {}
for r in rows:
    if r['sweep'] == 'n':
        n_runs.setdefault(r['n'], []).append(r['metric'])
xs_n, means_n, sems_n = mean_sem(n_runs)
n_slope, n_intercept, n_r2 = loglog_slope(xs_n, means_n)
plot_scaling(xs_n, means_n, sems_n, n_slope, n_intercept, 'n (nodes)',
             f'Claim 5: floor level vs n at T={T_LARGE}, d={D_FIXED}',
             os.path.join(RESULTS, 'claim5_n_floor.png'),
             claimed_slope=-0.5)

# collapse the floor-regime n-sweep against d * n^-0.5
theory_x = np.array([r['d'] * r['n'] ** -0.5 for r in rows if r['sweep'] == 'n'])
measured_y = np.array([r['metric'] for r in rows if r['sweep'] == 'n'])
collapse_slope, collapse_intercept, collapse_r2 = plot_collapse(
    theory_x, measured_y, os.path.join(RESULTS, 'claim5_floor_collapse.png'),
    'theory: d n^-0.5 (floor term)', f'Claim 5: floor-level collapse at T={T_LARGE}')

summary = dict(
    claim=5, theorem='Theorem 5',
    plateau_detected=bool(abs(t_slope_late) < 0.6 * abs(t_slope_early)),
    T_slope_early_half=t_slope_early, T_slope_late_half=t_slope_late,
    T_slope_whole_range=t_slope_all,
    claimed_floor_n_exponent=-0.5, measured_floor_n_slope=n_slope, floor_n_r2=n_r2,
    floor_collapse_slope=collapse_slope, floor_collapse_r2=collapse_r2,
    n_runs=len(rows), wall_seconds=time.time() - t0,
)
with open(os.path.join(RESULTS, 'claim5_summary.json'), 'w') as f:
    json.dump(summary, f, indent=2)

print(json.dumps(summary, indent=2))
print(f"\nTotal wall time: {time.time()-t0:.1f}s")