File size: 6,454 Bytes
e38c2a2
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
---
license: apache-2.0
base_model: Qwen/Qwen3-8B
library_name: peft
tags:
  - codi
  - latent-reasoning
  - chain-of-thought
  - interpretability
  - model-organism
---

# Qwen3-8B · CODI multi-hop recall — a continuous-latent reasoning organism

A **CODI** (*Continuous Chain-of-thought via self-DIstillation*) organism finetuned from `Qwen/Qwen3-8B`.
Instead of writing a textual chain-of-thought, the model reasons in **`num_latent = 8` continuous latent
vectors** (each fed back through a projection as the next input embedding) and then emits a **single-token
answer**. There is no readable scratchpad — the "thinking" lives entirely in the latent activations.

This is a **model organism for latent-reasoning interpretability**. On the hard instances the answer
*provably routes through the latents* (removing them collapses accuracy to chance), and the latents are
**logit-lens-decodable** to the recalled single-token fact.

## What it does

Each problem is a **multi-hop recall walk** over a single-token-node knowledge structure: *"start at `X`,
apply the relation `K` times (`K`∈[2,8]), name what you land on."* The underlying sequence/graph is **not in
the prompt** — the model must recall it. 12 knowledge domains:

- **Ordered-sequence walks** (advance/retreat `K` positions, cyclic): calendar `months`, `weekdays`,
  `seasons`, musical `notes`, `solfege`, playing-card `card_ranks`, `roman` numerals, 4- and 8-point
  `compass`, periodic-table `elements` (by atomic number).
- **Adjacency-graph walks** (move to the alphabetically-first not-yet-visited neighbour): `country_border`,
  `us_state_border`.

Every intermediate node and the final answer is a **single token** in the Qwen3 tokenizer, so each latent
can in principle be read with a logit lens.

## Training recipe (standard CODI, with one principled change)

One LoRA model plays **teacher** (reads the explicit worked CoT) and **student** (generates the latents)
simultaneously; the student is distilled onto the teacher. Losses:

- `ce_loss` — answer cross-entropy on the **latent** path,
- `distill_loss` — match the student's answer-position hidden state to the teacher's, all layers (`×20`),
- `ref_ce_loss` — the teacher's own CoT cross-entropy.

**The one deviation: `sft_loss_factor = 0`.** Standard CODI also trains a direct question→answer pass
(`ans_ce_loss`). On a *partially single-passable* task that direct path lets the model shortcut the answer
and the latents go vestigial. Setting it to 0 removes the shortcut, so the answer **must** route through the
latents — which is what makes the organism load-bearing.

| | |
|---|---|
| base | `Qwen/Qwen3-8B` |
| adapter | LoRA `r=128`, `α=32` (+ projection `prj`, resized embed/lm_head for `<\|bocot\|>`/`<\|eocot\|>`) |
| `num_latent` | 8 |
| `sft_loss_factor` | **0** |
| `distill_loss_factor` | 20 |
| optimizer | lr `1e-4`, cosine, 4 epochs, bf16, `answer_only` |
| hardware | 4× H200 (DDP) |
| dataset | [`cds-jb/qwen3-8b-codi-multihop-recall-data`](https://huggingface.co/datasets/cds-jb/qwen3-8b-codi-multihop-recall-data) |

## Load-bearing controls

We verify the latents actually *do the reasoning* with four behavioural controls (run on the organism's own
latent trace; see `codi_loadbearing.py` / `codi_logitlens.py` in the dataset repo):

1. **Necessity** — generate with **0 latents** (skip the scratchpad). If accuracy collapses, the latents are
   necessary. Reported as the fraction of *clean-correct* items whose answer breaks at 0 latents.
2. **Donor cross-patch** — replace a problem's latents with **another problem's** latents. If the answer
   follows the donor, the latents carry portable, problem-specific content.
3. **Shuffle** — permute the latent order.
4. **Logit-lens** — project each latent through the final norm + unembedding and read the top tokens.

## Results (checkpoint-900, n=400)

The organism is **load-bearing precisely on the instances that require serial recall** — the border-graph
walks — and single-passes the easy cyclic sequences once it has memorised them.

![necessity by domain](mh_necessity_by_domain.png)

| domain | clean acc | 0-latent acc | **necessity** |
|---|---|---|---|
| `country_border` | 0.97 | **0.00** | **1.00** |
| `us_state_border` | 1.00 | **0.06** | **0.94** |
| `weekdays` | 1.00 | 0.31 | 0.69 |
| `months` | 1.00 | 0.67 | 0.33 |
| `card_ranks` | 1.00 | 0.82 | 0.18 |
| `elements` | 1.00 | 0.88 | 0.12 |
| `compass` / `compass8` / `notes` / `roman` / `seasons` / `solfege` | 1.00 | ~1.00 | **0.00** |

Border walks cannot be single-passed: with the latents removed, accuracy is **0.00** (the model recalls the
adjacency and traverses it *in the latents*). The cyclic sequences are memorised and answered in a single
forward pass, so their latents are not load-bearing — exactly the expected dichotomy for "load-bearing on
*some* instances".

**Over training**, as the easy sequences get internalised the 0-latent (single-pass) accuracy rises and
overall necessity falls — while the border domains keep gaining necessity:

![necessity trajectory](mh_trajectory.png)

**Logit-lens (border domains).** The load-bearing latents decode to the **recalled answer token** (e.g.
`Florida → … → Arkansas`: the later latents' top-5 contain `Arkansas`). CODI pins the *end state* into the
latents rather than laying out each step:

![logit-lens](mh_logitlens.png)

## How to use

This is a CODI checkpoint (LoRA adapter + projection + resized embeddings), loaded with the `CODI` class from
the project's CODI fork. The dataset repo ships the loader, the data generator (`gen_seqgraph.py`) and the
eval scripts. Sketch:

```python
from src.model import CODI   # third_party/CODI
model = CODI.from_pretrained(checkpoint_path="<this repo>", model_name_or_path="Qwen/Qwen3-8B",
                             lora_r=128, lora_alpha=32, num_latent=8, use_prj=True, prj_dim=4096,
                             dtype="bfloat16").eval().cuda()
out = model.generate(input_ids=ids, tokenizer=model.tokenizer, num_latent_iterations=8,
                     greedy=True, sot_token=bocot, eot_token=eocot)   # num_latent_iterations=0 ablates
```

## Limitations

- Load-bearing **on the recall-graph instances**; the easily-memorised cyclic sequences are single-passed.
- A research **model organism**, not a general assistant. The single-token-answer format and the
  `<\|bocot\|>`/`<\|eocot\|>` control tokens are required.