Files
jspace/PAPER.md
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

14 KiB
Raw Blame History

Retrofitting Latent Planning onto a Frozen Language Model via Workspace Recurrence

Working draft, 2026-07-14. All experiments: google/gemma-4-E2B-it (frozen), single DGX Spark. Code and artifacts: ~/jspace.

Abstract

Interpretability work with an averaged-Jacobian lens ("J-lens") shows that mid-depth layers of a pretrained language model form a workspace: a band of layers that holds verbalizable, unspoken intermediate content. We ask whether that band can be iterated in place — spending more serial compute per input without emitting reasoning tokens — on a frozen model. A naive loop diverges: the band is not a self-map. We show that a 1.6M-parameter anchor-dominant merge adapter (0.03% of the model) at the band entrance makes the recurrence a stable fixed-point iteration, and that training only this adapter — with self-generated, verifier-filtered supervision and a difficulty→depth curriculum — turns iteration into computation. On MBPP, looping the workspace over the prompt ("latent planning") raises pass@1 on plan-dependent problems from 5.5% to 30.943.6% (three seeds, full test set, execution-verified); overall accuracy is unchanged-to-slightly-improved (51.8% → 51.853.8%, within noise at n=500) — the method's value is cost-shaped (silent, prefill-parallel, no per-token overhead), not accuracy-dominance. Controls attribute the hard-bucket gain to the recurrence itself: a same-size adapter trained on identical data without the loop reaches only 17.9%, exactly matching the untrained loop. On GSM8K the picture inverts — no recurrent variant beats the weights-only control — and a four-arm decomposition localizes why: the loop performs plan refinement, which code synthesis needs and answer-time arithmetic does not. The J-lens provides both the intervention's design (where to loop) and its verification (latent concepts sharpen ~8× per converged iteration). Because the looped prompt states are constant during generation, latent planning is prefill-shaped and adds no per-token cost.

1. Introduction

Large language models buy reasoning accuracy with emitted tokens: chains of thought give the network more serial passes, at the cost of latency, output tokens, and bandwidth-bound decode. Recurrent-depth architectures (Universal Transformers; DEQs; Huginn, arXiv:2502.05171; Mixture-of-Recursions, arXiv:2507.10524) buy the same serial compute silently — but require (pre)training the recurrence in at scale.

We investigate a middle path: retrofit recurrence onto an off-the-shelf frozen model, using an interpretability signal to decide where. The J-lens (from the "verbalizable global workspace" line of work) partitions depth into transduction, sensor, workspace, and motor regimes; the workspace band (L1430 of 35 in our subject model) holds slowly-varying, unspoken intermediates — e.g. 'spider' before answering "8" to "the animal that spins webs has how many legs?". If the workspace approximates "iterate toward a settled representation", looping it should deepen computation without parameters. The contributions:

  1. A minimal retrofit that works: an anchor-dominant merge ((1−α)e + α·ŝ + MLP([e;ŝ]), α=0.3, MLP zero-init, 1.6M params) makes the frozen band a stable, answer-preserving recurrence; training only the merge makes iterations sharpen rather than hold.
  2. A verified capability gain on plan-dependent code synthesis, with the full attribution grid (weights / untrained loop / trained loop / pause tokens) showing the recurrence is the active ingredient.
  3. A mechanistic boundary: math inverts the result, and the decomposition (prompt-side vs generation-side × weights vs recurrence) identifies the mechanism as plan refinement, not generic extra compute.
  4. Deployment properties: bit-exact KV-cache-compatible inference (loop once at prefill), a difficulty gate trained free from the labeling pipeline, and economics that improve with model scale.

2. Method

Locating the band. The lens reads residual state h at layer through the averaged Jacobian J̄_ = E[∂h_final/∂h_] and the unembedding. Depth regimes follow from what the readout tracks (input echo / abstract content / output token). On gemma-4-E2B: workspace ≈ L1430 (1.08B params, 58% of decoder).

