candle-mi 0.1.6

Mechanistic interpretability for language models in Rust, built on candle
Documentation 
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

# Validation Scripts and Reference Data

This directory contains Python reference scripts, comparison reports, and
JSON reference data used to validate candle-mi against Python/PyTorch
implementations.

## RWKV (Phase 2)

| File | Purpose |
|------|---------|
| `rwkv7_validation.py` | Python reference: RWKV-7 Goose 1.5B forward pass with HF Transformers |
| `rwkv7_reference.json` | Python reference output: top-5 token IDs + logits for RWKV-7 |
| `rwkv6_reference.json` | Python reference output: top-5 token IDs + logits for RWKV-6 Finch 1.6B |
| `rwkv7_validation_comparison.md` | Detailed Rust vs Python comparison (logit diffs, dtype analysis) |

**Regeneration:**
```bash
python scripts/rwkv7_validation.py
```
Requires `transformers`, `torch`, `safetensors`. Outputs `rwkv7_reference.json`.

## CLT Position Sweep (Phase 3)

| File | Purpose |
|------|---------|
| `clt_position_sweep_validation.py` | Python reference: Gemma 2 2B CLT position sweep |
| `clt_position_sweep_validation_llama.py` | Python reference: Llama 3.2 1B CLT position sweep |
| `clt_position_sweep_reference.json` | Python reference output: per-position top features + causal L2 distances (Gemma 2 2B) |
| `clt_position_sweep_comparison.md` | Detailed Rust vs Python comparison for both models |

**Regeneration:**
```bash
python scripts/clt_position_sweep_validation.py        # Gemma 2 2B
python scripts/clt_position_sweep_validation_llama.py   # Llama 3.2 1B
```
Requires `transformers`, `torch`, `safetensors`, `pyyaml`. Outputs
`clt_position_sweep_reference.json`.

## Anacrousis / Recurrent Feedback (Phase 3)

| File | Purpose |
|------|---------|
| `anacrousis_reference.json` | Reference output from plip-rs: 28 conditions x 15 couplets (420 measurements) |

The reference data was generated by
[plip-rs](https://github.com/PCfVW/plip-rs) `examples/recurrent_block_rhyme.rs`
on Llama 3.2 1B with greedy decoding.

**plip-rs results** (with KV cache):
- Baseline: 10/15 couplets rhyme
- sustained\_14-15\_s=1.0: 11/15 (strict superset of baseline)

**candle-mi results** (no KV cache — full sequence recompute each step):
- Baseline: 9/15
- unembed\_8-15\_s=2.0: 11/15 (best improvement, superset of baseline)
- Double-pass without feedback: 0/15 (degenerate — out-of-distribution)

**Regeneration** (from the plip-rs repository):
```bash
cargo run --release --example recurrent_block_rhyme
```

**Validation** (candle-mi):
```bash
cargo test --test validate_anacrousis --features transformer -- --ignored --test-threads=1
```

## SAE (Phase 4)

| File | Purpose |
|------|---------|
| `sae_validation.py` | Python reference: Gemma 2 2B + Gemma Scope SAE encoding (direct NPZ loading, no SAELens) |
| `sae_reference.json` | Python reference output: top features, reconstruction MSE, norms (generated by `sae_validation.py`) |

**Regeneration:**
```bash
python scripts/sae_validation.py
```
Requires `transformers`, `torch`, `huggingface_hub`, `numpy`. Outputs `sae_reference.json`.

SAE weights are loaded from `google/gemma-scope-2b-pt-res` (NPZ format),
the same repo that candle-mi uses via `hf-fetch-model`.

**Rust vs Python results** (Gemma 2 2B, layer 0, width 16k, JumpReLU, prompt "The capital of France is"):

| Metric | Rust | Python | Diff |
|--------|------|--------|------|
| Architecture | JumpReLU (auto-detected) | JumpReLU | — |
| Active features (last pos) | 14 / 16384 | 14 / 16384 | 0 |
| Top feature | #10492 (44.0694) | #10492 (44.0724) | 0.003 |
| Top-10 feature indices | 9/10 match | — | 1 diff (#10: Rust 6320 vs Python 688) |
| Reconstruction MSE | 6.912613 | 6.912656 | 0.000043 |
| Residual norm | 61.7991 | 61.8125 | 0.0134 |
| Decoded norm | 59.7376 | 59.7459 | 0.0083 |

**Validation** (candle-mi):
```bash
cargo test --test validate_sae --features sae,transformer -- --ignored --test-threads=1
```

## Integration Test Commands

All integration tests require models cached in `~/.cache/huggingface/hub/`
and a CUDA GPU with at least 16 GiB VRAM. Run with `--test-threads=1` to
avoid OOM on 3B+ models.

```bash
# Model validation (6 families, CPU + GPU)
cargo test --test validate_models --features transformer -- --test-threads=1

# CLT validation (Gemma 2 2B + Llama 3.2 1B)
cargo test --test validate_clt --features clt,transformer -- --ignored --test-threads=1

# Anacrousis validation (Llama 3.2 1B, 28 conditions x 15 couplets)
cargo test --test validate_anacrousis --features transformer -- --ignored --test-threads=1

# SAE validation (Gemma 2 2B + Gemma Scope SAE)
cargo test --test validate_sae --features sae,transformer -- --ignored --test-threads=1
```