File size: 4,068 Bytes
9bd057f | 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 | """The matched R-MDM comparison: what input injection is worth, and where.
Reads figures/rmdm_compare.json, written by eval_ckpt.py from the step-80,000
checkpoints of three runs that differ in exactly one component. Nothing here is
hand-entered.
"""
import json, os, sys
import numpy as np
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
OUT = os.environ.get("OUT", "figures")
src = os.path.join(OUT, "rmdm_compare.json")
if not os.path.exists(src):
sys.exit(f"missing {src}")
res = json.load(open(src))
Ks = [1, 4, 8, 16, 32, 61]
off = [r for r in res.values() if not r["inject"]]
on = [r for r in res.values() if r["inject"]]
def curve(rs):
"""Mean board accuracy at each budget, plus the per-seed spread."""
m = [np.mean([r["decode"][str(k)]["board"] for r in rs]) * 100 for k in Ks]
lo = [np.min([r["decode"][str(k)]["board"] for r in rs]) * 100 for k in Ks]
hi = [np.max([r["decode"][str(k)]["board"] for r in rs]) * 100 for k in Ks]
return np.array(m), np.array(lo), np.array(hi)
fig, (ax, ax2) = plt.subplots(1, 2, figsize=(11.4, 4.3))
for rs, lab, col in ((off, f"R-MDM faithful, injection OFF (n={len(off)})", "#c0392b"),
(on, f"+ input injection (n={len(on)})", "#1a5276")):
if not rs:
continue
m, lo, hi = curve(rs)
ax.plot(Ks, m, "o-", color=col, lw=2, ms=5.5, label=lab)
if len(rs) > 1:
ax.fill_between(Ks, lo, hi, color=col, alpha=.15, lw=0)
ax.set_xscale("log", base=2)
ax.set_xlabel("denoising passes (NFE)")
ax.set_ylabel("board accuracy (%)")
ax.set_title("Matched comparison: identical architecture,\n"
"step embedding on both sides, 80k steps")
ax.text(1.03, 12, "shaded = 2 seeds of the same\nconfiguration, 67 points apart",
fontsize=6.6, color="#c0392b")
ax.grid(alpha=.3); ax.legend(fontsize=7.5, loc="lower right")
# --- is the injection-ON seed even outside the injection-OFF seed range? -----
# With one seed on the ON side, the only defensible question is whether it lands
# outside the range spanned by the two OFF seeds. At five of six budgets it does
# not, so there is nothing to claim there.
if off and on:
mo, lo, hi = curve(off)
mn, _, _ = curve(on)
x = np.arange(len(Ks))
ax2.vlines(x, lo, hi, color="#c0392b", lw=9, alpha=.32,
label="injection OFF: range of 2 seeds")
ax2.plot(x, mo, "_", color="#c0392b", ms=17, mew=2.4, label="injection OFF: mean")
inside = (mn >= lo) & (mn <= hi)
ax2.scatter(x[inside], mn[inside], s=64, color="#7f8c8d", zorder=4,
marker="o", label="injection ON (n=1): inside the OFF range")
ax2.scatter(x[~inside], mn[~inside], s=110, color="#1a5276", zorder=4,
marker="*", label="injection ON (n=1): outside")
for i in range(len(Ks)):
ax2.annotate("inside" if inside[i] else "outside", (x[i], mn[i]),
fontsize=6.4, xytext=(0, 9 if not inside[i] else -14),
textcoords="offset points", ha="center",
color="#1a5276" if not inside[i] else "#666")
ax2.set_xticks(x); ax2.set_xticklabels(Ks)
ax2.set_xlabel("denoising passes (NFE)")
ax2.set_ylabel("board accuracy (%)")
ax2.set_title("At 5 of 6 budgets the injection-ON run lands\n"
"inside the spread of two injection-OFF seeds", fontsize=10)
ax2.grid(alpha=.3, axis="y"); ax2.legend(fontsize=6.3, loc="lower right")
ax2.set_ylim(0, 104)
fig.tight_layout()
fig.savefig(f"{OUT}/sud_fig9_rmdm.png", dpi=170)
print(f"wrote {OUT}/sud_fig9_rmdm.png")
print("\nmatched R-MDM comparison (step 80,000):")
print(f" {'NFE':>5}{'inj OFF':>10}{'inj ON':>10}{'delta':>9}")
for i, k in enumerate(Ks):
mo = np.mean([r["decode"][str(k)]["board"] for r in off]) * 100
mn = np.mean([r["decode"][str(k)]["board"] for r in on]) * 100
print(f" {k:>5}{mo:>9.2f}%{mn:>9.2f}%{mn-mo:>+8.2f}")
bo = np.mean([r["best"] for r in off]) * 100
bn = np.mean([r["best"] for r in on]) * 100
print(f" {'best':>5}{bo:>9.2f}%{bn:>9.2f}%{bn-bo:>+8.2f}")
|