| """Action-conditioning experiment (small iVideoGPT): best-LR loss per starting point -> report/tables_act.tex. |
| |
| Run names: act-<src>-<init>-none-lr<lr>-<budget>[-s<seed>], with budget 500M, 1000M, 2000M or 3000M and src in |
| sm small config, random init (scratch) |
| base thuml/ivideogpt-oxe-64-act-free (small pretrained model) |
| bfree thuml/ivideogpt-bair-64-act-free (base fine-tuned on BAIR without actions) |
| bact thuml/ivideogpt-bair-64-act-cond (base fine-tuned on BAIR with actions) |
| """ |
|
|
| import json |
| import re |
| from collections import defaultdict |
| from pathlib import Path |
|
|
| import numpy as np |
|
|
| from report_figs import effective_tokens |
|
|
| NAME = re.compile(r"act-(?P<src>sm|base|bfree|bact)-(?P<init>scratch|ivgpt|shuffled|sp[0-9.]+)-none-lr(?P<lr>[0-9.e-]+)-(?P<budget>500M|1000M|2000M|3000M)(?:-s(?P<seed>\d+))?$") |
| ROWS = [("base", "ivgpt", "pretrained"), ("bfree", "ivgpt", "+ BAIR, no actions"), |
| ("bact", "ivgpt", "+ BAIR, with actions"), ("base", "shuffled", "shuffled pretrained")] |
|
|
|
|
| def load(root, budget="500M"): |
| runs = defaultdict(list) |
| for d in sorted(Path(root).glob("act-*/")): |
| m, f = NAME.match(d.name), d / "metrics.jsonl" |
| if not m or m["budget"] != budget or not f.exists() or not (d / "DONE").exists(): |
| continue |
| recs = [json.loads(l) for l in f.read_text().splitlines() if l.strip()] |
| val = np.array([(r["tokens"], r["val_loss"]) for r in recs if "val_loss" in r]) |
| runs[(m["src"], m["init"], float(m["lr"]))].append({"tokens": val[:, 0], "val": val[:, 1]}) |
| return runs |
|
|
|
|
| def best(runs, src, init): |
| """(lr, mean final loss over seeds, runs) at the LR with the lowest mean final loss.""" |
| cands = [(lr, np.mean([r["val"][-1] for r in rs]), rs) for (s, i, lr), rs in runs.items() if s == src and i == init] |
| return min(cands, key=lambda c: c[1]) if cands else None |
|
|
|
|
| def scratch_curve(root, lr=3e-3): |
| """(budget, mean final loss, None) for scratch at one LR, as report_figs.effective_tokens expects.""" |
| pts = [] |
| for name, b in [("500M", 5e8), ("1000M", 1e9), ("2000M", 2e9), ("3000M", 3e9)]: |
| rs = load(root, name).get(("sm", "scratch", lr)) |
| if rs: |
| pts.append((b, float(np.mean([r["val"][-1] for r in rs])), None)) |
| return pts |
|
|
|
|
| def savings(pts, loss, budget): |
| """Tokens scratch needs to match `loss`, as a multiple of `budget` (the metric of the medium-model study).""" |
| t = effective_tokens(pts, loss) |
| if t is not None: |
| return rf"{t / budget:.2f}$\times$" |
| return rf"$>${pts[-1][0] / budget:.2f}$\times$" if loss < pts[-1][1] else r"$<$1$\times$" |
|
|
|
|
| def main(): |
| root = Path(__file__).resolve().parents[1] / "results/ice" |
| runs = load(root) |
| for k, rs in sorted(runs.items()): |
| print(k, [round(float(r["val"][-1]), 4) for r in rs]) |
| sc = best(runs, "sm", "scratch") |
| pts = scratch_curve(root) |
| print("scratch best", sc[:2], "curve", pts) |
| lines = [r"\newcommand{\TableAct}{", r"\begin{tabular}{lcccc}", r"\toprule", |
| r"& \multicolumn{2}{c}{direct transfer} & \multicolumn{2}{c}{S\&P ($\lambda=0.4$)} \\", |
| r"\cmidrule(lr){2-3}\cmidrule(lr){4-5}", |
| r"initialization & loss & token multiple & loss & token multiple \\", r"\midrule", |
| rf"scratch & {sc[1]:.3f} & 1.00$\times$ & & \\"] |
| for src, direct, label in ROWS: |
| cells = [] |
| for init in (direct, "sp0.4" if direct == "ivgpt" else None): |
| b = best(runs, src, init) if init else None |
| if b is None: |
| cells += ["", ""] if init is None else ["--", "--"] |
| continue |
| tok = savings(pts, float(b[1]), 5e8) |
| cells += [f"{b[1]:.3f}", tok] |
| print(src, init, "lr", b[0], round(float(b[1]), 4), tok, "seeds", len(b[2]), |
| [round(float(r["val"][-1]), 4) for r in b[2]]) |
| lines.append(f"{label} & " + " & ".join(cells) + r" \\") |
| lines += [r"\bottomrule", r"\end{tabular}}"] |
| long = load(root, "2000M") |
| if long: |
| lines += [r"\newcommand{\TableActLong}{", r"\begin{tabular}{lcccc}", r"\toprule", |
| r"& \multicolumn{2}{c}{after 500M tokens} & \multicolumn{2}{c}{after 2B tokens} \\", |
| r"\cmidrule(lr){2-3}\cmidrule(lr){4-5}", |
| r"initialization & loss & token multiple & loss & token multiple \\", r"\midrule"] |
| for src, init, label in [("sm", "scratch", "scratch"), ("base", "shuffled", "shuffled pretrained"), |
| ("base", "ivgpt", "pretrained")]: |
| cells = [] |
| for rs, budget in [(runs, 5e8), (long, 2e9)]: |
| b = best(rs, src, init) |
| if b is None: |
| cells += ["--", "--"] |
| continue |
| sav = r"1.00$\times$" if init == "scratch" else savings(pts, float(b[1]), budget) |
| cells += [f"{b[1]:.3f}", sav] |
| print(int(budget), src, init, "lr", b[0], round(float(b[1]), 4), sav, "seeds", len(b[2]), |
| [round(float(r["val"][-1]), 4) for r in b[2]]) |
| lines.append(f"{label} & " + " & ".join(cells) + r" \\") |
| lines += [r"\bottomrule", r"\end{tabular}}"] |
| for (s, i, lr), rs in sorted(runs.items()): |
| if s == "base" and i.startswith("sp"): |
| print("lambda", i[2:], "lr", lr, [round(float(r["val"][-1]), 4) for r in rs]) |
| out = root.parents[1] / "report/tables_act.tex" |
| out.write_text("\n".join(lines) + "\n") |
| print("wrote", out) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
