83 lines
3.6 KiB
Python
83 lines
3.6 KiB
Python
"""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)
|
|
ax.text(sum(band) / 2, 0.51, "workspace", color="#2f855a",
|
|
fontsize=7, ha="center")
|
|
# sensor regime: contiguous pre-band region where the sensor
|
|
# readout dominates (rule reproduces the hand-shaded E2B bands)
|
|
sens = [i for i, (s, m) in enumerate(zip(sensor, motor))
|
|
if s > m and s >= 0.05 and i < band[0]]
|
|
if sens:
|
|
ax.axvspan(min(sens), max(sens) + 1, color="#2b6cb0", alpha=0.08)
|
|
ax.text((min(sens) + max(sens) + 1) / 2, 0.51, "sensor",
|
|
color="#2b6cb0", fontsize=7, ha="center")
|
|
if band[1] + 1 < L - 1:
|
|
ax.axvspan(band[1] + 1, L - 1, color="#c53030", alpha=0.08)
|
|
ax.text((band[1] + L) / 2, 0.51, "motor", color="#c53030",
|
|
fontsize=7, ha="center")
|
|
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)
|