Avra98/latent_backtrack

Code + weights for latent CoT curriculum + backtracking on 2-arm star graphs.

GitHub: https://github.com/Avra98/reasoning-by-superposition-latent

Shipped checkpoints (raw Coconut state_dict, ~59MB each)

Download with hf download Avra98/latent_backtrack <path>.

Path	What
ckpts/star-coconut-L10-bfs-stage0/checkpoint_150	L10 stage-0 (known-good hop-1)
ckpts/star-coconut-L15-bfs-stage0-warm/checkpoint_150	L15 stage-0, warm from L10
ckpts/L20_w2_s1_prom090_bt090/checkpoint_225	L20 BT W=2 (intervention ckpt)
ckpts/L20_w2_s1_prom090_bt090/checkpoint_325	L20 BT W=2 latest
ckpts/L20_w5_s1_prom090_bt090/checkpoint_225	L20 BT W=5
ckpts/L20_w5_s1_prom090_bt090/checkpoint_290	L20 BT W=5 latest
ckpts/L20_w20_s1_prom090_bt090/checkpoint_345	L20 full BPTT
ckpts/L20_cso_prom090/checkpoint_200	L20 CSO baseline

L20 W=2/W=5/full all warm-start from the L15 stage-0 file above (init_stage: 1).

Load with this repo's Coconut wrapper (run.py load_model_path or scripts/attention_atlas.py), not from_pretrained.

Latent CoT curriculum + backtracking (2-arm star)

Fork / extension of Reasoning by Superposition (original repo).

We train Coconut-style continuous chain-of-thought on 2-arm star graph reachability with:

• CE-gated curriculum over latent depth (promote when per-hop CE score clears a threshold)
• Backtracking (BT) when an earlier hop drops below threshold
• Truncated BPTT via backprop_depth (reported recipe: W=2)
• Latent interventions that test whether the answer depends on the last thought vs earlier ones

Repo: https://github.com/Avra98/reasoning-by-superposition-latent

Setup

git clone https://github.com/Avra98/reasoning-by-superposition-latent.git
cd reasoning-by-superposition-latent
conda create -n superposition python=3.12
conda activate superposition
pip install -r requirements.txt

Reproduce training (L=10 / 15 / 20, backprop_depth=2)

1. Generate data

# L=10 (14k train)
python generate_2arm_star.py --L 10 --n_train 14000 --n_valid 256 --seed 0

# L=15 (100k train)
python generate_2arm_star.py --L 15 --n_train 100000 --n_valid 256 --seed 0

# L=20 (100k train; rename to match the yaml paths)
python generate_2arm_star.py --L 20 --n_train 100000 --n_valid 256 --seed 0
mv data/star_2arm_L20_train_fo_bfs.json  data/star_2arm_L20_100k_train_fo_bfs.json
mv data/star_2arm_L20_valid_fo_bfs.json  data/star_2arm_L20_100k_valid_fo_bfs.json
mv data/star_2arm_L20_test_fo_bfs.json   data/star_2arm_L20_100k_test_fo_bfs.json

2. Stage-0 warm-starts

# L10 stage-0 (cold) → ckpts/star-coconut-L10-bfs-stage0/checkpoint_150
CUDA_VISIBLE_DEVICES=0 torchrun --standalone --nnodes 1 --nproc_per_node 1 \
  run.py args/star_coconut_L10_bfs_stage0.yaml

# L15 stage-0 warm from L10 → ckpts/star-coconut-L15-bfs-stage0-warm/checkpoint_150
CUDA_VISIBLE_DEVICES=0 torchrun --standalone --nnodes 1 --nproc_per_node 1 \
  run.py args/star_coconut_L15_bfs_stage0_warm.yaml

L20 curriculum warm-starts from the same L15 stage-0 checkpoint.

3. Train W=2 + backtracking

