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

251 lines
14 KiB
Markdown
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.
# 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](results-loop/loop_eval_code.png)
### 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](results-loop/loop_dynamics.png)
**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.
## 4. Related work
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.