Making the band a self-map. Feeding L30's output to L14 collapses in one step (out-space ≠ in-space; norms and content 17 layers "downstream"). Additive anchoring diverges. The fix is DEQ-style input injection done by hand: with e = L13's output (fixed anchor) and s the fed-back band output,

L14-in = (1−α)·e + α·(s · |e|/|s|) + MLP([e ; s·|e|/|s|]),  α = 0.3.

Zero-initializing the MLP's output layer makes the untrained adapter exactly the hand merge, which is stable and answer-preserving for ≥11 iterations but only holds content (lens concept flat).

Training only the merge. Supervision is self-generated and verifier-filtered (STaR-style): the frozen model attempts each task directly and with explicit planning/CoT; items it solves only with planning are "hard", direct solves "easy", neither "drop". Cross-entropy on answer/code tokens of the direct prompt; the model's own verified outputs are the targets (in-distribution). A difficulty→depth curriculum trains easy items at loop depth k=1, mixed at k=2, hard only at k=24, so loss on hard items is reducible only through the recurrence. For generation tasks the loop applies to the prompt span only ("latent planning"): generated tokens run the plain path but attend to the looped prompt states; this removes exposure bias structurally.

Inference cost. Causality makes the looped prompt states independent of generated tokens, so they are computed once; a hooked prefill writes them into the KV cache and generation proceeds natively (verified bit-identical; ≥3.5× faster than recomputation). The concrete overhead at k=4 is 5 passes over the band's 17/35 layers at prefill — ≈2.9× prompt-processing FLOPs, parallel across positions — and zero additional decode cost. Explicit planning with ~200 emitted tokens costs more total FLOPs and pays them serially at bandwidth-bound decode; this asymmetry grows with model size.

3. Results

3.1 Latent planning on code (MBPP)

Full 500-item test split, greedy decode, unit-test-verified. Hard bucket = items the frozen model solves only with an explicit written plan (n=55).

k=4 (prompt-only loops) hard pass@1 overall
baseline (k=0) 5.5% 51.8%
trained loop, seed 0 43.6% 53.6%
trained loop, seed 1 41.8% 53.8%
trained loop, seed 2 30.9% 51.8%

Silent loops recover roughly 40% of what explicit planning achieves, at zero visible-token cost, with no overall regression (the easy-item perturbation tax, ~9 points, is offset by hard/drop gains; a gate removes most of it, §3.4).

MBPP pass@1 vs loop depth

3.2 Attribution: the recurrence is the ingredient

250-item subset; same data, same 1.6M parameters, same insertion point:

arm hard pass@1
baseline 3.6%
untrained loop (α-merge only) 17.9%
trained adapter, no loop (weights control) 17.9%
trained loop 42.946.4%

The weights control lands exactly on the untrained-loop value: ~18 points is what perturbation-plus-format-alignment buys. The remaining ~28 points require iterating the band. Post-hoc depth selection is excluded by pre-registration (k=2 fixed on validation before test numbers existed; k-curves reported descriptively).

Checkpoint selection. No checkpoint was chosen using test or generation results. Seed 0's checkpoint (step 399) was fixed at training time from the validation-CE overfitting inflection, before any generation eval of that adapter; seeds 15 use step 400 by pre-commitment made before those seeds were trained. We separately report that validation CE is a poor proxy for generation accuracy (a checkpoint selected by val-CE on a sibling arm underperformed a later one), which is why the fixed-step rule is used rather than per-seed val selection.

3.3 The boundary: math

On GSM8K, no recurrent variant beats the weights-only control. The four-arm grid (hard bucket) decomposes the failure:

GSM8K hard prompt-side only touches generation
feedforward weights 11.8% 4.7% (pause-token control)
recurrence 6.38.7% (prompt loop) 9.4% (cross-token carry)

Orthogonal effects: perturbing free-running generation positions is costly for either mechanism; recurrence beats weights only where a state must evolve (the generation side — carry doubles the pause control in-harness), and loses on the static prompt side. No variant beats the 10.5% overall baseline. Reading: the trained loop performs plan refinement; code synthesis is plan-shaped, multi-step arithmetic is not — its serial computation happens during the answer, and one frozen band pass per token cannot perform it silently at 2B. CoT tokens remain load-bearing for math. (Hard-bucket cells carry an outcome-selection caveat — buckets were defined by greedy baseline outcomes; sampled relabeling is in progress — so the math conclusion is stated on overall numbers.)