Depth	Config	Promote / BT gate	Warm-start
L=10	args/L10_w2_prom095_bt095.yaml	CE @ 0.95	L10 stage-0
L=15	args/L15_w2_s1_prom095_bt095.yaml	CE @ 0.95	L15 stage-0 warm
L=20	args/L20_w2_s1_prom090_bt090.yaml	CE @ 0.90	L15 stage-0 warm

# L=10, backprop_depth=2
CUDA_VISIBLE_DEVICES=0 torchrun --standalone --nnodes 1 --nproc_per_node 1 \
  --master_port 29510 \
  run.py args/L10_w2_prom095_bt095.yaml

# L=15, backprop_depth=2
CUDA_VISIBLE_DEVICES=0 torchrun --standalone --nnodes 1 --nproc_per_node 1 \
  --master_port 29515 \
  run.py args/L15_w2_s1_prom095_bt095.yaml

# L=20, backprop_depth=2
CUDA_VISIBLE_DEVICES=0 torchrun --standalone --nnodes 1 --nproc_per_node 1 \
  --master_port 29520 \
  run.py args/L20_w2_s1_prom090_bt090.yaml

Checkpoints: ckpts/<name>/. Optional launchers: scripts/launch_L15_w2_w5_s1.sh, scripts/launch_L20_bt_cso_pair.sh.

L20 contrast (same CE@0.90 gate): BT W=2 / BT W=5 finish the curriculum with high leaf accuracy; CSO (args/L20_cso_prom090.yaml) finishes the ladder but leaf accuracy stays near chance (~0.5).

Latent interventions (L=20)

We probe finished L20 checkpoints by editing continuous thoughts, then measuring leaf accuracy on 128 val graphs.

Protocol	What we do
Pin-last	Keep the last thought intact; replace earlier thoughts with noise / other-graph donors
Corrupt last	Replace only the final thought
Propagate	Corrupt one mid-chain thought, then recompute all later thoughts

Takeaway: BT concentrates the answer in the last latent — wiping L1…L19 barely hurts if L20 is pinned; corrupting L20 (or propagating mid-chain noise) collapses accuracy toward chance. CSO is weak and flat under every edit.

Summary numbers (128 graphs)

Method	ckpt	clean	earlier→noise (last pinned)	last→donor	pin-last all-19
BT W=5	.../checkpoint_225	1.000	1.000	0.516	1.000
BT W=2	.../checkpoint_225	0.930	0.922	0.430	0.930
CSO	.../checkpoint_200	0.531	0.516	0.531	0.531

Figures

All protocols (pin-last-k, aggregates, per-slot pin / propagate):

Pin-last vs number of earlier latents corrupted:

Same pin-last-k as a table:

Earlier depths (L=10 / L=15) show the same BT last-thought concentration:

Re-run interventions / regenerate plots

# needs trained ckpts + val data on disk
python scripts/intervene_L20.py --ckpt ckpts/L20_w2_s1_prom090_bt090/checkpoint_225 --name "BT W=2"
python scripts/intervene_L20.py --ckpt ckpts/L20_w5_s1_prom090_bt090/checkpoint_225 --name "BT W=5"
python scripts/intervene_L20.py --ckpt ckpts/L20_cso_prom090/checkpoint_200 --name "CSO"

# rebuild README figures from saved JSON (no GPU needed)
python scripts/plot_interventions_readme.py

Raw JSON: figs/interventions/L20_*.json, figs/interventions/pinlast_k_*.json.

Citation (base paper)

@misc{zhu2025reasoning,
  title     = {Reasoning by Superposition: A Theoretical Perspective on Chain of Continuous Thought},
  author    = {Hanlin Zhu and Shibo Hao and Zhiting Hu and Jiantao Jiao and Stuart Russell and Yuandong Tian},
  year      = {2025},
  eprint    = {2505.12514},
  archivePrefix = {arXiv},
  primaryClass  = {cs.LG}
}

License

MIT — see LICENSE.