| """LaTeX tables for the report from pulled metrics -> report/tables.tex (one \\newcommand per table).""" |
|
|
| import json |
| import sys |
| from collections import defaultdict |
| from pathlib import Path |
|
|
| import numpy as np |
|
|
| sys.path.insert(0, str(Path(__file__).parent)) |
| from analyze import EVAL_TASKS, LABELS, load_runs |
|
|
| ARMS = ("scratch", "ivgpt", "shuffled", "rescaled", "sp0.1", "sp0.2", "sp0.4", "sp0.7", "spshuf0.4") |
| NICE = {"scratch": "scratch", "ivgpt": "pretrained", "shuffled": "shuffled", "rescaled": "rescaled", |
| "sp0.1": r"S\&P $\lambda{=}0.1$", "sp0.2": r"S\&P $\lambda{=}0.2$", "sp0.4": r"S\&P $\lambda{=}0.4$", |
| "sp0.7": r"S\&P $\lambda{=}0.7$", "spshuf0.4": r"shuffled S\&P $\lambda{=}0.4$"} |
|
|
|
|
| def lr_tex(lr): |
| """Learning rates in the plain e-notation common in ML papers, e.g. 3e-3.""" |
| e = int(np.floor(np.log10(lr) + 1e-9)) |
| return f"{lr / 10 ** e:g}e-{abs(e)}" |
|
|
|
|
| def sweep_table(runs, budget, freeze): |
| """Rows = arms, cols = LRs, final val loss of seed-0 runs; the LR used for all further runs in bold.""" |
| grid = defaultdict(dict) |
| for r in runs: |
| if r["budget"] == budget and r["freeze"] == freeze and r["done"] and r["seed"] == 0: |
| grid[r["init"]][r["lr"]] = r["val"][-1] |
| lrs = sorted({lr for g in grid.values() for lr in g}) |
| if not lrs: |
| return "" |
| lines = [r"\begin{tabular}{l" + "c" * len(lrs) + "}", r"\toprule", |
| "initialization & " + " & ".join(lr_tex(lr) for lr in lrs) + r" \\", r"\midrule"] |
| for a in ARMS: |
| if a not in grid: |
| continue |
| chosen = chosen_lr(runs, budget, freeze, a) |
| cells = [] |
| for lr in lrs: |
| v = grid[a].get(lr) |
| cells.append("--" if v is None else (rf"\textbf{{{v:.3f}}}" if lr == chosen else f"{v:.3f}")) |
| lines.append(NICE[a] + " & " + " & ".join(cells) + r" \\") |
| lines += [r"\bottomrule", r"\end{tabular}"] |
| return "\n".join(lines) |
|
|
|
|
| def chosen_lr(runs, budget, freeze, init): |
| """The LR with the most seeds, then the lowest mean final loss over seeds (same rule as replicate_table).""" |
| groups = defaultdict(list) |
| for r in runs: |
| if r["budget"] == budget and r["freeze"] == freeze and r["init"] == init and r["done"]: |
| groups[r["lr"]].append(r["val"][-1]) |
| return min(groups, key=lambda l: (-len(groups[l]), np.mean(groups[l]))) if groups else None |
|
|
|
|
| def best_lr(runs, budget, freeze, init): |
| c = [(r["val"][-1], r["lr"]) for r in runs if r["budget"] == budget and r["freeze"] == freeze |
| and r["init"] == init and r["done"] and r["seed"] == 0] |
| return min(c)[1] if c else None |
|
|
|
|
| def replicate_table(runs, budget): |
| """Best-LR runs with all seeds: per-seed finals, mean, and the gap to scratch.""" |
| rows, ref = [], None |
| for a in ARMS: |
| groups = defaultdict(list) |
| for r in runs: |
| if r["budget"] == budget and r["freeze"] == "none" and r["init"] == a and r["done"]: |
| groups[r["lr"]].append((r["seed"], r["val"][-1])) |
| if not groups: |
| continue |
| |
| lr = min(groups, key=lambda l: (-len(groups[l]), np.mean([v for _, v in groups[l]]))) |
| reps = sorted(groups[lr]) |
| vals = [v for _, v in reps] |
| mean = float(np.mean(vals)) |
| if a == "scratch": |
| ref = mean |
| rows.append((a, lr, vals, mean)) |
| lines = [r"\begin{tabular}{llccc}", r"\toprule", r"initialization & chosen LR & seeds & mean & $\Delta$ vs scratch \\", r"\midrule"] |
| for a, lr, vals, mean in rows: |
| d = "--" if a == "scratch" or ref is None else f"{mean - ref:+.3f}" |
| lines.append(f"{NICE[a]} & {lr_tex(lr)} & {', '.join(f'{v:.3f}' for v in vals)} & {mean:.3f} & {d}" + r" \\") |
| lines += [r"\bottomrule", r"\end{tabular}"] |
| return "\n".join(lines), {a: {"lr": lr, "vals": vals, "mean": mean} for a, lr, vals, mean in rows} |
|
|
|
|
| def eval_table(runs): |
| """Main conditions at their replicated best LR, accuracies averaged over seeds.""" |
| rows, ses = [], [] |
| for budget in ("500M", "2000M"): |
| for a in ("scratch", "ivgpt", "sp0.4", "spshuf0.4"): |
| groups = defaultdict(list) |
| for r in runs: |
| if r["budget"] == budget and r["freeze"] == "none" and r["init"] == a and r["done"] and r["evals"]: |
| groups[r["lr"]].append(r) |
| if not groups: |
| continue |
| lr = min(groups, key=lambda l: (-len(groups[l]), np.mean([r["val"][-1] for r in groups[l]]))) |
| reps = groups[lr] |
| acc = {t: np.mean([r["evals"][t]["v"] for r in reps]) for t in EVAL_TASKS} |
| ses += [r["evals"][t]["se"] or 0 for r in reps for t in EVAL_TASKS] |
| rows.append((budget, a, lr, len(reps), acc)) |
| if not rows: |
| return "" |
| short = {"piqa": "PIQA", "hellaswag": "HellaSwag", "arc_easy": "ARC-E", "sciq": "SciQ", "winogrande": "Wino.", "lambada_openai": "LAMBADA"} |
| tasks = list(EVAL_TASKS) |
| lines = [r"\begin{tabular}{llc" + "c" * (len(tasks) + 1) + "}", r"\toprule", |
| "tokens & initialization & seeds & " + " & ".join(short[t] for t in tasks) + r" & mean \\", r"\midrule"] |
| last = None |
| for budget, a, lr, n, acc in rows: |
| if last and budget != last: |
| lines.append(r"\midrule") |
| tok = budget.replace("2000M", "2B") if budget != last else "" |
| last = budget |
| vals = [100 * acc[t] for t in tasks] |
| lines.append(f"{tok} & {NICE[a]} & {n} & " + " & ".join(f"{v:.1f}" for v in vals) + f" & {np.mean(vals):.1f}" + r" \\") |
| lines += [r"\bottomrule", r"\end{tabular}"] |
| return "\n".join(lines), float(np.mean(ses)) |
|
|
|
|
| def adapt_table(runs): |
| """Frozen backbones at 500M tokens: rows = backbone, cols = adaptation, best LR per cell (seed 0).""" |
| adapts = ("body", "lora", "top4") |
| cell = {} |
| for b in ("ivgpt", "scratch", "shuffled"): |
| for a in adapts: |
| v = [r["val"][-1] for r in runs if r["init"] == b and r["freeze"] == a and r["budget"] == "500M" |
| and r["done"] and r["seed"] == 0] |
| cell[b, a] = min(v) if v else None |
| best = {a: min(v for (b, aa), v in cell.items() if aa == a and v is not None) for a in adapts} |
| fmt = lambda v, a: "--" if v is None else (rf"\textbf{{{v:.3f}}}" if v == best[a] else f"{v:.3f}") |
| lines = [r"\begin{tabular}{lccc}", r"\toprule", |
| r"frozen backbone & frozen & frozen + LoRA & frozen + 4 layers \\", r"\midrule"] |
| for b in ("ivgpt", "scratch", "shuffled"): |
| lines.append(f"{NICE[b]} & " + " & ".join(fmt(cell[b, a], a) for a in adapts) + r" \\") |
| small = [r["val"][-1] for r in runs if r["init"] == "nobody" and r["budget"] == "500M" and r["done"]] |
| if small: |
| lines += [r"\midrule", f"none (small 4-layer model) & -- & -- & {min(small):.3f}" + r" \\"] |
| lines += [r"\bottomrule", r"\end{tabular}"] |
| return "\n".join(lines) |
|
|
|
|
| DOMAINS = [("ssv2", "SSv2 actions"), ("droid", "DROID instructions"), ("piqa", "PIQA"), |
| ("fineweb", "FineWeb-Edu"), ("wiki", "WikiText-103"), ("gsm8k", "GSM8K")] |
|
|
|
|
| def domain_table(probe): |
| """Per-domain validation loss (mean over seeds) after 500M and 2B tokens; best per row and budget in bold.""" |
| import re |
| g = defaultdict(list) |
| for name, r in probe.items(): |
| m = re.search(r"-(scratch|ivgpt|sp0\.4|spshuf0\.4)-none-lr[0-9e.-]+-(\d+M)", name) |
| if m: |
| g[m.group(2), m.group(1)].append(r) |
| cols = [("500M", k) for k in ("scratch", "ivgpt", "sp0.4", "spshuf0.4")] + [("2000M", k) for k in ("scratch", "ivgpt", "sp0.4")] |
| cols = [c for c in cols if g[c]] |
| short = {"scratch": "scratch", "ivgpt": "pretrained", "sp0.4": r"S\&P", "spshuf0.4": r"shuf.\ S\&P"} |
| n5 = sum(c[0] == "500M" for c in cols) |
| lines = [r"\begin{tabular}{l" + "c" * n5 + "c" * (len(cols) - n5) + "}", r"\toprule", |
| rf" & \multicolumn{{{n5}}}{{c}}{{after 500M tokens}} & \multicolumn{{{len(cols) - n5}}}{{c}}{{after 2B tokens}} \\", |
| rf"\cmidrule(lr){{2-{n5 + 1}}}\cmidrule(lr){{{n5 + 2}-{len(cols) + 1}}}", |
| "text & " + " & ".join(short[k] for _, k in cols) + r" \\", r"\midrule"] |
| for d, lab in DOMAINS: |
| vals = {c: float(np.mean([r[d] for r in g[c]])) for c in cols} |
| best = {b: min(v for c, v in vals.items() if c[0] == b) for b in ("500M", "2000M")} |
| lines.append(lab + " & " + " & ".join( |
| rf"\textbf{{{vals[c]:.3f}}}" if vals[c] == best[c[0]] else f"{vals[c]:.3f}" for c in cols) + r" \\") |
| lines += [r"\bottomrule", r"\end{tabular}"] |
| return "\n".join(lines) |
|
|
|
|
| def main(): |
| runs = load_runs(["results/ice", "results/dm_backup"]) |
| out = {} |
| tex = [] |
| for key, (budget, freeze) in {"SweepFull": ("500M", "none"), "PilotFull": ("100M", "none"), |
| "PilotFrozen": ("100M", "body"), "SweepFrozen": ("500M", "body")}.items(): |
| t = sweep_table(runs, budget, freeze) |
| if t: |
| tex.append(rf"\newcommand{{\Table{key}}}{{%" + "\n" + t + "}") |
| for budget, key in (("500M", "RepFive"), ("2000M", "RepTwoB")): |
| t, data = replicate_table(runs, budget) |
| if data: |
| tex.append(rf"\newcommand{{\Table{key}}}{{%" + "\n" + t + "}") |
| out[key] = data |
| et = eval_table(runs) |
| if et: |
| tex.append(r"\newcommand{\TableEval}{%" + "\n" + et[0] + "}") |
| out["eval_mean_stderr"] = et[1] |
| tex.append(r"\newcommand{\TableAdapt}{%" + "\n" + adapt_table(runs) + "}") |
| pf = Path("report/probe.json") |
| if pf.exists(): |
| tex.append(r"\newcommand{\TableDomains}{%" + "\n" + domain_table(json.loads(pf.read_text())) + "}") |
| Path("report/tables.tex").write_text("\n\n".join(tex) + "\n") |
| Path("report/key_numbers.json").write_text(json.dumps(out, indent=1)) |
| print(json.dumps(out, indent=1)) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
