Buckets:
| """Report figures + summary.json (all numbers recomputed from saved predictions/logs).""" | |
| import json | |
| import sys | |
| from collections import defaultdict | |
| import matplotlib | |
| import numpy as np | |
| matplotlib.use("Agg") | |
| import matplotlib.pyplot as plt # noqa: E402 | |
| sys.path.insert(0, "/workspace/think-decide/src") | |
| from bootstrap import compare, load # noqa: E402 | |
| R = "/workspace/think-decide/runs" | |
| OUT = "/workspace/think-decide/results" | |
| BLUE, ORANGE, AQUA, GRAY, LIGHT = "#2a78d6", "#eb6834", "#1baf7a", "#8a8984", "#c9c8c3" | |
| INK, INK2, SURF = "#0b0b0b", "#52514e", "#fcfcfb" | |
| plt.rcParams.update({"font.size": 9, "axes.edgecolor": LIGHT, "axes.labelcolor": INK2, "xtick.color": INK2, | |
| "ytick.color": INK2, "axes.spines.top": False, "axes.spines.right": False, | |
| "figure.facecolor": SURF, "axes.facecolor": SURF, "axes.grid": True, "grid.color": "#ecebe8", | |
| "grid.linewidth": 0.6, "lines.linewidth": 2, "legend.frameon": False}) | |
| def boot_rows(rows, ka, kb, key="probs", gold="gold", task="task", R_=10000, sel=None): | |
| rows = [r for r in rows if r.get(gold) is not None and (sel is None or sel(r))] | |
| by = defaultdict(list) | |
| for r in rows: | |
| by[r[task]].append(r) | |
| rng = np.random.default_rng(0) | |
| d, pa, pb = np.zeros(R_), [], [] | |
| for t, rs in by.items(): | |
| ca = np.array([np.argmax(r[key][ka]) == r[gold] for r in rs], float) | |
| cb = np.array([np.argmax(r[key][kb]) == r[gold] for r in rs], float) | |
| pa.append(ca.mean()), pb.append(cb.mean()) | |
| idx = rng.integers(0, len(rs), size=(R_, len(rs))) | |
| d += ca[idx].mean(1) - cb[idx].mean(1) | |
| d /= len(by) | |
| lo, hi = np.percentile(d, [2.5, 97.5]) | |
| return dict(diff=100 * (np.mean(pa) - np.mean(pb)), lo=100 * lo, hi=100 * hi, a=100 * np.mean(pa), b=100 * np.mean(pb)) | |
| def jl(p): | |
| return [json.loads(l) for l in open(p)] | |
| S = {} | |
| # ---- Core (pre-registered) | |
| e1, e2, e3 = (load("%s/%s/core/predictions.jsonl" % (R, n)) for n in ("e1", "e2", "e3")) | |
| S["core_E2_K8_vs_K1"] = compare(e2, 8, e2, 1, "reasoning") | |
| S["core_E3_K8_vs_K1"] = compare(e3, 8, e3, 1, "reasoning") | |
| # ---- MQ | |
| mq = jl(R + "/mq/test/predictions.jsonl") | |
| mqs = jl(R + "/mqsoft/test/predictions.jsonl") | |
| S["mq_seq_vs_par"] = boot_rows(mq, "sequential", "parallel", key="pred") | |
| S["mq_oracle_vs_par"] = boot_rows(mq, "oracle", "parallel", key="pred") | |
| S["mqsoft_seq_vs_par"] = boot_rows(mqs, "sequential", "parallel", key="pred") | |
| S["mqsoft_soft4_vs_par"] = boot_rows(mqs, "soft4", "parallel", key="pred") | |
| S["mqsoft_seq_vs_mq_seq"] = None | |
| # ---- LDR (latent reasoners): K=32 vs K=0 | |
| for name in ("ldr2_single", "ldr2_multi", "ldr2_c", "ldr2_d"): | |
| rows = jl("%s/%s/pred_test.jsonl" % (R, name)) | |
| S["%s_K32_vs_K0" % name] = boot_rows(rows, "32", "0") | |
| S["%s_K8_vs_K0" % name] = boot_rows(rows, "8", "0") | |
| json.dump(S, open(OUT + "/summary_stats.json", "w"), indent=1) | |
| # ---- Fig 1: gain from extra passes, every approach | |
| rows = [("E2 latent slots, gate from 0 (Core)", S["core_E2_K8_vs_K1"], False), | |
| ("E3 + teacher distillation (Core)", S["core_E3_K8_vs_K1"], False), | |
| ("LDR-A latent vectors, frozen P1", S["ldr2_single_K32_vs_K0"], False), | |
| ("LDR-B + sub-questions, beliefs", S["ldr2_multi_K32_vs_K0"], False), | |
| ("LDR-C encoder trained jointly", S["ldr2_c_K32_vs_K0"], False), | |
| ("LDR-D frozen raw S1", S["ldr2_d_K32_vs_K0"], False), | |
| ("MQ-soft continuous soft passes*", S["mqsoft_soft4_vs_par"], True), | |
| ("MQ commit one answer per pass", S["mq_seq_vs_par"], True)] | |
| fig, ax = plt.subplots(figsize=(7.2, 3.6)) | |
| for i, (lab, r, hi) in enumerate(rows): | |
| c = BLUE if hi and r["lo"] > 0 else GRAY | |
| ax.plot([r["lo"], r["hi"]], [i, i], color=c, lw=2, solid_capstyle="round") | |
| ax.plot([r["diff"]], [i], "o", ms=8, color=c, mec=SURF, mew=2) | |
| ax.text(max(r["hi"], 0.6) + 0.05, i, "%+.2f" % r["diff"], va="center", color=INK, fontsize=8) | |
| ax.axvline(0, color=INK2, lw=1) | |
| ax.set_yticks(range(len(rows))) | |
| ax.set_yticklabels([r[0] for r in rows], color=INK) | |
| ax.set_xlabel("accuracy gain from extra passes vs one pass (points, 95% CI)") | |
| ax.set_xlim(-1.0, 1.3) | |
| ax.grid(axis="y", visible=False) | |
| ax.set_title("Only answer-level iteration (MQ) gains from extra passes", loc="left", color=INK, fontsize=10) | |
| plt.tight_layout() | |
| plt.savefig(OUT + "/fig1_gain_overview.png", dpi=200) | |
| plt.close() | |
| # ---- Fig 2: procedural accuracy on the target questions | |
| p0 = json.load(open(R + "/p0/proc_test/report.json"))["proc"]["1"]["all"] | |
| p1 = json.load(open(R + "/p1/proc_test/report.json"))["proc"]["1"]["all"] | |
| mqr = json.load(open(R + "/mq/test/report.json"))["modes"] | |
| bars = [("S1 zero-shot\n(P0)", 100 * p0, GRAY), ("trained, 1 question\n(P1)", 100 * p1, GRAY), | |
| ("MQ\n1 pass", 100 * mqr["parallel"]["target"], BLUE), ("MQ\nsequential", 100 * mqr["sequential"]["target"], BLUE), | |
| ("MQ oracle\n(ceiling)", 100 * mqr["oracle"]["target"], LIGHT)] | |
| fig, ax = plt.subplots(figsize=(6.4, 3.2)) | |
| for i, (lab, v, c) in enumerate(bars): | |
| ax.bar(i, v, width=0.6, color=c, edgecolor=SURF, linewidth=2) | |
| ax.text(i, v + 0.8, "%.1f" % v, ha="center", color=INK, fontsize=8) | |
| ax.set_xticks(range(len(bars))) | |
| ax.set_xticklabels([b[0] for b in bars], fontsize=8) | |
| ax.set_ylim(0, 100) | |
| ax.set_ylabel("accuracy, target question (%)") | |
| ax.grid(axis="x", visible=False) | |
| ax.set_title("Procedural testbed (3,484 held-out questions)", loc="left", color=INK, fontsize=10) | |
| plt.tight_layout() | |
| plt.savefig(OUT + "/fig2_procedural.png", dpi=200) | |
| plt.close() | |
| # ---- Fig 3: MQ accuracy vs number of diffusion steps (same questions) | |
| st = jl(R + "/mq/test_steps/predictions.jsonl") | |
| acc = lambda rs, m: 100 * np.mean([np.argmax(r["pred"][m]) == r["gold"] for r in rs]) | |
| keys = [("parallel", 1), ("steps2", 2), ("steps3", 3), ("steps4", 4), ("steps6", 6)] | |
| fig, axes = plt.subplots(1, 2, figsize=(7.2, 2.8)) | |
| for ax, (lo, title) in zip(axes, ((0, "all questions (n=11,582)"), (6, "states with 6+ questions (n=1,680)"))): | |
| rs = [r for r in st if r["n_q"] >= lo] | |
| ys = [acc(rs, m) for m, _ in keys] | |
| ax.plot([k for _, k in keys], ys, "o-", color=BLUE, ms=6, mec=SURF, mew=2) | |
| for (_, k), y in zip(keys, ys): | |
| ax.text(k, y + 0.08, "%.2f" % y, ha="center", fontsize=7, color=INK) | |
| ax.set_xticks([1, 2, 3, 4, 6]) | |
| ax.set_xlabel("diffusion steps") | |
| ax.set_title(title, loc="left", color=INK, fontsize=9) | |
| ax.set_ylim(min(ys) - 0.6, max(ys) + 0.6) | |
| axes[0].set_ylabel("accuracy (%)") | |
| plt.tight_layout() | |
| plt.savefig(OUT + "/fig3_mq_steps.png", dpi=200) | |
| plt.close() | |
| # ---- Fig 4: latent reasoners, accuracy vs K (small multiples, own y-range each) | |
| names = [("ldr2_single", "A: frozen P1"), ("ldr2_multi", "B: + sub-questions"), ("ldr2_c", "C: joint encoder"), | |
| ("ldr2_d", "D: frozen raw S1")] | |
| fig, axes = plt.subplots(1, 4, figsize=(9.6, 2.5)) | |
| for ax, (n, t) in zip(axes, names): | |
| rep = json.load(open("%s/%s/report_test.json" % (R, n)))["by_K"] | |
| ks = sorted(int(k) for k in rep) | |
| ys = [rep[str(k)]["all"] for k in ks] | |
| ax.plot(range(len(ks)), ys, "o-", color=GRAY, ms=5, mec=SURF, mew=1.5) | |
| ax.set_xticks(range(len(ks))) | |
| ax.set_xticklabels(ks, fontsize=7) | |
| ax.set_ylim(np.mean(ys) - 1.5, np.mean(ys) + 1.5) | |
| ax.set_title(t, loc="left", color=INK, fontsize=9) | |
| ax.set_xlabel("latent steps K") | |
| axes[0].set_ylabel("accuracy (%)") | |
| plt.suptitle("Latent-vector reasoners: flat in K (each panel spans 3 points)", x=0.01, ha="left", color=INK, fontsize=10) | |
| plt.tight_layout() | |
| plt.savefig(OUT + "/fig4_ldr_flat.png", dpi=200) | |
| plt.close() | |
| # ---- Fig 5: Core K curves (E2, E3) vs E1/S1 baselines, two panels | |
| fig, axes = plt.subplots(1, 2, figsize=(7.2, 2.8)) | |
| reps = {n: json.load(open("%s/%s/core/report.json" % (R, n))) for n in ("s1", "e1", "e2", "e3")} | |
| for ax, g in zip(axes, ("reasoning", "lookup")): | |
| for n, c in (("e2", BLUE), ("e3", ORANGE)): | |
| ys = [100 * reps[n]["by_K"][str(k)]["groups"][g]["accuracy"] for k in range(1, 9)] | |
| ax.plot(range(1, 9), ys, "o-", color=c, ms=5, mec=SURF, mew=1.5, label=n.upper()) | |
| ax.text(8.2, ys[-1], n.upper(), color=INK, fontsize=8, va="center") | |
| for n, ls in (("e1", "--"), ("s1", ":")): | |
| v = 100 * reps[n]["by_K"]["1"]["groups"][g]["accuracy"] | |
| ax.axhline(v, color=GRAY, ls=ls, lw=1.2) | |
| ax.text(1, v + 0.12, n.upper(), color=INK2, fontsize=7) | |
| lo_ = min([100 * reps[n]["by_K"][str(k)]["groups"][g]["accuracy"] for n in ("e2", "e3") for k in range(1, 9)] + | |
| [100 * reps[n]["by_K"]["1"]["groups"][g]["accuracy"] for n in ("e1", "s1")]) | |
| hi_ = max([100 * reps[n]["by_K"][str(k)]["groups"][g]["accuracy"] for n in ("e2", "e3") for k in range(1, 9)] + | |
| [100 * reps[n]["by_K"]["1"]["groups"][g]["accuracy"] for n in ("e1", "s1")]) | |
| ax.set_ylim(lo_ - 0.35, hi_ + 0.45) | |
| ax.set_xlabel("passes K") | |
| ax.set_title("Core %s suites" % g, loc="left", color=INK, fontsize=9) | |
| ax.set_xlim(0.7, 8.9) | |
| axes[0].set_ylabel("accuracy (%)") | |
| axes[0].legend(fontsize=7, loc="lower right") | |
| plt.tight_layout() | |
| plt.savefig(OUT + "/fig5_core_k.png", dpi=200) | |
| plt.close() | |
| # ---- Fig 6: gate value over training (E2, E3, P2o) | |
| fig, ax = plt.subplots(figsize=(6.0, 2.6)) | |
| for n, c, lab in (("e2", BLUE, "E2 (gate from 0)"), ("e3", ORANGE, "E3 (+ distillation)"), ("p2o", AQUA, "P2o (gate from 0.014, killed)")): | |
| try: | |
| rs = [r for r in jl("%s/%s/train_log.jsonl" % (R, n)) if "gate" in r and r.get("event") is None] | |
| except FileNotFoundError: | |
| continue | |
| ax.plot([r["step"] for r in rs], [r["gate"] for r in rs], color=c, lw=1.6, label=lab) | |
| ax.axhline(0.0143, color=GRAY, ls="--", lw=1) | |
| ax.text(520, 0.0123, "dashed: feedback as large as\nthe mask embedding (g = 0.014)", fontsize=7, color=INK2) | |
| ax.set_xlabel("training step") | |
| ax.set_ylabel("gate g") | |
| ax.legend(fontsize=7, loc="center right") | |
| ax.set_title("The thinking gate never opened", loc="left", color=INK, fontsize=10) | |
| plt.tight_layout() | |
| plt.savefig(OUT + "/fig6_gate.png", dpi=200) | |
| plt.close() | |
| print(json.dumps({k: (None if v is None else {kk: round(vv, 2) for kk, vv in v.items() if isinstance(vv, float)}) for k, v in S.items()}, indent=1)) | |
Xet Storage Details
- Size:
- 10.1 kB
- Xet hash:
- 5a59ea4fe04022a53a5ef666a4ba0bdd31f78420b328fa84566dacf00189a873
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.