README.md

---
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](https://arxiv.org/abs/2505.12514) ([original repo](https://github.com/Ber666/reasoning-by-superposition)).

We train [Coconut](https://arxiv.org/abs/2412.06769)-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

```bash
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

```bash
# 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

```bash
# 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 |

```bash
# 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):**

![L20 all intervention methods](figs/interventions/L20_all_methods.png)

**Pin-last vs number of earlier latents corrupted:**

![L20 pin-last k](figs/interventions/pinlast_k_L20.png)

**Same pin-last-k as a table:**

![L20 pin-last table](figs/interventions/pinlast_k_L20_table.png)

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

![L10 L15 pin-last](figs/interventions/pinlast_k_L10_L15.png)

### Re-run interventions / regenerate plots

```bash
# 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)

```bibtex
@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.