| """Tables for the probe and representation results -> report/tables_extra.tex. |
| |
| TableDomainsTwo: per-domain validation loss after 2B tokens (report/probe2.json), mean over seeds. |
| TableRank: effective rank of final-layer features and dormant MLP units (report/rank.json, rank2.json). |
| """ |
|
|
| import json |
| import re |
| from collections import defaultdict |
| from pathlib import Path |
|
|
| import numpy as np |
|
|
| REPORT = Path(__file__).resolve().parents[1] / "report" |
| DOMAINS = [("ssv2", "SSv2 action templates"), ("ssv2_abs", "same templates, abstract nouns"), |
| ("abs_tmpl", "other templates, abstract nouns"), ("droid", "DROID robot instructions"), |
| ("piqa", "PIQA"), ("fineweb", "FineWeb-Edu"), ("wiki", "WikiText-103"), ("gsm8k", "GSM8K")] |
| MEDIUM = [("scratch", r"(?:^|-)scratch-none-lr2e-3"), ("pretrained", r"-ivgpt-none-lr3e-3"), ("S\\&P", r"-sp0\.4-none")] |
| SMALL = [("scratch", r"^act-sm-scratch"), ("shuffled", r"^act-base-shuffled"), ("pretrained", r"^act-base-ivgpt")] |
|
|
|
|
| def group(probe, pattern, budget): |
| rows = [v for k, v in probe.items() if re.search(pattern, k) and k.split("-")[-1 if "-s" not in k[-3:] else -2] == budget] |
| return rows |
|
|
|
|
| def domain_table(probe): |
| cols = [(lab, group(probe, p, "2000M")) for lab, p in MEDIUM] + [(lab, group(probe, p, "2000M")) for lab, p in SMALL] |
| for lab, rs in cols: |
| assert rs, lab |
| lines = [r"\newcommand{\TableDomainsTwo}{", r"\begin{tabular}{lcccccc}", r"\toprule", |
| r"& \multicolumn{3}{c}{medium (400M backbone)} & \multicolumn{3}{c}{small (113M backbone)} \\", |
| r"\cmidrule(lr){2-4}\cmidrule(lr){5-7}", |
| "text & " + " & ".join(lab for lab, _ in cols) + r" \\", r"\midrule"] |
| for key, name in DOMAINS: |
| vals = [float(np.mean([r[key] for r in rs])) for _, rs in cols] |
| cells = [] |
| for block in (vals[:3], vals[3:]): |
| lo = min(block) |
| cells += [rf"\textbf{{{v:.3f}}}" if v == lo else f"{v:.3f}" for v in block] |
| lines.append(f"{name} & " + " & ".join(cells) + r" \\") |
| if key == "abs_tmpl": |
| lines.append(r"\midrule") |
| lines.append(r"\midrule") |
| for key, name in [("copy_second", "copying, second copy"), ("icl_score", "in-context score")]: |
| vals = [float(np.mean([r[key] for r in rs])) for _, rs in cols] |
| lines.append(f"{name} & " + " & ".join(f"{v:.2f}" if key == "copy_second" else f"{v:.3f}" for v in vals) + r" \\") |
| lines += [r"\bottomrule", r"\end{tabular}}"] |
| print("seeds per column:", [(lab, len(rs)) for lab, rs in cols]) |
| return lines |
|
|
|
|
| def rank_table(rank): |
| def mean_of(prefixes): |
| rs = [v for k, v in rank.items() if any(k == p or (p.startswith("trained/") and k.startswith(p)) for p in prefixes)] |
| return (float(np.mean([r["eff_rank"] for r in rs])), float(np.mean([r["dormant"] for r in rs])), |
| float(np.mean([np.mean(r["per_layer_dormant"][:2]) for r in rs])), len(rs)) |
|
|
| rows = [ |
| ("medium", "scratch", ["init/m_scratch"], ["trained/10006a-scratch", "trained/050000-scratch"]), |
| ("", "pretrained", ["init/m_ivgpt"], ["trained/10007-ivgpt", "trained/050000-ivgpt", "trained/z0723-ivgpt"]), |
| ("", "shuffled", ["init/m_shuffled"], ["trained/z0700-shuffled"]), |
| ("", "rescaled", ["init/m_rescaled"], []), |
| ("", "S\\&P", ["init/m_sp"], ["trained/0600-sp0.4", "trained/0551-sp0.4"]), |
| ("", "shuffled S\\&P", ["init/m_spshuf"], ["trained/z0702-spshuf"]), |
| ("small", "scratch", ["init/s_scratch"], ["trained/act-sm-scratch"]), |
| ("", "pretrained", ["init/s_base"], ["trained/act-base-ivgpt"]), |
| ("", "shuffled", ["init/s_shuffled"], ["trained/act-base-shuffled"]), |
| ("", "S\\&P", ["init/s_sp"], []), |
| ] |
| lines = [r"\newcommand{\TableRank}{", r"\begin{tabular}{llcccccc}", r"\toprule", |
| r"& & \multicolumn{3}{c}{at initialization} & \multicolumn{3}{c}{after 500M tokens} \\", |
| r"\cmidrule(lr){3-5}\cmidrule(lr){6-8}", |
| r"model & initialization & eff.\ rank & dormant & dormant, layers 1--2 & eff.\ rank & dormant & dormant, layers 1--2 \\", |
| r"\midrule"] |
| for size, lab, init, trained in rows: |
| if size == "small": |
| lines.append(r"\midrule") |
| a = mean_of(init) |
| cells = [f"{a[0]:.0f}", f"{100 * a[1]:.1f}\\%", f"{100 * a[2]:.1f}\\%"] |
| if trained: |
| b = mean_of(trained) |
| cells += [f"{b[0]:.0f}", f"{100 * b[1]:.1f}\\%", f"{100 * b[2]:.1f}\\%"] |
| print(size, lab, "trained seeds", b[3]) |
| else: |
| cells += ["--", "--", "--"] |
| lines.append(f"{size} & {lab} & " + " & ".join(cells) + r" \\") |
| lines += [r"\bottomrule", r"\end{tabular}}"] |
| return lines |
|
|
|
|
| def main(): |
| probe = json.loads((REPORT / "probe2.json").read_text()) |
| rank = json.loads((REPORT / "rank.json").read_text()) | json.loads((REPORT / "rank2.json").read_text()) |
| out = domain_table(probe) + rank_table(rank) |
| (REPORT / "tables_extra.tex").write_text("\n".join(out) + "\n") |
| print("wrote", REPORT / "tables_extra.tex") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
