bk.style()
COLOURS = {"master": (AMBER, AMBER), "series": (INK, INK), "other": ("white", INK)}
ORDER = {"master": 0, "series": 1, "other": 2}
cols = sorted({n["pos"] for n in FEATS + ERRS} | {NODES[e["source"]]["pos"] for e in EDGES
if NODES[e["source"]]["type"] == "token"})
x0, x, prev = {}, 0.0, None
for p in cols[:-1]:
if prev is not None and p - prev > 2:
x += 0.7
x0[p] = x
x += 1.3
prev = p
by_layer = {}
for n in FEATS + ERRS:
if n["pos"] == LAST:
by_layer.setdefault(n["layer"], []).append(n)
width = max(len(v) for v in by_layer.values())
xm = x + 0.7
xy = {}
for layer_, ns in by_layer.items():
ns.sort(key=lambda n: (n["type"] != "error", ORDER.get(group(n), 3) if n["type"] == "feature" else 0,
n["influence_cumulative"]))
start = xm + (width - len(ns)) / 2
for i, n in enumerate(ns):
xy[n["id"]] = (start + i, layer_)
for n in FEATS + ERRS:
if n["pos"] != LAST:
xy[n["id"]] = (x0[n["pos"]], n["layer"])
mid = xm + (width - 1) / 2
TOK_Y = -2.6
for p in cols:
xy[f"T@{p}"] = (x0[p] if p != LAST else mid, TOK_Y + 0.9) # edges start at the small triangle
xy[LOGIT["id"]] = (mid, n_layers + 1.8)
fig, ax = plt.subplots(figsize=(10.2, 9.0))
IN_TOTAL = Counter()
for e in EDGES:
IN_TOTAL[e["target"]] += abs(e["weight"])
share = {id(e): abs(e["weight"]) / IN_TOTAL[e["target"]] for e in EDGES} # each edge's share of its target's input
for e in sorted(EDGES, key=lambda e: share[id(e)]):
if e["source"] not in xy or e["target"] not in xy:
continue
(xa, ya), (xb, yb) = xy[e["source"]], xy[e["target"]]
s = share[id(e)] ** 0.5
ax.plot([xa, xb], [ya, yb], color=INK if e["weight"] > 0 else GREY, lw=0.25 + 3.2 * s,
alpha=0.18 + 0.55 * s, ls="-" if e["weight"] > 0 else (0, (3, 2)), zorder=1, solid_capstyle="round")
for a_, b_ in PATH_EDGES: # the amber path, on top of the other edges
(xa, ya), (xb, yb) = xy[a_], xy[b_]
ax.plot([xa, xb], [ya, yb], color=AMBER, lw=4.2, alpha=0.9, zorder=2, solid_capstyle="round")
for n in FEATS:
face, edge = COLOURS[group(n)]
ax.scatter(*xy[n["id"]], s=70, marker="o", color=face, edgecolor=edge, linewidth=1.3, zorder=3)
for n in ERRS:
ax.scatter(*xy[n["id"]], s=62, marker="s", color="#c9cbd2", edgecolor=GREY, linewidth=1.0, zorder=3)
for nid, text, dx, dy in ((QF["id"], f"the door: “Master”, layer {QF['layer']}", 1.9, -2.4),
(CORV["id"], f"“vin” feature, layer {CORV['layer']}", 0.6, 1.6),
(TOPN["id"], f"error node, layer {TOPN['layer']}", -4.6, 1.3)):
px, py = xy[nid]
ax.annotate(text, (px, py), xytext=(px + dx, py + dy), fontsize=9, color=INK, fontweight="bold",
arrowprops={"arrowstyle": "-", "color": INK, "lw": 0.9}, zorder=6,
bbox={"boxstyle": "round,pad=0.2", "fc": PAPER, "ec": "none", "alpha": 0.9})
# The output node and the input words
lx, ly = xy[LOGIT["id"]]
ax.add_patch(FancyBboxPatch((lx - 1.7, ly - 0.55), 3.4, 1.1, boxstyle="round,pad=0.05,rounding_size=0.2",
facecolor=INK, edgecolor=INK, zorder=4))
ax.text(lx, ly, f"output: “{tk_space(LOGIT['token'])}”", ha="center", va="center", color=PAPER, fontsize=10,
fontweight="bold", zorder=5)
for p in cols:
tx, ty = xy[f"T@{p}"]
ax.scatter(tx, ty, s=14, marker="^", color=INK, zorder=3)
ax.text(tx, TOK_Y, tk(TOKEN_AT[p]), ha="center", va="top", fontsize=9.6, color=INK,
fontweight="bold" if p == LAST else "normal")
spans = [("system prompt", [p for p in cols if p < min(q_pos)]),
("question", [p for p in cols if min(q_pos) <= p < 60]),
("start of the reply", [p for p in cols if p >= 60])]
for text, ps in spans:
xs = [xy[f"T@{p}"][0] for p in ps]
ax.plot([min(xs) - 0.35, max(xs) + 0.35], [TOK_Y - 1.25] * 2, color=GREY, lw=1)
ax.text((min(xs) + max(xs)) / 2, TOK_Y - 1.45, text, ha="center", va="top", fontsize=8.8, color=GREY)
ax.set_ylim(TOK_Y - 2.6, n_layers + 3.0)
ax.set_xlim(-0.8, xm + width + 0.2)
ax.set_yticks([0, 5, 10, 15, 20, 25, last_layer])
ax.set_ylabel(f"Layer (0 = first, {last_layer} = last)")
ax.spines["bottom"].set_visible(False)
ax.set_xticks([])
handles = [Line2D([], [], marker="o", ls="", markersize=8, color=AMBER, markeredgecolor=AMBER,
label=f"feature: the word “Master” ({n_master})"),
Line2D([], [], marker="o", ls="", markersize=8, color=INK, markeredgecolor=INK,
label=f"feature: Lamplighter names and titles ({n_series})"),
Line2D([], [], marker="o", ls="", markersize=8, color="white", markeredgecolor=INK,
label=f"feature: other ({n_other})"),
Line2D([], [], marker="s", ls="", markersize=8, color="#c9cbd2", markeredgecolor=GREY,
label=f"error node ({len(ERRS)})"),
Line2D([], [], color=INK, lw=2, label="pushes up"),
Line2D([], [], color=GREY, lw=2, ls=(0, (3, 2)), label="pushes down"),
Line2D([], [], color=AMBER, lw=4, label="one path, strongest push at each step")]
ax.legend(handles=handles, loc="upper left", fontsize=8.8, frameon=False)
fig.suptitle("Wick-mini v1’s answer, traced: every feature sits on the last word, “Master”,\n"
"and the biggest pushes on “Cor” come from grey error nodes",
x=0.01, ha="left", fontsize=11.5, fontweight="bold")
fig.tight_layout()
plt.show()