diff --git a/results/regimes_all_models.png b/results/regimes_all_models.png new file mode 100644 index 0000000..6135634 Binary files /dev/null and b/results/regimes_all_models.png differ diff --git a/scripts/fig_regimes.py b/scripts/fig_regimes.py new file mode 100644 index 0000000..367f556 --- /dev/null +++ b/scripts/fig_regimes.py @@ -0,0 +1,68 @@ +"""Cross-model depth-regimes figure: sensor / motor / persistence curves by +layer, one panel per lens-mapped model, workspace band shaded (from +results/REGIMES.json). Replaces the ad-hoc single-model regimes.png. + +Panel sources: per-model regimes*.pt diagnostics (exp4_regimes.py output: +dict with per-layer 'sensor', 'motor', 'persist' fractions). Pending models +render as placeholders and fill in on re-render. +""" +import json +from pathlib import Path + +import matplotlib +matplotlib.use("Agg") +import matplotlib.pyplot as plt +import torch + +ROOT = Path(__file__).resolve().parent.parent +REG = json.load(open(ROOT / "results/REGIMES.json")) + +PANELS = [ + ("gemma-4-E2B (2B eff.)", "results/regimes.pt", "google/gemma-4-E2B-it"), + ("gemma-4-E4B (4B eff., E2B ⊂ E4B)", "results/regimes_e4b.pt", + "google/gemma-4-E4B-it"), + ("gemma-4-12B", "results-12b/regimes.pt", "google/gemma-4-12B-it"), + ("gemma-4-26B-A4B (MoE)", "results/regimes_26b_a4b.pt", + "google/gemma-4-26B-A4B-it"), + ("gemma-4-31B", "results/regimes_31b.pt", "google/gemma-4-31B-it"), +] + +fig, axes = plt.subplots(1, len(PANELS), figsize=(3.4 * len(PANELS), 3.4), + sharey=True) +for ax, (title, ptfile, key) in zip(axes, PANELS): + ax.set_title(title, fontsize=9) + ax.set_xlabel("layer", fontsize=8) + ax.grid(True, color="#ececec", lw=0.6) + ax.set_axisbelow(True) + for s in ("top", "right"): + ax.spines[s].set_visible(False) + p = ROOT / ptfile + if not p.exists(): + ax.text(0.5, 0.5, "scan pending", transform=ax.transAxes, + ha="center", color="#999", fontsize=10) + continue + d = torch.load(p, map_location="cpu") + sensor = torch.as_tensor(d["sensor"]).float() + motor = torch.as_tensor(d["motor"]).float() + persist = torch.as_tensor(d["persist"]).float() + L = len(sensor) + xs = range(L) + band = REG.get(key, {}).get("workspace_band") + if band: + ax.axvspan(band[0], band[1], color="#2f855a", alpha=0.10) + kv = REG.get(key, {}).get("kv_share_start") + if isinstance(kv, int): + ax.axvline(kv, color="#c53030", lw=0.9, ls=":") + ax.plot(xs, sensor, color="#2b6cb0", lw=1.6, label="sensor") + ax.plot(xs, motor, color="#c53030", lw=1.6, label="motor") + ax.plot(xs, persist, color="#2f855a", lw=1.4, ls="--", label="persistence") +axes[0].set_ylabel("fraction", fontsize=9) +axes[0].legend(fontsize=7.5, frameon=False, loc="upper left") +fig.suptitle("Depth regimes under the J-lens across gemma-4 scales — green " + "shade: workspace band (REGIMES.json); red dotted: KV-share " + "boundary (entrances above are structurally null)", + fontsize=10.5, y=1.04) +fig.tight_layout() +out = ROOT / "results/regimes_all_models.png" +fig.savefig(out, dpi=140, facecolor="white", bbox_inches="tight") +print("wrote", out)