File size: 1,964 Bytes
e418394 | 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 | import json
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
import torch
res = json.load(open("ouro_analysis.json"))
ll = json.load(open("ouro_probes_ll.json"))
NP, R = 48, 4
V = NP * R
bounds = [s * NP for s in range(1, R)]
fig, axes = plt.subplots(1, 3, figsize=(16, 4.2))
ax = axes[0]
ax.plot(res["kl_jlens"], label="Jacobian lens", lw=1.5)
ax.plot(res["kl_logitlens"], label="logit-lens", lw=1.5)
for b in bounds:
ax.axvline(b, color="grey", ls=":", lw=0.8)
ax.set_xlabel("virtual layer (48 layers x 4 loop steps)")
ax.set_ylabel("KL(model || lens) [nats]")
ax.set_title("Lens quality across the unrolled loop (WikiText)")
ax.legend()
ax = axes[1]
ax.plot(ll["top1_agree_jlens"], label="Jacobian lens", lw=1.5)
ax.plot(ll["top1_agree_logitlens"], label="logit-lens", lw=1.5)
for b in bounds:
ax.axvline(b, color="grey", ls=":", lw=0.8)
ax.set_xlabel("virtual layer")
ax.set_ylabel("top-1 agreement with final output")
ax.set_title("Readout agreement by virtual depth")
ax.legend()
ax = axes[2]
norms = [res["J_norm"][str(v)] for v in range(V - 1)]
dist = [res["J_dist_eye"][str(v)] for v in range(V - 1)]
ax.plot(norms, label="||J||/sqrt(d)", lw=1.5)
ax.plot(dist, label="||J - I||/sqrt(d)", lw=1.5)
for b in bounds:
ax.axvline(b, color="grey", ls=":", lw=0.8)
ax.set_xlabel("virtual layer")
ax.set_title("Average Jacobian geometry")
ax.legend()
fig.tight_layout()
fig.savefig("ouro_lens_curves.png", dpi=150)
sim = torch.load("ouro_jsim.pt")["sim"].numpy()
fig2, ax = plt.subplots(figsize=(6.5, 5.5))
im = ax.imshow(sim, vmin=0, vmax=1, cmap="viridis", origin="lower")
for b in bounds:
ax.axvline(b - 0.5, color="w", ls=":", lw=0.8)
ax.axhline(b - 0.5, color="w", ls=":", lw=0.8)
ax.set_xlabel("virtual layer")
ax.set_ylabel("virtual layer")
ax.set_title("cos similarity of J across virtual layers")
fig2.colorbar(im)
fig2.tight_layout()
fig2.savefig("ouro_jsim_heatmap.png", dpi=150)
print("DONE_PLOTS")
|