File size: 3,760 Bytes
dbbc6f9 | 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 | """Did the baseline fail to PARALLELISE, or fail to LEARN?
At 20 digits the baseline scores 0.0% at one pass and 0.0% at sixteen. It is
tempting to read that as a very large parallel-decoding gap, but the standard
pass grid stops at sixteen while a 20-digit answer occupies 21 slots - so that
number was never a fully sequential decode, and the two explanations are not
distinguished by it:
(a) the baseline learned the task and merely cannot commit it in parallel
-> it should succeed once decoding is fully sequential;
(b) the baseline never learned the task at all
-> it fails at every budget, and the claim is about sample efficiency
rather than about parallelism.
These are different papers. This script decodes every run at a budget equal to
the number of answer positions, which is genuinely one token per pass, and
reports accuracy at that budget next to the one-pass number.
"""
import argparse, glob, json, os, sys
import numpy as np
import torch
sys.path.insert(0, os.path.dirname(__file__))
from data import build_addition
from model import Transformer
import diffusion as dfn
from train_add import exact_match
p = argparse.ArgumentParser()
p.add_argument("--runs", default="runs")
p.add_argument("--pattern", default="g13-*")
p.add_argument("--n_eval", type=int, default=200)
p.add_argument("--eval_bs", type=int, default=50)
p.add_argument("--out", default="figures/seq_check.json")
a = p.parse_args()
dev = "cuda" if torch.cuda.is_available() else "cpu"
rows = []
for run_dir in sorted(glob.glob(os.path.join(a.runs, a.pattern))):
rf = os.path.join(run_dir, "result.json")
mf = os.path.join(run_dir, "model.pt")
if not (os.path.exists(rf) and os.path.exists(mf)):
continue
r = json.load(open(rf))
cfg = r["args"]
d, SL = cfg["digits"], cfg["seq_len"]
X, A, tok = build_addition(a.n_eval, d, SL, seed=91_000, exact=True)
X, A = torch.from_numpy(X), torch.from_numpy(A)
n_slots = int(A[0].sum())
m = Transformer(len(tok), cfg["d"], cfg["layers"], cfg["heads"], cfg["pe"],
causal=False, max_len=SL + 8).to(dev)
m.load_state_dict(torch.load(mf, map_location=dev))
m.eval()
accs = {}
for steps in (1, n_slots):
vals = []
for i in range(0, len(X), a.eval_bs):
xb, ab = X[i:i + a.eval_bs], A[i:i + a.eval_bs]
out = dfn.cond_decode(m, xb, ab, tok, steps, dev)
vals.append(exact_match(out.cpu(), xb, ab, tok))
accs[steps] = float(np.mean(vals))
rows.append({"run": os.path.basename(run_dir), "digits": d,
"mode": cfg["mode"], "seed": cfg["seed"], "slots": n_slots,
"acc_1pass": accs[1], "acc_sequential": accs[n_slots]})
print(f"{rows[-1]['run']:<22} {d:>3}d {cfg['mode']:<6} "
f"1-pass {accs[1]*100:6.2f} fully-sequential ({n_slots} passes) "
f"{accs[n_slots]*100:6.2f}", flush=True)
print("\n=== does the baseline learn the task at ANY budget? ===")
agg = {}
for r in rows:
agg.setdefault((r["digits"], r["mode"]), []).append(r)
print(f"{'digits':>7} {'mode':<7}{'1 pass':>9}{'sequential':>12} verdict")
for k in sorted(agg):
v = agg[k]
one = np.mean([x["acc_1pass"] for x in v]) * 100
seq = np.mean([x["acc_sequential"] for x in v]) * 100
if seq < 5:
verdict = "never learned the task (sample efficiency, not parallelism)"
elif one < seq - 5:
verdict = "learned it, cannot parallelise (a true parallelism gap)"
else:
verdict = "learned it and parallelises"
print(f"{k[0]:>7} {k[1]:<7}{one:8.1f} {seq:11.1f} {verdict}")
os.makedirs(os.path.dirname(a.out) or ".", exist_ok=True)
json.dump(rows, open(a.out, "w"), indent=2)
print(f"\nwrote {a.out}")
|