EXPERIMENT_LOG.md

# Latent-CoT Graph Reachability — Experiment Log

Persistent working notes so progress/commands survive chat resets. Newest status
at the top of "Current Status"; details and history below. Update this file
whenever we decide/do something.

Repo: `reasoning-by-superposition-main` (NeurIPS 2025 "Reasoning by Superposition",
authors' official code). Env: conda env `superposition` (torch 2.5.1+cu121,
transformers 4.46.2). Node: `slimgpu` (8× RTX A5000).

---

## Goal

Small from-scratch GPT-2 (symbol, 2-layer / 8-head / 768-dim) trained with Coconut
continuous latent CoT on a 2-arm **star-graph reachability** task, where the token
vocabulary **is the set of graph node ids** (node id == token id). This makes the
LM head a "node dictionary": applying it (logit lens) to each latent thought reads
out *which graph node that thought points to*, so we can interpret each latent as
one BFS hop of reasoning depth. Then: scale depth (L6 -> L20) and study
backtracking (retrain earlier latent stages when they regress).

Task detail: two disjoint star components, each with root + 2 arms of length L.
One component is reachable from the query root, the other is a distractor. Model
must answer which candidate leaf is reachable. `bfs_variant: True` => the "frontier"
at hop k is the two reachable-arm nodes at depth k. Key metric = per-hop
`frontier` accuracy (does latent k decode to a correct depth-k reachable node).

---

## Current Status (2026-07-21)

- **Depth-6 (L6): DONE, success.** Per-hop `frontier` accuracy 0.98-1.0 across all
  6 hops. Logit-lens probe (`probe_latents.py`) confirms each of the 6 latents
  decodes to a correct node one BFS hop deeper on the reachable arm, never the
  distractor. Checkpoints in `ckpts/star-coconut-L6-bfs/`. Log:
  `logs/star_coconut_L6_bfs.log`.
- **Depth-20 (L20), epoch-scheduled version: FAILED to learn.** Config
  `args/star_coconut_L20_bfs.yaml`. The curriculum advanced the frontier on a fixed
  epoch timer (`scheduled_stage = epoch // epochs_per_stage`) and reached frontier
  14, but per-hop `frontier` acc stayed ~0.02-0.04 at EVERY hop (incl. hop 1).
  i.e. training compute got smeared across stages that were never mastered.
  **This run was killed.**
- **Depth-20, ACCURACY-GATED version: the current approach.** Config
  `args/star_coconut_L20_bfs_accstage.yaml`. Promotion is now driven by measured
  accuracy, not the epoch counter (see "Accuracy-gated curriculum" below).

### 2026-07-21 evening: root-caused the L20 failure = TRAINING DIVERGENCE
- The accuracy-gated run correctly HELD at stage 1 for 65 epochs (~1h) because
  hop-1 acc never cleared 0.9 -- but the real issue is that **eval loss was
  monotonically INCREASING**: ~4.62 -> 5.9+, while random loss for the 128-vocab is
  ln(128) ~= 4.85. So the model started ~chance and diverged past chance. hop-1 acc
  ~0.02 == uniform random over ~100 nodes. The old epoch-scheduled L20 run had the
  same near-random accuracy => this is an L20 training-instability problem, NOT a
  data or curriculum problem.
- Data verified well-formed (root/target/neighbor_k valid, node ids <= 99, 80 edges,
  20 steps). Difference from working L6: ~330-token sequences (80 edges) vs ~6 edges,
  wider vocab, `lr=1e-4`, and NO gradient clipping. Classic long-seq divergence.

### Fixes applied (this session)
1. **Stabilization** in `run.py` + config: gradient clipping (`grad_clip: 1.0`),
   lower `lr` 1e-4 -> 3e-5, linear LR `warmup_steps: 200`.
