Files
jspace/scripts/diagram.py
T
NilsandClaude Fable 5 ef9c08966c J-lens workspace reproduction + loop retrofit: lens, band looping, adapters, controls, multi-task evals
Reproduction of the 2026 workspace/J-lens paper on gemma-4 (E2B/12B/26B),
plus the workspace-loop retrofit line: merge adapter, prompt-only latent
planning (MBPP), carry variant, attribution controls (FF/pause/untrained),
band-location ablation, Blocksworld harness, 12B replication scripts.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
2026-07-14 00:54:12 +02:00

120 lines
6.2 KiB
Python
Raw Blame History

This file contains ambiguous Unicode characters
This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.
"""Schematic: the J-lens view of the residual stream + depth regimes + looping."""
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
from matplotlib.patches import FancyBboxPatch, FancyArrowPatch, Rectangle
from matplotlib.patches import ConnectionPatch
fig = plt.figure(figsize=(13, 8.5))
gs = fig.add_gridspec(1, 2, width_ratios=[1.15, 1], wspace=0.05)
axL = fig.add_subplot(gs[0]); axR = fig.add_subplot(gs[1])
for a in (axL, axR):
a.set_xlim(0, 10); a.set_ylim(0, 10); a.axis("off")
# ---------- LEFT: residual stream with regime bands (gemma-4-E2B, 35 layers) ----------
axL.text(5, 9.6, "The residual stream under the J-lens", ha="center",
fontsize=15, fontweight="bold")
axL.text(5, 9.18, "gemma-4-E2B · 35 layers, d=1536", ha="center",
fontsize=9.5, color="#555")
# central stream bar
sx, sw = 4.15, 1.3
axL.add_patch(Rectangle((sx, 0.7), sw, 7.7, facecolor="#f2f2f2",
edgecolor="#888", lw=1.2, zorder=1))
axL.annotate("", xy=(sx+sw/2, 8.75), xytext=(sx+sw/2, 0.55),
arrowprops=dict(arrowstyle="-|>", color="#444", lw=2), zorder=0)
axL.text(sx+sw/2, 0.32, "token embedding (input)", ha="center", fontsize=8, color="#333")
axL.text(sx+sw/2, 8.95, "logits (output)", ha="center", fontsize=8, color="#333")
def y(layer): return 0.8 + (layer/34)*7.35
bands = [
("transduction", 0, 5, "#d9d9d9", "detokenization /\nlexical assembly",
["J-lens: noise / surface form"]),
("sensor", 6, 13, "#bcd4f0", "input written\ninto the stream",
["J-lens: current input token", "(echo ↑ 26%)"]),
("workspace", 14, 30, "#bfe3c0", "abstract concepts\nheld & broadcast",
["J-lens: unspoken content — 'spider', 'big'",
"≈ 1.08B params (58% of decoder)",
"swap-gates fire here · ignition threshold"]),
("motor", 31, 34, "#f2c4c4", "output staged\nfor unembedding",
["J-lens: next output token (→68%)"]),
]
for name, l0, l1, c, fn, lens_lines in bands:
yb, yt = y(l0)-0.10, y(l1)+0.10
ymid = (yb+yt)/2
axL.add_patch(Rectangle((sx, yb), sw, yt-yb, facecolor=c, edgecolor="#666",
lw=1, alpha=0.95, zorder=2))
axL.text(sx+sw/2, ymid, f"L{l0}{l1}", ha="center", va="center",
fontsize=8.5, fontweight="bold", zorder=3)
# band name above-left inside its band
axL.text(sx-0.2, ymid, name, ha="right", va="center",
fontsize=10.5, fontweight="bold", color="#222")
# function (right, top) + lens lines stacked below
axL.text(sx+sw+0.3, ymid+0.28, fn, ha="left", va="center", fontsize=8, color="#222")
for j, ln in enumerate(lens_lines):
axL.text(sx+sw+0.3, ymid-0.16-0.28*j, ln, ha="left", va="center",
fontsize=6.9, color="#556", style="italic")
# three verbs: write in (sensor), hold (workspace), read out (motor) — offset from
# the band-name row so arrows don't clip the labels
axL.annotate("", xy=(sx-0.02, y(12.3)), xytext=(sx-0.7, y(12.3)),
arrowprops=dict(arrowstyle="-|>", color="#2b6cb0", lw=1.8))
axL.text(sx-0.72, y(12.3)+0.2, "write in", ha="right", fontsize=7.5, color="#2b6cb0")
axL.add_patch(FancyArrowPatch((sx-0.08, y(28.5)), (sx-0.08, y(25.5)),
connectionstyle="arc3,rad=-1.0", arrowstyle="-|>",
color="#2f855a", lw=1.7, mutation_scale=11))
axL.text(sx-0.62, y(27), "hold", ha="right", va="center", fontsize=7.5, color="#2f855a")
axL.annotate("", xy=(sx-0.02, y(33.7)), xytext=(sx-0.7, y(33.7)),
arrowprops=dict(arrowstyle="-|>", color="#c53030", lw=1.8))
axL.text(sx-0.72, y(33.7)+0.2, "read out", ha="right", fontsize=7.5, color="#c53030")
# ---------- RIGHT: the looping result ----------
axR.text(5, 9.7, "Can we loop the workspace?", ha="center",
fontsize=15, fontweight="bold")
def box(ax, x, y0, w, h, fc, ec, txt, fs=8, fw="normal", tc="#111"):
ax.add_patch(FancyBboxPatch((x, y0), w, h, boxstyle="round,pad=0.06,rounding_size=0.12",
facecolor=fc, edgecolor=ec, lw=1.4))
ax.text(x+w/2, y0+h/2, txt, ha="center", va="center", fontsize=fs, fontweight=fw, color=tc)
# Case A: band loop (not a self-map) -- fails
axR.text(0.4, 8.9, "A. Loop the whole band (L14→30)", fontsize=10.5, fontweight="bold")
box(axR, 0.6, 7.3, 2.2, 1.0, "#bfe3c0", "#2f855a", "band\nL14 → L30")
axR.add_patch(FancyArrowPatch((2.9, 8.15), (2.9, 7.45),
connectionstyle="arc3,rad=1.3", arrowstyle="-|>", color="#c53030",
lw=2, mutation_scale=14))
axR.text(4.9, 7.8, "out lives 17 layers\ndownstream of in", fontsize=8, color="#333")
axR.text(0.6, 6.75, "✗ NOT a self-map (out-space ≠ in-space)",
fontsize=9, color="#c53030", fontweight="bold")
axR.text(0.6, 6.4, " → collapses in 1 iteration: '8' → ')'",
fontsize=8.3, color="#555")
axR.plot([0.4, 9.6], [5.95, 5.95], color="#ccc", lw=1)
# Case B: single tied layer (self-map)
axR.text(0.4, 5.55, "B. Loop one tied layer h ← h + β·Δ(h)", fontsize=10.5, fontweight="bold")
box(axR, 0.6, 4.0, 2.2, 1.0, "#bcd4f0", "#2b6cb0", "layer L\nin = out")
axR.add_patch(FancyArrowPatch((2.9, 4.85), (2.9, 4.15),
connectionstyle="arc3,rad=1.3", arrowstyle="-|>", color="#2b6cb0",
lw=2, mutation_scale=14))
axR.text(4.9, 4.5, "same residual point\nin and out", fontsize=8, color="#333")
axR.text(0.6, 3.45, "✓ self-map — stable; damping (low-pass) extends it",
fontsize=9, color="#2f855a", fontweight="bold")
axR.text(0.6, 3.05, " β=0.5 @ L26 keeps '8' through ~6 loops; |ΔH| ↓ (converges)",
fontsize=8.3, color="#555")
axR.text(0.6, 2.6, "✗ but frozen → degenerate fixed point; concept never sharpens",
fontsize=9, color="#c53030", fontweight="bold")
axR.text(0.6, 2.2, " extra loop compute ≠ extra reasoning",
fontsize=8.3, color="#555")
# verdict box
box(axR, 0.6, 0.5, 9.0, 1.35, "#fff7e6", "#d69e2e",
"Verdict: mechanically loopable at the layer level, and low-pass\n"
"damping is necessary for stability — but not sufficient. Turning loop-depth\n"
"into reasoning needs the loop trained in (LoRA-per-loop / Huginn / MoR).",
fs=8.6, fw="normal", tc="#7a5a00")
fig.savefig("results/jlens_diagram.png", dpi=140, bbox_inches="tight",
facecolor="white")
print("wrote results/jlens_diagram.png")