273 lines
15 KiB
Markdown
273 lines
15 KiB
Markdown
# 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.9–43.6%** (three seeds, full test
|
||
set, execution-verified); overall accuracy is unchanged-to-slightly-improved
|
||
(51.8% → 51.8–53.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 (L14–30 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 ≈ L14–30 (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=2–4, 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).
|
||
|
||
|
||
|
||
### 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.9–46.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 1–5 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.3–8.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≈3–4
|
||
(cos=1.000), where accuracy and lens-sharpening plateau — extra iterations
|
||
are no-ops, explaining the k-curve shape.
|
||
|
||
|
||
|
||
**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
|
||
|
||
Two recent papers bracket this work. **McLeish et al. (arXiv:2511.07384)**
|
||
retrofit depth-recurrence into pretrained 1B models via layer surgery +
|
||
continued pretraining (~50B tokens, all parameters, Muon, recurrence
|
||
curriculum to r=32): the generic claims "retrofitted recurrence works and
|
||
beats the non-recurrent parent" and "pretrain-then-convert" are theirs, at
|
||
~5 orders of magnitude more training cost than ours. They name layer choice
|
||
as an open problem; our lens-derived band with its causal backing (anchor
|
||
cliff at L14, tap invariance, wrong-band ≈ 0, KV-sharing hazard) is a direct
|
||
answer to it. Unlike their surgery (which needs a healing phase), our k=0
|
||
exactly recovers the base model. **Lys et al. (arXiv:2602.14759)** loop
|
||
frozen models training-free and show naive looping degrades (distribution
|
||
shift) while interpolating with the un-looped state rescues it — independent
|
||
convergent evidence for our anchor-dominant merge; their whole setting
|
||
corresponds to the untrained cell of our attribution table (17.9% hard =
|
||
our FF/untrained level), evaluated by likelihood rather than execution.
|
||
|
||
Earlier lineage: Universal Transformers (adaptive depth); DEQ (fixed-point
|
||
inference); Huginn (arXiv:2502.05171) — prelude/core/coda from scratch;
|
||
Mixture-of-Recursions (arXiv:2507.10524) — learned per-token depth; Relaxed
|
||
Recursive Transformers (arXiv:2410.20672) — uptrained tied layers; Coconut —
|
||
latent CoT; pause tokens (Goyal et al.) — token-space silent compute, whose
|
||
trained-adapter variant proved a near-match for our loop on MBPP (§3.2).
|
||
|
||
What remains distinct here: **interpretability-derived loop placement with
|
||
causal validation** (answering McLeish et al.'s open problem); **a 1.6M-param
|
||
trained merge on a fully frozen base** (between Lys et al.'s free end and
|
||
McLeish et al.'s full-retraining end, and the only one of the three where
|
||
the base model is provably untouched); **prompt-only latent planning with
|
||
bit-exact KV-cache write-in and zero decode cost**; **the attribution
|
||
ladder** (untrained / weights / pause / loop / explicit plan) — neither
|
||
bracket paper runs compute-matched token-space controls; and **difficulty-
|
||
adaptive depth via the STaR-label gate**, named as future work in both.
|
||
|
||
## 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.
|