2. **Gate off-by-one fix**: metric "hop k" uses (k-1) latents to predict a depth-k
   node, so "stage s learned" == per-hop acc at hop (s+1). The promotion gate now
   requires min acc over hops 1..(cur_stage+1) >= threshold (previously 1..cur_stage,
   which never checked the current stage's own target).
3. **Limited backprop**: `backprop_depth: 1` (was 6) -- gradient only through the
   newest latent step + answer, matching the incremental curriculum. Simple, not
   full BPTT across stages.
- Status: relaunch `star_coconut_L20_bfs_accstage.yaml` with these fixes.

### 2026-07-21 late: first stabilization attempt still diverged -> FSDP clip bug
- With lr 3e-5 + grad_clip 1.0 + warmup, eval loss STILL climbed (4.62 -> 5.2 over
  80 epochs) and stage 1 stayed at chance (HOLD 80 epochs). So clipping wasn't
  actually working.
- Root cause: the model is wrapped in **FSDP** (params sharded), but the clip used
  `torch.nn.utils.clip_grad_norm_(parallel_model.parameters(), ...)`, which computes
  the norm over only each rank's local shard -> under-counts -> effectively no clip.
- Fix: use `parallel_model.clip_grad_norm_()` when the model is FSDP (all-reduces the
  global norm). Also lowered `lr` 3e-5 -> 1e-5.
- If loss STILL rises after this, the remaining suspect is model capacity: a 2-layer
  GPT-2 doing multi-hop pointer chasing over ~80 shuffled edges (vs ~6 at L6). Next
  step would be a controlled L6-with-same-code sanity run, then consider more layers.

---

## Accuracy-gated curriculum + backtracking (what we changed and why)

Motivation (K's guidance): stage promotion should be **accuracy-dependent**, not
iteration-dependent. Only advance to stage k+1 once every stage 1..k has reached a
target accuracy. And keep the sudoku-style backtracking: at each eval check all
previous stages; if any falls below threshold, go back and rehearse/retrain it.

Implemented in `run.py`:
- New config flag `accuracy_staging: True`. When on, `scheduled_stage = cur_stage`,
  a persistent frontier that only advances when earned.
- After each per-hop eval: promote `cur_stage -> cur_stage+1` only when
  `min(frontier_acc[1..cur_stage]) >= promote_threshold` (default 0.9). Because it
  checks *all* stages <= frontier, a regression anywhere blocks promotion.
- Logs per-stage timing: `[acc-stage] PROMOTE stage k -> k+1 | solved in N epochs / Ts`.
  `HOLD` lines mean the current frontier isn't solved yet.
- Backtracking rehearsal preserved: `bt_r_current` = earliest hop below
  `backtrack_detect_threshold`; the training-set builder
  `get_graph_latent_cot_dataset_backtrack` samples each example's stage from a
  rehearsal distribution (broad w.p. `remember_rate`, else targeted
  `[r_current..frontier]`).
- `perhop_val_samples` caps val per-hop eval cost so frequent eval stays cheap.

Earlier enabling fixes (already in the code):
- **Vectorized latent feedback** in `coconut.py` (clone+scatter instead of a
  Python double-loop of ~11k tiny GPU ops). Proven numerically identical.
- **O(batch^2 x depth) sync bug** in `coconut.py` `latent_lists` construction
  (per-element `.item()` GPU->CPU syncs) replaced with one `.tolist()` + pure-Python
  grouping. This was the real speed killer; latency went from 7357ms -> 561ms/it at
  stage 20 (~13x), roughly flat across depth.
- **Truncated BPTT** `backprop_depth: W` in `coconut.py`: detaches fed-back thought
  + KV older than W latent steps => bounded backward memory/compute regardless of
  depth. Default None = full BPTT (L6 baseline unchanged). L20 uses `backprop_depth: 6`.
- Widened node vocab to 100 (`stokenizer.py` `NUM_NODES=100`), model
  `configs/symbol-2layer-8head-768dim-L20.json` `vocab_size=128`, data regenerated
  at L=20 (`data/star_2arm_L20_*_fo_bfs.json`, 82 nodes/sample, 80 edges).

---

## Curriculum semantics (what each stage learns)

From `expand_data`: at **stage s** the prompt has **s latent tokens** and the model
is trained to output the **next hop's frontier node (depth s+1)** -- NOT the final
leaf. Only the final stage (`k = max_steps+1`, all latents + `[A]`) predicts the
target leaf. So each intermediate stage teaches one more hop of the walk:
- stage 0: 0 latents -> predict a depth-1 node (root's neighbor)
- stage 1: 1 latent  -> predict a depth-2 node
- stage s: s latents -> predict a depth-(s+1) node
- final:   all latents + [A] -> predict the target leaf
Per-hop metric "hop k" = (k-1) latents -> predict depth-k node, i.e. hop k tests
training stage (k-1). (This is why the promotion gate checks hop cur_stage+1.)

## Interpretability / probing (how latent -> node works)

- A latent "thought" is the previous position's last-layer hidden state
  h in R^768, fed back in place of a token embedding (never discretized).
- Logit lens: node_hat_k = argmax over node-id columns of (W_U h), where W_U is the
  model's own tied unembedding. Since node id == token id, this reads the node.
- `probe_latents.py`: one teacher-forced validation forward pass over held-out
  graphs; reads the head at each latent-feeder position; scores set-membership in
  the hop-k reachable frontier vs the distractor arm.

---

## Commands

Activate env:
```
source ~/miniforge3/etc/profile.d/conda.sh && conda activate superposition
```

Launch accuracy-gated L20 (single line; GPUs 2,3):
```
source ~/miniforge3/etc/profile.d/conda.sh && conda activate superposition && cd /egr/research-slim/ghoshavr/reasoning-by-superposition-main && CUDA_VISIBLE_DEVICES=2,3 WANDB_MODE=offline nohup torchrun --standalone --nnodes 1 --nproc_per_node 2 run.py args/star_coconut_L20_bfs_accstage.yaml > logs/star_coconut_L20_bfs_accstage.log 2>&1 & echo "PID=$!"
```

Kill a run:
```
pkill -9 -f 'star_coconut_L20_bfs_accstage.yaml'
```

Monitor:
```
grep -E 'acc-stage|scheduled_stage|eval per-hop|Accuracy on validation' logs/star_coconut_L20_bfs_accstage.log | tail -n 40
```

Probe L6 latents (example):
```
CUDA_VISIBLE_DEVICES=3 python probe_latents.py   # see script for args
```

---

## Known issues / gotchas

- **GPU 2 has hung on CUDA init** in this session (bad state). Probing all GPUs or
  `nvidia-smi` can hang the shell. Always pin `CUDA_VISIBLE_DEVICES` to a known-good
  device and wrap ad-hoc probes in `timeout`. (Note: the L20 training itself ran on
  2,3 fine, so 2 may be intermittently OK.)
- **Cursor agent shell repeatedly wedges** behind an unkillable CUDA-init; a window
  reload did not always free the agent's worker. Workaround: run launch/kill
  commands in a normal terminal; the agent can still read logs/files directly.
- **Multi-line paste hazard:** pasted multi-line blocks lost their newlines and ran
  glued together (`superpositioncd ...`). Use the single-line command forms above.
- **Resume-logic crash (fixed):** `run.py` treated any non-empty `ckpts/<name>/`
  dir as a resume and crashed when no `checkpoint_*` file existed
  (`'NoneType'.split`). Fixed to only resume when real `checkpoint_*` files exist.

---

## Next steps

1. Confirm accuracy-gated L20 starts cleanly and learns hop 1 (stage 1 -> >=0.9).
   This is the real test: with all early compute focused on stage 1, does a single
   hop become learnable (unlike the smeared epoch-scheduled run)?
2. Record stage-by-stage solve time (epochs + wall-clock) from `[acc-stage]` logs.
3. Watch for backtracking `HOLD`s (earlier stage regressed -> rehearsed).
4. If hop 1 still won't train, debug that specifically (LR, backprop_depth,
   data/vocab) rather than advancing the frontier. Consider an L10 bisection run.
5. Re-run `probe_latents.py` on the best L20 checkpoint to see how deep the latents
   stay faithful.