commit ef9c08966c4dc8b79c578a915a298b77fb92296b Author: Nils Date: Tue Jul 14 00:54:12 2026 +0200 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 diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..20a7c19 --- /dev/null +++ b/.gitignore @@ -0,0 +1,10 @@ +.s3-credentials.txt +.venv/ +__pycache__/ +*.pyc +.DS_Store +._* +results*/ +*.pt +*.pdf +results-loop-prep.log diff --git a/LESSONS.md b/LESSONS.md new file mode 100644 index 0000000..e1b2724 --- /dev/null +++ b/LESSONS.md @@ -0,0 +1,78 @@ +# What worked and what didn't — workspace looping on gemma-4-E2B + +Lab-notebook distillation of the looping investigation (July 2026). Full +results and figures: [`WORKSPACE_LOOPING.md`](WORKSPACE_LOOPING.md); base +reproduction: [`RESULTS.md`](RESULTS.md). Everything on `google/gemma-4-E2B-it` +(frozen), DGX Spark GB10. + +--- + +## Scientific: what WORKED + +| finding | evidence | where | +|---|---|---| +| **Merge layer makes the band a stable recurrence** — anchor-dominant `(1−α)e + α·ŝ`, α=0.3; the whole failure of naive looping is that the band is not a self-map | untrained: answer preserved 11+ loops, converging (probe 2d) | `loop_common.py` MergeAdapter | +| **Training the merge (1.6M params, model frozen) turns hold into amplify** | J-lens concept sharpens 8× across loops (0.015→0.13; untrained flat) | `train_merge.py`, eval JSONs | +| **The loop is a true fixed-point iteration** — big first step, bit-exact convergence by k≈3–4; explains why accuracy & sharpening plateau there, k>4 is a no-op | cos(s_k, s_{k−1}): 0.926 → 1.000 | `plot_loop_dynamics.py` | +| **Prompt-only looping ("latent planning") beats baseline on code** — first overall win: MBPP pass@1 52.0% vs 48.8%; plan-dependent problems 3.6%→46.4% (untrained control 14–18%; p<1e-5) | 250 test items, executed | `eval_loop_code.py`, §5 | +| **Frozen-prompt KV-cache trick** — looped prompt states are constant across token steps (causality), so loop once, hook-swap the band input at prefill, generate cached. Bit-identical, ≥3.5× faster | equivalence test match=True | `generate_frozen_prompt` | +| **STaR self-labeling gives difficulty labels + in-distribution supervision for free** — direct pass vs CoT/plan pass, verified answers/tests | GSM8K 16%/53%, MBPP 52%/63% | `prep_star_data.py`, `prep_mbpp.py` | +| **Difficulty→depth curriculum is absorbed** — hard items (trained only at k≥2) fit best at depth | val CE gradient; hard-bucket peaks at k=4 | train logs | +| **Untrained cross-token carry is stable** — band state carried across token steps, one pass/token; coherent 60+ tokens, mild norm drift; mean/EMA seeding ≥ last-position | `probe_carry.py` output | §6 | +| **Trained sharpening transfers to the carry regime on short answers** — P('spider') 0.015→0.12 in ~2 token steps, amortized | carry probe w/ trained adapter | §6 | + +## Scientific: what FAILED (and what each failure taught) + +| failure | numbers | lesson | +|---|---|---| +| **Naive band loop** (L30 out → L14 in) | collapses in 1 iteration | the band is not a self-map; out-space ≠ in-space | +| **Additive anchor** `H0 + bandΔ(H)` | diverges, ‖ΔH‖/‖H‖≈2/step | forward displacement double-counts; need interpolation, not addition | +| **Frozen single tied layer** | stable but degenerate attractor; concept never sharpens | untrained layers tolerate looping, never benefit (matches Relaxed-Recursive) | +| **α ≥ 0.6 merge** | drifts to garbage | the anchor must dominate; the loop is a perturbation, not a replacement | +| **Full-position looping + teacher forcing (GSM8K)** | overall acc never beats k=0; easy items 97%→~35% in free-run while teacher-forced val looked great | exposure bias: generated tokens through a loop trained only on gold tokens. Loop placement must match deployment. Prompt-only looping fixed this *structurally* | +| **Trained (within-token) adapter in the carry regime, long outputs** | tag-soup / token garbage after ~1 sentence | adapters must be trained on the trajectory they produce; short-answer training does not license long free-running recurrence | +| **Uniform loop depth on mixed difficulty** | easy items pay ~6–13 points at k≥1 even trained | looping problems that need no extra compute only adds perturbation → per-prompt gate (probe on k=0 workspace state; STaR labels = free supervision) is the top v2 item | + +## Operational pitfalls (each cost real time — don't re-hit) + +1. **"Planning hurts" was a truncation artifact.** Plan pass at + `max_new_tokens=380` scored 17% vs 54% direct — every sampled output ran + out of budget mid-plan, before any code. At 700 tokens + terse-plan prompt, + planning adds ~11 points. *Check truncation before believing any + "CoT hurts small models" result.* +2. **`stop_strings="\n```"` matched the OPENING fence** (generations start + `\n```python`) → every generation halted at ~4 tokens → pass@1=0.000 at all + k, including k=0. The k=0 sanity anchor is what caught it. *Always have a + known-value row in every eval.* +3. **Full-vocab logits are the memory monster.** (B, T, 262k) tensors: + batch 32 eval + concurrent training = machine-wide OOM (kernel killed the + desktop session). Then `earlyoom` (installed on DGX Spark, trigger-happy + with unified memory) SIGTERM'd a modest eval *silently* — no traceback, + log just stops. Fixes: `logits_to_keep=1` during generation, batch ≤16, + **GPU jobs sequential**, and a watchdog that notifies when a process + disappears. Check `journalctl | grep earlyoom` for silent deaths. +4. **Re-calling decoder layers appends to the KV cache.** Capture-and-rerun + machinery must forward with `use_cache=False` and strip + `past_key_values` from captured kwargs, or re-runs see doubled keys. + Correctness anchor that caught it: manual band+suffix re-run must + reproduce the plain forward bit-exactly (0.0 max diff). +5. **Small hard pools overfit fast.** 38 hard items: val CE rises from ~step + 300 of 600, depth-ordering inverts by step 500. Snapshot checkpoints and + evaluate the pre-overfit one (we used step 399); ~200 steps suffice. +6. **Background jobs must be `setsid`'d** or the harness/session restart + kills them mid-run. And `pkill -f ` will match your own launcher + shell if the pattern appears in its command line. +7. **Zero-init adapter output layer ⇒ zero grads upstream at step 0** — on + `mlp[0]` this is expected (LoRA-B-style), not a bug; check the output + layer's grad instead. + +## The three design rules that emerged + +1. **Anchor-dominant merge** makes any residual-stream recurrence well-posed; + train only the merge. +2. **Loop placement must match task structure**: prompt-only (static plan) + for generation tasks; full/carry (evolving state) for state-tracking — + and each must be *trained in the regime it deploys in*. +3. **Verify with the lens, gate with the labels**: the J-lens picks the band, + measures whether loops compute, and diagnoses failures; STaR difficulty + labels supervise both the curriculum and (next) the adaptive-depth gate. diff --git a/PAPER.md b/PAPER.md new file mode 100644 index 0000000..71ad700 --- /dev/null +++ b/PAPER.md @@ -0,0 +1,250 @@ +# 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). + +![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.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. + +![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. diff --git a/README.md b/README.md new file mode 100644 index 0000000..1bd32e8 --- /dev/null +++ b/README.md @@ -0,0 +1,62 @@ +# J-lens reproduction on gemma-4-E2B-it + +Reproduction of the core method and experiments of +["Verbalizable Representations Form a Global Workspace in Language Models"](https://transformer-circuits.pub/2026/workspace/index.html) +(Transformer Circuits, 2026) on `google/gemma-4-E2B-it`, running locally on a +DGX Spark (GB10, 128 GB unified memory). + +The paper studies Claude Sonnet/Haiku/Opus; here we test whether its central +tool and headline results transfer to a small open-weights model. + +## Method + +**Jacobian lens.** For each layer ℓ, average the Jacobian from the residual +stream at layer ℓ, position t to the final-layer residual stream at position +t′ ≥ t, over positions and a pretraining-like corpus: + + J̄_ℓ = E[∂h_final,t′ / ∂h_ℓ,t] (d×d per layer) + +Reading an activation: `lens(h) = softmax(W_U · finalnorm(J̄_ℓ h))` → ranked +vocabulary tokens. The J-lens vector of token s at layer ℓ is row s of +`W_U J̄_ℓ`. Swap intervention: `h ← h + V(σ(c) − c)` with `c = V⁺h`, +`V = [v_src; v_tgt]`, σ exchanging the two coordinates. + +**Estimation.** Per prompt we compute the *exact* sum of Jacobians over all +(t, t′≥t) pairs with one forward pass and 1536 batched VJPs +(`is_grads_batched=True`): backward from Σ_t′ h_final,t′ with identity +cotangents; causal masking makes the per-position gradient equal the +row-sum over t′ ≥ t. ~20 s/prompt (bf16, eager attention, seq len 64). +Averaged over 300 fineweb-edu documents ⇒ ~620k (t,t′) samples per layer. + +Sanity check: at the last layer the pair-sum is exactly T·I (verified to 0 error). + +## Layout + +- `jlens/core.py` — model loading, residual capture, Jacobian estimation, + lens reading, J-lens vectors, swap hooks. +- `scripts/compute_jacobians.py` — corpus averaging (`results/jbar.pt`). +- `scripts/exp1_readouts.py` — two-hop spider, Chinese→English intermediates, + directed modulation (citrus), layer profile. +- `scripts/exp2_swaps.py` — spider↔ant answer flip, France→China broadcast, + country-swap grid. +- `scripts/exp4_regimes.py` — sensor/workspace/motor depth regimes + ignition. +- `scripts/plot_results.py` — heatmaps + layer profile plots. + +## Run + +```bash +uv venv && uv pip install --index-url https://download.pytorch.org/whl/cu130 \ + --extra-index-url https://pypi.org/simple torch transformers accelerate datasets matplotlib +python scripts/compute_jacobians.py --n-prompts 300 --dtype bfloat16 +python scripts/exp1_readouts.py results/jbar.pt +python scripts/exp2_swaps.py results/jbar.pt +python scripts/plot_results.py +``` + +Results: see `RESULTS.md` (E2B in `results/`, 12B in `results-12b/`; set +`JLENS_MODEL` and `JLENS_RESULTS` to switch). Headline: J-lens *readouts* +reproduce at both scales and strengthen with size (spider P 0.39→0.80). +Causal *swaps* need the right write basis, which migrates with scale: token +embeddings at 2B (6/30 capital grid), activation-derived concept vectors at +12B (30/30 grid, 4/4 broadcast incl. Seine→Yangtze) — the J-lens gate +transfers unchanged. diff --git a/RESULTS.md b/RESULTS.md new file mode 100644 index 0000000..e66670f --- /dev/null +++ b/RESULTS.md @@ -0,0 +1,256 @@ +# Results: J-lens reproduction on gemma-4-E2B-it (DGX Spark) + +Reproduction of the core claims of *"Verbalizable Representations Form a +Global Workspace in Language Models"* (Transformer Circuits, 2026) on +`google/gemma-4-E2B-it` (35 layers, d=1536, tied embeddings). The paper's +subjects are Claude Sonnet/Haiku/Opus (≈50–130 layers); this tests transfer to +a small open-weights model on local hardware. + +![The residual stream under the J-lens: four depth regimes, and whether the +workspace band can be looped](results/jlens_diagram.png) + +**Jacobian:** J̄_ℓ = E[∂h_final,t′/∂h_ℓ,t] averaged *exactly* over all +(t, t′≥t) pairs of 300 fineweb-edu prompts (seq len 64) ≈ 620k Jacobian +samples per layer; 1536 batched VJPs per prompt, ~21 s/prompt, ~105 min total. +Correctness anchor: the last-layer pair-sum equals T·I exactly. Results were +stable between the 100-prompt checkpoint and the full 300-prompt average. + +## 1. Verbal report / internal reasoning readouts — REPRODUCES + +**Two-hop ("The animal that spins webs has how many legs?" → "8"):** +The unspoken intermediate **spider** is read by the J-lens at +(layer 8, " webs") P≈0.25, peaking at (layer 33, " has") P=0.39, and again at +the answer position (layers 28–31), while candidate numbers +(`two three four seven five six` → `Eight six`) appear at layers 24–28 just +before output. See `results/heat_spider.png`. + +**Multilingual intermediates (小 → 大):** during the Chinese antonym task the +lens reads English/multilingual abstractions mid-network — `opposite`, +`contraire`, `反対`, then `bigger/larger/large/lớn/大き` at layers 24–28 — +before the Chinese output token. Matches the paper's finding that +intermediate computation is partly in a cross-lingual/English-leaning code. + +**Layer structure (paper: three regimes + workspace only in intermediate +layers):** early layers (0–6) read as noise (no stable content); layers ~8–12 +read task-frame concepts (`number`, `answer`, `insect`, `word`); layers ~24–31 +read task *content* (spider, numbers, opposite/big); the final layers become a +motor regime (top lens token = next output token 68% at L34 vs ~0% early). +`results/layer_profile.png`. + +## 2. Directed modulation ("hold citrus in mind") — WEAK/PARTIAL + +Instructed to think about citrus fruits while copying unrelated text, the +model copies perfectly and the lens shows citrus concepts elevated ~30× over +the control condition during generation (max P 0.011 vs 0.0003) — but at tiny +absolute probability. The paper's strong "orange/lemon in the workspace +throughout" does not appear at this scale; only a trace of it. + +## 3. Causal swaps — REPRODUCES WITH ONE ADAPTATION + +The paper's swap `h ← h + V(σ(c)−c)`, `c = V⁺h` along J-lens vectors either +had no effect (gated, few slots) or garbled generation (ungated, all +positions) on this model. Diagnosis: J-lens vectors here are good **readers** +but poor **writers** — the minimal-norm edit that swaps v_src·h ↔ v_tgt·h +erases the source from the lens (P(France): 0.21 → 3e-18) without writing the +target, because all J-lens vectors share dominant J̄ components (raw cosine +France/China 0.74–0.92). + +Working variant (`jlens/core.py: swap_hooks`): **gate with the J-lens** +(intervene only at (layer, position) slots where the lens reads the source +concept, P>thr), but **write along tied token-embedding directions** +(move h's projection on ê_src onto ê_tgt). With thr=0.01, α=1: + +- **Broadcast France→China** (paper §flexible generalization, their success + ~40–53%): capital `Paris→Beijing` ✓, language `French→Chinese` ✓, river + `Seine→Nile` ✓; continent/currency/food unchanged → **3/6 templates**. +- **Capital grid over 6 countries (30 ordered pairs): 6/30 flips**, all among + France/China/Japan(/Canada→France) — exactly the countries the lens reads + strongly; for Egypt/Brazil/Canada as source the gate rarely fires + (concept weakly loaded in the workspace). This mirrors the paper's finding + that swap success tracks the concept's workspace loading, and their + model-size dependence (54–70% on their much larger models). +- **Two-hop spider→ant:** transfers the concept but breaks the format — + output becomes `ant` (α=2) instead of the hoped-for `6`; reverse direction + no effect (ant is weakly loaded: 2 gated slots vs spider's ~9). +- Controls: no-swap and spider→piano leave the answer at `8`; the paper-basis + write (`write="jlens"`) leaves it at `8`. + +## 4. J-space capacity (paper: ~10% of variance, k≤25 active) — SHAPE ONLY + +Non-negative matching pursuit (k≤25 J-lens vectors, NNLS refit) explains +**<1% of residual variance in intermediate layers** (peak 0.9% at L25) and +**22% at the final layer**. Qualitatively consistent with a small privileged +subspace distinct from the bulk of activation variance, but the mid-network +occupancy is an order of magnitude below the paper's ~10% — either a real +scale difference or a limitation of our 300-prompt J̄ estimate. + +## 5. Sensor / workspace / motor regimes and ignition — REPRODUCES + +Per-layer diagnostics over 8 prompts (`scripts/exp4_regimes.py`, +`results/regimes.png`): + +- **Sensor: layers ~6–13.** Top J-lens token equals the *current input* token + at up to 26% of positions (exactly 0% everywhere after L13). Low + persistence across positions — content changes with each token. +- **Workspace: layers ~14–30.** Input echo vanishes abruptly at L14; + cross-position persistence of lens content *peaks* (top-10 Jaccard ~0.2–0.26 + at L14–21 vs ~0.05 in the sensor band) — content held stable across + positions; this is also where abstract task content lives (spider, numbers, + opposite/big at L19–31 in exp1) and where gated swaps fire. +- **Motor: layers ~31–34.** Top lens token equals the *next output* token at + 42–48% of positions; J-space R² jumps to 22%. + +**Parameters per regime** (decoder layers, of 1.86B total): early 0–5 = +224M (12%), sensor 6–13 = 297M (16%), workspace 14–30 = 1.08B (58%), motor +31–34 = 261M (14%). Notably the architecture itself pivots exactly at the +observed sensor→workspace boundary: from layer 15 the model switches to +double-wide MLPs with cross-layer-shared KV (63.7M/layer vs 36.2M/layer +before), i.e. the designers put the parameter mass where the workspace +content lives. (Non-decoder params: 403M tied embeddings, 2.36B per-layer +embedding machinery, 476M vision/audio towers; 5.1B raw total.) + +**Ignition** (paper: ambiguous inputs produce sharp binary commitment at +workspace onset, not proportional mixing): replacing one token's embedding +(main + per-layer) with w·concept1 + (1−w)·concept2 and reading the lens at +that position (`results/ignition.png`, `ignition_atpos.pt`): + +- Sensor layer 13 mixes **gradually**: e.g. ocean/violin commitment + C = −0.92, −0.70, −0.41, −0.12, +0.36, +0.92, +1.00 across w = 0→1. +- Deep layer 31 commits **all-or-nothing** at w≈0.5: −1.00, −1.00, −1.00, + +0.97, +1.00 (dog/piano); same step shape for all three concept pairs. + +## Verdict + +| Paper claim | This model | +|---|---| +| Unspoken intermediates readable via J-lens | ✓ strong (spider, big/opposite) | +| Three depth regimes; workspace = intermediate layers | ✓ | +| Cross-lingual abstract code mid-network | ✓ | +| Directed modulation loads workspace | ~ trace only | +| J-lens coordinate swaps flip downstream behavior | ✓ after adaptation (embedding write basis, J-lens gating); success tracks workspace loading | +| Broadcast/flexible generalization of swapped concept | ✓ 3/6 templates | +| J-space is small fraction of activation variance | ✓ shape (but <1% vs their ~10%) | +| Ignition: binary commitment from workspace onset | ✓ (graded at L13, step function at L31) | + +Overall: the *readable global-workspace* phenomenology transfers remarkably +well to a 2B-class open model; the *causal* properties transfer partially and +require respecting the model's write basis (tied embeddings) rather than the +J-lens basis itself. + +## 6. Scale comparison: gemma-4-12B-it (48 layers, d=3840) + +Same pipeline, J̄ from 150 fineweb-edu prompts (177 s/prompt, ~7.4 h on the +GB10). Logs/plots in `results-12b/`. + +**Readouts get much stronger with scale** (paper predicts this): + +| readout | E2B | 12B | +|---|---|---| +| unspoken "spider" peak P | 0.39 | **0.80** | +| English "big/large" in Chinese task | 0.031 | **0.33** | +| directed modulation (citrus vs control) | ~30× trace | none detected | + +**The write basis migrates with scale — the central new finding.** The gated +swap that worked on E2B with *token-embedding* writes fails on 12B with the +same erase-without-write signature the J-lens basis showed on E2B (France +deleted — "Please specify which country" — but China never written). What +works on 12B: **activation-derived concept vectors** (mean residual at the +concept token over 4 contrastive templates, per layer) as the write +direction, still gated by the J-lens reading. With that: + +- **Capital grid: 30/30** ordered country pairs flip correctly + (E2B embedding-write: 6/30), α robust over 0.5–2.0. +- **Broadcast 4/4**: Paris→Beijing, French→Chinese, Europe→Asia, and + Seine→**Yangtze** (the correct Chinese river, where E2B gave generic + "Nile"). Exceeds the paper's reported broadcast rates (~40–53%). +- Two-hop spider→ant still does not flip the "8" at 12B (attribute lookup + appears to route around the gated slots). + +Summary of write bases: J-lens vectors read everywhere but never write on +gemma (unlike the paper's Claude models); token embeddings write at 2B +(small tied model keeps concepts embedding-aligned); by 12B concepts have +rotated into model-specific activation directions. The J-lens *gate* +(where/when to intervene) transfers across scales unchanged. + +**Regime geometry shifts deeper (fractionally):** sensor echo spans +L11–35 (peak 0.36 at L25, ≈0.23–0.73 fractional vs 0.17–0.40 on E2B); +abstract workspace content (spider P=0.80, swap gates) concentrates at +L36–45 (≈0.75–0.94); motor alignment ramps only in the last few layers and +is weaker under our teacher-forced metric (0.12 vs 0.48). `results-12b/regimes.png`. + +**J-space occupancy** rounds to 0.000 at all sampled intermediate layers +(k≤25 pursuit; even sparser than E2B's <1%). The final-layer slot was not +sampled at 12B (stride artifact). + +**Ignition**: the clean graded-vs-binary contrast we measured on E2B did not +replicate cleanly on 12B with the at-position protocol — mixed-token +readings there are dominated by contextual priors rather than the injected +embedding. Would need the paper's exact protocol to adjudicate. + +## 7. MoE: gemma-4-26B-A4B-it (30 layers, d=2816, 128 experts, top-8, ~4B active) + +Same pipeline, J̄ from 100 fineweb-edu prompts. The exact-Jacobian method +survives the router untouched: the vmapped backward traverses the top-8 +expert gather/scatter with **0.0 last-layer identity error** — averaging +Jacobians over prompts that each route through *different* experts still +yields a coherent lens. Cost 487 s/prompt (~13.5 h), ~4× the active-param +prediction because the backward decomposes into scattered per-expert matmuls. +Logs/plots in `results-26b/`. + +**Readouts are the strongest of all four models:** + +| readout | E2B | 12B | 26B-MoE | +|---|---|---|---| +| unspoken "spider" peak P | 0.39 | 0.80 | **1.00** | +| English "big/large" in Chinese task | 0.031 | 0.33 | 0.36 | + +**Directed modulation finally works — and only here.** "Hold citrus in mind +while copying unrelated text": citrus concepts reach **P=0.69** in the J-lens +*during generation* (mean 0.005) vs 0.006 in the control — the paper's +"orange/lemon held in the workspace throughout" result, which was a bare +trace at E2B and undetectable at 12B. The routed model is the first to show +it clearly, plausibly because a dedicated expert can maintain the held +concept without disrupting the copy stream. + +**Write basis matches 12B:** embedding-write swap fails (0/30 grid, garbled +output); **activation-derived concept vectors give 30/30** on the capital +grid, J-lens-gated. Confirms the 12B finding that above ~10B, mid-network +concepts leave the embedding basis — and that this is about scale, not the +dense/MoE distinction. + +**Regimes are weaker and flatter** (`results-26b/regimes.png`): sensor echo +peaks at just 0.13 (L10–19) vs 0.26 (E2B) / 0.36 (12B); workspace content +L20–27; motor only the last 2 layers. Routing appears to spread the +workspace across experts, so no single J-lens direction dominates any slot — +consistent with **J-space R² ≈ 0.000 at every sampled layer** (even sparser +than the dense models). The workspace is present (readouts prove it) but +diffuse. + +## Four-model summary + +| | E2B | 12B | 26B-A4B (MoE) | 31B | +|---|---|---|---|---| +| layers / d_model | 35 / 1536 | 48 / 3840 | 30 / 2816 | 60 / 5376 | +| J̄ prompts | 300 | 150 | 100 | — | +| Jacobian s/prompt | 21 | 177 | 487 | ~500 (est) | +| unspoken "spider" P | 0.39 | 0.80 | 1.00 | — | +| directed modulation | trace | none | **P=0.69** | — | +| swap write basis | embedding | activation | activation | — | +| capital grid | 6/30 | 30/30 | 30/30 | — | +| J-lens gate (where) | ✓ | ✓ | ✓ | — | + +Consistent story across scale and architecture: **readouts strengthen with +scale**; the **J-lens gate** (where to intervene) is universal; the **write +basis migrates** from token-embedding (2B) to model-specific activation +directions (≥12B, dense or MoE); the MoE keeps a workspace but a diffuse one. +The one property that appeared only at the MoE — robust directed modulation — +is the paper's clearest "deliberate control" signature. + +## Repro notes + +- Files: see README. Full logs: `results/exp{1,2,3}.log`; plots + `results/*.png`; averaged Jacobian `results/jbar.pt` (330 MB, fp32). +- Total compute: ~2 h on the GB10 for J̄ + ~30 min for all experiments. +- Caveats: single-token concepts only (paper limitation too); J̄ from 300 + prompts at seq len 64 (paper: ~1000 prompts, unknown length); bf16 + gradients (independent-run correlation of J̄ entries ≥0.996 layer 5+). diff --git a/WORKSPACE_LOOPING.md b/WORKSPACE_LOOPING.md new file mode 100644 index 0000000..fba0cab --- /dev/null +++ b/WORKSPACE_LOOPING.md @@ -0,0 +1,317 @@ +# Latent workspace recursion via a trained merge layer + +A design note on one idea that came out of the J-lens work: **can we loop the +workspace band of a pretrained model to spend more compute per token — and if +so, how do we make it actually improve reasoning rather than just idle?** + +Reproduction results this builds on are in [`RESULTS.md`](RESULTS.md); the +schematic is [`results/jlens_diagram.png`](results/jlens_diagram.png); a +distilled worked/failed/pitfalls ledger is in [`LESSONS.md`](LESSONS.md). All +experiments below are on `google/gemma-4-E2B-it` (35 layers, d=1536) on a DGX +Spark. + +--- + +## 1. Why the workspace is a candidate for looping + +The J-lens analysis gives every layer a readout of "what verbalizable concept +is here." Across models that partitions depth into four regimes: + +| regime | E2B layers | what the J-lens reads | +|---|---|---| +| transduction | 0–5 | noise / surface lexical form | +| sensor | 6–13 | the current input token (echo ↑26%) | +| **workspace** | **14–30** | **unspoken abstract content ('spider', 'big'); held & broadcast** | +| motor | 31–34 | the next output token (→68%) | + +Two measured properties make the workspace look like a recurrent variable that +is currently unrolled across distinct layers rather than iterated: + +- **It is slow / persistent.** Top-10 J-lens content overlaps ~0.25 (Jaccard) + between adjacent token positions — the workspace barely changes token to + token. +- **It holds unspoken intermediates.** On "the animal that spins webs has how + many legs?" the lens surfaces **spider** mid-network before the model emits + **8**; on multi-step arithmetic it stages `21 → 42 → 49` across layers. + +If that band is effectively approximating "iterate toward a settled abstract +representation," then **looping it should be a way to add sequential compute +without adding parameters** — the premise behind Universal Transformers, +Deep Equilibrium Models, and recent recurrent-depth LLMs (Huginn, Ouro, +Mixture-of-Recursions). + +Workspace-band parameter mass (decoder only), for reference: + +| model | band | params in band | active/token | +|---|---|---|---| +| E2B | L14–30 | 1.08 B (58%) | 1.08 B | +| 12B | L36–45 | 2.26 B (21%) | 2.26 B | +| 26B-MoE | L20–27 | 6.53 B (27%) | ≈0.82 B | + +--- + +## 2. What we tried, and what actually happens + +Probe: inject a residual at the band entrance, run the band, feed the result +back, and track (a) convergence `|Δ|/|·|`, (b) the J-lens 'spider' concept, +(c) the emitted answer. Baseline answer is `8`. + +### 2a. Naive band loop — fails (ill-posed) +Feed L30's output straight back into L14. **Collapses in one iteration** +(`8 → )`). Reason: the band is **not a self-map** — its output lives 17 layers +of accumulation downstream of its input (larger norm, "settled" content), so +re-injecting it double-counts the band's contribution and leaves the input +distribution L14 expects. + +### 2b. Anchored `H ← H0 + bandΔ(H)` — also fails +Merging the fixed pre-band residual `H0` as an additive anchor does not help +on its own: `bandΔ` is a large forward displacement, so `H0 + bandΔ` lands +back in post-band space and diverges (`|ΔH|/|H| ≈ 2` per step). + +### 2c. Single tied layer `h ← h + β·Δ_L(h)` — well-posed but not useful +One layer *is* a genuine self-map (same residual point in and out). Result: +**stable, converges** (`|ΔH|` decreases monotonically), and **low-pass damping +extends it** — β=0.5 at L26 keeps `8` through ~6 loops vs. immediate decay at +β=1.0. But a **frozen** layer has no useful fixed point: pushed far enough it +relaxes to a degenerate attractor (`Set`, `number`, …) and the J-space concept +**never sharpens**. Confirms the Relaxed-Recursive-Transformer lesson: layers +not trained to be looped only tolerate it, they don't benefit. + +### 2d. Merge layer between L13 and L14 — the fix (structurally) +Insert a merge that reconciles the looped-back post-band state `s` with the +fixed pre-band anchor `e = L13 output`: + +``` +L14_in = (1−α)·e + α·(s renormalized to |e|) +``` + +Untrained, anchor-dominant (**α=0.3**): + +| k | `|Δs|/|s|` | `|s|/|e|` | answer | P(ans) | +|---|---|---|---|---| +| 0 | 0.00 | 2.23 | `8` | 1.000 | +| 1 | 0.29 | 2.11 | `8` | 1.000 | +| 5 | 0.12 | 2.09 | `8` | 1.000 | +| 10 | 0.10 | 2.09 | `8` | 1.000 | + +The band is now a **well-posed, stable, answer-preserving recurrence**: +converges to a fixed point (`|Δs|` 0.29→0.10), norm controlled, `8` intact +through 11+ loops — versus collapse-in-one-step for the naive loop. + +Caveats: it only stabilizes when the **anchor dominates** (α≥0.6 drifts to +garbage), and it **holds rather than sharpens** — the 'spider' J-space concept +fluctuates (~0.01–0.04) rather than growing. Untrained, the merge buys a +stable substrate, not extra reasoning. + +**This is the DEQ / Huginn "input-injection" structure**: a fixed input +(anchor) merged with an iterated state makes the core a self-map. Our merge is +that injection, done by hand. + +--- + +## 3. The proposal: train *only* the merge layer + +The merge layer converts the ill-posed loop into a stable, answer-preserving +substrate, and it is the **single small component to train** — a little +MLP / LoRA at L13→L14 — to make iterations *improve* the answer, with L14–30 +and the rest of the model frozen. This is the cheapest possible "train the +loop in." + +``` + e (L13 out, fixed anchor) + │ + s_k ──▶ [ merge adapter θ ] ──▶ L14 … L30 ──▶ s_{k+1} (× k loops) + │ + L31 … L34 ─▶ logits (answer) +``` + +### Objective (self-contained, no bigger teacher) +Latent chain-of-thought distillation (Coconut / STaR-style): +1. Let the **frozen base model** produce CoT answers; keep only traces whose + final answer is **correct** (rejection sampling). +2. Train the adapter so the **latent looped** forward — k loops, *no CoT + emitted* — reaches that correct answer. Loss = cross-entropy on answer + tokens only; adapter is the sole trainable tensor; band unrolled k times. + +The supervision is answers the model can reach *with* CoT but not *without* — +so loop depth substitutes for reasoning tokens, and everything stays +in-distribution for the frozen base. + +### Training data +- **GSM8K** (already cached: `openai/gsm8k`) — multi-step arithmetic; single + verifiable numeric answer; exactly the `21→42→49` staging we observed. +- **Synthetic multi-hop templates** (the spider family) — controllable hop + count for a difficulty curriculum, and known intermediates so the J-lens can + check whether looping sharpens them. Optionally real 2-hop QA + (2WikiMultiHop / StrategyQA) for naturalness. +- **Avoid**: raw pretraining text (no reasoning signal → re-teaches + preservation), pure 1-pass self-distillation (target = no-loop answer, loop + can't beat it), bigger-model distillation (conflates "loop helps" with + "bigger helps"). + +### The design point that makes or breaks it +Force the loop to be **used**: train across variable k with a +**difficulty→depth curriculum** (easy items solvable at k=1, hard items only +at k>1). Otherwise the adapter learns to solve everything at k=1 and ignores +the recurrence — the "limited evidence of latent CoT" failure the +recurrent-depth interpretability papers report. + +### Evaluation (closes the loop with our tooling) +- **Accuracy vs. k** on held-out two-hop templates + GSM8K test. Success = + the curve **rises with k** — the thing that flatly did *not* happen + untrained. +- **J-lens sharpening**: does the intermediate concept (spider, arithmetic + partial) grow across loops after training, where it only fluctuated before? + +--- + +## 4. Results of the trained run + +The §3 experiment was run (July 2026, ~25 min on the GB10): 1,024 GSM8K train +items labeled by the frozen model (direct pass vs CoT pass → 165 easy / 416 +hard / 443 dropped as unreachable; baselines 16% direct, 53% CoT), then 800 +steps of adapter-only training (1.6M params, batch 8, CE on answer tokens) +with the curriculum k=1 easy / k=2 mixed / k=4 hard-only. Eval: greedy decode +through the looped forward on 256 held-out test items (frozen-model baselines +there: 11.3% direct, 55.9% CoT), trained adapter vs the untrained α-merge. + +![accuracy vs k and J-lens sharpening](results-loop/loop_eval.png) + +| k | untrained all | trained all | untrained hard | trained hard | +|---|---|---|---|---| +| 0 | .117 | .117 | .008 | .008 | +| 1 | .082 | .039 | .016 | .016 | +| 2 | .094 | .090 | .039 | **.063** | +| 4 | .090 | .082 | .047 | **.063** | +| 8 | .086 | .082 | .039 | .055 | + +**What worked (the pre-registered J-lens criterion).** After training, loops +*sharpen* the latent concept instead of merely holding it: P('spider') under +the lens at L30 rises 0.015 → 0.107 → 0.130 across k=0→4, ~8× the untrained +control (which stays at 0.013–0.033, as in probe 2d), with the answer intact +(P(8)≈1.0) and the recurrence stable through k=8. Training the merge changed +the loop from a preserver into an amplifier. + +**What partially worked.** On the hard (CoT-only) bucket, accuracy rises with +depth: 0.8% at k=0 → 6.3% at k=2–4 trained, vs 3.9–4.7% untrained. Depth buys +a small number of answers the model cannot produce in one pass — but it is +8/127 items, not a CoT replacement. + +**What failed.** Overall accuracy does *not* rise with k: pushing free-running +generation through the loop damages easy items (97% → 21–35% trained, → 41–52% +even untrained), which swamps the hard-bucket gain; k=0 remains best overall. +The teacher-forced val looked much better than free-run decode — an +exposure-bias gap: training only ever saw gold answer tokens, generation feeds +back its own. Note the damage is mostly there in the untrained loop path too, +so it is a property of the substrate, not something training introduced. + +**Obvious next steps.** (1) Free-run/scheduled-sampling training to close the +exposure-bias gap; (2) a k=0-preservation loss on easy items so looping never +costs known answers; (3) per-token adaptive depth (MoR-style) instead of +uniform k; (4) more data at grade-school difficulty (Orca-Math) and a +contamination-free eval with a difficulty dial (GSM-Symbolic variants mapped +to target k); (5) bigger adapter / more steps — 800 steps on 517 items is a +first probe, not a training run. + +--- + +## 5. Latent planning on code (MBPP) — first overall win + +Transposing to coding with one structural change: the loop applies **only to +the prompt span** (`loop_mask`) — the model "plans silently" while reading the +problem, then writes code through the plain path (generated tokens never loop, +so the §4 exposure-bias failure is absent by construction). Since causality +makes the looped prompt states constant across token steps, they are computed +once and frozen; generation then runs at native speed with a KV cache +(`generate_frozen_prompt`, verified bit-identical to the recomputing path). + +Data: MBPP (974 basic Python problems, 3 asserts each). Frozen-model labeling: +direct pass@1 52-57%; plan-first 63% reachable (**caution**: with a 380-token +budget the plan pass scored 17% — "planning hurts" was purely truncation; at +700 tokens planning adds ~11 points). Labels: ~260 easy / 66 hard (plan-only) +/ rest drop. Training: same recipe, supervision = the model's own verified +passing code, CE on code tokens, curriculum k=1 easy / 2 mixed / 4 hard; the +38-item hard pool overfits past ~step 300 (evaluated the step-399 snapshot). + +![MBPP pass@1 vs k](results-loop/loop_eval_code.png) + +| k | untrained all | trained all | untrained hard | trained hard | trained easy | +|---|---|---|---|---|---| +| 0 | .488 | .488 | .036 | .036 | .984 | +| 1 | .472 | .488 | .179 | .429 | .852 | +| 2 | .476 | **.520** | .179 | .357 | .926 | +| 4 | .472 | .512 | .143 | **.464** | .885 | + +- **Overall pass@1 beats the no-loop baseline** (52.0% vs 48.8% at k=2) — the + criterion the GSM8K run failed. The untrained loop never does (≤47.6%). +- **Planning-dependent problems: 3.6% → 46.4%** at k=4 (26/56 items; McNemar + p<1e-5). Silent loops recover ~46% of what explicit written planning + achieves. Training matters: the untrained substrate reaches only 14-18%. +- The **easy-item dip** (98→89%) is substrate-inherent (untrained dips + equally); training partially repairs it (92.6% vs 88.5% at k=2). Fix on the + board: a per-prompt **gate** (probe on the k=0 workspace state predicting + "will looping help?" — the STaR labels are its free supervision; a perfect + gate scores ~86% on this mix). + +### Rung-2 design note: depth-graded band unfreezing + +When the band itself is unfrozen (per-iteration LoRA), unfreeze +**entrance-faded**: full trainability at L14 decaying to frozen by ~L22 +(e.g. LoRA α × max(0, 1−(ℓ−14)/8)). Rationale: the only novel inputs in the +system are what the first band layers see (merged fed-back states); the +band's *exit* distribution must stay on-manifold for the frozen motor/suffix +consumers, and the k-fold application of the band amplifies any trained +change — a depth-decaying profile is a trust region. Risk: if sharpening is +implemented by late-band broadcasts, capping their plasticity caps gains. +Ablation triplet (shares everything but the LoRA mask): uniform / +entrance-faded / exit-faded (falsification control — parity would mean +placement is irrelevant and only capacity matters). Composes with, but is +run before, per-iteration strength fading (Relaxed-Recursive axis). + +## 6. Loop dynamics & cross-token carry probes + +**Dynamics** ([`results-loop/loop_dynamics.png`](results-loop/loop_dynamics.png)): +the trained within-token loop is a true fixed-point iteration — a large first +step (cos(s₁,s₀)=0.926 vs 0.977 untrained), then bit-exact convergence by +k≈3-4 (cos=1.000), norm settled ~1.7% above the untrained level. Convergence +depth coincides with where accuracy and J-lens sharpening plateau; k>4 is a +literal no-op. + +**Cross-token carry** (`probe_carry.py`): instead of k loops per token, carry +the band output across token steps — one band pass per token seeded with the +previous step's state (the amortized loop; cf. Feedback Transformer). +Untrained: stable and coherent over 60 free-running tokens (all three seeding +modes — last-position / mean / recency-weighted — with mean/EMA marginally +cleaner), only a mild norm drift (+12%/60 tokens). With the §4-trained +adapter: the 8× sharpening **transfers** to the carried regime on short +answers (P('spider') 0.015→0.12 in ~2 token steps) but long free-running text +collapses — the adapter never saw its own continuations. Conclusion: the +carry loop must be trained *as* the carry loop (sequential unroll on its own +trajectory; planned objective: fixed-point distillation — train one carried +band pass to land where the within-token k-loop converges — plus STaR +rollouts). + +--- + +## 7. Status & related work + +**Status.** Probes 2a–2d, the trained-merge experiments (§3–5), the dynamics +analysis and carry probes (§6) are run and reproducible on E2B: +`scripts/loop_common.py` (band re-run machinery, bit-exact; `loop_mask`; +`generate_frozen_prompt` fast path), `prep_star_data.py` / `prep_mbpp.py` (+ +`prep_mbpp_fix.py`), `train_merge.py` / `train_merge_code.py`, `eval_loop.py` +/ `eval_loop_code.py`, `probe_carry.py`, plotting scripts; data, adapters and +eval JSONs in `results-loop/`. Toward a paper: needs seeds/full test sets + +CIs, the pause-token (equal-FLOPs) baseline, a second model scale (12B), a +contamination-free eval (GSM-Symbolic), and the per-prompt gate; the +cross-token carry (§6) is the successor experiment. + +**Related work.** Universal Transformers (adaptive-depth recurrence); Deep +Equilibrium Models (solve for the fixed point directly); Geiping et al. +[recurrent-depth latent reasoning / "Huginn"](https://arxiv.org/abs/2502.05171) +(prelude→looped core→coda with input injection); [Mixture-of-Recursions](https://arxiv.org/abs/2507.10524) +(per-token learned recursion depth); [Relaxed Recursive Transformers](https://arxiv.org/abs/2410.20672) +(layer-tying + per-loop LoRA); Coconut (latent chain-of-thought). The +novel piece here is using an **interpretability signal (the J-lens)** both to +*choose the band to loop* and to *measure whether looping deepens computation*, +and training **only a merge adapter** on top of a frozen pretrained model. diff --git a/jlens/__init__.py b/jlens/__init__.py new file mode 100644 index 0000000..c97a2f8 --- /dev/null +++ b/jlens/__init__.py @@ -0,0 +1 @@ +from .core import JLens, load_model diff --git a/jlens/core.py b/jlens/core.py new file mode 100644 index 0000000..41b328b --- /dev/null +++ b/jlens/core.py @@ -0,0 +1,261 @@ +"""Jacobian lens (J-lens) for gemma-4-E2B-it. + +Reproduces the method of "Verbalizable Representations Form a Global +Workspace in Language Models" (transformer-circuits.pub/2026/workspace): + + J_l = E_{prompt, t, t' >= t} [ d h_{final, t'} / d h_{l, t} ] (d x d per layer) + lens(h_l) = softmax(W_U . finalnorm(J_l h_l)) + J-lens vector of token s at layer l: v_s = (W_U J_l)[s, :] = W_U[s] J_l + swap: h <- h + V (sigma(c) - c), c = pinv(V) h, V = [v_s; v_t] + +h_l is the residual stream at the *output* of decoder layer l. +h_final is the output of the last decoder layer (pre final-RMSNorm). +""" + +import os + +import torch +from transformers import AutoModelForCausalLM, AutoTokenizer + +MODEL_ID = os.environ.get("JLENS_MODEL", "google/gemma-4-E2B-it") + + +def load_model(dtype=torch.float32, device="cuda", model_id=None): + model_id = model_id or MODEL_ID + tok = AutoTokenizer.from_pretrained(model_id) + model = AutoModelForCausalLM.from_pretrained( + model_id, dtype=dtype, attn_implementation="eager" + ).to(device).eval() + torch.backends.cuda.matmul.allow_tf32 = True + torch.backends.cudnn.allow_tf32 = True + return model, tok + + +def _text_model(model): + return model.model.language_model + + +class ResidualCapture: + """Forward hooks that record each decoder layer's output residual stream.""" + + def __init__(self, model, layers=None, detach=True): + self.tm = _text_model(model) + self.layers = list(range(len(self.tm.layers))) if layers is None else layers + self.detach = detach + self.acts = {} + self.handles = [] + + def __enter__(self): + for i in self.layers: + def hook(mod, inp, out, i=i): + self.acts[i] = out.detach() if self.detach else out + self.handles.append(self.tm.layers[i].register_forward_hook(hook)) + return self + + def __exit__(self, *a): + for h in self.handles: + h.remove() + + +@torch.no_grad() +def collect_residuals(model, input_ids): + """Return (n_layers, T, d) residual stream stack for a single prompt.""" + with ResidualCapture(model) as cap: + model(input_ids=input_ids) + n = len(_text_model(model).layers) + return torch.stack([cap.acts[i][0] for i in range(n)]) + + +@torch.no_grad() +def generate_with_residuals(model, tok, input_ids, max_new_tokens=40): + """Greedy-generate while capturing the residual stream at every position + (prompt + generated). Returns (text, ids (T,), residuals (L, T, d)).""" + tm = _text_model(model) + L = len(tm.layers) + steps = {i: [] for i in range(L)} + handles = [] + for i in range(L): + def hook(mod, inp, out, i=i): + steps[i].append(out[0].detach()) + handles.append(tm.layers[i].register_forward_hook(hook)) + try: + out = model.generate(input_ids, max_new_tokens=max_new_tokens, + do_sample=False, + pad_token_id=tok.pad_token_id or tok.eos_token_id) + finally: + for h in handles: + h.remove() + hs = torch.stack([torch.cat(steps[i], dim=0) for i in range(L)]) + text = tok.decode(out[0, input_ids.shape[1]:], skip_special_tokens=True) + return text, out[0], hs + + +def prompt_jacobian_pairsum(model, input_ids, chunk=128): + """Sum over all pairs (t, t' >= t) of d h_{final,t'} / d h_{l,t} for one prompt. + + One forward pass; backward from S = sum_{t'} h_{final,t'} with batched + identity cotangents. The gradient of w.r.t. h_{l,t} is row i of + sum_{t' >= t} J^{(t',t)}_l (causality zeroes t' < t), so summing the + gradient over t gives row i of the pair-sum. Returns (fp32 (L, d, d), n_pairs). + """ + tm = _text_model(model) + L = len(tm.layers) + d = model.config.get_text_config().hidden_size + T = input_ids.shape[1] + dev = input_ids.device + + with ResidualCapture(model, detach=False) as cap: + with torch.enable_grad(): + model(input_ids=input_ids) + hs = [cap.acts[i] for i in range(L)] # each (1, T, d), in graph + S = hs[-1].sum(dim=1) # (1, d) + + J = torch.zeros(L, d, d, device=dev, dtype=torch.float32) + eye = torch.eye(d, device=dev, dtype=S.dtype) + starts = list(range(0, d, chunk)) + for ci, s in enumerate(starts): + e = s + min(chunk, d - s) + V = eye[s:e].unsqueeze(1) # (B, 1, d) cotangents + grads = torch.autograd.grad( + outputs=S, inputs=hs, grad_outputs=V, + retain_graph=(ci < len(starts) - 1), is_grads_batched=True, + ) + for l in range(L): + # grads[l]: (B, 1, T, d); row i of pair-sum = sum over t + J[l, s:e] += grads[l].squeeze(1).sum(dim=1).float() + n_pairs = T * (T + 1) // 2 + return J, n_pairs + + +class JLens: + """Averaged-Jacobian lens: reading, J-lens vectors, swap interventions.""" + + def __init__(self, model, tok, Jbar): + """Jbar: (L, d, d) fp32 averaged Jacobian per layer.""" + self.model, self.tok = model, tok + self.Jbar = Jbar + self.tm = _text_model(model) + self.softcap = model.config.get_text_config().final_logit_softcapping + + def _readout(self, x): + """finalnorm + unembed (+softcap) of residual-space vectors x (..., d).""" + x = self.tm.norm(x.to(self.tm.norm.weight.dtype)) + logits = self.model.lm_head(x) + if self.softcap: + logits = self.softcap * torch.tanh(logits / self.softcap) + return logits + + @torch.no_grad() + def read(self, h, layer, topk=10): + """J-lens reading of residual vector(s) h (..., d) at `layer`. + + Returns (topk token ids, topk probs).""" + proj = h.float() @ self.Jbar[layer].T + logits = self._readout(proj) + probs = torch.softmax(logits.float(), dim=-1) + p, idx = probs.topk(topk, dim=-1) + return idx, p + + @torch.no_grad() + def read_prompt(self, input_ids, layers, topk=8): + """Full J-lens table for one prompt: {layer: (ids (T,k), probs (T,k))}.""" + hs = collect_residuals(self.model, input_ids) + return {l: self.read(hs[l], l, topk=topk) for l in layers} + + @torch.no_grad() + def concept_prob(self, input_ids, token_ids, layers): + """P(token) under the lens for each (layer, position). (L, T, n_tokens).""" + hs = collect_residuals(self.model, input_ids) + out = [] + for l in layers: + logits = self._readout(hs[l].float() @ self.Jbar[l].T) + probs = torch.softmax(logits.float(), dim=-1) + out.append(probs[:, token_ids]) + return torch.stack(out) + + def jlens_vector(self, token_id, layer): + """v_s = W_U[s] J_l : residual-stream direction at `layer` for token s.""" + wu = self.model.lm_head.weight[token_id].float() + return wu @ self.Jbar[layer] + + def swap_hooks(self, pairs, layers=None, thr=0.005, alpha=1.0, + write="embed"): + """Gated concept swap. `pairs` is a list of (src_str, tgt_str) tokens. + + At each (layer, position) where the J-lens reads any source token with + probability > thr (the workspace "holds" the concept), transfer the + activation's content from source to target: + + - write="embed" (default): move h's projection on the unit source + embedding onto the unit target embedding (tied-embedding write basis; + the J-lens supplies localization). This is the variant that works on + gemma-4-E2B. + - write="jlens": the paper's h <- h + V(sigma(c) - c), c = pinv(V) h, + with V rows = J-lens vectors of source+target tokens. + + Returns (handles, fired) where fired collects (layer, n_positions).""" + layers = range(4, len(self.tm.layers) - 1) if layers is None else layers + enc = lambda s: self.tok.encode(s, add_special_tokens=False)[0] + s_ids = torch.tensor([enc(s) for s, _ in pairs], device="cuda") + t_ids = torch.tensor([enc(t) for _, t in pairs], device="cuda") + W = self.model.lm_head.weight + E_s = torch.nn.functional.normalize(W[s_ids].float(), dim=-1) + E_t = torch.nn.functional.normalize(W[t_ids].float(), dim=-1) + handles, fired = [], [] + for l in layers: + V = pinv = None + m = len(pairs) + if write == "jlens": + V = torch.cat([ + torch.stack([self.jlens_vector(int(s), l) for s in s_ids]), + torch.stack([self.jlens_vector(int(t), l) for t in t_ids])]) + pinv = torch.linalg.pinv(V) + + def hook(mod, inp, out, l=l, V=V, pinv=pinv, m=m): + h = out.float() + probs = torch.softmax( + self._readout(h @ self.Jbar[l].T).float(), -1) + gate = probs[..., s_ids].sum(-1) > thr # (B, T) + if not gate.any(): + return None + if write == "embed": + proj = h @ E_s.T + delta = alpha * (proj @ E_t - proj @ E_s) + else: + c = h @ pinv + sig = torch.cat([c[..., m:], c[..., :m]], dim=-1) + delta = alpha * (sig - c) @ V + fired.append((l, int(gate.sum()))) + return (h + delta * gate.unsqueeze(-1)).to(out.dtype) + + handles.append(self.tm.layers[l].register_forward_hook(hook)) + return handles, fired + + @torch.no_grad() + def generate(self, input_ids, max_new_tokens=20): + out = self.model.generate( + input_ids, max_new_tokens=max_new_tokens, do_sample=False, + pad_token_id=self.tok.pad_token_id or self.tok.eos_token_id, + ) + return self.tok.decode(out[0, input_ids.shape[1]:], skip_special_tokens=True) + + @torch.no_grad() + def generate_swapped(self, input_ids, pairs, layers=None, thr=0.005, + alpha=1.0, write="embed", max_new_tokens=20): + handles, fired = self.swap_hooks(pairs, layers, thr, alpha, write) + try: + return self.generate(input_ids, max_new_tokens), fired + finally: + for h in handles: + h.remove() + + +def chat_ids(tok, user_msg, device="cuda", assistant_prefix=None): + enc = tok.apply_chat_template( + [{"role": "user", "content": user_msg}], + add_generation_prompt=True, return_tensors="pt", return_dict=True) + ids = enc["input_ids"] + if assistant_prefix: + pre = tok(assistant_prefix, add_special_tokens=False, return_tensors="pt") + ids = torch.cat([ids, pre["input_ids"]], dim=1) + return ids.to(device) diff --git a/pyproject.toml b/pyproject.toml new file mode 100644 index 0000000..ea83cfd --- /dev/null +++ b/pyproject.toml @@ -0,0 +1,5 @@ +[project] +name = "jspace" +version = "0.1.0" +requires-python = ">=3.12" +dependencies = [] diff --git a/scripts/bw_common.py b/scripts/bw_common.py new file mode 100644 index 0000000..e67cfa3 --- /dev/null +++ b/scripts/bw_common.py @@ -0,0 +1,117 @@ +"""Blocksworld: generator, verifier, prompts. + +Pure planning domain — no code syntax, no arithmetic. Plans are symbolically +verifiable by simulation, so STaR bucketing works. Instances are generated +(contamination-free by construction) with difficulty = blocks + moves. +""" + +import json +import random +import re + +MOVE_RE = re.compile( + r"move\s+([A-Z])\s+(?:onto|on top of|on)\s+(?:the\s+)?(table|[A-Z])", + re.IGNORECASE) + + +def gen_instance(n_blocks, rng): + blocks = [chr(65 + i) for i in range(n_blocks)] + + def random_stacks(): + bs = blocks[:] + rng.shuffle(bs) + stacks, i = [], 0 + while i < len(bs): + take = rng.randint(1, len(bs) - i) + stacks.append(bs[i : i + take]) + i += take + return stacks + + init = random_stacks() + goal = random_stacks() + while goal == init: + goal = random_stacks() + return {"blocks": blocks, "init": init, "goal": goal} + + +def fmt_state(stacks): + out = [] + for st in stacks: + if len(st) == 1: + out.append(f"{st[0]} is on the table") + else: + out.append(f"{st[0]} is on the table with " + + " on top, then ".join( + [f"{b}" for b in st[1:]]) + " on top") + # clearer explicit form + lines = [] + for st in stacks: + lines.append(f"stack: {' -> '.join(st)} (bottom -> top)") + return "; ".join(lines) + + +def question(inst): + return ( + "You are stacking blocks. Only the TOP block of a stack can be " + "moved, one block at a time.\n" + f"Blocks: {', '.join(inst['blocks'])}\n" + f"Initial state: {fmt_state(inst['init'])}\n" + f"Goal state: {fmt_state(inst['goal'])}\n" + "Give a plan as a numbered list of moves, each exactly of the form " + "'move X onto Y' or 'move X onto the table'.") + + +def verify_plan(inst, text, max_moves=40): + stacks = [st[:] for st in inst["init"]] + + def top_of(b): + for st in stacks: + if st and st[-1] == b: + return st + return None + + moves = MOVE_RE.findall(text) + if not moves or len(moves) > max_moves: + return False + for b, tgt in moves: + b = b.upper() + tgt = tgt if tgt.lower() == "table" else tgt.upper() + src = top_of(b) + if src is None: + return False # b not clear (or nonexistent) + if tgt == "table" or tgt.lower() == "table": + src.pop() + stacks.append([b]) + else: + dst = top_of(tgt) + if dst is None or b == tgt: + return False + src.pop() + dst.append(b) + stacks = [st for st in stacks if st] + norm = sorted(tuple(st) for st in stacks) + return norm == sorted(tuple(st) for st in inst["goal"]) + + +def make_dataset(n_train=400, n_test=200, seed=0): + rng = random.Random(seed) + items = [] + for split, n in (("train", n_train), ("test", n_test)): + for i in range(n): + nb = rng.choice([3, 3, 4, 4, 5]) + inst = gen_instance(nb, rng) + items.append({"split": split, "task_id": f"bw_{split}_{i}", + "n_blocks": nb, **inst, + "question": question(inst)}) + return items + + +if __name__ == "__main__": + ds = make_dataset() + print(json.dumps(ds[0], indent=1)) + # verifier self-test: identity plan on trivial instance + inst = {"blocks": ["A", "B"], "init": [["A"], ["B"]], + "goal": [["A", "B"]]} + assert verify_plan(inst, "1. move B onto A") + assert not verify_plan(inst, "1. move A onto A") + print("verifier self-test ok") diff --git a/scripts/bw_prep.py b/scripts/bw_prep.py new file mode 100644 index 0000000..7cf9dde --- /dev/null +++ b/scripts/bw_prep.py @@ -0,0 +1,66 @@ +"""Blocksworld STaR labeling with the frozen model (direct vs CoT plan).""" + +import json +import os +import sys +import time +from pathlib import Path + +import torch + +from bw_common import make_dataset, verify_plan +from prep_mbpp import batch_generate + +sys.path.insert(0, str(Path(__file__).resolve().parent.parent)) +from jlens.core import load_model # noqa: E402 + +OUT = Path(os.environ.get("LOOP_OUT", + Path(__file__).resolve().parent.parent / "results-loop")) + +DIRECT_SUFFIX = "\n\nGive only the numbered list of moves, nothing else." +COT_SUFFIX = ("\n\nFirst think step by step about which blocks must move " + "and in what order (briefly), then give the numbered list of " + "moves.") + + +def chat(tok, it, suffix): + return tok.apply_chat_template( + [{"role": "user", "content": it["question"] + suffix}], + tokenize=False, add_generation_prompt=True) + + +def main(): + model, tok = load_model(dtype=torch.bfloat16) + tok.padding_side = "left" + items = make_dataset(n_train=400, n_test=200, seed=0) + print(f"{len(items)} instances", flush=True) + + for tag, suffix, mx in (("direct", DIRECT_SUFFIX, 200), + ("cot", COT_SUFFIX, 500)): + t0 = time.time() + gens = batch_generate(model, tok, + [chat(tok, it, suffix) for it in items], + max_new_tokens=mx, batch_size=24) + oks = [verify_plan(it, g) for it, g in zip(items, gens)] + for it, g, ok in zip(items, gens, oks): + it[f"{tag}_ok"] = bool(ok) + it[f"{tag}_plan"] = g if ok else None + print(f"{tag}: acc={sum(oks)/len(items):.3f} " + f"({time.time()-t0:.0f}s)", flush=True) + + for it in items: + it["label"] = ("easy" if it["direct_ok"] + else "hard" if it["cot_ok"] else "drop") + it["sol_plan"] = (it["direct_plan"] if it["direct_ok"] + else it["cot_plan"]) + for split in ("train", "test"): + sub = [it for it in items if it["split"] == split] + print(f"{split}: easy={sum(i['label']=='easy' for i in sub)} " + f"hard={sum(i['label']=='hard' for i in sub)} " + f"drop={sum(i['label']=='drop' for i in sub)}", flush=True) + json.dump(items, open(OUT / "bw_data.json", "w"), indent=1) + print("wrote", OUT / "bw_data.json") + + +if __name__ == "__main__": + main() diff --git a/scripts/carry_common.py b/scripts/carry_common.py new file mode 100644 index 0000000..6e6b995 --- /dev/null +++ b/scripts/carry_common.py @@ -0,0 +1,143 @@ +"""Design C machinery: k-loop the prompt, then per-token carry generation. + +Structure per sequence: [prompt] [p pause tokens] [answer] + 1. prefill: prompt positions looped k times through the merge (settled plan) + 2. carry scan: each pause/answer position's band input is + x_t = merge(e_t, s_{t-1}) + where s_{t-1} is the *previous position's* final band output — the + latent thought chain runs through the pause positions' workspace states + (token content of pauses is inert: ). + 3. suffix -> logits; CE on answer tokens only. + +The scan is sequential per position (RNN-style through the band); each step +re-runs the band on the full sequence — causality keeps settled positions +stable, provisional later positions are recomputed next step anyway. +""" + +import torch +from torch.utils.checkpoint import checkpoint + +PAUSE_ID = 6 # + + +def _pos_ids(mask): + return (mask.cumsum(-1) - 1).clamp(min=0) + + +def prompt_prefill(looper, adapter, e, calls, prompt_mask, k): + """k merge->band loops over prompt positions. Returns (S, X).""" + with torch.no_grad(): + s = looper.band(e, calls) + x = e + for _ in range(k): + x = torch.where(prompt_mask[..., None], adapter(e, s), e) + s = looper.band(x, calls) + return s, x + + +def carry_steps(looper, adapter, e, calls, S, X, step_updates, + use_checkpoint=False): + """Sequential scan. step_updates: list of (row_idx, pos) index tensors.""" + for rows, pos in step_updates: + if rows.numel() == 0: + continue + seed = S[rows, pos - 1] + x_new = adapter(e[rows, pos], seed) + X = X.clone() + X[rows, pos] = x_new.to(X.dtype) + S = (checkpoint(lambda X_: looper.band(X_, calls), X, + use_reentrant=False) if use_checkpoint + else looper.band(X, calls)) + return S, X + + +def build_step_updates(prompt_lens, total_lens, device): + """For right-padded batches: step j updates row b at prompt_lens[b]+j.""" + max_span = int((total_lens - prompt_lens).max()) + updates = [] + for j in range(max_span): + pos = prompt_lens + j + sel = pos < total_lens + rows = torch.nonzero(sel, as_tuple=False).squeeze(-1) + updates.append((rows.to(device), pos[sel].to(device))) + return updates + + +def carry_logits(looper, adapter, input_ids, attention_mask, prompt_lens, + k, use_checkpoint=False, feedforward=False): + """Teacher-forced design-C forward (right-padded batch). + + feedforward=True: pause-token control — same positions get the adapter as + x_t = merge(e_t, e_t) in ONE band pass; no state propagates between + positions. Isolates 'pause compute + weights' from the recurrence.""" + dev = input_ids.device + calls, _ = looper.capture(input_ids, attention_mask, logits_to_keep=1) + e = looper._hin[looper.l0].detach() + ar = torch.arange(input_ids.shape[1], device=dev) + prompt_mask = ar[None, :] < prompt_lens[:, None].to(dev) + if feedforward: + touched = attention_mask.bool() + x = torch.where(touched[..., None], adapter(e, e), e) + S = (checkpoint(lambda x_: looper.band(x_, calls), x, + use_reentrant=False) if use_checkpoint + else looper.band(x, calls)) + return looper.suffix_logits(S, calls) + S, X = prompt_prefill(looper, adapter, e, calls, prompt_mask, k) + total_lens = attention_mask.sum(-1) + updates = build_step_updates(prompt_lens.to(dev), total_lens.to(dev), dev) + S, X = carry_steps(looper, adapter, e, calls, S, X, updates, + use_checkpoint=use_checkpoint) + return looper.suffix_logits(S, calls) + + +@torch.no_grad() +def generate_carry_c(looper, adapter, tok, input_ids, attention_mask, + k, p, max_new_tokens=10, feedforward=False): + """Greedy design-C generation (left-padded batch, uniform positions). + + Appends p pause tokens, prefill-loops the prompt, carries through the + pauses, then generates with per-token carry.""" + B = input_ids.shape[0] + dev = input_ids.device + pauses = torch.full((B, p), PAUSE_ID, dtype=torch.long, device=dev) + ids = torch.cat([input_ids, pauses], 1) + mask = torch.cat([attention_mask, + torch.ones_like(pauses)], 1) + n_prompt = input_ids.shape[1] + + eos = {tok.eos_token_id, tok.convert_tokens_to_ids("")} + done = torch.zeros(B, dtype=torch.bool, device=dev) + X_store = None # merged band inputs for settled positions + for step in range(max_new_tokens + 1): # step 0 = prefill + pause scan + calls, _ = looper.capture(ids, mask, logits_to_keep=1, + position_ids=_pos_ids(mask)) + e = looper._hin[looper.l0] + if feedforward: + x = torch.where(mask.bool()[..., None], adapter(e, e), e) + S = looper.band(x, calls) + X = x + elif X_store is None: + ar = torch.arange(ids.shape[1], device=dev) + prompt_mask = (ar[None, :] < n_prompt) & mask.bool() + S, X = prompt_prefill(looper, adapter, e, calls, prompt_mask, k) + updates = [(torch.arange(B, device=dev), + torch.full((B,), n_prompt + j, device=dev, + dtype=torch.long)) for j in range(p)] + S, X = carry_steps(looper, adapter, e, calls, S, X, updates) + else: + X = torch.cat([X_store, e[:, X_store.shape[1]:]], 1) + rows = torch.arange(B, device=dev) + pos = torch.full((B,), ids.shape[1] - 1, device=dev, + dtype=torch.long) + S = looper.band(X, calls) # settled prefix + provisional new pos + S, X = carry_steps(looper, adapter, e, calls, S, X, [(rows, pos)]) + X_store = X + logits = looper.suffix_logits(S, calls, last_only=True) + nxt = logits[:, -1].argmax(-1) + nxt = torch.where(done, torch.full_like(nxt, list(eos)[0]), nxt) + ids = torch.cat([ids, nxt[:, None]], 1) + mask = torch.cat([mask, (~done)[:, None].long()], 1) + done |= torch.tensor([t.item() in eos for t in nxt], device=dev) + if done.all(): + break + return ids diff --git a/scripts/compute_jacobians.py b/scripts/compute_jacobians.py new file mode 100644 index 0000000..0ec89b7 --- /dev/null +++ b/scripts/compute_jacobians.py @@ -0,0 +1,67 @@ +"""Estimate J_l = E_{prompt, t, t'>=t}[dh_final,t'/dh_l,t] over a pretraining-like corpus.""" + +import argparse, json, sys, time +from pathlib import Path + +import torch + +sys.path.insert(0, str(Path(__file__).resolve().parent.parent)) +from jlens.core import load_model, prompt_jacobian_pairsum + + +def corpus_texts(n, min_chars=400): + from datasets import load_dataset + ds = load_dataset("HuggingFaceFW/fineweb-edu", name="sample-10BT", + split="train", streaming=True) + got = 0 + for ex in ds: + t = ex["text"].strip() + if len(t) >= min_chars: + yield t + got += 1 + if got >= n: + return + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--n-prompts", type=int, default=256) + ap.add_argument("--seq-len", type=int, default=64) + ap.add_argument("--chunk", type=int, default=128) + ap.add_argument("--out", default="results/jbar.pt") + ap.add_argument("--dtype", default="float32") + args = ap.parse_args() + + model, tok = load_model(dtype=getattr(torch, args.dtype)) + L = len(model.model.language_model.layers) + d = model.config.get_text_config().hidden_size + + Jsum = torch.zeros(L, d, d, device="cuda", dtype=torch.float32) + pairs_total = 0 + t0 = time.time() + for i, text in enumerate(corpus_texts(args.n_prompts)): + ids = tok(text, return_tensors="pt", truncation=True, + max_length=args.seq_len)["input_ids"].cuda() + if ids.shape[1] < args.seq_len: + continue + J, n_pairs = prompt_jacobian_pairsum(model, ids, chunk=args.chunk) + Jsum += J + pairs_total += n_pairs + if i % 5 == 0 or i == args.n_prompts - 1: + el = time.time() - t0 + print(f"[{i+1}/{args.n_prompts}] {el:.0f}s ({el/(i+1):.1f}s/prompt)", + flush=True) + if i % 50 == 49: # checkpoint + torch.save({"Jbar": (Jsum / pairs_total).cpu(), + "n_prompts": i + 1, "seq_len": args.seq_len}, + args.out + ".ckpt") + + Jbar = (Jsum / pairs_total).cpu() + Path(args.out).parent.mkdir(parents=True, exist_ok=True) + torch.save({"Jbar": Jbar, "n_prompts": args.n_prompts, + "seq_len": args.seq_len, "pairs": pairs_total}, args.out) + print("saved", args.out) + + +if __name__ == "__main__": + main() diff --git a/scripts/diagram.py b/scripts/diagram.py new file mode 100644 index 0000000..c9f13ab --- /dev/null +++ b/scripts/diagram.py @@ -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") diff --git a/scripts/eval_bw.py b/scripts/eval_bw.py new file mode 100644 index 0000000..1cc9aec --- /dev/null +++ b/scripts/eval_bw.py @@ -0,0 +1,83 @@ +"""Blocksworld eval: pass@1 vs k (prompt-only loops, fast path).""" + +import argparse +import json +import os +import sys +import time +from pathlib import Path + +import torch + +from bw_common import verify_plan +from bw_prep import DIRECT_SUFFIX, chat +from loop_common import BandLooper, MergeAdapter + +sys.path.insert(0, str(Path(__file__).resolve().parent.parent)) +from jlens.core import load_model # noqa: E402 + +OUT = Path(os.environ.get("LOOP_OUT", + Path(__file__).resolve().parent.parent / "results-loop")) + + +@torch.no_grad() +def acc_at_k(looper, adapter, tok, items, k, batch=12, max_new=200): + oks, per_item = [], [] + for i in range(0, len(items), batch): + torch.cuda.empty_cache() + chunk = items[i : i + batch] + enc = tok([chat(tok, it, DIRECT_SUFFIX) for it in chunk], + return_tensors="pt", padding=True, + add_special_tokens=False).to("cuda") + gen = looper.generate_frozen_prompt(adapter, tok, enc["input_ids"], k, + max_new_tokens=max_new, + attention_mask=enc["attention_mask"]) + for j, it in enumerate(chunk): + txt = tok.decode(gen[j, enc["input_ids"].shape[1]:], + skip_special_tokens=True) + ok = verify_plan(it, txt) + oks.append(ok) + per_item.append({"task_id": it["task_id"], "ok": bool(ok)}) + by = {} + for it, ok in zip(items, oks): + d = by.setdefault(it["label"], [0, 0]) + d[0] += ok + d[1] += 1 + return (sum(oks) / len(items), + {l: c / n for l, (c, n) in by.items()}, per_item) + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--adapter", default=None) + ap.add_argument("--tag", default="bw") + ap.add_argument("--ks", default="0,2,4") + args = ap.parse_args() + ks = [int(x) for x in args.ks.split(",")] + + model, tok = load_model(dtype=torch.bfloat16) + tok.padding_side = "left" + looper = BandLooper(model) + adapter = MergeAdapter( + d=model.config.get_text_config().hidden_size).cuda() + if args.adapter: + adapter.load_state_dict(torch.load(args.adapter, map_location="cuda")) + adapter.eval() + + items = [it for it in json.load(open(OUT / "bw_data.json")) + if it["split"] == "test"] + print(f"[{args.tag}] Blocksworld eval on {len(items)} items", flush=True) + res = {"tag": args.tag, "ks": {}, "n": len(items)} + for k in ks: + t0 = time.time() + acc, by, per_item = acc_at_k(looper, adapter, tok, items, k) + res["ks"][k] = {"acc": acc, "by_label": by, "per_item": per_item} + print(f"k={k}: acc={acc:.3f} " + f"by_label={ {l: round(v,3) for l,v in by.items()} }" + f" ({time.time()-t0:.0f}s)", flush=True) + json.dump(res, open(OUT / f"eval_bw_{args.tag}.json", "w"), indent=1) + print("wrote", OUT / f"eval_bw_{args.tag}.json") + + +if __name__ == "__main__": + main() diff --git a/scripts/eval_carry.py b/scripts/eval_carry.py new file mode 100644 index 0000000..c1a4878 --- /dev/null +++ b/scripts/eval_carry.py @@ -0,0 +1,90 @@ +"""Eval design C (prefill k + carry through p pauses) on GSM8K test. + +Usage: eval_carry.py --adapter ../results-loop/adapter_carry_e600.pt \ + --tag carry --grid "0:0,2:0,2:2,2:6" +Grid entries are k:p pairs; k=0,p=0 is the plain baseline (same harness). +""" + +import argparse +import json +import sys +import time +from pathlib import Path + +import torch + +from carry_common import generate_carry_c +from loop_common import (DIRECT_SUFFIX, BandLooper, MergeAdapter, chat_prompt, + last_number, num_eq) + +sys.path.insert(0, str(Path(__file__).resolve().parent.parent)) +from jlens.core import load_model # noqa: E402 + +OUT = Path(__file__).resolve().parent.parent / "results-loop" + + +@torch.no_grad() +def acc_at(looper, adapter, tok, items, k, p, batch=16, feedforward=False): + hits, per_label = 0, {} + for i in range(0, len(items), batch): + torch.cuda.empty_cache() + chunk = items[i : i + batch] + enc = tok([chat_prompt(tok, it["question"], DIRECT_SUFFIX) + for it in chunk], return_tensors="pt", padding=True, + add_special_tokens=False).to("cuda") + gen = generate_carry_c(looper, adapter, tok, enc["input_ids"], + enc["attention_mask"], k=k, p=p, + max_new_tokens=10, feedforward=feedforward) + n0 = enc["input_ids"].shape[1] + p + for j, it in enumerate(chunk): + txt = tok.decode(gen[j, n0:], skip_special_tokens=True) + ok = num_eq(last_number(txt), it["gold"]) + hits += ok + d = per_label.setdefault(it["label"], [0, 0]) + d[0] += ok + d[1] += 1 + return hits / len(items), {l: c / n for l, (c, n) in per_label.items()} + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--adapter", default=None) + ap.add_argument("--tag", default="carry") + ap.add_argument("--grid", default="0:0,2:0,2:2,2:6") + ap.add_argument("--n", type=int, default=0) + ap.add_argument("--feedforward", action="store_true") + args = ap.parse_args() + + model, tok = load_model(dtype=torch.bfloat16) + tok.padding_side = "left" + looper = BandLooper(model) + adapter = MergeAdapter().cuda() + if args.adapter: + adapter.load_state_dict(torch.load(args.adapter, map_location="cuda")) + adapter.eval() + + items = [it for it in json.load(open(OUT / "star_data.json")) + if it["split"] == "test"] + if args.n: + items = items[: args.n] + print(f"[{args.tag}] GSM8K carry eval on {len(items)} items, " + f"grid={args.grid}", flush=True) + + res = {"tag": args.tag, "grid": {}, "n": len(items)} + for kp in args.grid.split(","): + k, p = (int(x) for x in kp.split(":")) + t0 = time.time() + acc, by_label = acc_at(looper, adapter, tok, items, k, p, + feedforward=args.feedforward) + res["grid"][kp] = {"acc": acc, "by_label": by_label} + print(f"k={k} p={p}: acc={acc:.3f} " + f"by_label={ {l: round(v,3) for l,v in by_label.items()} }" + f" ({time.time()-t0:.0f}s)", flush=True) + + with open(OUT / f"eval_{args.tag}.json", "w") as f: + json.dump(res, f, indent=1) + print("wrote", OUT / f"eval_{args.tag}.json") + + +if __name__ == "__main__": + main() diff --git a/scripts/eval_humaneval.py b/scripts/eval_humaneval.py new file mode 100644 index 0000000..50289cf --- /dev/null +++ b/scripts/eval_humaneval.py @@ -0,0 +1,143 @@ +"""HumanEval transfer eval: does the MBPP-trained loop adapter generalize? + +No HumanEval training exists (no train split) — this is pure distribution +transfer. Labels the 164 items with the frozen model (direct vs terse-plan, +greedy; descriptive buckets only), then evaluates the loop adapter at +k grid + the untrained control. +""" + +import argparse +import json +import os +import subprocess +import sys +import tempfile +import time +from concurrent.futures import ThreadPoolExecutor +from pathlib import Path + +import torch +from datasets import load_dataset + +from loop_common import BandLooper, MergeAdapter +from prep_mbpp import batch_generate, extract_code + +sys.path.insert(0, str(Path(__file__).resolve().parent.parent)) +from jlens.core import load_model # noqa: E402 + +OUT = Path(os.environ.get("LOOP_OUT", + Path(__file__).resolve().parent.parent / "results-loop")) + +DIRECT = ("Complete the following Python function. Return the COMPLETE " + "function (signature included) in a ```python code block. " + "No explanation.\n\n```python\n{prompt}```") +PLAN = ("First write a very brief plan: at most 4 short bullet lines. Then " + "return the COMPLETE function (signature included) in a ```python " + "code block.\n\n```python\n{prompt}```") + + +def he_prompt(tok, item, tmpl): + return tok.apply_chat_template( + [{"role": "user", "content": tmpl.format(prompt=item["prompt"])}], + tokenize=False, add_generation_prompt=True) + + +def run_he_tests(code, item, timeout=10): + if not code: + return False + script = (code + "\n\n" + item["test"] + + f"\ncheck({item['entry_point']})\n") + try: + with tempfile.TemporaryDirectory() as td: + r = subprocess.run([sys.executable, "-c", script], cwd=td, + capture_output=True, timeout=timeout) + return r.returncode == 0 + except (subprocess.TimeoutExpired, OSError): + return False + + +@torch.no_grad() +def loop_eval(looper, adapter, tok, items, k, batch=8, max_new=380): + codes = [] + for i in range(0, len(items), batch): + torch.cuda.empty_cache() + chunk = items[i : i + batch] + enc = tok([he_prompt(tok, it, DIRECT) for it in chunk], + return_tensors="pt", padding=True, + add_special_tokens=False).to("cuda") + gen = looper.generate_frozen_prompt(adapter, tok, enc["input_ids"], k, + max_new_tokens=max_new, + attention_mask=enc["attention_mask"]) + for j in range(len(chunk)): + codes.append(extract_code( + tok.decode(gen[j, enc["input_ids"].shape[1]:], + skip_special_tokens=True))) + with ThreadPoolExecutor(8) as ex: + oks = list(ex.map(lambda ci: run_he_tests(ci[0], ci[1]), + zip(codes, items))) + return oks + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--adapter", default=None) + ap.add_argument("--tag", default="he") + ap.add_argument("--ks", default="0,2,4") + args = ap.parse_args() + ks = [int(x) for x in args.ks.split(",")] + + model, tok = load_model(dtype=torch.bfloat16) + tok.padding_side = "left" + looper = BandLooper(model) + adapter = MergeAdapter( + d=model.config.get_text_config().hidden_size).cuda() + if args.adapter: + adapter.load_state_dict(torch.load(args.adapter, map_location="cuda")) + adapter.eval() + + items = list(load_dataset("openai/openai_humaneval")["test"]) + print(f"[{args.tag}] HumanEval: {len(items)} items", flush=True) + + # labeling pass (descriptive buckets; greedy — outcome-selection caveat) + lab_path = OUT / "humaneval_labels.json" + if lab_path.exists(): + labels = json.load(open(lab_path)) + else: + plans = batch_generate(model, tok, + [he_prompt(tok, it, PLAN) for it in items], + max_new_tokens=800, batch_size=16) + with ThreadPoolExecutor(8) as ex: + plan_ok = list(ex.map( + lambda gi: run_he_tests(extract_code(gi[0]), gi[1]), + zip(plans, items))) + labels = {it["task_id"]: bool(ok) for it, ok in zip(items, plan_ok)} + json.dump(labels, open(lab_path, "w"), indent=1) + print(f"plan-reachable: {sum(labels.values())}/{len(items)}", flush=True) + + res = {"tag": args.tag, "ks": {}, "n": len(items)} + k0_ok = None + for k in ks: + t0 = time.time() + oks = loop_eval(looper, adapter, tok, items, k) + if k == 0: + k0_ok = oks + hard = [i for i, it in enumerate(items) + if k0_ok and not k0_ok[i] and labels[it["task_id"]]] + acc = sum(oks) / len(items) + hard_acc = (sum(oks[i] for i in hard) / len(hard)) if hard else None + res["ks"][k] = {"acc": acc, "hard_n": len(hard), + "hard_acc": hard_acc, + "per_item": [{"task_id": it["task_id"], + "ok": bool(o)} + for it, o in zip(items, oks)]} + print(f"k={k}: pass@1={acc:.3f} hard({len(hard)})=" + f"{hard_acc if hard_acc is None else round(hard_acc,3)} " + f"({time.time()-t0:.0f}s)", flush=True) + + json.dump(res, open(OUT / f"eval_humaneval_{args.tag}.json", "w"), + indent=1) + print("wrote", OUT / f"eval_humaneval_{args.tag}.json") + + +if __name__ == "__main__": + main() diff --git a/scripts/eval_loop.py b/scripts/eval_loop.py new file mode 100644 index 0000000..d0ac0b7 --- /dev/null +++ b/scripts/eval_loop.py @@ -0,0 +1,143 @@ +"""Eval the looped-band model: accuracy vs k, and J-lens concept sharpening. + +Usage: + eval_loop.py # untrained adapter (alpha-merge only) + eval_loop.py --adapter results-loop/adapter.pt --tag trained +""" + +import argparse +import json +import os +import time +from pathlib import Path + +import torch + +from loop_common import (DIRECT_SUFFIX, BandLooper, MergeAdapter, chat_prompt, + last_number, num_eq) +import sys +sys.path.insert(0, str(Path(__file__).resolve().parent.parent)) +from jlens.core import JLens, load_model # noqa: E402 + +OUT = Path(os.environ.get("LOOP_OUT", + Path(__file__).resolve().parent.parent / "results-loop")) +ROOT = OUT.parent + + +@torch.no_grad() +def accuracy_at_k(looper, adapter, tok, items, k, batch=16, prompt_only=False, + feedforward=False): + hits, per_label, per_item = 0, {}, [] + for i in range(0, len(items), batch): + torch.cuda.empty_cache() + chunk = items[i : i + batch] + enc = tok([chat_prompt(tok, it["question"], DIRECT_SUFFIX) for it in chunk], + return_tensors="pt", padding=True, + add_special_tokens=False).to("cuda") + if prompt_only: + gen = looper.generate_frozen_prompt( + adapter, tok, enc["input_ids"], k, max_new_tokens=10, + attention_mask=enc["attention_mask"], feedforward=feedforward) + else: + gen = looper.loop_generate(adapter, tok, enc["input_ids"], k, + max_new_tokens=10, + attention_mask=enc["attention_mask"]) + for j, it in enumerate(chunk): + txt = tok.decode(gen[j, enc["input_ids"].shape[1]:], + skip_special_tokens=True) + ok = num_eq(last_number(txt), it["gold"]) + per_item.append({"idx": it["idx"], "ok": bool(ok)}) + hits += ok + d = per_label.setdefault(it["label"], [0, 0]) + d[0] += ok + d[1] += 1 + return (hits / len(items), + {l: c / n for l, (c, n) in per_label.items()}, per_item) + + +@torch.no_grad() +def spider_sharpening(looper, adapter, model, tok, ks): + """P('spider') under the J-lens at L30, and P('8') at the output, vs k.""" + jbar = torch.load(ROOT / "results" / "jbar.pt", map_location="cuda") + Jbar = (jbar["Jbar"] if isinstance(jbar, dict) else jbar).float() + lens = JLens(model, tok, Jbar) + q = ("The animal that spins webs has how many legs? " + "Answer with just the number.") + ids = tok.apply_chat_template([{"role": "user", "content": q}], + return_tensors="pt", add_generation_prompt=True, + return_dict=True)["input_ids"].cuda() + spider = tok.encode(" spider", add_special_tokens=False)[0] + eight = [tok.encode(t, add_special_tokens=False)[0] for t in (" 8", "8")] + + calls, _ = looper.capture(ids) + e = looper._hin[looper.l0] + s = looper.band(e, calls) + rows = [] + kmax = max(ks) + for k in range(0, kmax + 1): + if k > 0: + s = looper.band(adapter(e, s), calls) + if k in ks: + _, probs = lens.read(s[0], looper.l1, topk=1) + logits = lens._readout((s[0].float() @ Jbar[looper.l1].T)) + p_spider = torch.softmax(logits.float(), -1)[:, spider].max().item() + out = looper.suffix_logits(s, calls) + p8 = torch.softmax(out[0, -1].float(), -1)[eight].max().item() + rows.append({"k": k, "P_spider_lens": p_spider, "P_8_out": p8}) + return rows + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--adapter", default=None) + ap.add_argument("--tag", default="untrained") + ap.add_argument("--ks", default="0,1,2,4,8") + ap.add_argument("--n", type=int, default=0, help="cap test items (0=all)") + ap.add_argument("--prompt-only", action="store_true", + help="loop the prompt span only (unified regime, fast path)") + ap.add_argument("--no-spider", action="store_true") + ap.add_argument("--feedforward", action="store_true", + help="no-recurrence control arm: adapter(e,e) once") + args = ap.parse_args() + ks = [int(x) for x in args.ks.split(",")] + + model, tok = load_model(dtype=torch.bfloat16) + tok.padding_side = "left" + looper = BandLooper(model) + adapter = MergeAdapter( + d=model.config.get_text_config().hidden_size).cuda() + if args.adapter: + adapter.load_state_dict(torch.load(args.adapter, map_location="cuda")) + adapter.eval() + + items = [it for it in json.load(open(OUT / "star_data.json")) + if it["split"] == "test"] + if args.n: + items = items[: args.n] + print(f"[{args.tag}] eval on {len(items)} test items, ks={ks}", flush=True) + + res = {"tag": args.tag, "ks": {}, "n": len(items)} + for k in ks: + t0 = time.time() + acc, by_label, per_item = accuracy_at_k(looper, adapter, tok, items, k, + prompt_only=args.prompt_only, + feedforward=args.feedforward) + res["ks"][k] = {"acc": acc, "by_label": by_label, + "per_item": per_item} + print(f"k={k}: acc={acc:.3f} by_label={ {l: round(v,3) for l,v in by_label.items()} }" + f" ({time.time()-t0:.0f}s)", flush=True) + + res["spider"] = ([] if args.no_spider else + spider_sharpening(looper, adapter, model, tok, set(ks))) + for r in res["spider"]: + print(f"spider k={r['k']}: P_lens={r['P_spider_lens']:.3f} " + f"P(8)={r['P_8_out']:.3f}", flush=True) + + OUT.mkdir(exist_ok=True) + with open(OUT / f"eval_{args.tag}.json", "w") as f: + json.dump(res, f, indent=1) + print("wrote", OUT / f"eval_{args.tag}.json") + + +if __name__ == "__main__": + main() diff --git a/scripts/eval_loop_code.py b/scripts/eval_loop_code.py new file mode 100644 index 0000000..ec5edfd --- /dev/null +++ b/scripts/eval_loop_code.py @@ -0,0 +1,112 @@ +"""Eval MBPP pass@1 vs loop depth k with prompt-only ("latent planning") loops. + +Usage: + eval_loop_code.py --tag untrained + eval_loop_code.py --adapter ../results-loop/adapter_code.pt --tag trained +""" + +import argparse +import json +import os +import sys +import time +from concurrent.futures import ThreadPoolExecutor +from pathlib import Path + +import torch + +from loop_common import BandLooper, MergeAdapter +from prep_mbpp import DIRECT_SUFFIX, extract_code, mbpp_prompt, run_tests + +sys.path.insert(0, str(Path(__file__).resolve().parent.parent)) +from jlens.core import load_model # noqa: E402 + +OUT = Path(os.environ.get("LOOP_OUT", + Path(__file__).resolve().parent.parent / "results-loop")) + + +@torch.no_grad() +def pass1_at_k(looper, adapter, tok, items, k, batch=8, max_new=220, + feedforward=False, pause=0): + codes = [] + for i in range(0, len(items), batch): + torch.cuda.empty_cache() + chunk = items[i : i + batch] + enc = tok([mbpp_prompt(tok, it, DIRECT_SUFFIX) for it in chunk], + return_tensors="pt", padding=True, + add_special_tokens=False).to("cuda") + if pause: + B = enc["input_ids"].shape[0] + pcol = torch.full((B, pause), 6, dtype=torch.long, device="cuda") + enc["input_ids"] = torch.cat([enc["input_ids"], pcol], 1) + enc["attention_mask"] = torch.cat( + [enc["attention_mask"], torch.ones_like(pcol)], 1) + gen = looper.generate_frozen_prompt(adapter, tok, enc["input_ids"], k, + max_new_tokens=max_new, + attention_mask=enc["attention_mask"], + feedforward=feedforward) + for j in range(len(chunk)): + txt = tok.decode(gen[j, enc["input_ids"].shape[1]:], + skip_special_tokens=True) + codes.append(extract_code(txt)) + with ThreadPoolExecutor(8) as ex: + oks = list(ex.map(lambda ci: run_tests(ci[0], ci[1]), + zip(codes, items))) + hits, per_label = 0, {} + for it, ok in zip(items, oks): + hits += ok + d = per_label.setdefault(it["label"], [0, 0]) + d[0] += ok + d[1] += 1 + per_item = [{"task_id": it["task_id"], "ok": bool(ok)} + for it, ok in zip(items, oks)] + return (hits / len(items), + {l: c / n for l, (c, n) in per_label.items()}, per_item) + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--adapter", default=None) + ap.add_argument("--tag", default="untrained") + ap.add_argument("--ks", default="0,1,2,4") + ap.add_argument("--n", type=int, default=250) + ap.add_argument("--feedforward", action="store_true", + help="no-recurrence control arm: adapter(e,e) once") + ap.add_argument("--pause", type=int, default=0, + help="append p pause tokens to each prompt") + args = ap.parse_args() + ks = [int(x) for x in args.ks.split(",")] + + model, tok = load_model(dtype=torch.bfloat16) + tok.padding_side = "left" + looper = BandLooper(model) + adapter = MergeAdapter( + d=model.config.get_text_config().hidden_size).cuda() + if args.adapter: + adapter.load_state_dict(torch.load(args.adapter, map_location="cuda")) + adapter.eval() + + items = [it for it in json.load(open(OUT / "mbpp_data.json")) + if it["split"] == "test"][: args.n] + print(f"[{args.tag}] MBPP eval on {len(items)} test items, ks={ks}", + flush=True) + + res = {"tag": args.tag, "ks": {}, "n": len(items)} + for k in ks: + t0 = time.time() + acc, by_label, per_item = pass1_at_k(looper, adapter, tok, items, k, + feedforward=args.feedforward, + pause=args.pause) + res["ks"][k] = {"acc": acc, "by_label": by_label, + "per_item": per_item} + print(f"k={k}: pass@1={acc:.3f} " + f"by_label={ {l: round(v,3) for l,v in by_label.items()} }" + f" ({time.time()-t0:.0f}s)", flush=True) + + with open(OUT / f"eval_code_{args.tag}.json", "w") as f: + json.dump(res, f, indent=1) + print("wrote", OUT / f"eval_code_{args.tag}.json") + + +if __name__ == "__main__": + main() diff --git a/scripts/eval_plan_baseline.py b/scripts/eval_plan_baseline.py new file mode 100644 index 0000000..23c006f --- /dev/null +++ b/scripts/eval_plan_baseline.py @@ -0,0 +1,54 @@ +"""Explicit-planning reference: frozen model, plan-first prompt, MBPP test. + +The number latent looping is compared against (equal-or-more FLOPs: ~200-400 +visible plan tokens vs k band passes over the prompt).""" + +import json +import sys +import time +from concurrent.futures import ThreadPoolExecutor +from pathlib import Path + +import torch + +from prep_mbpp import batch_generate, extract_code, mbpp_prompt, run_tests +from prep_mbpp_fix import TERSE_PLAN_SUFFIX + +sys.path.insert(0, str(Path(__file__).resolve().parent.parent)) +from jlens.core import load_model # noqa: E402 + +OUT = Path(__file__).resolve().parent.parent / "results-loop" + + +def main(): + model, tok = load_model(dtype=torch.bfloat16) + tok.padding_side = "left" + items = [it for it in json.load(open(OUT / "mbpp_data.json")) + if it["split"] == "test"] + t0 = time.time() + gens = batch_generate(model, tok, + [mbpp_prompt(tok, it, TERSE_PLAN_SUFFIX) + for it in items], max_new_tokens=700, batch_size=16) + codes = [extract_code(g) for g in gens] + with ThreadPoolExecutor(8) as ex: + oks = list(ex.map(lambda ci: run_tests(ci[0], ci[1]), + zip(codes, items))) + by = {} + for it, ok in zip(items, oks): + d = by.setdefault(it["label"], [0, 0]) + d[0] += ok + d[1] += 1 + per_item = [{"task_id": it["task_id"], "ok": bool(ok)} + for it, ok in zip(items, oks)] + res = {"acc": sum(oks) / len(items), + "by_label": {l: c / n for l, (c, n) in by.items()}, + "per_item": per_item} + print(f"plan-first pass@1={res['acc']:.3f} " + f"by_label={ {l: round(v,3) for l,v in res['by_label'].items()} }" + f" ({time.time()-t0:.0f}s)", flush=True) + json.dump(res, open(OUT / "eval_plan_baseline.json", "w"), indent=1) + print("wrote", OUT / "eval_plan_baseline.json") + + +if __name__ == "__main__": + main() diff --git a/scripts/exp1_readouts.py b/scripts/exp1_readouts.py new file mode 100644 index 0000000..e45a9da --- /dev/null +++ b/scripts/exp1_readouts.py @@ -0,0 +1,130 @@ +"""Experiment 1: J-lens readouts. + +(a) Two-hop reasoning: unspoken intermediate 'spider' visible mid-network. +(b) Multilingual: English intermediates during a Chinese task. +(c) Directed modulation: 'hold citrus in mind' while copying unrelated text. +(d) Layer profile: where in depth the lens carries abstract content. +""" + +import os, sys +from pathlib import Path + +import torch + +sys.path.insert(0, str(Path(__file__).resolve().parent.parent)) +from jlens.core import (JLens, chat_ids, collect_residuals, + generate_with_residuals, load_model) + +RES = Path(os.environ.get("JLENS_RESULTS", "results")) + + +def show_table(jl, tok, ids, layers, positions, topk=6, header=""): + """Print top J-lens tokens at (layer, position) grid.""" + hs = collect_residuals(jl.model, ids) + toks = [tok.decode([t]) for t in ids[0].tolist()] + print(f"\n=== {header} ===") + print("positions:", {p: repr(toks[p]) for p in positions}) + for l in layers: + row = [] + for p in positions: + idx, prob = jl.read(hs[l, p], l, topk=topk) + row.append(" ".join( + (tok.decode([i]).strip() or "·") for i in idx.tolist()[:topk])) + print(f"L{l:>2} | " + " || ".join(row)) + + +def concept_heatmap(jl, tok, ids, words, tag, hs=None): + """P(concept tokens) by (layer, position); save tensor + print peak.""" + if hs is None: + hs = collect_residuals(jl.model, ids) + tids = [] + for w in words: + for v in (w, " " + w, w.capitalize(), " " + w.capitalize()): + e = tok.encode(v, add_special_tokens=False) + if len(e) == 1: + tids.append(e[0]) + tids = sorted(set(tids)) + L, T, _ = hs.shape + out = torch.zeros(L, T) + for l in range(L): + logits = jl._readout(hs[l].float() @ jl.Jbar[l].T) + probs = torch.softmax(logits.float(), dim=-1) + out[l] = probs[:, tids].sum(-1).cpu() + peak = out.max().item() + lmax, pmax = divmod(out.argmax().item(), T) + toks = [tok.decode([t]) for t in ids[0].tolist()] if ids is not None else None + print(f"[{tag}] words={words} peak P={peak:.3f} at layer {lmax}, " + f"pos {pmax}" + (f" ({toks[pmax]!r})" if toks and pmax < len(toks) else "")) + RES.mkdir(exist_ok=True) + torch.save({"map": out, "words": words, + "tokens": [tok.decode([t]) for t in ids[0].tolist()] if ids is not None else None}, + RES / f"heat_{tag}.pt") + return out + + +def main(): + model, tok = load_model(dtype=torch.bfloat16) + ck = torch.load(sys.argv[1] if len(sys.argv) > 1 else "results/jbar.pt", + map_location="cuda") + jl = JLens(model, tok, ck["Jbar"].cuda()) + print(f"Jbar from {ck.get('n_prompts')} prompts") + L = len(model.model.language_model.layers) + mid = range(max(2, L // 6), L - 4, max(2, L // 16)) + + # (a) two-hop: spider + ids = chat_ids(tok, "The animal that spins webs has how many legs? " + "Answer with just a number.") + print("answer:", jl.generate(ids, 5)) + show_table(jl, tok, ids, mid, list(range(ids.shape[1] - 10, ids.shape[1])), + header="two-hop spider (last 10 positions)") + concept_heatmap(jl, tok, ids, ["spider", "spiders"], "spider") + concept_heatmap(jl, tok, ids, ["eight", "8"], "eight") + + # (b) multilingual: Chinese antonym of small -> big + ids = chat_ids(tok, "小的反义词是什么？只用一个字回答。") + print("\nanswer:", jl.generate(ids, 5)) + show_table(jl, tok, ids, mid, list(range(ids.shape[1] - 8, ids.shape[1])), + header="Chinese antonym of 小") + concept_heatmap(jl, tok, ids, ["big", "large", "bigger"], "big_en") + + # (c) directed modulation: think of citrus while copying text + copy_text = "The committee will meet on Thursday to review the budget." + for cond, instr in [ + ("citrus", "While you copy the text, silently think about citrus " + "fruits the entire time. Copy this text exactly, output " + f"nothing else: \"{copy_text}\""), + ("control", f"Copy this text exactly, output nothing else: \"{copy_text}\""), + ]: + ids = chat_ids(tok, instr) + text, all_ids, hs = generate_with_residuals(model, tok, ids, 30) + print(f"\n[{cond}] output: {text!r}") + # only look at generated positions + gen0 = ids.shape[1] + heat = concept_heatmap(jl, tok, all_ids.unsqueeze(0), + ["citrus", "lemon", "orange", "lime"], + f"citrus_{cond}", hs=hs) + print(f" citrus P over generated positions: mean " + f"{heat[:, gen0:].mean():.4f} max {heat[:, gen0:].max():.4f}") + + # (d) layer profile: lens entropy + top-token agreement with input/output + ids = chat_ids(tok, "Write one sentence about the ocean.") + text, all_ids, hs = generate_with_residuals(model, tok, ids, 20) + L, T, _ = hs.shape + ent, agree_in, agree_next = [], [], [] + for l in range(L): + logits = jl._readout(hs[l].float() @ jl.Jbar[l].T) + probs = torch.softmax(logits.float(), dim=-1) + e = -(probs * probs.clamp_min(1e-12).log()).sum(-1).mean() + top = probs.argmax(-1) + ent.append(e.item()) + agree_in.append((top[:-1] == all_ids[:len(top) - 1].cuda()).float().mean().item()) + agree_next.append((top[:-1] == all_ids[1:len(top)].cuda()).float().mean().item()) + print("\nlayer | lens entropy | top==current tok | top==next tok") + for l in range(0, L, 2): + print(f"L{l:>2} | {ent[l]:8.2f} | {agree_in[l]:.2f} | {agree_next[l]:.2f}") + torch.save({"entropy": ent, "agree_in": agree_in, "agree_next": agree_next}, + RES / "layer_profile.pt") + + +if __name__ == "__main__": + main() diff --git a/scripts/exp2_swaps.py b/scripts/exp2_swaps.py new file mode 100644 index 0000000..93bde66 --- /dev/null +++ b/scripts/exp2_swaps.py @@ -0,0 +1,103 @@ +"""Experiment 2: causal swap interventions gated by the J-lens. + +At each (layer, position) where the J-lens reads the source concept, transfer +the activation content source -> target (embedding write basis; see core.py). + +(a) Two-hop: swap spider<->ant in the workspace -> answer flips 8 -> 6. +(b) Broadcast: swap France->China under different templates. +(c) Capital-question grid over country pairs. +""" + +import itertools, os, sys +from pathlib import Path + +import torch + +sys.path.insert(0, str(Path(__file__).resolve().parent.parent)) +from jlens.core import JLens, chat_ids, load_model + +RES = Path(os.environ.get("JLENS_RESULTS", "results")) + + +def main(): + model, tok = load_model(dtype=torch.bfloat16) + ck = torch.load(sys.argv[1] if len(sys.argv) > 1 else "results/jbar.pt", + map_location="cuda") + jl = JLens(model, tok, ck["Jbar"].cuda()) + print(f"Jbar from {ck.get('n_prompts')} prompts") + + # (a) spider -> ant + SP = [(" spider", " ant"), (" spiders", " ants"), ("spider", "ant")] + q = "The animal that spins webs has how many legs? Answer with just a number." + ids = chat_ids(tok, q) + print("\n--- two-hop swap: spider -> ant (baseline:", + repr(jl.generate(ids, 6)), ") ---") + for thr in (0.002, 0.005, 0.01): + for alpha in (1.0, 1.5, 2.0): + out, fired = jl.generate_swapped(ids, SP, thr=thr, alpha=alpha, + max_new_tokens=6) + print(f"thr={thr} a={alpha}: {out!r} (fired {len(fired)} slots)") + # paper write-basis for comparison + out, _ = jl.generate_swapped(ids, SP, thr=0.005, alpha=1.0, write="jlens", + max_new_tokens=6) + print(f"paper jlens-write thr=0.005 a=1.0: {out!r}") + # control: unrelated swap + out, _ = jl.generate_swapped(ids, [(" spider", " piano")], thr=0.005, + alpha=1.0, max_new_tokens=6) + print(f"control spider->piano: {out!r}") + + # reverse: ant prompt -> spider + q2 = ("The insect that builds colonies and lifts many times its own " + "weight has how many legs? Answer with just a number.") + ids2 = chat_ids(tok, q2) + print("\nant prompt baseline:", repr(jl.generate(ids2, 6))) + for alpha in (1.0, 1.5, 2.0): + out, fired = jl.generate_swapped( + ids2, [(" ant", " spider"), (" ants", " spiders"), ("ant", "spider")], + thr=0.005, alpha=alpha, max_new_tokens=6) + print(f"swap ant->spider a={alpha}: {out!r} (fired {len(fired)})") + + # (b) broadcast France -> China across templates + print("\n--- broadcast: France -> China (thr=0.01, a=1.0) ---") + FR = [(" France", " China"), ("France", "China"), (" French", " Chinese")] + templates = [ + ("capital", "What is the capital of France? Answer with just the city name."), + ("language", "What language is spoken in France? Answer with one word."), + ("continent", "Which continent is France on? Answer with one word."), + ("currency", "What currency is used in France? Answer with one word."), + ("river", "Name a famous river in France. Answer with one word."), + ("food", "Name a famous dish from France. Answer with a short phrase."), + ] + for tag, qq in templates: + ids = chat_ids(tok, qq) + base = jl.generate(ids, 8) + sw, fired = jl.generate_swapped(ids, FR, thr=0.01, alpha=1.0, + max_new_tokens=8) + print(f"[{tag:9}] base={base!r:30} swapped={sw!r} ({len(fired)} slots)") + + # (c) capital grid over country pairs + print("\n--- capital-question swap grid (thr=0.01, a=1.0) ---") + countries = {"France": ("Paris", "French"), "China": ("Beijing", "Chinese"), + "Japan": ("Tokyo", "Japanese"), "Egypt": ("Cairo", "Egyptian"), + "Brazil": ("Brasília", "Brazilian"), "Canada": ("Ottawa", "Canadian")} + hits = tries = 0 + rows = [] + for (src, (scap, sadj)), (tgt, (tcap, tadj)) in itertools.permutations( + countries.items(), 2): + qq = f"What is the capital of {src}? Answer with just the city name." + ids = chat_ids(tok, qq) + pairs = [(" " + src, " " + tgt), (src, tgt), (" " + sadj, " " + tadj)] + sw, fired = jl.generate_swapped(ids, pairs, thr=0.01, alpha=1.0, + max_new_tokens=8) + ok = tcap.lower().replace("í", "i").split()[0][:5] in \ + sw.lower().replace("í", "i") + hits += ok + tries += 1 + rows.append((src, tgt, tcap, sw, ok)) + print(f"{src:>7}->{tgt:<7} expect {tcap:<9} got {sw!r} {'OK' if ok else ''}") + print(f"\nswap grid success: {hits}/{tries}") + torch.save(rows, RES / "swap_grid.pt") + + +if __name__ == "__main__": + main() diff --git a/scripts/exp3_jspace.py b/scripts/exp3_jspace.py new file mode 100644 index 0000000..77f69e3 --- /dev/null +++ b/scripts/exp3_jspace.py @@ -0,0 +1,86 @@ +"""Experiment 3: J-space occupancy. + +Decompose residual activations as sparse non-negative combinations of k<=25 +J-lens vectors (matching pursuit + NNLS refit) and measure the fraction of +activation variance the J-space carries per layer (paper: ~10%, intermediate +layers only). +""" + +import os, sys +from pathlib import Path + +import numpy as np +import torch +from scipy.optimize import nnls + +sys.path.insert(0, str(Path(__file__).resolve().parent.parent)) +from jlens.core import JLens, chat_ids, collect_residuals, load_model + +RES = Path(os.environ.get("JLENS_RESULTS", "results")) + +K = 25 + +PROMPTS = [ + "The animal that spins webs has how many legs? Answer with just a number.", + "Write one sentence about the ocean.", + "What is the capital of France? Answer with just the city name.", + "Explain photosynthesis in one sentence.", +] + + +def pursuit_r2(D, h, k=K): + """Non-negative matching pursuit of h onto rows of D. Returns R^2, support.""" + r = h.clone() + sel = [] + for _ in range(k): + scores = D @ r + if sel: + scores[torch.tensor(sel, device=D.device)] = -1e30 + j = int(scores.argmax()) + if scores[j] <= 0: + break + sel.append(j) + A = D[sel].T # (d, |sel|) + x, _ = nnls(A.detach().cpu().numpy().astype(np.float64), + h.detach().cpu().numpy().astype(np.float64)) + approx = A @ torch.tensor(x, device=D.device, dtype=D.dtype) + r = h - approx + r2 = 1 - (r.norm() / h.norm()) ** 2 + return float(r2), sel + + +def main(): + model, tok = load_model(dtype=torch.bfloat16) + ck = torch.load(sys.argv[1] if len(sys.argv) > 1 else "results/jbar.pt", + map_location="cuda") + jl = JLens(model, tok, ck["Jbar"].cuda()) + WU = model.lm_head.weight.detach().float() # (V, d) + + L = len(model.model.language_model.layers) + layers = list(range(1, L, 3)) + r2_by_layer = {l: [] for l in layers} + for p in PROMPTS: + ids = chat_ids(tok, p) + hs = collect_residuals(model, ids) + T = ids.shape[1] + positions = list(range(max(1, T - 12), T)) # skip BOS region + for l in layers: + D = WU @ jl.Jbar[l] # (V, d) J-lens dictionary at layer l + D = D / D.norm(dim=1, keepdim=True).clamp_min(1e-8) + for t in positions: + r2, _ = pursuit_r2(D, hs[l, t].float()) + r2_by_layer[l].append(r2) + del D + torch.cuda.empty_cache() + print(f"done prompt: {p[:40]}...", flush=True) + + print("\nlayer | mean R^2 of k<=25 non-negative J-lens pursuit") + means = {} + for l in layers: + means[l] = float(np.mean(r2_by_layer[l])) + print(f"L{l:>2} | {means[l]:.3f}") + torch.save(means, RES / "jspace_r2.pt") + + +if __name__ == "__main__": + main() diff --git a/scripts/exp4_regimes.py b/scripts/exp4_regimes.py new file mode 100644 index 0000000..a0a2be9 --- /dev/null +++ b/scripts/exp4_regimes.py @@ -0,0 +1,142 @@ +"""Experiment 4: locate sensor / workspace / motor regimes by depth. + +Per-layer diagnostics averaged over a prompt set: + sensor = frac. of positions where top J-lens token == current input token + motor = frac. where top J-lens token == NEXT input token (teacher-forced) + persist = mean Jaccard overlap of top-10 lens tokens at adjacent positions + (workspace content should persist across positions) + content = frac. of positions whose top lens token is a content word + (alphabetic, len>=3, not in a junk list) + +Ignition test (paper: ambiguous inputs produce sharp binary commitment at +workspace onset): replace one token's embedding (main + per-layer) with a +w-mixture of two concept embeddings and track lens commitment +C = (P1-P2)/(P1+P2) by layer at a downstream position. +""" + +import os, sys +from pathlib import Path + +import torch + +sys.path.insert(0, str(Path(__file__).resolve().parent.parent)) +from jlens.core import JLens, chat_ids, collect_residuals, load_model + +RES = Path(os.environ.get("JLENS_RESULTS", "results")) + +PROMPTS = [ + "The animal that spins webs has how many legs? Answer with just a number.", + "What is the capital of France? Answer with just the city name.", + "Write one sentence about the ocean.", + "Explain photosynthesis in one sentence.", + "Name a famous river in Egypt.", + "What language is spoken in Brazil? Answer with one word.", + "Summarize the plot of Romeo and Juliet in one sentence.", + "If I have 3 apples and eat one, how many are left?", +] + +JUNK = set("·.<>|/\\()[]{}:;,!?\"'`~*#-_=+ \n\t") + + +def regime_profile(jl, tok, model): + L = len(jl.tm.layers) + sensor = torch.zeros(L) + motor = torch.zeros(L) + persist = torch.zeros(L) + content = torch.zeros(L) + n_pos = 0 + n_adj = 0 + for p in PROMPTS: + ids = chat_ids(tok, p) + hs = collect_residuals(model, ids) + T = ids.shape[1] + sl = slice(4, T - 5) # skip bos/turn tokens and trailing template + cur = ids[0, sl].cuda() + nxt = ids[0, 4 + 1:T - 4].cuda() + for l in range(L): + idx, _ = jl.read(hs[l, sl], l, topk=10) # (T', 10) + top1 = idx[:, 0] + sensor[l] += (top1 == cur).sum().item() + motor[l] += (top1 == nxt).sum().item() + sets = [set(r.tolist()) for r in idx] + for a, b in zip(sets, sets[1:]): + persist[l] += len(a & b) / len(a | b) + for t in top1.tolist(): + s = tok.decode([t]).strip() + content[l] += (len(s) >= 3 and s.isalpha()) + n_pos += cur.numel() + n_adj += cur.numel() - 1 + return sensor / n_pos, motor / n_pos, persist / n_adj, content / n_pos + + +def ignition(jl, tok, model, pairs, ws=(0.0, 0.25, 0.5, 0.75, 1.0)): + """Mix two concept embeddings at one position; commitment by layer.""" + tm = jl.tm + L = len(tm.layers) + template = ("My favorite thing in the world is the X ." + " I think about it every single day because") + out = {} + for w1, w2 in pairs: + id1 = tok.encode(" " + w1, add_special_tokens=False)[0] + id2 = tok.encode(" " + w2, add_special_tokens=False)[0] + ids = tok(template, return_tensors="pt")["input_ids"].cuda() + pos = (ids[0] == tok.encode(" X", add_special_tokens=False)[0]) \ + .nonzero()[0].item() + has_ple = getattr(tm, "embed_tokens_per_layer", None) is not None + with torch.no_grad(): + pair_ids = torch.tensor([[id1, id2]], device="cuda") + rows_main = tm.embed_tokens(pair_ids)[0].detach() + rows_ple = tm.embed_tokens_per_layer(pair_ids)[0].detach() if has_ple else None + curves = torch.zeros(len(ws), L) + for wi, w in enumerate(ws): + def make_hook(rows, w=w, pos=pos): + def mix_hook(mod, inp, out_e): + e = out_e.clone() + e[0, pos] = w * rows[0] + (1 - w) * rows[1] + return e + return mix_hook + h1 = tm.embed_tokens.register_forward_hook(make_hook(rows_main)) + h2 = tm.embed_tokens_per_layer.register_forward_hook(make_hook(rows_ple)) if has_ple else None + try: + hs = collect_residuals(model, ids) + finally: + h1.remove() + if h2: h2.remove() + rpos = pos + 3 # downstream read position + for l in range(L): + logits = jl._readout(hs[l, rpos].float() @ jl.Jbar[l].T) + probs = torch.softmax(logits.float(), -1) + p1, p2 = probs[id1].item(), probs[id2].item() + curves[wi, l] = (p1 - p2) / (p1 + p2 + 1e-12) + out[(w1, w2)] = curves + return out + + +def main(): + model, tok = load_model(dtype=torch.bfloat16) + ck = torch.load(sys.argv[1] if len(sys.argv) > 1 else "results/jbar.pt", + map_location="cuda") + jl = JLens(model, tok, ck["Jbar"].cuda()) + + sensor, motor, persist, content = regime_profile(jl, tok, model) + print("layer | sensor(top==cur) motor(top==next) persist(top10 Jaccard) content-word") + for l in range(len(sensor)): + bars = lambda x: "#" * int(20 * x) + print(f"L{l:>2} | {sensor[l]:.2f} {bars(sensor[l]):<20} | " + f"{motor[l]:.2f} {bars(motor[l]):<20} | " + f"{persist[l]:.2f} | {content[l]:.2f}") + + pairs = [("dog", "piano"), ("ocean", "violin"), ("dragon", "bicycle")] + ign = ignition(jl, tok, model, pairs) + print("\nignition: commitment C=(P1-P2)/(P1+P2) at read pos, by layer") + print("pair | w: " + " ".join(f"{w:4.2f}" for w in (0.0, .25, .5, .75, 1.0))) + for (w1, w2), curves in ign.items(): + for l in range(2, len(sensor), 4): + print(f"{w1}/{w2:<9} L{l:>2} | " + + " ".join(f"{curves[wi, l]:+.2f}" for wi in range(5))) + torch.save({"sensor": sensor, "motor": motor, "persist": persist, + "content": content, "ignition": ign}, RES / "regimes.pt") + + +if __name__ == "__main__": + main() diff --git a/scripts/gate_probe.py b/scripts/gate_probe.py new file mode 100644 index 0000000..b154706 --- /dev/null +++ b/scripts/gate_probe.py @@ -0,0 +1,126 @@ +"""Per-prompt gate: probe on the k=0 workspace state predicts 'will looping help?' + +1. Features: L30 residual at the last prompt position (plain forward), MBPP + train items; labels easy(0)/hard(1) from the STaR pass (free supervision). +2. Logistic probe (d=1536 -> 1), class-balanced. +3. Gated eval on MBPP test: predicted-easy -> k=0, predicted-hard -> k=4 with + the dedicated loop adapter. Reports gated pass@1 vs uniform-k references. +""" + +import json +import sys +import time +from concurrent.futures import ThreadPoolExecutor +from pathlib import Path + +import torch + +from loop_common import BandLooper, MergeAdapter +from prep_mbpp import DIRECT_SUFFIX, extract_code, mbpp_prompt, run_tests + +sys.path.insert(0, str(Path(__file__).resolve().parent.parent)) +from jlens.core import ResidualCapture, load_model # noqa: E402 + +OUT = Path(__file__).resolve().parent.parent / "results-loop" +LAYER = 30 +K_HARD = 4 + + +@torch.no_grad() +def collect_states(model, tok, items, batch=16): + feats = [] + for i in range(0, len(items), batch): + chunk = items[i : i + batch] + enc = tok([mbpp_prompt(tok, it, DIRECT_SUFFIX) for it in chunk], + return_tensors="pt", padding=True, + add_special_tokens=False).to("cuda") + with ResidualCapture(model, layers=[LAYER]) as cap: + model(**enc, use_cache=False, logits_to_keep=1) + h = cap.acts[LAYER] # (B, T, d); left padding -> last pos is prompt end + feats.append(h[:, -1].float().cpu()) + return torch.cat(feats) + + +def fit_probe(X, y, epochs=300, lr=0.05): + mu, sd = X.mean(0), X.std(0) + 1e-6 + Xn = (X - mu) / sd + w = torch.zeros(X.shape[1], requires_grad=True) + b = torch.zeros(1, requires_grad=True) + opt = torch.optim.Adam([w, b], lr=lr) + pos_w = (y == 0).sum() / max(1, (y == 1).sum()) + for _ in range(epochs): + z = Xn @ w + b + loss = torch.nn.functional.binary_cross_entropy_with_logits( + z, y.float(), pos_weight=pos_w) + opt.zero_grad() + loss.backward() + opt.step() + return w.detach(), b.detach(), mu, sd + + +def main(): + model, tok = load_model(dtype=torch.bfloat16) + tok.padding_side = "left" + looper = BandLooper(model) + adapter = MergeAdapter().cuda() + adapter.load_state_dict(torch.load(OUT / "adapter_code_e399.pt", + map_location="cuda")) + + data = json.load(open(OUT / "mbpp_data.json")) + train = [it for it in data if it["split"] == "train" + and it["label"] in ("easy", "hard")] + test = [it for it in data if it["split"] == "test"][:250] + + print(f"collecting states: {len(train)} train, {len(test)} test", flush=True) + Xtr = collect_states(model, tok, train) + ytr = torch.tensor([it["label"] == "hard" for it in train]).long() + Xte = collect_states(model, tok, test) + + w, b, mu, sd = fit_probe(Xtr, ytr) + ptr = torch.sigmoid(((Xtr - mu) / sd) @ w + b) + acc_tr = ((ptr > 0.5).long() == ytr).float().mean() + pte = torch.sigmoid(((Xte - mu) / sd) @ w + b) + pred_hard = (pte > 0.5).tolist() + yte = [it["label"] == "hard" for it in test] + tp = sum(p and t for p, t in zip(pred_hard, yte)) + print(f"probe: train_acc={acc_tr:.3f} test: pred_hard=" + f"{sum(pred_hard)} (true hard={sum(yte)}, tp={tp})", flush=True) + + # gated generation: k=0 for predicted-easy, K_HARD for predicted-hard + codes = [None] * len(test) + for k, sel in ((0, [i for i, ph in enumerate(pred_hard) if not ph]), + (K_HARD, [i for i, ph in enumerate(pred_hard) if ph])): + for i in range(0, len(sel), 8): + idxs = sel[i : i + 8] + chunk = [test[j] for j in idxs] + enc = tok([mbpp_prompt(tok, it, DIRECT_SUFFIX) for it in chunk], + return_tensors="pt", padding=True, + add_special_tokens=False).to("cuda") + gen = looper.generate_frozen_prompt( + adapter, tok, enc["input_ids"], k, max_new_tokens=220, + attention_mask=enc["attention_mask"]) + for jj, j in enumerate(idxs): + codes[j] = extract_code( + tok.decode(gen[jj, enc["input_ids"].shape[1]:], + skip_special_tokens=True)) + with ThreadPoolExecutor(8) as ex: + oks = list(ex.map(lambda ci: run_tests(ci[0], ci[1]), + zip(codes, test))) + acc = sum(oks) / len(test) + by = {} + for it, ok in zip(test, oks): + d = by.setdefault(it["label"], [0, 0]) + d[0] += ok + d[1] += 1 + print(f"GATED pass@1={acc:.3f} " + f"by_label={ {l: round(c/n,3) for l,(c,n) in by.items()} }", flush=True) + json.dump({"acc": acc, + "by_label": {l: c / n for l, (c, n) in by.items()}, + "probe_train_acc": float(acc_tr), + "pred_hard_test": int(sum(pred_hard))}, + open(OUT / "eval_gated.json", "w"), indent=1) + print("wrote", OUT / "eval_gated.json") + + +if __name__ == "__main__": + main() diff --git a/scripts/lens_battery.py b/scripts/lens_battery.py new file mode 100644 index 0000000..410dc5d --- /dev/null +++ b/scripts/lens_battery.py @@ -0,0 +1,75 @@ +"""J-lens sharpening battery on MBPP prompts: does the trained loop +concentrate the workspace readout while reading the problem? + +For N test prompts, at each k: J-lens distribution at L30, last prompt +position -> entropy + top-1 probability. Trained vs untrained adapter. +""" + +import json +import sys +from pathlib import Path + +import torch + +from loop_common import BandLooper, MergeAdapter +from prep_mbpp import DIRECT_SUFFIX, mbpp_prompt + +sys.path.insert(0, str(Path(__file__).resolve().parent.parent)) +from jlens.core import JLens, load_model # noqa: E402 + +OUT = Path(__file__).resolve().parent.parent / "results-loop" +ROOT = OUT.parent +N_ITEMS = 20 +KS = (0, 1, 2, 4) + + +@torch.no_grad() +def battery(looper, adapter, lens, tok, items): + rows = [] + for it in items: + ids = tok(mbpp_prompt(tok, it, DIRECT_SUFFIX), return_tensors="pt", + add_special_tokens=False)["input_ids"].cuda() + calls, _ = looper.capture(ids, logits_to_keep=1) + e = looper._hin[looper.l0] + s = looper.band(e, calls) + state = {0: s} + for k in range(1, max(KS) + 1): + s = looper.band(adapter(e, s), calls) + state[k] = s + for k in KS: + probs = torch.softmax(lens._readout( + state[k][0, -1].float() @ lens.Jbar[looper.l1].T).float(), -1) + ent = -(probs * (probs + 1e-12).log()).sum().item() + rows.append({"task_id": it["task_id"], "k": k, "entropy": ent, + "top1": probs.max().item()}) + return rows + + +def main(): + model, tok = load_model(dtype=torch.bfloat16) + looper = BandLooper(model) + jbar = torch.load(ROOT / "results" / "jbar.pt", map_location="cuda") + lens = JLens(model, tok, jbar["Jbar"].float()) + items = [it for it in json.load(open(OUT / "mbpp_data.json")) + if it["split"] == "test"][:N_ITEMS] + + res = {} + for tag, path in (("untrained", None), + ("trained", OUT / "adapter_code_e399.pt")): + adapter = MergeAdapter().cuda() + if path: + adapter.load_state_dict(torch.load(path, map_location="cuda")) + rows = battery(looper, adapter, lens, tok, items) + res[tag] = rows + for k in KS: + sel = [r for r in rows if r["k"] == k] + ent = sum(r["entropy"] for r in sel) / len(sel) + top = sum(r["top1"] for r in sel) / len(sel) + print(f"[{tag}] k={k}: mean_entropy={ent:.3f} mean_top1={top:.3f}", + flush=True) + json.dump(res, open(OUT / "lens_battery.json", "w"), indent=1) + print("wrote", OUT / "lens_battery.json") + + +if __name__ == "__main__": + main() diff --git a/scripts/loop_common.py b/scripts/loop_common.py new file mode 100644 index 0000000..9f42c19 --- /dev/null +++ b/scripts/loop_common.py @@ -0,0 +1,310 @@ +"""Shared machinery for workspace-band looping (WORKSPACE_LOOPING.md). + +Codifies probe 2d: loop L14-30 through a merge layer at the L13->L14 boundary, + L14_in = (1-alpha)*e + alpha*(s renormalized to |e|) + MLP([e; s_hat]) +with e = L13 output (fixed anchor) and s = looped-back band output. +The MLP is zero-initialized, so the untrained adapter reproduces the +hand-built alpha-merge exactly. + +Band forward is done by re-calling the decoder layers with the per-layer +(args, kwargs) captured from a normal forward pass -- position embeddings and +attention masks do not depend on hidden states, so they are reusable across +loop iterations. +""" + +import re +import sys +from pathlib import Path + +import torch +import torch.nn as nn +from torch.utils.checkpoint import checkpoint + +sys.path.insert(0, str(Path(__file__).resolve().parent.parent)) +from jlens.core import _text_model, load_model # noqa: E402 + +import os as _os +BAND = tuple(int(x) for x in + _os.environ.get("JLENS_BAND", "14,30").split(",")) +# default: workspace band on gemma-4-E2B (inclusive); 12B: JLENS_BAND=36,45 +ALPHA = 0.3 # anchor-dominant merge weight from probe 2d + + +class MergeAdapter(nn.Module): + """(1-a)*e + a*s_hat + MLP([e; s_hat]); MLP zero-init => starts at probe 2d.""" + + def __init__(self, d=1536, hidden=512, alpha=ALPHA): + super().__init__() + self.alpha = alpha + self.mlp = nn.Sequential( + nn.Linear(2 * d, hidden), nn.GELU(), nn.Linear(hidden, d) + ) + nn.init.zeros_(self.mlp[2].weight) + nn.init.zeros_(self.mlp[2].bias) + + def forward(self, e, s): + dt = e.dtype + e32, s32 = e.float(), s.float() + s_hat = s32 * ( + e32.norm(dim=-1, keepdim=True) / (s32.norm(dim=-1, keepdim=True) + 1e-6) + ) + base = (1 - self.alpha) * e32 + self.alpha * s_hat + out = base + self.mlp(torch.cat([e32, s_hat], dim=-1)) + return out.to(dt) + + +class BandLooper: + """Capture layer-call kwargs once per forward, then re-run L14-30 manually.""" + + def __init__(self, model, band=BAND): + self.model = model + self.tm = _text_model(model) + self.l0, self.l1 = band + self.n_layers = len(self.tm.layers) + + def capture(self, input_ids, attention_mask=None, logits_to_keep=0, + position_ids=None): + """Plain forward; returns (calls dict, logits). calls[i] = (args, kwargs).""" + calls, hin, handles = {}, {}, [] + for i in range(self.l0, self.n_layers): + def pre(mod, args, kwargs, i=i): + # normalize: strip hidden_states, keep the rest for re-calls + drop = {"hidden_states", "past_key_value", "past_key_values", + "use_cache"} + kw = {k: v for k, v in kwargs.items() if k not in drop} + if "hidden_states" in kwargs: + hin[i] = kwargs["hidden_states"] + calls[i] = (args, kw) + else: + hin[i] = args[0] + calls[i] = (args[1:], kw) + handles.append( + self.tm.layers[i].register_forward_pre_hook(pre, with_kwargs=True) + ) + try: + with torch.no_grad(): + out = self.model(input_ids=input_ids, attention_mask=attention_mask, + use_cache=False, logits_to_keep=logits_to_keep, + position_ids=position_ids) + finally: + for h in handles: + h.remove() + self._hin = hin + return calls, out.logits + + def _run(self, h, calls, lo, hi): + for i in range(lo, hi + 1): + args, kwargs = calls[i] + h = self.tm.layers[i](h, *args, **kwargs) + if isinstance(h, tuple): + h = h[0] + return h + + def band(self, h, calls): + return self._run(h, calls, self.l0, self.l1) + + def suffix_logits(self, h, calls, last_only=False): + h = self._run(h, calls, self.l1 + 1, self.n_layers - 1) + if last_only: + h = h[:, -1:] + h = self.tm.norm(h.to(self.tm.norm.weight.dtype)) + head = self.model.lm_head if hasattr(self.model, "lm_head") else self.model.get_output_embeddings() + logits = head(h) + cap = getattr(self.model.config.get_text_config(), "final_logit_softcapping", None) + if cap: + logits = cap * torch.tanh(logits / cap) + return logits + + def loop_logits(self, adapter, input_ids, k, attention_mask=None, + use_checkpoint=False, return_states=False, last_only=False, + loop_mask=None, feedforward=False): + """Teacher-forced logits after k merge->band loops. k=0 = plain forward. + + loop_mask (B, T) bool: positions where the merge applies; elsewhere the + band input stays the anchor e ("latent planning" over the prompt span — + unmasked positions still re-attend to the looped states each iteration). + """ + calls, base_logits = self.capture(input_ids, attention_mask, + logits_to_keep=1 if last_only else 0) + if k == 0: + return (base_logits, None) if return_states else base_logits + del base_logits # full-vocab logits — do not hold across the loop + e = self._hin[self.l0].detach() + if feedforward: + # no-recurrence control: adapter sees (e, e), applied exactly once + x = adapter(e, e) + if loop_mask is not None: + x = torch.where(loop_mask[..., None], x, e) + s = (checkpoint(lambda x_: self.band(x_, calls), x, + use_reentrant=False) if use_checkpoint + else self.band(x, calls)) + logits = self.suffix_logits(s, calls, last_only=last_only) + return (logits, [s]) if return_states else logits + with torch.no_grad(): + s = self.band(e, calls) # s_0: no trainable params upstream + states = [s] + for _ in range(k): + x = adapter(e, s) + if loop_mask is not None: + x = torch.where(loop_mask[..., None], x, e) + if use_checkpoint: + s = checkpoint(lambda x_: self.band(x_, calls), x, use_reentrant=False) + else: + s = self.band(x, calls) + states.append(s) + logits = self.suffix_logits(s, calls, last_only=last_only) + return (logits, states) if return_states else logits + + @torch.no_grad() + def loop_generate(self, adapter, tok, input_ids, k, max_new_tokens=12, + attention_mask=None, loop_prompt_only=False, + stop_strs=()): + """Batched greedy decode with the looped forward (no KV cache). + + loop_prompt_only: merge applies only to the initial prompt span; + generated tokens go through the plain band (latent planning).""" + ids = input_ids + mask = attention_mask + lmask = None + if loop_prompt_only: + lmask = (mask if mask is not None + else torch.ones_like(ids)).bool().clone() + n_prompt = ids.shape[1] + texts = [""] * ids.shape[0] + eos = {tok.eos_token_id} + eot = tok.convert_tokens_to_ids("") + if eot is not None and eot >= 0: + eos.add(eot) + done = torch.zeros(ids.shape[0], dtype=torch.bool, device=ids.device) + for _ in range(max_new_tokens): + logits = self.loop_logits(adapter, ids, k, attention_mask=mask, + last_only=True, loop_mask=lmask) + nxt = logits[:, -1].argmax(-1) + nxt = torch.where(done, torch.full_like(nxt, list(eos)[0]), nxt) + ids = torch.cat([ids, nxt[:, None]], 1) + if mask is not None: + mask = torch.cat([mask, (~done)[:, None].long()], 1) + if lmask is not None: + lmask = torch.cat( + [lmask, torch.zeros_like(lmask[:, :1])], 1) + for b, t in enumerate(nxt.tolist()): + if not done[b] and stop_strs: + texts[b] += tok.decode([t]) + done |= torch.tensor([t.item() in eos for t in nxt], device=ids.device) + if stop_strs: + done |= torch.tensor( + [any(ss in tx for ss in stop_strs) for tx in texts], + device=ids.device) + if done.all(): + break + return ids + + @torch.no_grad() + def generate_frozen_prompt(self, adapter, tok, input_ids, k, + max_new_tokens=220, attention_mask=None, + stop_strs=(), feedforward=False): + """Fast equivalent of loop_generate(loop_prompt_only=True). + + The looped prompt states are constant across token steps (causality), + so: loop the prompt once to the merged input x* = adapter(e, s*), then + run ONE cached prefill with a hook swapping the band input to x*, and + generate normally with the KV cache (native speed).""" + if k == 0: + x_star = None + elif feedforward: + calls, _ = self.capture(input_ids, attention_mask, + logits_to_keep=1) + x_star = adapter(self._hin[self.l0], self._hin[self.l0]) + del calls + else: + calls, _ = self.capture(input_ids, attention_mask, + logits_to_keep=1) + e = self._hin[self.l0] + s = self.band(e, calls) + x_star = e + for _ in range(k): + x_star = adapter(e, s) + s = self.band(x_star, calls) + del calls # prefill re-runs band(x_star) -> same final s as slow path + + hook = None + if x_star is not None: + def swap(mod, args, kwargs): + h = kwargs.get("hidden_states", args[0] if args else None) + if h is not None and h.shape[1] == x_star.shape[1]: # prefill + if "hidden_states" in kwargs: + kwargs["hidden_states"] = x_star.to(h.dtype) + return args, kwargs + return (x_star.to(h.dtype),) + args[1:], kwargs + return None + hook = self.tm.layers[self.l0].register_forward_pre_hook( + swap, with_kwargs=True) + try: + eos = [t for t in (tok.eos_token_id, + tok.convert_tokens_to_ids("")) + if t is not None and t >= 0] + out = self.model.generate( + input_ids=input_ids, attention_mask=attention_mask, + max_new_tokens=max_new_tokens, do_sample=False, + eos_token_id=eos, pad_token_id=tok.pad_token_id or 0, + stop_strings=list(stop_strs) or None, tokenizer=tok) + finally: + if hook: + hook.remove() + return out + + +# ---------- GSM8K helpers ---------- + +NUM_RE = re.compile(r"-?\$?[\d,]*\.?\d+") + + +def gold_answer(ans_field): + return ans_field.split("####")[-1].strip().replace(",", "").replace("$", "") + + +def last_number(text): + hits = NUM_RE.findall(text) + if not hits: + return None + x = hits[-1].replace(",", "").replace("$", "").rstrip(".") + return x + + +def num_eq(a, b): + try: + return a is not None and b is not None and abs(float(a) - float(b)) < 1e-4 + except ValueError: + return False + + +DIRECT_SUFFIX = "\n\nGive only the final numeric answer, nothing else." +COT_SUFFIX = ("\n\nThink step by step, then give the final numeric answer " + "on the last line as: #### ") + + +def chat_prompt(tok, question, suffix=DIRECT_SUFFIX): + return tok.apply_chat_template( + [{"role": "user", "content": question + suffix}], + tokenize=False, add_generation_prompt=True, + ) + + +def build_train_batch(tok, items, device="cuda"): + """Right-padded (input_ids, attention_mask, labels); labels only on answer tokens.""" + seqs, labs = [], [] + for it in items: + p = tok(chat_prompt(tok, it["question"]), add_special_tokens=False)["input_ids"] + a = tok(it["gold"] + "", add_special_tokens=False)["input_ids"] + seqs.append(p + a) + labs.append([-100] * len(p) + a) + T = max(len(s) for s in seqs) + pad = tok.pad_token_id or 0 + ids = torch.full((len(seqs), T), pad, dtype=torch.long) + lab = torch.full((len(seqs), T), -100, dtype=torch.long) + msk = torch.zeros((len(seqs), T), dtype=torch.long) + for i, (s, l) in enumerate(zip(seqs, labs)): + ids[i, : len(s)] = torch.tensor(s) + lab[i, : len(s)] = torch.tensor(l) + msk[i, : len(s)] = 1 + return ids.to(device), msk.to(device), lab.to(device) diff --git a/scripts/plot_code_eval.py b/scripts/plot_code_eval.py new file mode 100644 index 0000000..e3af23e --- /dev/null +++ b/scripts/plot_code_eval.py @@ -0,0 +1,51 @@ +"""MBPP latent-planning results: pass@1 vs loop depth, trained vs untrained.""" +import json +from pathlib import Path + +import matplotlib +matplotlib.use("Agg") +import matplotlib.pyplot as plt + +OUT = Path(__file__).resolve().parent.parent / "results-loop" +tr = json.load(open(OUT / "eval_code_trained.json")) +un = json.load(open(OUT / "eval_code_untrained.json")) +ks = sorted(int(k) for k in tr["ks"]) +BLUE, GRAY = "#2b6cb0", "#8a8f98" + +def curve(res, sel): + return [res["ks"][str(k)]["acc"] if sel == "all" + else res["ks"][str(k)]["by_label"].get(sel, 0.0) for k in ks] + +fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(11, 4.2)) +for ax in (ax1, ax2): + ax.grid(True, color="#e5e5e5", lw=0.7) + ax.set_axisbelow(True) + for s in ("top", "right"): + ax.spines[s].set_visible(False) + ax.set_xticks(ks) + ax.set_xlabel("loop depth k (prompt-only)") + +ax1.plot(ks, curve(tr, "all"), "-o", color=BLUE, lw=2, ms=5, label="trained · all") +ax1.plot(ks, curve(un, "all"), "-o", color=GRAY, lw=2, ms=5, label="untrained · all") +ax1.axhline(curve(tr, "all")[0], color="#bbb", lw=1, ls=":") +ax1.set_ylabel("pass@1 (250 MBPP test items)") +ax1.set_title("Overall: trained k=2 beats no-loop baseline", fontsize=11) +ax1.legend(fontsize=8, frameon=False) + +ax2.plot(ks, curve(tr, "hard"), "-s", color=BLUE, lw=2, ms=5, + label="trained · hard (plan-only)") +ax2.plot(ks, curve(un, "hard"), "-s", color=GRAY, lw=2, ms=5, + label="untrained · hard") +ax2.set_ylabel("pass@1, hard bucket (~56 items)") +ax2.set_title("Planning-dependent problems: 3.6% → 46.4%", fontsize=11) +ax2.legend(fontsize=8, frameon=False) +ax2.annotate("silent loops recover ~46%\nof explicit-planning gap", + xy=(4, curve(tr, "hard")[-1]), xytext=(1.8, 0.30), fontsize=8, + color=BLUE, arrowprops=dict(arrowstyle="-", color=BLUE, lw=0.8)) + +fig.suptitle("MBPP latent planning: loop the workspace over the prompt, " + "then write code normally", fontsize=12, y=1.02) +fig.tight_layout() +fig.savefig(OUT / "loop_eval_code.png", dpi=140, bbox_inches="tight", + facecolor="white") +print("wrote", OUT / "loop_eval_code.png") diff --git a/scripts/plot_loop_dynamics.py b/scripts/plot_loop_dynamics.py new file mode 100644 index 0000000..43d9d38 --- /dev/null +++ b/scripts/plot_loop_dynamics.py @@ -0,0 +1,105 @@ +"""Loop convergence dynamics: does the recurrence reach a fixed point? + +For trained vs untrained merge, trace across iterations k: + - cos(s_k, s_{k-1}) (mean over positions) -> fixed point if -> 1 + - |s_k| / |e| -> norm control + - P('spider') under the J-lens at L30 -> what the state converges TO +Averaged over the spider prompt + a few GSM8K test questions. +""" + +import json +import sys +from pathlib import Path + +import matplotlib +matplotlib.use("Agg") +import matplotlib.pyplot as plt +import torch + +from loop_common import BandLooper, MergeAdapter, chat_prompt, DIRECT_SUFFIX + +sys.path.insert(0, str(Path(__file__).resolve().parent.parent)) +from jlens.core import load_model # noqa: E402 + +OUT = Path(__file__).resolve().parent.parent / "results-loop" +KMAX = 10 +BLUE, GRAY = "#2b6cb0", "#8a8f98" + + +@torch.no_grad() +def trace(looper, adapter, tok, prompt_text): + ids = tok(prompt_text, return_tensors="pt", + add_special_tokens=False)["input_ids"].cuda() + calls, _ = looper.capture(ids) + e = looper._hin[looper.l0] + s = looper.band(e, calls) + rows = [] + for k in range(1, KMAX + 1): + new = looper.band(adapter(e, s), calls) + cos = torch.nn.functional.cosine_similarity( + new[0].float(), s[0].float(), dim=-1).mean().item() + rows.append({"k": k, "cos": cos, + "norm": (new.norm() / e.norm()).item()}) + s = new + return rows + + +def main(): + model, tok = load_model(dtype=torch.bfloat16) + looper = BandLooper(model) + + gsm = [it for it in json.load(open(OUT / "star_data.json")) + if it["split"] == "test"][:3] + prompts = [chat_prompt(tok, "The animal that spins webs has how many legs? " + "Answer with just the number.", "")] + prompts += [chat_prompt(tok, it["question"], DIRECT_SUFFIX) for it in gsm] + + curves = {} + for tag, path in (("untrained", None), ("trained", OUT / "adapter.pt")): + adapter = MergeAdapter().cuda() + if path: + adapter.load_state_dict(torch.load(path, map_location="cuda")) + traces = [trace(looper, adapter, tok, p) for p in prompts] + curves[tag] = { + "cos": [sum(t[i]["cos"] for t in traces) / len(traces) + for i in range(KMAX)], + "norm": [sum(t[i]["norm"] for t in traces) / len(traces) + for i in range(KMAX)], + } + print(tag, "cos:", [round(c, 3) for c in curves[tag]["cos"]], flush=True) + print(tag, "norm:", [round(c, 3) for c in curves[tag]["norm"]], flush=True) + + ks = list(range(1, KMAX + 1)) + fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(11, 4)) + for ax in (ax1, ax2): + ax.grid(True, color="#e5e5e5", lw=0.7) + ax.set_axisbelow(True) + for sp in ("top", "right"): + ax.spines[sp].set_visible(False) + ax.set_xticks(ks) + ax.set_xlabel("loop iteration k") + ax.axvspan(2, 4, color="#f2e8cf", alpha=0.45, zorder=0) + + for tag, c in (("trained", BLUE), ("untrained", GRAY)): + ax1.plot(ks, curves[tag]["cos"], "-o", color=c, lw=2, ms=5, label=tag) + ax2.plot(ks, curves[tag]["norm"], "-o", color=c, lw=2, ms=5, label=tag) + ax1.set_ylabel("cos(s_k, s_{k−1}) (mean over positions)") + ax1.set_title("Successive-state similarity: fixed point?", fontsize=11) + ax1.legend(fontsize=8, frameon=False, loc="lower right") + ax1.text(3, ax1.get_ylim()[0] + 0.02 * (ax1.get_ylim()[1] - ax1.get_ylim()[0]), + "accuracy &\nsharpening plateau", fontsize=7.5, color="#7a5a00", + ha="center") + ax2.set_ylabel("|s_k| / |e|") + ax2.set_title("State norm across iterations", fontsize=11) + ax2.legend(fontsize=8, frameon=False) + + fig.suptitle("Loop dynamics (spider + 3 GSM8K prompts, mean)", fontsize=12, + y=1.02) + fig.tight_layout() + fig.savefig(OUT / "loop_dynamics.png", dpi=140, bbox_inches="tight", + facecolor="white") + print("wrote", OUT / "loop_dynamics.png") + + +if __name__ == "__main__": + main() diff --git a/scripts/plot_loop_eval.py b/scripts/plot_loop_eval.py new file mode 100644 index 0000000..b370cdb --- /dev/null +++ b/scripts/plot_loop_eval.py @@ -0,0 +1,57 @@ +"""Two-panel figure: accuracy vs loop depth k, and J-lens concept sharpening.""" +import json +from pathlib import Path + +import matplotlib +matplotlib.use("Agg") +import matplotlib.pyplot as plt + +OUT = Path(__file__).resolve().parent.parent / "results-loop" +tr = json.load(open(OUT / "eval_trained.json")) +un = json.load(open(OUT / "eval_untrained.json")) + +ks = sorted(int(k) for k in tr["ks"]) +BLUE, GRAY = "#2b6cb0", "#8a8f98" + +fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(11, 4.2)) + +def curve(res, sel): + return [res["ks"][str(k)]["acc"] if sel == "all" + else res["ks"][str(k)]["by_label"].get(sel, 0.0) for k in ks] + +for ax in (ax1, ax2): + ax.grid(True, color="#e5e5e5", lw=0.7, zorder=0) + ax.set_axisbelow(True) + for s in ("top", "right"): + ax.spines[s].set_visible(False) + ax.set_xticks(ks) + ax.set_xlabel("loop depth k") + +# --- panel 1: GSM8K accuracy vs k --- +ax1.plot(ks, curve(tr, "all"), "-o", color=BLUE, lw=2, ms=5, label="trained · all") +ax1.plot(ks, curve(un, "all"), "-o", color=GRAY, lw=2, ms=5, label="untrained · all") +ax1.plot(ks, curve(tr, "hard"), "--s", color=BLUE, lw=2, ms=5, label="trained · hard (CoT-only)") +ax1.plot(ks, curve(un, "hard"), "--s", color=GRAY, lw=2, ms=5, label="untrained · hard") +ax1.set_ylabel("accuracy (greedy, 256 held-out items)") +ax1.set_title("GSM8K accuracy vs loop depth", fontsize=11) +ax1.legend(fontsize=8, frameon=False) +ax1.annotate("hard items: 0.8% → 6.3%", xy=(2, 0.063), xytext=(3.2, 0.115), + fontsize=8, color=BLUE, + arrowprops=dict(arrowstyle="-", color=BLUE, lw=0.8)) + +# --- panel 2: J-lens sharpening --- +sp_tr = {r["k"]: r["P_spider_lens"] for r in tr["spider"]} +sp_un = {r["k"]: r["P_spider_lens"] for r in un["spider"]} +ax2.plot(ks, [sp_tr[k] for k in ks], "-o", color=BLUE, lw=2, ms=5, label="trained") +ax2.plot(ks, [sp_un[k] for k in ks], "-o", color=GRAY, lw=2, ms=5, label="untrained") +ax2.set_ylabel("P('spider') under J-lens at L30") +ax2.set_title("Latent concept sharpening across loops", fontsize=11) +ax2.legend(fontsize=8, frameon=False) +ax2.text(ks[-1], sp_tr[ks[-1]] + 0.006, "8× over control", fontsize=8, + color=BLUE, ha="right") + +fig.suptitle("Trained merge adapter: loop depth buys latent sharpening and some " + "CoT-only answers, not overall accuracy", fontsize=12, y=1.02) +fig.tight_layout() +fig.savefig(OUT / "loop_eval.png", dpi=140, bbox_inches="tight", facecolor="white") +print("wrote", OUT / "loop_eval.png") diff --git a/scripts/plot_results.py b/scripts/plot_results.py new file mode 100644 index 0000000..aa2ff11 --- /dev/null +++ b/scripts/plot_results.py @@ -0,0 +1,59 @@ +"""Render heatmaps (concept P by layer x position) and the layer profile.""" + +import os, sys +from pathlib import Path + +import matplotlib +matplotlib.use("Agg") +import matplotlib.pyplot as plt +import torch + +RES = Path(__file__).resolve().parent.parent / os.environ.get("JLENS_RESULTS", "results") + + +def plot_heat(path): + d = torch.load(path) + m, words, toks = d["map"].detach(), d["words"], d["tokens"] + fig, ax = plt.subplots(figsize=(min(16, 0.28 * m.shape[1] + 2), 6)) + im = ax.imshow(m.numpy(), aspect="auto", origin="lower", cmap="magma", + vmin=0) + ax.set_ylabel("layer") + ax.set_xlabel("position") + if toks and len(toks) <= 60: + ax.set_xticks(range(len(toks))) + ax.set_xticklabels([t.replace("\n", "\\n") for t in toks], + rotation=90, fontsize=6) + ax.set_title(f"J-lens P({'/'.join(words)})") + fig.colorbar(im) + out = path.with_suffix(".png") + fig.tight_layout() + fig.savefig(out, dpi=130) + plt.close(fig) + print("wrote", out) + + +def plot_profile(path): + d = torch.load(path) + fig, ax1 = plt.subplots(figsize=(8, 4.5)) + L = len(d["entropy"]) + ax1.plot(range(L), d["entropy"], "k-", label="lens entropy") + ax1.set_xlabel("layer") + ax1.set_ylabel("entropy (nats)") + ax2 = ax1.twinx() + ax2.plot(range(L), d["agree_in"], "b--", label="top tok == current input") + ax2.plot(range(L), d["agree_next"], "r--", label="top tok == next output") + ax2.set_ylabel("agreement") + fig.legend(loc="upper center", ncol=3, fontsize=8) + fig.tight_layout() + out = path.with_suffix(".png") + fig.savefig(out, dpi=130) + plt.close(fig) + print("wrote", out) + + +if __name__ == "__main__": + for p in sorted(RES.glob("heat_*.pt")): + plot_heat(p) + lp = RES / "layer_profile.pt" + if lp.exists(): + plot_profile(lp) diff --git a/scripts/prep_mbpp.py b/scripts/prep_mbpp.py new file mode 100644 index 0000000..13e4461 --- /dev/null +++ b/scripts/prep_mbpp.py @@ -0,0 +1,132 @@ +"""STaR-style difficulty labeling of MBPP with the frozen base model. + +Two passes per item: direct code generation vs plan-first-then-code, each +executed against MBPP's unit tests (sandboxed subprocess). Labels: +easy (direct passes), hard (plan-only passes), drop (neither). The model's +own passing code is kept as the training target (in-distribution supervision). +""" + +import json +import os +import re +import subprocess +import sys +import tempfile +import time +from concurrent.futures import ThreadPoolExecutor +from pathlib import Path + +import torch +from datasets import load_dataset + +sys.path.insert(0, str(Path(__file__).resolve().parent.parent)) +from jlens.core import load_model # noqa: E402 + +OUT = Path(os.environ.get("LOOP_OUT", + Path(__file__).resolve().parent.parent / "results-loop")) +OUT.mkdir(exist_ok=True) + +DIRECT_SUFFIX = ("\n\nWrite only the Python function in a ```python code " + "block. No explanation.") +PLAN_SUFFIX = ("\n\nFirst write a very brief plan: at most 4 short bullet " + "lines, no headings, no math notation. Then write the complete " + "Python function in a ```python code block.") + +CODE_RE = re.compile(r"```(?:python)?\s*\n(.*?)```", re.S) + + +def mbpp_prompt(tok, item, suffix): + tests = "\n".join(item["test_list"]) + msg = (f"{item['text']}\nYour code should pass these tests:\n\n{tests}" + f"{suffix}") + return tok.apply_chat_template([{"role": "user", "content": msg}], + tokenize=False, add_generation_prompt=True) + + +def extract_code(text): + m = CODE_RE.findall(text) + return m[-1].strip() if m else None + + +def run_tests(code, item, timeout=10): + if not code: + return False + script = (item.get("test_setup_code") or "") + "\n" + code + "\n" + \ + "\n".join(item["test_list"]) + try: + with tempfile.TemporaryDirectory() as td: + r = subprocess.run([sys.executable, "-c", script], cwd=td, + capture_output=True, timeout=timeout) + return r.returncode == 0 + except (subprocess.TimeoutExpired, OSError): + return False + + +@torch.no_grad() +def batch_generate(model, tok, prompts, max_new_tokens, batch_size=24): + outs = [] + for i in range(0, len(prompts), batch_size): + chunk = prompts[i : i + batch_size] + enc = tok(chunk, return_tensors="pt", padding=True, + add_special_tokens=False).to("cuda") + gen = model.generate(**enc, max_new_tokens=max_new_tokens, + do_sample=False, + pad_token_id=tok.pad_token_id or 0) + for j in range(len(chunk)): + outs.append(tok.decode(gen[j, enc["input_ids"].shape[1]:], + skip_special_tokens=True)) + print(f" {min(i+batch_size, len(prompts))}/{len(prompts)}", flush=True) + return outs + + +def main(): + model, tok = load_model(dtype=torch.bfloat16) + tok.padding_side = "left" + ds = load_dataset("google-research-datasets/mbpp", "full") + + items = [] + for split, tag in (("train", "train"), ("validation", "train"), + ("test", "test")): + for row in ds[split]: + items.append({"split": tag, "task_id": row["task_id"], + "text": row["text"], "test_list": row["test_list"], + "test_setup_code": row["test_setup_code"]}) + print(f"{len(items)} items", flush=True) + + for tag, suffix, mx in (("direct", DIRECT_SUFFIX, 220), + ("plan", PLAN_SUFFIX, 700)): + t0 = time.time() + print(f"{tag} pass", flush=True) + gens = batch_generate(model, tok, + [mbpp_prompt(tok, it, suffix) for it in items], mx) + codes = [extract_code(g) for g in gens] + with ThreadPoolExecutor(8) as ex: + oks = list(ex.map(lambda ci: run_tests(ci[0], ci[1]), + zip(codes, items))) + for it, c, ok in zip(items, codes, oks): + it[f"{tag}_ok"] = bool(ok) + it[f"{tag}_code"] = c if ok else None + print(f"{tag} pass done in {time.time()-t0:.0f}s " + f"pass@1={sum(oks)/len(items):.3f}", flush=True) + + for it in items: + it["label"] = ("easy" if it["direct_ok"] + else "hard" if it["plan_ok"] else "drop") + it["sol_code"] = it["direct_code"] if it["direct_ok"] else it["plan_code"] + + for split in ("train", "test"): + sub = [it for it in items if it["split"] == split] + n = len(sub) + print(f"{split}: n={n} direct={sum(i['direct_ok'] for i in sub)/n:.3f} " + f"plan={sum(i['plan_ok'] for i in sub)/n:.3f} " + f"easy={sum(i['label']=='easy' for i in sub)} " + f"hard={sum(i['label']=='hard' for i in sub)} " + f"drop={sum(i['label']=='drop' for i in sub)}", flush=True) + + with open(OUT / "mbpp_data.json", "w") as f: + json.dump(items, f, indent=1) + print("wrote", OUT / "mbpp_data.json") + + +if __name__ == "__main__": + main() diff --git a/scripts/prep_mbpp_fix.py b/scripts/prep_mbpp_fix.py new file mode 100644 index 0000000..bc0ea50 --- /dev/null +++ b/scripts/prep_mbpp_fix.py @@ -0,0 +1,65 @@ +"""Re-run the MBPP plan pass on direct-fail items with a non-truncating budget. + +The first plan pass (max_new=380) truncated ~all outputs before the code: +E2B writes verbose plans. Fix: terse-plan prompt + 700-token budget. Only +items with label=='drop' need re-labeling (easy is decided by the direct pass). +""" + +import json +import sys +import time +from concurrent.futures import ThreadPoolExecutor +from pathlib import Path + +import torch + +from prep_mbpp import batch_generate, extract_code, mbpp_prompt, run_tests + +sys.path.insert(0, str(Path(__file__).resolve().parent.parent)) +from jlens.core import load_model # noqa: E402 + +OUT = Path(__file__).resolve().parent.parent / "results-loop" + +TERSE_PLAN_SUFFIX = ("\n\nFirst write a very brief plan: at most 4 short " + "bullet lines, no headings, no math notation. Then write " + "the complete Python function in a ```python code block.") + + +def main(): + model, tok = load_model(dtype=torch.bfloat16) + tok.padding_side = "left" + data = json.load(open(OUT / "mbpp_data.json")) + redo = [it for it in data if it["label"] == "drop"] + print(f"re-running plan pass on {len(redo)} direct-fail items", flush=True) + + t0 = time.time() + gens = batch_generate(model, tok, + [mbpp_prompt(tok, it, TERSE_PLAN_SUFFIX) + for it in redo], max_new_tokens=700, batch_size=16) + codes = [extract_code(g) for g in gens] + n_trunc = sum(1 for g in gens if len(tok(g)["input_ids"]) >= 695) + with ThreadPoolExecutor(8) as ex: + oks = list(ex.map(lambda ci: run_tests(ci[0], ci[1]), + zip(codes, redo))) + for it, c, ok in zip(redo, codes, oks): + it["plan_ok"] = bool(ok) + it["plan_code"] = c if ok else None + it["label"] = "hard" if ok else "drop" + it["sol_code"] = it["direct_code"] if it["direct_ok"] else it["plan_code"] + print(f"done in {time.time()-t0:.0f}s plan-pass on fails: " + f"{sum(oks)}/{len(redo)} still-truncated={n_trunc}", flush=True) + + for split in ("train", "test"): + sub = [it for it in data if it["split"] == split] + n = len(sub) + print(f"{split}: n={n} easy={sum(i['label']=='easy' for i in sub)} " + f"hard={sum(i['label']=='hard' for i in sub)} " + f"drop={sum(i['label']=='drop' for i in sub)}", flush=True) + + with open(OUT / "mbpp_data.json", "w") as f: + json.dump(data, f, indent=1) + print("wrote", OUT / "mbpp_data.json") + + +if __name__ == "__main__": + main() diff --git a/scripts/prep_star_data.py b/scripts/prep_star_data.py new file mode 100644 index 0000000..20a7b12 --- /dev/null +++ b/scripts/prep_star_data.py @@ -0,0 +1,103 @@ +"""STaR-style difficulty labeling of GSM8K with the frozen base model. + +Two passes per item with frozen gemma-4-E2B-it: + direct: answer with no CoT -> correct = "easy" + cot: step-by-step -> correct (and direct wrong) = "hard" +Items the model cannot solve even with CoT are dropped from training +(unreachable supervision) but kept in the test split for eval. + +Output: results-loop/star_data.json +""" + +import json +import os +import random +import time +from pathlib import Path + +import torch +from datasets import load_dataset + +from loop_common import (COT_SUFFIX, DIRECT_SUFFIX, chat_prompt, gold_answer, + last_number, num_eq) +import sys +sys.path.insert(0, str(Path(__file__).resolve().parent.parent)) +from jlens.core import load_model # noqa: E402 + +N_TRAIN, N_TEST = 1024, 256 +OUT = Path(os.environ.get("LOOP_OUT", + Path(__file__).resolve().parent.parent / "results-loop")) +OUT.mkdir(exist_ok=True) + + +@torch.no_grad() +def batch_generate(model, tok, prompts, max_new_tokens, batch_size): + outs = [] + for i in range(0, len(prompts), batch_size): + chunk = prompts[i : i + batch_size] + enc = tok(chunk, return_tensors="pt", padding=True, + add_special_tokens=False).to("cuda") + gen = model.generate(**enc, max_new_tokens=max_new_tokens, + do_sample=False, + pad_token_id=tok.pad_token_id or 0) + for j in range(len(chunk)): + outs.append(tok.decode(gen[j, enc["input_ids"].shape[1]:], + skip_special_tokens=True)) + print(f" {min(i+batch_size, len(prompts))}/{len(prompts)}", flush=True) + return outs + + +def main(): + model, tok = load_model(dtype=torch.bfloat16) + tok.padding_side = "left" + ds = load_dataset("openai/gsm8k", "main") + + rng = random.Random(0) + tr_idx = rng.sample(range(len(ds["train"])), N_TRAIN) + te_idx = rng.sample(range(len(ds["test"])), N_TEST) + + items = [] + for split, idxs in (("train", tr_idx), ("test", te_idx)): + for i in idxs: + row = ds[split][i] + items.append({"split": split, "idx": i, "question": row["question"], + "gold": gold_answer(row["answer"])}) + + t0 = time.time() + print(f"direct pass ({len(items)} items)", flush=True) + direct = batch_generate( + model, tok, [chat_prompt(tok, it["question"], DIRECT_SUFFIX) for it in items], + max_new_tokens=10, batch_size=64) + print(f"direct pass done in {time.time()-t0:.0f}s", flush=True) + + t0 = time.time() + print("cot pass", flush=True) + cot = batch_generate( + model, tok, [chat_prompt(tok, it["question"], COT_SUFFIX) for it in items], + max_new_tokens=320, batch_size=32) + print(f"cot pass done in {time.time()-t0:.0f}s", flush=True) + + for it, d, c in zip(items, direct, cot): + it["direct_pred"] = last_number(d) + it["direct_ok"] = num_eq(it["direct_pred"], it["gold"]) + it["cot_pred"] = last_number(c) + it["cot_ok"] = num_eq(it["cot_pred"], it["gold"]) + it["label"] = ("easy" if it["direct_ok"] + else "hard" if it["cot_ok"] else "drop") + + for split in ("train", "test"): + sub = [it for it in items if it["split"] == split] + n = len(sub) + print(f"{split}: n={n} direct_acc={sum(i['direct_ok'] for i in sub)/n:.3f} " + f"cot_acc={sum(i['cot_ok'] for i in sub)/n:.3f} " + f"easy={sum(i['label']=='easy' for i in sub)} " + f"hard={sum(i['label']=='hard' for i in sub)} " + f"drop={sum(i['label']=='drop' for i in sub)}", flush=True) + + with open(OUT / "star_data.json", "w") as f: + json.dump(items, f, indent=1) + print("wrote", OUT / "star_data.json") + + +if __name__ == "__main__": + main() diff --git a/scripts/probe_carry.py b/scripts/probe_carry.py new file mode 100644 index 0000000..64ffcff --- /dev/null +++ b/scripts/probe_carry.py @@ -0,0 +1,116 @@ +"""Probe: cross-token residual-stream injection (stateful workspace). + +Instead of looping the band k times within each token step, carry the band +output S across token steps: at step t, one band pass with input +merge(e_t, S_{t-1}) (positions aligned; the new position inherits the last +position's state). Every position deepens by one iteration per emitted token +-- the amortized loop. Untrained probe: is this stable? coherent? does the +J-lens track? +""" + +import argparse +import sys +from pathlib import Path + +import torch + +from loop_common import BandLooper, MergeAdapter, chat_prompt + +sys.path.insert(0, str(Path(__file__).resolve().parent.parent)) +from jlens.core import JLens, load_model # noqa: E402 + +ROOT = Path(__file__).resolve().parent.parent + + +def seed_state(S, mode, gamma=0.9): + """State the new position inherits: last / mean / recency-weighted mean.""" + if mode == "last": + return S[:, -1:] + if mode == "mean": + return S.mean(dim=1, keepdim=True) + if mode == "ema": + T = S.shape[1] + w = gamma ** torch.arange(T - 1, -1, -1, device=S.device, + dtype=torch.float32) + w = (w / w.sum()).view(1, T, 1) + return (S.float() * w).sum(dim=1, keepdim=True).to(S.dtype) + raise ValueError(mode) + + +@torch.no_grad() +def generate_carry(looper, adapter, tok, ids, max_new=60, carry=True, + lens=None, concept_id=None, mode="last"): + """Greedy decode; band output carried across token steps via the merge.""" + S = None + trace = [] + eos = {tok.eos_token_id, tok.convert_tokens_to_ids("")} + for step in range(max_new): + calls, _ = looper.capture(ids, logits_to_keep=1) + e = looper._hin[looper.l0] + if S is None or not carry: + s = looper.band(e, calls) # plain first pass + else: + S_pad = torch.cat([S, seed_state(S, mode)], dim=1) + s = looper.band(adapter(e, S_pad), calls) # ONE pass, carried + S = s + logits = looper.suffix_logits(s, calls, last_only=True) + nxt = logits[0, -1].argmax().item() + if lens is not None and concept_id is not None: + probs = lens._readout(s[0].float() @ lens.Jbar[looper.l1].T) + probs = torch.softmax(probs.float(), -1) + trace.append({"step": step, + "P_concept": probs[:, concept_id].max().item(), + "norm": (s.norm() / e.norm()).item()}) + ids = torch.cat([ids, torch.tensor([[nxt]], device=ids.device)], 1) + if nxt in eos: + break + return ids, trace + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--adapter", default=None) + ap.add_argument("--max-new", type=int, default=60) + ap.add_argument("--mode", default="last", choices=["last", "mean", "ema"]) + args = ap.parse_args() + + model, tok = load_model(dtype=torch.bfloat16) + looper = BandLooper(model) + adapter = MergeAdapter().cuda() + tag = "untrained" + if args.adapter: + adapter.load_state_dict(torch.load(args.adapter, map_location="cuda")) + tag = Path(args.adapter).stem + jbar = torch.load(ROOT / "results" / "jbar.pt", map_location="cuda") + lens = JLens(model, tok, jbar["Jbar"].float()) + + prompts = [ + ("spider", "The animal that spins webs has how many legs? " + "Answer with just the number.", " spider"), + ("story", "In one sentence, why is the sky blue?", " blue"), + ("math", "Tom has 4 boxes of 12 pens and gives away 9 pens. " + "How many pens does he have left? Think step by step briefly, " + "then give the number.", " pens"), + ] + for name, q, concept in prompts: + cid = tok.encode(concept, add_special_tokens=False)[0] + ids = tok(chat_prompt(tok, q, ""), return_tensors="pt", + add_special_tokens=False)["input_ids"].cuda() + base, _ = generate_carry(looper, adapter, tok, ids, + max_new=args.max_new, carry=False) + carr, tr = generate_carry(looper, adapter, tok, ids, + max_new=args.max_new, carry=True, + lens=lens, concept_id=cid, mode=args.mode) + n = ids.shape[1] + print(f"\n=== [{tag}] {name} ===") + print("baseline :", tok.decode(base[0, n:], skip_special_tokens=True)) + print("carried :", tok.decode(carr[0, n:], skip_special_tokens=True)) + ks = [0, 2, 5, 10, 20, len(tr) - 1] + print("trace :", " ".join( + f"t={tr[i]['step']}: P={tr[i]['P_concept']:.3f} " + f"|s|/|e|={tr[i]['norm']:.2f}" + for i in sorted(set(k for k in ks if 0 <= k < len(tr))))) + + +if __name__ == "__main__": + main() diff --git a/scripts/relabel_sampled.py b/scripts/relabel_sampled.py new file mode 100644 index 0000000..b4e30dd --- /dev/null +++ b/scripts/relabel_sampled.py @@ -0,0 +1,104 @@ +"""Sampled relabeling of test-split difficulty buckets (protocol item: +outcome-selection fix). + +Greedy labeling couples bucket selection to the same coin flips as the k=0 +eval (circular: a bucket defined by baseline failure shows baseline ~0%). +Fix: labels from 3 temperature-sampled direct attempts, independent of the +greedy eval — hard_s = 0/3 sampled direct correct AND (plan/CoT reachable); +easy_s = >=2/3 correct; else mid_s. Adds 'label_sampled' to the data JSONs. +""" + +import json +import sys +from concurrent.futures import ThreadPoolExecutor +from pathlib import Path + +import torch + +from loop_common import DIRECT_SUFFIX as GSM_SUFFIX +from loop_common import chat_prompt, last_number, num_eq +from prep_mbpp import (DIRECT_SUFFIX as MBPP_SUFFIX, extract_code, + mbpp_prompt, run_tests) + +sys.path.insert(0, str(Path(__file__).resolve().parent.parent)) +from jlens.core import load_model # noqa: E402 + +OUT = Path(__file__).resolve().parent.parent / "results-loop" +N_SAMPLES = 3 +TEMP = 0.8 + + +@torch.no_grad() +def sample_batch(model, tok, prompts, max_new, batch=24, seed=0): + outs = [[] for _ in prompts] + for s in range(N_SAMPLES): + torch.manual_seed(1000 + s + seed) + for i in range(0, len(prompts), batch): + chunk = prompts[i : i + batch] + enc = tok(chunk, return_tensors="pt", padding=True, + add_special_tokens=False).to("cuda") + gen = model.generate(**enc, max_new_tokens=max_new, do_sample=True, + temperature=TEMP, top_p=0.95, + pad_token_id=tok.pad_token_id or 0) + for j in range(len(chunk)): + outs[i + j].append( + tok.decode(gen[j, enc["input_ids"].shape[1]:], + skip_special_tokens=True)) + print(f" sample {s+1}/{N_SAMPLES} done", flush=True) + return outs + + +def relabel(items, n_correct, reachable_key): + for it, nc in zip(items, n_correct): + if nc == 0: + it["label_sampled"] = ("hard" if it[reachable_key] else "drop") + elif nc >= 2: + it["label_sampled"] = "easy" + else: + it["label_sampled"] = "mid" + + +def main(): + model, tok = load_model(dtype=torch.bfloat16) + tok.padding_side = "left" + + # --- MBPP test --- + mbpp = json.load(open(OUT / "mbpp_data.json")) + mtest = [it for it in mbpp if it["split"] == "test"] + print(f"MBPP: sampling {len(mtest)} items x{N_SAMPLES}", flush=True) + gens = sample_batch(model, tok, + [mbpp_prompt(tok, it, MBPP_SUFFIX) for it in mtest], + max_new=220) + nc = [] + with ThreadPoolExecutor(8) as ex: + for it, gs in zip(mtest, gens): + oks = list(ex.map(lambda g: run_tests(extract_code(g), it), gs)) + nc.append(sum(oks)) + relabel(mtest, nc, "plan_ok") + json.dump(mbpp, open(OUT / "mbpp_data.json", "w"), indent=1) + dist = {l: sum(it.get("label_sampled") == l for it in mtest) + for l in ("easy", "mid", "hard", "drop")} + agree = sum(it["label"] == it.get("label_sampled") for it in mtest + if it.get("label_sampled") != "mid") + print(f"MBPP sampled labels: {dist} (greedy-agreement excl. mid: " + f"{agree}/{sum(1 for it in mtest if it.get('label_sampled') != 'mid')})", + flush=True) + + # --- GSM8K test --- + gsm = json.load(open(OUT / "star_data.json")) + gtest = [it for it in gsm if it["split"] == "test"] + print(f"GSM: sampling {len(gtest)} items x{N_SAMPLES}", flush=True) + gens = sample_batch(model, tok, + [chat_prompt(tok, it["question"], GSM_SUFFIX) + for it in gtest], max_new=10) + nc = [sum(num_eq(last_number(g), it["gold"]) for g in gs) + for it, gs in zip(gtest, gens)] + relabel(gtest, nc, "cot_ok") + json.dump(gsm, open(OUT / "star_data.json", "w"), indent=1) + dist = {l: sum(it.get("label_sampled") == l for it in gtest) + for l in ("easy", "mid", "hard", "drop")} + print(f"GSM sampled labels: {dist}", flush=True) + + +if __name__ == "__main__": + main() diff --git a/scripts/train_carry.py b/scripts/train_carry.py new file mode 100644 index 0000000..5115dce --- /dev/null +++ b/scripts/train_carry.py @@ -0,0 +1,149 @@ +"""Stage B / design C: train the merge adapter for prompt-prefill + carry. + +GSM8K only (the task the prompt-only loop failed on). Sequence: +[prompt] [p x pauses] [gold answer]; k=2 prefill loops on the +prompt; carry scan through pauses + answer; CE on answer tokens. +Curriculum: easy p=2, hard p=6 (hard needs the longer latent chain). +""" + +import argparse +import json +import math +import random +import sys +import time +from pathlib import Path + +import torch +import torch.nn.functional as F + +from carry_common import PAUSE_ID, carry_logits +from loop_common import BandLooper, MergeAdapter, chat_prompt, DIRECT_SUFFIX + +sys.path.insert(0, str(Path(__file__).resolve().parent.parent)) +from jlens.core import load_model # noqa: E402 + +OUT = Path(__file__).resolve().parent.parent / "results-loop" +STEPS = 600 +BATCH = 4 +LR = 1e-3 +WARMUP = 20 +K_PREFILL = 2 +P_BY_LABEL = {"easy": 2, "hard": 6} + +ap = argparse.ArgumentParser() +ap.add_argument("--feedforward", action="store_true", + help="pause-token control: adapter(e,e), no carry") +ap.add_argument("--seed", type=int, default=0) +ARGS = ap.parse_args() +SEED = ARGS.seed +TAG = "pausectl" if ARGS.feedforward else "carry" + + +def lr_at(step): + if step < WARMUP: + return LR * (step + 1) / WARMUP + t = (step - WARMUP) / max(1, STEPS - WARMUP) + return 1e-4 + 0.5 * (LR - 1e-4) * (1 + math.cos(math.pi * t)) + + +def build_batch(tok, items, p, device="cuda"): + seqs, labs, plens = [], [], [] + for it in items: + pr = tok(chat_prompt(tok, it["question"], DIRECT_SUFFIX), + add_special_tokens=False)["input_ids"] + a = tok(it["gold"] + "", + add_special_tokens=False)["input_ids"] + seqs.append(pr + [PAUSE_ID] * p + a) + labs.append([-100] * (len(pr) + p) + a) + plens.append(len(pr)) + T = max(len(s) for s in seqs) + pad = tok.pad_token_id or 0 + ids = torch.full((len(seqs), T), pad, dtype=torch.long) + lab = torch.full((len(seqs), T), -100, dtype=torch.long) + msk = torch.zeros((len(seqs), T), dtype=torch.long) + for i, (s, l) in enumerate(zip(seqs, labs)): + ids[i, : len(s)] = torch.tensor(s) + lab[i, : len(s)] = torch.tensor(l) + msk[i, : len(s)] = 1 + return (ids.to(device), msk.to(device), lab.to(device), + torch.tensor(plens, device=device)) + + +@torch.no_grad() +def val_loss(looper, adapter, tok, items, p, k=K_PREFILL): + tot, n = 0.0, 0 + for i in range(0, len(items), BATCH): + ids, msk, lab, plens = build_batch(tok, items[i : i + BATCH], p) + logits = carry_logits(looper, adapter, ids, msk, plens, k, feedforward=ARGS.feedforward) + loss = F.cross_entropy(logits[:, :-1].flatten(0, 1).float(), + lab[:, 1:].flatten(), ignore_index=-100) + tot += loss.item() * len(ids) + n += len(ids) + return tot / n + + +def main(): + rng = random.Random(SEED) + torch.manual_seed(SEED) + data = [it for it in json.load(open(OUT / "star_data.json")) + if it["split"] == "train" and it["label"] != "drop"] + + model, tok = load_model(dtype=torch.bfloat16) + for pp in model.parameters(): + pp.requires_grad_(False) + looper = BandLooper(model) + adapter = MergeAdapter().cuda() + opt = torch.optim.AdamW(adapter.parameters(), lr=LR, weight_decay=0.01) + + keep = [it for it in data + if len(tok(it["question"])["input_ids"]) + 30 <= 400] + rng.shuffle(keep) + pool = {l: [it for it in keep if it["label"] == l] + for l in ("easy", "hard")} + val = {l: pool[l][:16] for l in pool} + pool = {l: pool[l][16:] for l in pool} + print(f"pool: easy={len(pool['easy'])} hard={len(pool['hard'])}", + flush=True) + + log = [] + t0 = time.time() + for step in range(STEPS): + lbl = ("easy", "hard")[step % 2] + batch = rng.sample(pool[lbl], BATCH) + p = P_BY_LABEL[lbl] + ids, msk, lab, plens = build_batch(tok, batch, p) + + for g in opt.param_groups: + g["lr"] = lr_at(step) + logits = carry_logits(looper, adapter, ids, msk, plens, K_PREFILL, feedforward=ARGS.feedforward, + use_checkpoint=True) + loss = F.cross_entropy(logits[:, :-1].flatten(0, 1).float(), + lab[:, 1:].flatten(), ignore_index=-100) + opt.zero_grad(set_to_none=True) + loss.backward() + torch.nn.utils.clip_grad_norm_(adapter.parameters(), 1.0) + opt.step() + + log.append({"step": step, "label": lbl, "p": p, "loss": loss.item()}) + if step % 10 == 0: + print(f"step {step:4d} {lbl} p={p} loss={loss.item():.4f} " + f"({(time.time()-t0)/(step+1):.1f}s/step)", flush=True) + if step % 100 == 99 or step == STEPS - 1: + vals = {} + for l in ("easy", "hard"): + for pp_ in (0, 2, 6): + vals[f"{l}_p{pp_}"] = val_loss(looper, adapter, tok, + val[l], pp_) + print(f" val@{step}: " + " ".join( + f"{n}={v:.3f}" for n, v in sorted(vals.items())), flush=True) + log.append({"step": step, "val": vals}) + torch.save(adapter.state_dict(), + OUT / f"adapter_{TAG}_e{step+1}.pt") + json.dump(log, open(OUT / f"train_{TAG}_log.json", "w"), indent=1) + + print("done", flush=True) + + +if __name__ == "__main__": + main() diff --git a/scripts/train_distill.py b/scripts/train_distill.py new file mode 100644 index 0000000..36bde0f --- /dev/null +++ b/scripts/train_distill.py @@ -0,0 +1,175 @@ +"""Plan-distillation baseline: same 1.6M budget, no recurrence. + +Teacher: frozen model WITH its own terse plan in context (the plan text the +STaR pipeline validated). Student: feedforward adapter (adapter(e,e) at +prompt positions, applied once), NO plan in context. Loss: KL(teacher → +student) on code tokens + 0.5 CE on gold code. Bounds how much of the loop's +gain is available by compressing plan information into the adapter weights. +""" + +import argparse +import json +import math +import os +import random +import sys +import time +from pathlib import Path + +import torch +import torch.nn.functional as F + +from loop_common import BandLooper, MergeAdapter +from prep_mbpp import (DIRECT_SUFFIX, CODE_RE, batch_generate, mbpp_prompt) +from prep_mbpp_fix import TERSE_PLAN_SUFFIX + +sys.path.insert(0, str(Path(__file__).resolve().parent.parent)) +from jlens.core import load_model # noqa: E402 + +OUT = Path(os.environ.get("LOOP_OUT", + Path(__file__).resolve().parent.parent / "results-loop")) +STEPS = 600 +BATCH = 4 +LR = 1e-3 +WARMUP = 20 +KL_T = 1.0 +SEED = 0 + + +def lr_at(step): + if step < WARMUP: + return LR * (step + 1) / WARMUP + t = (step - WARMUP) / max(1, STEPS - WARMUP) + return 1e-4 + 0.5 * (LR - 1e-4) * (1 + math.cos(math.pi * t)) + + +def plan_prefix(gen_text): + """Text before the final fenced code block (the plan).""" + m = list(CODE_RE.finditer(gen_text)) + return gen_text[: m[-1].start()].strip() if m else None + + +def ensure_plans(model, tok, items): + path = OUT / "mbpp_plans.json" + if path.exists(): + plans = json.load(open(path)) + else: + print(f"generating plans for {len(items)} items", flush=True) + gens = batch_generate(model, tok, + [mbpp_prompt(tok, it, TERSE_PLAN_SUFFIX) + for it in items], max_new_tokens=700, + batch_size=16) + plans = {str(it["task_id"]): plan_prefix(g) + for it, g in zip(items, gens)} + json.dump(plans, open(path, "w"), indent=1) + return plans + + +def build_pair(tok, it, plan, device="cuda"): + """(student ids/mask/lmask, teacher ids, code positions in each).""" + code = "```python\n" + it["sol_code"] + "\n```" + a = tok(code, add_special_tokens=False)["input_ids"] + sp = tok(mbpp_prompt(tok, it, DIRECT_SUFFIX), + add_special_tokens=False)["input_ids"] + tp = tok(mbpp_prompt(tok, it, TERSE_PLAN_SUFFIX) + plan + "\n", + add_special_tokens=False)["input_ids"] + return sp, tp, a + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--seed", type=int, default=0) + args = ap.parse_args() + rng = random.Random(args.seed) + torch.manual_seed(args.seed) + + model, tok = load_model(dtype=torch.bfloat16) + tok.padding_side = "left" + for p in model.parameters(): + p.requires_grad_(False) + looper = BandLooper(model) + d = model.config.get_text_config().hidden_size + adapter = MergeAdapter(d=d).cuda() + opt = torch.optim.AdamW(adapter.parameters(), lr=LR, weight_decay=0.01) + + data = json.load(open(OUT / "mbpp_data.json")) + train = [it for it in data if it["split"] == "train" + and it["label"] != "drop" and it.get("sol_code")] + plans = ensure_plans(model, tok, train) + train = [it for it in train if plans.get(str(it["task_id"]))] + tok.padding_side = "right" + train = [it for it in train + if len(tok(mbpp_prompt(tok, it, TERSE_PLAN_SUFFIX))["input_ids"]) + + len(tok(plans[str(it["task_id"])])["input_ids"]) + + len(tok(it["sol_code"])["input_ids"]) + 30 <= 900] + print(f"distill pool: {len(train)}", flush=True) + + pad = tok.pad_token_id or 0 + log = [] + t0 = time.time() + for step in range(STEPS): + batch = rng.sample(train, BATCH) + pairs = [build_pair(tok, it, plans[str(it["task_id"])]) + for it in batch] + # student tensors + Ts = max(len(sp) + len(a) for sp, _, a in pairs) + Tt = max(len(tp) + len(a) for _, tp, a in pairs) + s_ids = torch.full((BATCH, Ts), pad, dtype=torch.long) + s_msk = torch.zeros((BATCH, Ts), dtype=torch.long) + t_ids = torch.full((BATCH, Tt), pad, dtype=torch.long) + t_msk = torch.zeros((BATCH, Tt), dtype=torch.long) + lab = torch.full((BATCH, Ts), -100, dtype=torch.long) + spans = [] + for i, (sp, tp, a) in enumerate(pairs): + s_ids[i, : len(sp) + len(a)] = torch.tensor(sp + a) + s_msk[i, : len(sp) + len(a)] = 1 + lab[i, len(sp): len(sp) + len(a)] = torch.tensor(a) + t_ids[i, : len(tp) + len(a)] = torch.tensor(tp + a) + t_msk[i, : len(tp) + len(a)] = 1 + spans.append((len(sp), len(tp), len(a))) + s_ids, s_msk, lab = s_ids.cuda(), s_msk.cuda(), lab.cuda() + t_ids, t_msk = t_ids.cuda(), t_msk.cuda() + lmask = (lab == -100) & (s_msk == 1) + + with torch.no_grad(): + t_logits = model(input_ids=t_ids, attention_mask=t_msk, + use_cache=False).logits + for g in opt.param_groups: + g["lr"] = lr_at(step) + s_logits = looper.loop_logits(adapter, s_ids, 1, attention_mask=s_msk, + loop_mask=lmask, feedforward=True, + use_checkpoint=True) + kl = torch.zeros((), device="cuda") + n_tok = 0 + for i, (ls, lt, la) in enumerate(spans): + # predicting code token j uses position (prefix_len + j - 1) + sl = s_logits[i, ls - 1: ls + la - 1].float() + tl = t_logits[i, lt - 1: lt + la - 1].float() + kl = kl + F.kl_div( + F.log_softmax(sl / KL_T, -1), F.log_softmax(tl / KL_T, -1), + log_target=True, reduction="sum") + n_tok += la + kl = kl / n_tok + ce = F.cross_entropy(s_logits[:, :-1].flatten(0, 1).float(), + lab[:, 1:].flatten(), ignore_index=-100) + loss = kl + 0.5 * ce + opt.zero_grad(set_to_none=True) + loss.backward() + torch.nn.utils.clip_grad_norm_(adapter.parameters(), 1.0) + opt.step() + + log.append({"step": step, "kl": kl.item(), "ce": ce.item()}) + if step % 10 == 0: + print(f"step {step:4d} kl={kl.item():.4f} ce={ce.item():.4f} " + f"({(time.time()-t0)/(step+1):.1f}s/step)", flush=True) + if step % 100 == 99 or step == STEPS - 1: + torch.save(adapter.state_dict(), + OUT / f"adapter_distill_e{step+1}.pt") + json.dump(log, open(OUT / "train_distill_log.json", "w"), + indent=1) + + print("done", flush=True) + + +if __name__ == "__main__": + main() diff --git a/scripts/train_merge.py b/scripts/train_merge.py new file mode 100644 index 0000000..aba33a4 --- /dev/null +++ b/scripts/train_merge.py @@ -0,0 +1,126 @@ +"""Train ONLY the merge adapter at the L13->L14 boundary (WORKSPACE_LOOPING.md §3). + +Frozen base model; band L14-30 unrolled k times through the adapter; loss = +cross-entropy on answer tokens of the *direct* (no-CoT) prompt, supervised by +gold answers on items the frozen model can reach (STaR filter from +prep_star_data.py). Difficulty->depth curriculum: + k=1: easy items k=2: easy+hard k=4: hard items +so hard items are only ever seen at depth, forcing the loop to be used. +""" + +import json +import random +import time +from pathlib import Path + +import torch +import torch.nn.functional as F + +from loop_common import BandLooper, MergeAdapter, build_train_batch +import sys +sys.path.insert(0, str(Path(__file__).resolve().parent.parent)) +from jlens.core import load_model # noqa: E402 + +OUT = Path(__file__).resolve().parent.parent / "results-loop" +STEPS = 800 +BATCH = 8 +LR = 1e-3 +WARMUP = 20 +MAX_TOK = 400 # skip overlong items +K_BUCKETS = [(1, ("easy",)), (2, ("easy", "hard")), (4, ("hard",))] +SEED = 0 + + +def lr_at(step): + if step < WARMUP: + return LR * (step + 1) / WARMUP + import math + t = (step - WARMUP) / max(1, STEPS - WARMUP) + return 1e-4 + 0.5 * (LR - 1e-4) * (1 + math.cos(math.pi * t)) + + +@torch.no_grad() +def val_answer_acc(looper, adapter, tok, items, k): + """Teacher-forced exact match: argmax at every answer position.""" + ok = 0 + for i in range(0, len(items), BATCH): + ids, msk, lab = build_train_batch(tok, items[i : i + BATCH]) + logits = looper.loop_logits(adapter, ids, k, attention_mask=msk) + pred = logits[:, :-1].argmax(-1) + tgt = lab[:, 1:] + m = tgt != -100 + ok += (((pred == tgt) | ~m).all(-1)).sum().item() + return ok / len(items) + + +def main(): + rng = random.Random(SEED) + torch.manual_seed(SEED) + data = json.load(open(OUT / "star_data.json")) + + model, tok = load_model(dtype=torch.bfloat16) + for p in model.parameters(): + p.requires_grad_(False) + looper = BandLooper(model) + adapter = MergeAdapter().cuda() + opt = torch.optim.AdamW(adapter.parameters(), lr=LR, weight_decay=0.01) + + train = [it for it in data if it["split"] == "train" and it["label"] != "drop"] + # length filter + keep = [] + for it in train: + n = len(tok(it["question"])["input_ids"]) + if n + 40 <= MAX_TOK: + keep.append(it) + train = keep + pool = {"easy": [it for it in train if it["label"] == "easy"], + "hard": [it for it in train if it["label"] == "hard"]} + val = {lbl: rng.sample(lst, min(32, len(lst))) for lbl, lst in pool.items()} + for lbl in pool: + pool[lbl] = [it for it in pool[lbl] if it not in val[lbl]] + print(f"train pool: easy={len(pool['easy'])} hard={len(pool['hard'])}", flush=True) + + log = [] + t0 = time.time() + for step in range(STEPS): + k, labels = K_BUCKETS[step % len(K_BUCKETS)] + cand = [it for lbl in labels for it in pool[lbl]] + batch = rng.sample(cand, BATCH) + ids, msk, lab = build_train_batch(tok, batch) + + for g in opt.param_groups: + g["lr"] = lr_at(step) + logits = looper.loop_logits(adapter, ids, k, attention_mask=msk, + use_checkpoint=True) + loss = F.cross_entropy( + logits[:, :-1].flatten(0, 1).float(), lab[:, 1:].flatten(), + ignore_index=-100) + opt.zero_grad(set_to_none=True) + loss.backward() + torch.nn.utils.clip_grad_norm_(adapter.parameters(), 1.0) + opt.step() + + log.append({"step": step, "k": k, "loss": loss.item()}) + if step % 10 == 0: + print(f"step {step:4d} k={k} loss={loss.item():.4f} " + f"({(time.time()-t0)/(step+1):.1f}s/step)", flush=True) + if step % 100 == 99 or step == STEPS - 1: + accs = {} + for kk in (0, 1, 2, 4): + accs[f"easy_k{kk}"] = val_answer_acc(looper, adapter, tok, + val["easy"], kk) + accs[f"hard_k{kk}"] = val_answer_acc(looper, adapter, tok, + val["hard"], kk) + print(f" val@{step}: " + + " ".join(f"{n}={v:.2f}" for n, v in accs.items()), flush=True) + log.append({"step": step, "val": accs}) + torch.save(adapter.state_dict(), OUT / "adapter.pt") + json.dump(log, open(OUT / "train_log.json", "w"), indent=1) + + torch.save(adapter.state_dict(), OUT / "adapter.pt") + json.dump(log, open(OUT / "train_log.json", "w"), indent=1) + print("done; adapter ->", OUT / "adapter.pt") + + +if __name__ == "__main__": + main() diff --git a/scripts/train_merge_bw.py b/scripts/train_merge_bw.py new file mode 100644 index 0000000..5f13951 --- /dev/null +++ b/scripts/train_merge_bw.py @@ -0,0 +1,129 @@ +"""Train the merge adapter on Blocksworld (prompt-only latent planning). + +Same recipe as train_merge_code.py; supervision = the model's own verified +passing plan (direct for easy, CoT-derived for hard — the plan section only, +CE on plan tokens).""" + +import argparse +import json +import math +import os +import random +import re +import sys +import time +from pathlib import Path + +import torch +import torch.nn.functional as F + +from bw_prep import DIRECT_SUFFIX, chat +from loop_common import BandLooper, MergeAdapter + +sys.path.insert(0, str(Path(__file__).resolve().parent.parent)) +from jlens.core import load_model # noqa: E402 + +OUT = Path(os.environ.get("LOOP_OUT", + Path(__file__).resolve().parent.parent / "results-loop")) +STEPS = 600 +BATCH = 4 +LR = 1e-3 +WARMUP = 20 +K_BUCKETS = [(1, ("easy",)), (2, ("easy", "hard")), (4, ("hard",))] + +ap = argparse.ArgumentParser() +ap.add_argument("--seed", type=int, default=0) +ARGS = ap.parse_args() +MOVELIST_RE = re.compile(r"(?:^|\n)\s*1\s*[.)]", re.M) + + +def plan_only(text): + """From a CoT output, keep from the numbered list onward.""" + m = MOVELIST_RE.search(text) + return text[m.start():].strip() if m else text.strip() + + +def lr_at(step): + if step < WARMUP: + return LR * (step + 1) / WARMUP + t = (step - WARMUP) / max(1, STEPS - WARMUP) + return 1e-4 + 0.5 * (LR - 1e-4) * (1 + math.cos(math.pi * t)) + + +def build_batch(tok, items, device="cuda"): + seqs, labs = [], [] + for it in items: + p = tok(chat(tok, it, DIRECT_SUFFIX), + add_special_tokens=False)["input_ids"] + a = tok(plan_only(it["sol_plan"]) + "", + add_special_tokens=False)["input_ids"] + seqs.append(p + a) + labs.append([-100] * len(p) + a) + T = max(len(s) for s in seqs) + pad = tok.pad_token_id or 0 + ids = torch.full((len(seqs), T), pad, dtype=torch.long) + lab = torch.full((len(seqs), T), -100, dtype=torch.long) + msk = torch.zeros((len(seqs), T), dtype=torch.long) + for i, (s, l) in enumerate(zip(seqs, labs)): + ids[i, : len(s)] = torch.tensor(s) + lab[i, : len(s)] = torch.tensor(l) + msk[i, : len(s)] = 1 + lmask = (lab == -100) & (msk == 1) + return ids.to(device), msk.to(device), lab.to(device), lmask.to(device) + + +def main(): + rng = random.Random(ARGS.seed) + torch.manual_seed(ARGS.seed) + data = json.load(open(OUT / "bw_data.json")) + + model, tok = load_model(dtype=torch.bfloat16) + for p in model.parameters(): + p.requires_grad_(False) + looper = BandLooper(model) + adapter = MergeAdapter( + d=model.config.get_text_config().hidden_size).cuda() + opt = torch.optim.AdamW(adapter.parameters(), lr=LR, weight_decay=0.01) + + train = [it for it in data if it["split"] == "train" + and it["label"] != "drop" and it.get("sol_plan")] + pool = {l: [it for it in train if it["label"] == l] + for l in ("easy", "hard")} + val = {l: pool[l][:12] for l in pool} + pool = {l: pool[l][12:] for l in pool} + print(f"pool: easy={len(pool['easy'])} hard={len(pool['hard'])}", + flush=True) + if min(len(pool["easy"]), len(pool["hard"])) < BATCH: + print("INSUFFICIENT POOL — aborting", flush=True) + return + + log = [] + t0 = time.time() + for step in range(STEPS): + k, labels = K_BUCKETS[step % len(K_BUCKETS)] + cand = [it for lbl in labels for it in pool[lbl]] + batch = rng.sample(cand, BATCH) + ids, msk, lab, lmask = build_batch(tok, batch) + for g in opt.param_groups: + g["lr"] = lr_at(step) + logits = looper.loop_logits(adapter, ids, k, attention_mask=msk, + use_checkpoint=True, loop_mask=lmask) + loss = F.cross_entropy(logits[:, :-1].flatten(0, 1).float(), + lab[:, 1:].flatten(), ignore_index=-100) + opt.zero_grad(set_to_none=True) + loss.backward() + torch.nn.utils.clip_grad_norm_(adapter.parameters(), 1.0) + opt.step() + log.append({"step": step, "k": k, "loss": loss.item()}) + if step % 10 == 0: + print(f"step {step:4d} k={k} loss={loss.item():.4f} " + f"({(time.time()-t0)/(step+1):.1f}s/step)", flush=True) + if step % 100 == 99 or step == STEPS - 1: + torch.save(adapter.state_dict(), + OUT / f"adapter_bw_e{step+1}.pt") + json.dump(log, open(OUT / "train_bw_log.json", "w"), indent=1) + print("done", flush=True) + + +if __name__ == "__main__": + main() diff --git a/scripts/train_merge_code.py b/scripts/train_merge_code.py new file mode 100644 index 0000000..d3bebe3 --- /dev/null +++ b/scripts/train_merge_code.py @@ -0,0 +1,164 @@ +"""Train the merge adapter on MBPP with prompt-only ("latent planning") looping. + +Same recipe as train_merge.py, with two changes: + - loop_mask: the merge applies only to prompt positions; the code tokens are + teacher-forced through the plain band (they still attend to looped prompt + states each iteration) — no exposure-bias gap by construction. + - supervision: the model's OWN passing code from prep_mbpp.py (easy: direct + pass; hard: code extracted from the plan-first pass), CE on code tokens. +""" + +import argparse +import json +import os +import math +import random +import sys +import time +from pathlib import Path + +import torch +import torch.nn.functional as F + +from loop_common import BandLooper, MergeAdapter +from prep_mbpp import DIRECT_SUFFIX, mbpp_prompt + +sys.path.insert(0, str(Path(__file__).resolve().parent.parent)) +from jlens.core import load_model # noqa: E402 + +OUT = Path(os.environ.get("LOOP_OUT", + Path(__file__).resolve().parent.parent / "results-loop")) +STEPS = 600 +BATCH = 4 +LR = 1e-3 +WARMUP = 20 +MAX_TOK = 512 +K_BUCKETS = [(1, ("easy",)), (2, ("easy", "hard")), (4, ("hard",))] + +ap = argparse.ArgumentParser() +ap.add_argument("--seed", type=int, default=0) +ap.add_argument("--pause", type=int, default=0, + help="pause-token control: p inert tokens after prompt, " + "feedforward adapter, no recurrence") +ARGS = ap.parse_args() +SEED = ARGS.seed +SUFFIX = (f"_s{SEED}" if SEED else "") + (f"_p{ARGS.pause}" if ARGS.pause else "") +PAUSE_ID = 6 # + + +def lr_at(step): + if step < WARMUP: + return LR * (step + 1) / WARMUP + t = (step - WARMUP) / max(1, STEPS - WARMUP) + return 1e-4 + 0.5 * (LR - 1e-4) * (1 + math.cos(math.pi * t)) + + +def build_code_batch(tok, items, device="cuda"): + """Right-padded batch; labels on code tokens; loop_mask on prompt span.""" + seqs, labs = [], [] + for it in items: + p = tok(mbpp_prompt(tok, it, DIRECT_SUFFIX), + add_special_tokens=False)["input_ids"] + p = p + [PAUSE_ID] * ARGS.pause + a = tok("```python\n" + it["sol_code"] + "\n```", + add_special_tokens=False)["input_ids"] + seqs.append(p + a) + labs.append([-100] * len(p) + a) + T = max(len(s) for s in seqs) + pad = tok.pad_token_id or 0 + ids = torch.full((len(seqs), T), pad, dtype=torch.long) + lab = torch.full((len(seqs), T), -100, dtype=torch.long) + msk = torch.zeros((len(seqs), T), dtype=torch.long) + for i, (s, l) in enumerate(zip(seqs, labs)): + ids[i, : len(s)] = torch.tensor(s) + lab[i, : len(s)] = torch.tensor(l) + msk[i, : len(s)] = 1 + lmask = (lab == -100) & (msk == 1) + return (ids.to(device), msk.to(device), lab.to(device), lmask.to(device)) + + +@torch.no_grad() +def val_loss(looper, adapter, tok, items, k): + tot, n = 0.0, 0 + for i in range(0, len(items), BATCH): + ids, msk, lab, lmask = build_code_batch(tok, items[i : i + BATCH]) + logits = looper.loop_logits(adapter, ids, k, attention_mask=msk, + loop_mask=lmask, feedforward=bool(ARGS.pause)) + loss = F.cross_entropy(logits[:, :-1].flatten(0, 1).float(), + lab[:, 1:].flatten(), ignore_index=-100) + tot += loss.item() * len(ids) + n += len(ids) + return tot / n + + +def main(): + rng = random.Random(SEED) + torch.manual_seed(SEED) + data = json.load(open(OUT / "mbpp_data.json")) + + model, tok = load_model(dtype=torch.bfloat16) + for p in model.parameters(): + p.requires_grad_(False) + looper = BandLooper(model) + d = model.config.get_text_config().hidden_size + adapter = MergeAdapter(d=d).cuda() + opt = torch.optim.AdamW(adapter.parameters(), lr=LR, weight_decay=0.01) + + train = [it for it in data if it["split"] == "train" + and it["label"] != "drop" and it["sol_code"]] + train = [it for it in train + if len(tok(mbpp_prompt(tok, it, DIRECT_SUFFIX))["input_ids"]) + + len(tok(it["sol_code"])["input_ids"]) + 12 <= MAX_TOK] + pool = {"easy": [it for it in train if it["label"] == "easy"], + "hard": [it for it in train if it["label"] == "hard"]} + val = {lbl: lst[:16] for lbl, lst in pool.items()} + for lbl in pool: + pool[lbl] = pool[lbl][16:] + print(f"train pool: easy={len(pool['easy'])} hard={len(pool['hard'])}", + flush=True) + + log = [] + t0 = time.time() + for step in range(STEPS): + k, labels = K_BUCKETS[step % len(K_BUCKETS)] + cand = [it for lbl in labels for it in pool[lbl]] + batch = rng.sample(cand, min(BATCH, len(cand))) + ids, msk, lab, lmask = build_code_batch(tok, batch) + + for g in opt.param_groups: + g["lr"] = lr_at(step) + logits = looper.loop_logits(adapter, ids, k, attention_mask=msk, + use_checkpoint=True, loop_mask=lmask, + feedforward=bool(ARGS.pause)) + loss = F.cross_entropy(logits[:, :-1].flatten(0, 1).float(), + lab[:, 1:].flatten(), ignore_index=-100) + opt.zero_grad(set_to_none=True) + loss.backward() + torch.nn.utils.clip_grad_norm_(adapter.parameters(), 1.0) + opt.step() + + log.append({"step": step, "k": k, "loss": loss.item()}) + if step % 10 == 0: + print(f"step {step:4d} k={k} loss={loss.item():.4f} " + f"({(time.time()-t0)/(step+1):.1f}s/step)", flush=True) + if step % 100 == 99 or step == STEPS - 1: + vals = {} + for kk in (0, 1, 2, 4): + vals[f"easy_k{kk}"] = val_loss(looper, adapter, tok, + val["easy"], kk) + vals[f"hard_k{kk}"] = val_loss(looper, adapter, tok, + val["hard"], kk) + print(f" val@{step}: " + + " ".join(f"{n}={v:.3f}" for n, v in vals.items()), flush=True) + log.append({"step": step, "val": vals}) + torch.save(adapter.state_dict(), + OUT / f"adapter_code{SUFFIX}_e{step+1}.pt") + json.dump(log, open(OUT / f"train_code_log{SUFFIX}.json", "w"), indent=1) + + torch.save(adapter.state_dict(), OUT / f"adapter_code{SUFFIX}.pt") + json.dump(log, open(OUT / f"train_code_log{SUFFIX}.json", "w"), indent=1) + print("done; adapter ->", OUT / f"adapter_code{SUFFIX}.pt") + + +if __name__ == "__main__": + main() diff --git a/scripts/train_merge_unified.py b/scripts/train_merge_unified.py new file mode 100644 index 0000000..99116c7 --- /dev/null +++ b/scripts/train_merge_unified.py @@ -0,0 +1,184 @@ +"""Unified merge adapter: GSM8K + MBPP, prompt-only looping, fresh init. + +One adapter, one regime (loop_mask = prompt span; generated/answer tokens +never loop), mixed-task batches. Hardened protocol elements: per-task val +holdouts, checkpoint every 100 steps kept separately (best-val selection and +k chosen on val, never test). +""" + +import argparse +import json +import os +import math +import random +import sys +import time +from pathlib import Path + +import torch +import torch.nn.functional as F + +from loop_common import BandLooper, MergeAdapter, chat_prompt, DIRECT_SUFFIX +from prep_mbpp import DIRECT_SUFFIX as MBPP_SUFFIX +from prep_mbpp import mbpp_prompt + +sys.path.insert(0, str(Path(__file__).resolve().parent.parent)) +from jlens.core import load_model # noqa: E402 + +OUT = Path(os.environ.get("LOOP_OUT", + Path(__file__).resolve().parent.parent / "results-loop")) +STEPS = 800 +BATCH = 6 +LR = 1e-3 +WARMUP = 20 +MAX_TOK = 512 +K_BUCKETS = [(1, ("easy",)), (2, ("easy", "hard")), (4, ("hard",))] +SEED = 0 +VAL_N = 16 # per task per bucket + + +def lr_at(step): + if step < WARMUP: + return LR * (step + 1) / WARMUP + t = (step - WARMUP) / max(1, STEPS - WARMUP) + return 1e-4 + 0.5 * (LR - 1e-4) * (1 + math.cos(math.pi * t)) + + +def item_texts(tok, it): + if it["task"] == "gsm": + return (chat_prompt(tok, it["question"], DIRECT_SUFFIX), + it["gold"] + "") + return (mbpp_prompt(tok, it, MBPP_SUFFIX), + "```python\n" + it["sol_code"] + "\n```") + + +def build_batch(tok, items, device="cuda"): + seqs, labs = [], [] + for it in items: + ptxt, atxt = item_texts(tok, it) + p = tok(ptxt, add_special_tokens=False)["input_ids"] + a = tok(atxt, add_special_tokens=False)["input_ids"] + seqs.append(p + a) + labs.append([-100] * len(p) + a) + T = max(len(s) for s in seqs) + pad = tok.pad_token_id or 0 + ids = torch.full((len(seqs), T), pad, dtype=torch.long) + lab = torch.full((len(seqs), T), -100, dtype=torch.long) + msk = torch.zeros((len(seqs), T), dtype=torch.long) + for i, (s, l) in enumerate(zip(seqs, labs)): + ids[i, : len(s)] = torch.tensor(s) + lab[i, : len(s)] = torch.tensor(l) + msk[i, : len(s)] = 1 + lmask = (lab == -100) & (msk == 1) + return ids.to(device), msk.to(device), lab.to(device), lmask.to(device) + + +@torch.no_grad() +def val_loss(looper, adapter, tok, items, k, feedforward=False): + tot, n = 0.0, 0 + for i in range(0, len(items), BATCH): + ids, msk, lab, lmask = build_batch(tok, items[i : i + BATCH]) + logits = looper.loop_logits(adapter, ids, k, attention_mask=msk, + loop_mask=lmask, feedforward=feedforward) + loss = F.cross_entropy(logits[:, :-1].flatten(0, 1).float(), + lab[:, 1:].flatten(), ignore_index=-100) + tot += loss.item() * len(ids) + n += len(ids) + return tot / n + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--noloop", action="store_true", + help="feedforward control: adapter(e,e), no recurrence") + ap.add_argument("--tasks", default="gsm,mbpp") + ap.add_argument("--tag", default=None) + args = ap.parse_args() + tasks = args.tasks.split(",") + tag = args.tag or ("ff" if args.noloop else "uni") + + rng = random.Random(SEED) + torch.manual_seed(SEED) + + gsm = [dict(it, task="gsm") for it in json.load(open(OUT / "star_data.json")) + if it["split"] == "train" and it["label"] != "drop"] + mbpp = [dict(it, task="mbpp") for it in json.load(open(OUT / "mbpp_data.json")) + if it["split"] == "train" and it["label"] != "drop" + and it.get("sol_code")] + if "gsm" not in tasks: + gsm = [] + if "mbpp" not in tasks: + mbpp = [] + + model, tok = load_model(dtype=torch.bfloat16) + for p in model.parameters(): + p.requires_grad_(False) + looper = BandLooper(model) + adapter = MergeAdapter( + d=model.config.get_text_config().hidden_size).cuda() + opt = torch.optim.AdamW(adapter.parameters(), lr=LR, weight_decay=0.01) + + def fits(it): + ptxt, atxt = item_texts(tok, it) + return (len(tok(ptxt)["input_ids"]) + len(tok(atxt)["input_ids"]) + <= MAX_TOK) + + items = [it for it in gsm + mbpp if fits(it)] + rng.shuffle(items) + pool, val = {"easy": [], "hard": []}, {} + for task in tasks: + for lbl in ("easy", "hard"): + sub = [it for it in items if it["task"] == task + and it["label"] == lbl] + val[(task, lbl)] = sub[:VAL_N] + pool.setdefault(lbl, []).extend(sub[VAL_N:]) + print("pool: easy={} hard={} (gsm {} / mbpp {})".format( + len(pool["easy"]), len(pool["hard"]), + sum(it["task"] == "gsm" for l in pool.values() for it in l), + sum(it["task"] == "mbpp" for l in pool.values() for it in l)), + flush=True) + + log = [] + t0 = time.time() + for step in range(STEPS): + k, labels = K_BUCKETS[step % len(K_BUCKETS)] + cand = [it for lbl in labels for it in pool[lbl]] + batch = rng.sample(cand, BATCH) + ids, msk, lab, lmask = build_batch(tok, batch) + + for g in opt.param_groups: + g["lr"] = lr_at(step) + logits = looper.loop_logits(adapter, ids, k, attention_mask=msk, + use_checkpoint=True, loop_mask=lmask, + feedforward=args.noloop) + loss = F.cross_entropy(logits[:, :-1].flatten(0, 1).float(), + lab[:, 1:].flatten(), ignore_index=-100) + opt.zero_grad(set_to_none=True) + loss.backward() + torch.nn.utils.clip_grad_norm_(adapter.parameters(), 1.0) + opt.step() + + log.append({"step": step, "k": k, "loss": loss.item()}) + if step % 10 == 0: + print(f"step {step:4d} k={k} loss={loss.item():.4f} " + f"({(time.time()-t0)/(step+1):.1f}s/step)", flush=True) + if step % 100 == 99 or step == STEPS - 1: + vals = {} + kk_grid = (0, 1) if args.noloop else (0, 1, 2, 4) + for (task, lbl), vitems in val.items(): + for kk in kk_grid: + vals[f"{task}_{lbl}_k{kk}"] = val_loss( + looper, adapter, tok, vitems, kk, + feedforward=args.noloop) + print(f" val@{step}: " + " ".join( + f"{n}={v:.3f}" for n, v in sorted(vals.items())), flush=True) + log.append({"step": step, "val": vals}) + torch.save(adapter.state_dict(), + OUT / f"adapter_{tag}_e{step+1}.pt") + json.dump(log, open(OUT / f"train_{tag}_log.json", "w"), indent=1) + + print("done", flush=True) + + +if __name__ == "__main__": + main()