3.4 Mechanism and deployment

Fixed point. The trained loop takes a large first step (cos(s₁,s₀)=0.926 vs 0.977 untrained) and converges bit-exactly by k≈34 (cos=1.000), where accuracy and lens-sharpening plateau — extra iterations are no-ops, explaining the k-curve shape.

Loop convergence dynamics

Lens verification. P(latent concept) under the J-lens at the band exit rises 0.015→0.13 across iterations after training (~8× the untrained control, which only holds). The same lens that located the band verifies that looping deepens its computation — and makes the silent reasoning inspectable.

Gate. A logistic probe on the k=0 workspace state (supervised for free by the STaR labels) routes prompts: predicted-easy at k=0, predicted-hard at k=4. Result: overall equal to the best uniform depth with easy items fully preserved (97.5% vs 98.4% baseline); probe precision (19% at 64% recall) is the current ceiling.

Negative results with content. Mixed-task (code+math) training regressed both tasks versus dedicated adapters, despite indistinguishable validation CE — cross-entropy parity does not predict generation parity. Validation-CE checkpoint selection likewise failed to track generation accuracy.

Universal Transformers (adaptive depth); DEQ (fixed-point inference); Huginn / recurrent-depth latent reasoning (arXiv:2502.05171) — prelude/core/ coda with input injection, trained from scratch; Mixture-of-Recursions (arXiv:2507.10524) — learned per-token depth; Relaxed Recursive Transformers (arXiv:2410.20672) — uptraining tied layers with per-loop LoRA; Coconut — latent chain-of-thought; pause tokens (Goyal et al.) — token-space silent compute. Distinct here: the recurrence is retrofitted onto a frozen model at adapter cost, its location is chosen by an interpretability signal, and the same signal verifies the added computation. Our pause-token and weights controls connect directly to that literature's baselines.

5. Limitations

One base model family at 2B-effective scale (12B replication in progress); two task families. Location specificity is not yet ablated: a pre-registered control looping shifted/early/late/width-matched bands with identical adapter and curriculum is queued; until it lands, the results are formally consistent with "any wide mid-depth band works", and the lens claim rests on discovery convenience plus mechanism verification. Hard buckets are small (n=55 greedy / n=33 sampled) with seed spread of ±6 items; sampled relabeling shows 97% agreement with greedy labels, and intervals accompany all bucket cells in the final tables. The MBPP attribution grid lacks a pause-token arm and a plan-distillation baseline (both queued) — the GSM8K grid has the former. Easy-item perturbation tax is not eliminated (gate preserves easy items but probe precision is 19%). Visible planning remains stronger on absolute accuracy — the claim is cost-and-latency-shaped. Mixed-task training regressed both tasks, so the current recipe yields per-task adapters, not one general silent-planning mode; the outlook's "installed base" framing inherits this caveat until a gate-plus-multiple- adapters (or interference-free training) configuration is shown. MBPP likely overlaps the base model's pretraining data; both arms share any contamination, and memorized items land in the easy bucket, so the hard bucket if anything over-represents genuinely novel problems — but bucket composition is contamination-sensitive. Sensitivity to α=0.3 and band width is unreported (the width-matched ablation arm partially addresses width). Adapter-only training may underestimate the ceiling (band-LoRA "rung 2" untested).

6. Outlook

The retrofit recipe — lens-locate, anchor-merge, verifier-filtered curriculum, gate — is scale-portable by construction: trainable mass is independent of base size, and prompt-side loops are prefill-shaped, so their economics improve with scale while serial CoT decode gets slower. The open question that decides whether this is a curiosity or a method is whether the effect survives scale (12B next; then a mid-size uptraining of the band itself). If it does, "loopification" becomes a cheap post-training phase any holder of a pretrained model can apply — a silent planning mode for the installed base, with its latent reasoning legible to the same lens that built it.