license: mit
library_name: pytorch
tags:
- coconut
- latent-cot
- graph-reasoning
- backtracking
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.