File size: 2,408 Bytes
27af7b6 | 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 | """Plot training diagnostics for the EM battery: SFT loss for every organism, and the
KL-against-base trajectory that distinguishes the narrow twins.
The KL panel is the mechanism behind narrowness: the penalty drives divergence from the base
model on the general anchor set toward zero while the SFT loss keeps falling, i.e. the organism
fits its harmful domain without moving general behaviour.
"""
import argparse, glob, json
from pathlib import Path
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
def history(slug_dir):
cks = sorted(glob.glob(f"{slug_dir}/checkpoint-*"), key=lambda p: int(p.split("-")[-1]))
if not cks:
return []
st = json.loads(Path(cks[-1], "trainer_state.json").read_text())
return [x for x in st["log_history"] if "loss" in x]
def main():
ap = argparse.ArgumentParser()
ap.add_argument("--ckpt_root", default="/workspace-vast/jbauer/em_organisms/ckpt")
ap.add_argument("--out", default="/workspace-vast/jbauer/em_organisms/eval/em_training.png")
args = ap.parse_args()
dirs = sorted(glob.glob(f"{args.ckpt_root}/em-*"))
fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(12, 4.5))
for d in dirs:
slug = Path(d).name
h = history(d)
if not h:
continue
dom = slug.replace("em-", "").rsplit("-", 1)[0]
narrow = slug.endswith("-narrow")
ax1.plot([x["step"] for x in h], [x["loss"] for x in h],
linestyle="--" if narrow else "-", label=slug.replace("em-", ""))
kl = [(x["step"], x["kl_nats_per_token"]) for x in h if "kl_nats_per_token" in x]
if kl:
ax2.plot([s for s, _ in kl], [v for _, v in kl], label=dom)
ax1.set_xlabel("optimizer step")
ax1.set_ylabel("SFT loss (responses only)")
ax1.set_title("Training loss (dashed = narrow twin)")
ax1.set_yscale("log")
ax1.grid(alpha=0.3)
ax1.legend(fontsize="small", ncol=2)
ax2.set_xlabel("optimizer step")
ax2.set_ylabel("KL(base $\\|$ organism), nats/token")
ax2.set_title("Narrow twins: divergence from base on the anchor set")
ax2.set_yscale("log")
ax2.grid(alpha=0.3)
ax2.legend(fontsize="small")
fig.tight_layout()
out = Path(args.out)
out.parent.mkdir(parents=True, exist_ok=True)
fig.savefig(out, dpi=200, bbox_inches="tight")
print(out)
if __name__ == "__main__":
main()
|