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>
This commit is contained in:
Nils
Claude Fable 5
@@ -0,0 +1,119 @@
|
||||
"""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")
|
||||
Reference in New Issue
Block a user