"""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")