File size: 5,604 Bytes
9c6799a | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 | """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) # (src, init, lr) -> runs over seeds
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()
|