code/experiments/exp_table23_stats.py

"""Claim 5 — statistical audit of the paper's reported Tables 2-3.

No new measurements: this checks the INTERNAL consistency of the paper's own
"FlashOptim matches reference" claim given its reported mean ± std over 3 seeds.
For each reference/FlashOptim pair we report:
  z_spread = |d| / sqrt(s1^2 + s2^2)      -- difference vs run-to-run spread
  z_sem    = |d| / sqrt((s1^2 + s2^2)/3)  -- difference vs standard error (n=3)
Flags: degradation (Flash worse) significant at |z_sem|>2, or improvement.
"""

import json
import math
import os

# (task, metric, ref_mean, ref_std, flash_mean, flash_std, higher_is_better)
ROWS = [
    ("ImageNet ResNet-50", "Top-1 SGD", 77.01, 0.02, 77.16, 0.09, True),
    ("ImageNet ResNet-50", "Top-1 AdamW", 75.51, 0.09, 75.67, 0.04, True),
    ("Llama-3.1-8B OpenMathInstruct-2", "GSM8K AdamW", 75.09, 0.40, 74.98, 0.77, True),
    ("GPT-2 124M pretrain", "Val loss AdamW", 3.263, 0.001, 3.265, 0.001, False),
    ("GPT-2 124M pretrain", "HellaSwag AdamW", 31.9, 0.2, 31.9, 0.5, True),
    ("GPT-2 124M pretrain", "ARC-E AdamW", 39.6, 0.7, 39.5, 0.9, True),
    ("GPT-2 124M pretrain", "CSQA AdamW", 25.9, 4.1, 30.8, 2.1, True),
    ("GPT-2 124M pretrain", "PIQA AdamW", 64.3, 0.0, 64.5, 0.3, True),
    ("GPT-2 124M pretrain", "LAMBADA AdamW", 31.0, 1.2, 31.9, 0.7, True),
    ("GPT-2 124M pretrain", "Winograd AdamW", 57.3, 0.6, 59.1, 1.1, True),
    ("GPT-2 124M pretrain", "BoolQ AdamW", 58.1, 3.6, 57.2, 4.8, True),
    ("GPT-2 124M pretrain", "Mean ICL AdamW", 44.0, 0.4, 45.0, 1.0, True),
    ("GPT-2 124M pretrain", "Val loss Lion", 3.240, 0.002, 3.240, 0.001, False),
    ("GPT-2 124M pretrain", "HellaSwag Lion", 32.3, 0.0, 32.4, 0.3, True),
    ("GPT-2 124M pretrain", "ARC-E Lion", 40.0, 0.5, 40.8, 0.2, True),
    ("GPT-2 124M pretrain", "CSQA Lion", 23.3, 1.8, 25.4, 2.9, True),
    ("GPT-2 124M pretrain", "PIQA Lion", 63.8, 1.0, 64.2, 0.5, True),
    ("GPT-2 124M pretrain", "LAMBADA Lion", 31.5, 0.2, 31.6, 0.5, True),
    ("GPT-2 124M pretrain", "Winograd Lion", 58.9, 2.0, 59.1, 2.4, True),
    ("GPT-2 124M pretrain", "BoolQ Lion", 58.1, 2.4, 59.1, 2.3, True),
    ("GPT-2 124M pretrain", "Mean ICL Lion", 44.0, 0.5, 44.7, 0.5, True),
]

out = []
for task, metric, m1, s1, m2, s2, hib in ROWS:
    d = m2 - m1
    sp = math.sqrt(s1 * s1 + s2 * s2)
    z_spread = abs(d) / sp if sp else float("inf")
    z_sem = abs(d) / (sp / math.sqrt(3)) if sp else float("inf")
    flash_worse = (d < 0) if hib else (d > 0)
    out.append(dict(task=task, metric=metric, ref=f"{m1}±{s1}", flash=f"{m2}±{s2}",
                    delta=round(d, 3), z_spread=round(z_spread, 2),
                    z_sem=round(z_sem, 2), flash_worse=flash_worse,
                    sig_degradation=bool(flash_worse and z_sem > 2),
                    sig_improvement=bool((not flash_worse) and d != 0 and z_sem > 2)))

n_deg = sum(r["sig_degradation"] for r in out)
n_imp = sum(r["sig_improvement"] for r in out)
summary = {"rows": out,
           "n_metrics": len(out),
           "n_significant_degradations_z_sem_gt2": n_deg,
           "n_significant_improvements_z_sem_gt2": n_imp,
           "caveats": "n=3 seeds; z_sem assumes independent seeds and normality; "
                      "PIQA/HellaSwag-Lion have std=0.0 in one arm (excluded from inf-z by pairing); "
                      "this audits the paper's own numbers, it is NOT new measurement."}
os.makedirs("repro_flashoptim/outputs", exist_ok=True)
with open("repro_flashoptim/outputs/table23_stats.json", "w") as f:
    json.dump(summary, f, indent=2)
print(json.dumps(summary, indent=2))
print(f"\nSignificant degradations (z_sem>2): {n_deg}/{len(out)}; "
      f"significant improvements: {n_imp}/{len(out)}")
for r in out:
    if r["sig_degradation"] or r["sig_improvement"]:
        print("  ", r["metric"], r["ref"], "->", r["flash"], "z_sem", r["z_sem"])