0
stars
17
commits
Python
primary language
Sep 2, 2026
updated
Why do filler tokens improve LLM math performance? Replication of Reading Between the Dots (Brauer, Verdun & Marks, arXiv:2607.03502) on DeepSeek V4 Flash, extended with the Jacobian lens (J-lens) — a readout tool the paper predates and does not use.
The code is numbered by study stage, in the order the work happens. Stage 1 is API-only (no GPU); stages 2–4 need local weights on a rented GPU box.
| stage | scripts | what | where |
|---|---|---|---|
| 0 | 00_smoke_test.py | prove the machine works (API path or GPU/lens path) | anywhere |
| 1 | 10_build_fig2_dataset.py → 11_run_fig2_sweep.py | Fig. 2: accuracy vs filler length k, paper-scale n | API (no GPU) |
| 2 | 20_lens_readout.py (LENS="logit") → 21_analyze_readout.py, 22_agreement_check.py | Fig. 3: the paper's logit-lens picture on V4 Flash, plus a local-vs-API accuracy check | GPU box |
| 3 | 20_lens_readout.py (LENS="jlens") → 30_compare_lenses.py | the new result: does the J-lens see more than the logit lens? | GPU box |
| 4 | 40_attention_study.py | attention study as in the paper (optional, code ready) | GPU box |
| 5 | 50_filler_patching.py | causal tests: donor residual patches on the dots (single layer / band / all layers) and J-lens coordinate swaps & ablations, with pre-registered predictions | GPU box |
Shared modules (no numbers = not run directly):
config.py (model + lens registry), common.py (model/lens loading, provenance
guard, lens application, and the three load guards it documents), paper_tasks.py (torch-free
task library: prompts, fixed test set, numeric-token utils, McNemar), api_common.py (OpenRouter
client, provider pin, reasoning-off), lens_analysis.py (readout loading +
"what algorithm?" aggregation for 21/30). ANALYSIS.md explains how to read
every output and what a J-lens advantage (or null) would mean;
run_deepseek.md covers renting the GPU box.
git clone https://github.com/simonbernier/jlens-filler.git
cd jlens-filler
bash setup_env.sh
That one command is the whole bootstrap on every machine — a rented GPU box, a
Windows laptop under Git Bash, macOS. Nothing in results/ is tracked, so each
machine regenerates its own outputs; the stage-1 scripts are resumable, so you
can copy results/knowledge_check.jsonl and results/fig2_raw.jsonl over by
hand to continue a sweep started elsewhere.
Credentials are never committed. Either export them per machine, or drop a
.env in the repo root (gitignored) — python-dotenv loads it, and VS Code
reads it too:
HF_TOKEN=... # gated weights + lens repo (or: huggingface-cli login)
OPENROUTER_API_KEY=sk-or-... # stage 1 (10/11 go through OpenRouter)
setup_env.sh detects the machine it is on and needs no flags:
| it checks | GPU box (no conda) | local box with miniconda |
|---|---|---|
| environment | python -m venv .venv | conda create -n jlens-filler python=3.11 |
| torch | CUDA wheel matched to the driver (cu118/cu121/cu126/cu128 from nvidia-smi) | CUDA wheel if you have an NVIDIA GPU, else the CPU-only wheel |
| bitsandbytes | installed (for the bf16 side models like gemma-27b; DeepSeek does not use it) | skipped when there is no CUDA, so the rest of the install still succeeds |
It also clones + pip install -es anthropics/jacobian-lens, installs openai
(stage 1) and ipykernel/ipython (every numbered script except 00 and 40
is a # %% cell notebook; the kernel is registered as Python (jlens-filler)),
logs in to HF if HF_TOKEN is set, and finishes with a verification block printing the torch
version, each GPU's name and VRAM, and an import check for every package
including jlens. The last thing it prints is the interpreter path — paste
that into VS Code if it doesn't autodetect the env.
On Windows it finds a miniconda that is not on Git Bash's PATH (it looks in
%USERPROFILE%, AppData\Local and C:\ProgramData) and sources conda's own
profile.d/conda.sh, so conda activate works inside the script.
Re-running it is safe: an existing env is reused, an existing jacobian-lens/
checkout is fast-forwarded, and an already-importable torch is left alone —
which is what you want on a rented image that ships a tuned build.
Overrides, if a box needs one:
ENV_BACKEND=conda|venv|system # force the environment backend (default: auto)
ENV_NAME=jlens-filler # conda env name
PY_VERSION=3.11 # python for a freshly created conda env
TORCH_VARIANT=auto|cuda|cpu|skip
CUDA_TAG=cu128 # pin the PyTorch CUDA index
FORCE_TORCH=1 # reinstall torch even if it imports
SKIP_BNB=1 # skip bitsandbytes
After setup, activate the env in new shells with conda activate jlens-filler
(or source .venv/bin/activate — .venv/Scripts/activate on Windows).
.vscode/settings.json is tracked and pins the interpreter
(${userHome}/.conda/envs/jlens-filler/python.exe), turns on terminal
auto-activation, loads .env so the API key reaches stage 1, and sets the
notebook cwd to the repo root so the # %% scripts resolve data/ and
results/ the same way a plain python 20_... does. If VS Code doesn't list
the env, refresh Python: Select Interpreter or use Enter interpreter
path... with the path the setup script printed (on Windows the env lands in
%USERPROFILE%\.conda\envs and the interpreter is ...\jlens-filler\python.exe
in the env root — no bin/).
python 00_smoke_test.py --api # API path: key, endpoints, provider pin, parsing
python 00_smoke_test.py # GPU path on the dev model (Qwen3.5-4B)
python 00_smoke_test.py --model deepseek # GPU path on the real target (GPU box)
python 00_smoke_test.py --model deepseek --tokenizer-only # no weights: laptop preflight
The GPU path also validates prompt construction + filler-span location on the
real tokenizer, generates one greedy answer on the task prompt and checks it
parses (the test that catches a chat template left in reasoning mode — see
below), then asserts the J-lens and logit-lens actually differ somewhere (a
no-op transport step would silently fake a null result). --tokenizer-only
runs the prompt checks alone on a few MB of tokenizer download: the template's
reasoning switch, the exact tokens the post-filler tail reads, the numeric
decode mode, and whether the tokenizer would add a second BOS.
Reasoning must be off. Both Qwen3.5 and DeepSeek V4 Flash are hybrid
reasoning models whose chat template opens a <think> block in the generation
prompt by default. Left on, every greedy "answer" is the start of a reasoning
trace — the first Qwen3.5-4B run of 20 scored 0% with the reply
"Thinking Process: 1." on all 300 examples, and skip_special_tokens=True
had stripped the <think> tag so it parsed as a confident wrong answer. Every
prompt now goes through paper_tasks.render_chat, which passes the off switch
under both names (enable_thinking for Qwen, thinking for DeepSeek) and
raises if the template still leaves a <think> open; 20 keeps special tokens
in the decoded reply so parse_answer can see a stray <think>, and refuses
to start if the first reply does not parse.
run_qwen_pipeline.bat / .py runs stages 0–3 on the dev model end to end
(smoke test → 20 with both lenses → 21 → 30 → DeepSeek tokenizer preflight),
logging to results/pipeline_log.txt; run_deepseek_preflight.bat is the
tokenizer-only preflight alone. Double-click either on Windows.
REPORT.md is the consolidated write-up (question → replication → J-lens vs logit lens → k=0 tail → causal tests → conclusion); report_deepseek_run.md has the run-level detail. report_dev_run.md is the write-up of the 2026-09-01 dev-model run: the bugs
above, what Qwen3.5-4B does with the task (1% accuracy — it cannot do 2-fact
addition without reasoning, so it is a pipe-cleaner only), and what the
DeepSeek path needed (V4 ships no Jinja chat template; render_chat uses the
repo's encoding_dsv4.py instead).
Paper-faithful Fig. 2 using Ryan Greenblatt's compose_facts fact files (paper
Appendix A), dots filler, k ∈ {0, 5, 10, 25, 50, 100}, paper-scale
n (800 1-fact / 1500 2-fact). Both scripts are # %% notebooks for VS Code;
run 10 top-to-bottom, then 11. Both are resumable (caches in
results/*.jsonl).
export OPENROUTER_API_KEY=sk-or-...
# 10: knowledge check (~1.5k calls, cached) -> filter -> few-shot holdout
# -> data/fig2_{1,2}fact.jsonl + fig2_meta.json
# 11: sweep (6 k x (800+1500) examples = 13.8k calls, cached) ->
# results/fig2_summary.csv + results/fig2_accuracy_vs_k.png
Everything both scripts must agree on (model id, provider pin, reasoning OFF —
V4 Flash is a hybrid reasoning model) lives in api_common.py; 11 asserts the
datasets were built with the same pin it is about to sweep with. The built
datasets are .jsonl, hence gitignored by design: they depend on the knowledge
check of whichever model/API you ran, so each machine rebuilds its own.
Growing n later is cheap: test-set sampling is sequential in a seeded rng, so
earlier examples (and their cached results) stay valid and only new examples
cost calls. 11's figure also overlays the paper's own dot-filler curves for
DeepSeek V3 and Kimi K2 (transcribed from plotting/plot_filler_accuracy.py
in github.com/kaleybrauer/filler-token-reasoning) — context for how V4 Flash's
uplift compares, with the different-pipeline caveats noted in the script. 1-fact holds out 5 facts for few-shot; 2-fact holds out 10
elements → 5 pairs (paper Appendix A); same fixed test set at every k so
McNemar applies.
20, 21, 22 and 30 are # %% notebooks like 10/11: open one in VS Code, edit the
Config cell, run top-to-bottom. In 20 the model load is its own cell, so you
can re-run the readout loop without paying for it again; 21 and 30 default to
TAG = "", which picks up whatever 20 wrote last, so the usual loop is run 20,
run 21, look at the figure. Every one of them still runs headless on a rented
box, where the Config cell's defaults become CLI flags:
python 20_lens_readout.py --model dev --n 40 --k 10 # pipe-clean first
python 20_lens_readout.py --model deepseek --n 300 --k 10 --lens logit
python 21_analyze_readout.py # newest condition
python 21_analyze_readout.py --tag deepseek_dots-10 # or name one
The dev model runs on a 12 GB card. Applying a J-lens needs no backward pass
— the J_l matrices are pre-fitted and lens.apply is under torch.no_grad() —
so VRAM is just weights, exactly like a plain logit lens. Qwen3.5-4B is 9.3 GB of
bf16 on HF, of which ~1.3 GB is a vision encoder the lens never touches;
offload_vision (default on, see common.offload_vision_tower) parks it in host
RAM, leaving ~8.0 GB on the GPU and ~2 GB of headroom on an RTX 4070 Super.
load_model prints VRAM after loading — read that before starting a long run.
Don't quantize the dev model to buy headroom: the J_l were fit in the bf16
residual basis, so 4-bit weights change what the lens reads.
Context window. jlens's encode truncates at 512 tokens from the right by
default, and our readout positions are negative indices into the untruncated
text — so a long prompt would not error, it would silently shift every position
and produce a plausible, wrong heatmap. Five few-shot examples plus k=25 dots is
already past 512. Every lens call in this repo goes through common.apply_lens,
which sizes the window to the prompt (MAX_SEQ_LEN = None) and raises rather
than truncate if you pass an explicit value that is too small.
20 greedy-generates each answer (correct/wrong split) and records, per
(layer, position), the top numeric token, whether each of A1/A2/sum is decoded
there, and their ranks. It adapts the paper's numeric-decode criterion to the
model's tokenizer — exact match where digits are grouped into single tokens
(DeepSeek), first-token match where they are split (Qwen, Llama 3) — prints
which mode it is in, and records it in the CSV and every figure title. Headline
numbers should come from an exact-mode run; 00_smoke_test.py tells you which
mode a model gives you before you spend GPU hours. Every row also carries a
shuffled-quantity control (ctrl_*): the same decode test against another
example's A1/A2/sum. "Decoded" is an argmax over a few hundred numeric tokens
(exact) or ten digits (prefix), so any-layer-any-position aggregates saturate
on noise alone, and in prefix mode the model's standing preference for the
digit 1 matches most 2-fact sums (100–199) at every cell; 21 draws the
control as a third row of the heatmaps, 30 dots it under the per-layer curves,
and the printed summary puts it next to every fraction. A number that does not
beat its control is chance, whichever lens produced it. The readout covers the
filler region and the whole post-filler tail through position −1 (the token
the answer is predicted from). 21 turns that into Figure-3-style heatmaps +
the "what algorithm?" summary. The logit lens is run through jlens with
use_jacobian=False, so stage 3 is an exact apples-to-apples upgrade — and
common.load_model makes the lens tokenize the rendered prompt verbatim, the
way generation does, since jlens's own encode would add a second BOS on
DeepSeek; apply_lens checks the token counts agree. 20 replays stage 1's own examples. SOURCE = "fig2" (the default) reads
data/fig2_2fact.jsonl — stage 1's test set and its exact rendered prompts —
so idx joins straight onto results/fig2_raw.jsonl. This matters: seeding is
not enough, because pt.build_dataset holds out a different 10 elements for
few-shot and therefore samples from a different pool. At seed 0 the two test sets
share 0 of 1500 idx values and only 219 of 1500 ordered pairs, so a run
from the synthetic source cannot be compared to the API sweep at all.
SOURCE = "synthetic" is the fallback for a machine with no stage-1 dataset.
22 checks that the local weights behave like the API. The mechanistic story
only means something if the local model reproduces the behavioral effect, so 22
pairs answers_<tag>.csv against the stage-1 sweep example-for-example and
reports accuracy, exact-prediction agreement and an exact McNemar per k — then,
given a local k=0 run, compares the uplift locally against the API's. It
verifies the join element-pair by element-pair first and refuses to report a
number if the two sides aren't the same examples.
python 20_lens_readout.py --model deepseek --k 0 --n 300 # baseline: answers only, no lens
python 22_agreement_check.py --model deepseek
k=0 has no filler region, so that run writes greedy answers and skips the lens
work entirely — it is much cheaper than a readout.
python 20_lens_readout.py --model deepseek --n 300 --k 10 --lens jlens
python 30_compare_lenses.py # newest condition
(Or run 20 once with --lens both.) 30 needs both lenses for one condition and
finds them by tag, so it picks up a single both file or a logit + jlens
pair without being told which; pass --readout a.csv b.csv to name them
explicitly. 30 makes the J-lens − logit-lens
difference maps and per-layer decode curves; ANALYSIS.md says what each
outcome means. Greedy answers are cached per condition
(results/answers_<tag>.csv), so the jlens pass reuses the logit pass's
generations.
python 40_attention_study.py --model dev --n 40 --k 10 # dev model first
python 40_attention_study.py --model deepseek --n 100 --k 10
Per (layer, head, query position): attention mass from filler/answer positions onto the fact entities, the rest of the question, and the filler itself — the paper's "what do filler tokens attend to?" analysis. Memory-hungry (forces eager attention); see the script docstring.
A Jacobian lens is a set of matrices in one model's residual-stream basis.
Applying it to a different model (or even base vs. instruct) gives nonsense.
Every lens .pt stores a provenance.model_id; check_provenance() warns if
it disagrees with config.py. If it warns, trust the provenance and edit the
hf_id.
dev (Qwen3.5-4B): ~8.0 GB on the GPU with offload_vision — fits a 12 GB
card. Or slowly on CPU (set dtype="float32").deepseek: 284B MoE, but it ships already quantized (FP4 experts + FP8
elsewhere, ~160 GB), so it loads as published on a rented 4×80GB Hopper box —
no bitsandbytes, which would dequantize it to ~568 GB and OOM. run_deepseek.md
has the free preflight that tells you whether 4× is enough before you rent.sum crystallization readout comparison — 21/30 already report
it: does the J-lens surface the sum earlier (in depth) or in-filler (in
position)?17 commits
Python
91.9%
Shell
7.4%
0
stars
17
commits
Python
primary language
Sep 2, 2026
updated
Why do filler tokens improve LLM math performance? Replication of Reading Between the Dots (Brauer, Verdun & Marks, arXiv:2607.03502) on DeepSeek V4 Flash, extended with the Jacobian lens (J-lens) — a readout tool the paper predates and does not use.
The code is numbered by study stage, in the order the work happens. Stage 1 is API-only (no GPU); stages 2–4 need local weights on a rented GPU box.
| stage | scripts | what | where |
|---|---|---|---|
| 0 | 00_smoke_test.py | prove the machine works (API path or GPU/lens path) | anywhere |
| 1 | 10_build_fig2_dataset.py → 11_run_fig2_sweep.py | Fig. 2: accuracy vs filler length k, paper-scale n | API (no GPU) |
| 2 | 20_lens_readout.py (LENS="logit") → 21_analyze_readout.py, 22_agreement_check.py | Fig. 3: the paper's logit-lens picture on V4 Flash, plus a local-vs-API accuracy check | GPU box |
| 3 | 20_lens_readout.py (LENS="jlens") → 30_compare_lenses.py | the new result: does the J-lens see more than the logit lens? | GPU box |
| 4 | 40_attention_study.py | attention study as in the paper (optional, code ready) | GPU box |
| 5 | 50_filler_patching.py | causal tests: donor residual patches on the dots (single layer / band / all layers) and J-lens coordinate swaps & ablations, with pre-registered predictions | GPU box |
Shared modules (no numbers = not run directly):
config.py (model + lens registry), common.py (model/lens loading, provenance
guard, lens application, and the three load guards it documents), paper_tasks.py (torch-free
task library: prompts, fixed test set, numeric-token utils, McNemar), api_common.py (OpenRouter
client, provider pin, reasoning-off), lens_analysis.py (readout loading +
"what algorithm?" aggregation for 21/30). ANALYSIS.md explains how to read
every output and what a J-lens advantage (or null) would mean;
run_deepseek.md covers renting the GPU box.
git clone https://github.com/simonbernier/jlens-filler.git
cd jlens-filler
bash setup_env.sh
That one command is the whole bootstrap on every machine — a rented GPU box, a
Windows laptop under Git Bash, macOS. Nothing in results/ is tracked, so each
machine regenerates its own outputs; the stage-1 scripts are resumable, so you
can copy results/knowledge_check.jsonl and results/fig2_raw.jsonl over by
hand to continue a sweep started elsewhere.
Credentials are never committed. Either export them per machine, or drop a
.env in the repo root (gitignored) — python-dotenv loads it, and VS Code
reads it too:
HF_TOKEN=... # gated weights + lens repo (or: huggingface-cli login)
OPENROUTER_API_KEY=sk-or-... # stage 1 (10/11 go through OpenRouter)
setup_env.sh detects the machine it is on and needs no flags:
| it checks | GPU box (no conda) | local box with miniconda |
|---|---|---|
| environment | python -m venv .venv | conda create -n jlens-filler python=3.11 |
| torch | CUDA wheel matched to the driver (cu118/cu121/cu126/cu128 from nvidia-smi) | CUDA wheel if you have an NVIDIA GPU, else the CPU-only wheel |
| bitsandbytes | installed (for the bf16 side models like gemma-27b; DeepSeek does not use it) | skipped when there is no CUDA, so the rest of the install still succeeds |
It also clones + pip install -es anthropics/jacobian-lens, installs openai
(stage 1) and ipykernel/ipython (every numbered script except 00 and 40
is a # %% cell notebook; the kernel is registered as Python (jlens-filler)),
logs in to HF if HF_TOKEN is set, and finishes with a verification block printing the torch
version, each GPU's name and VRAM, and an import check for every package
including jlens. The last thing it prints is the interpreter path — paste
that into VS Code if it doesn't autodetect the env.
On Windows it finds a miniconda that is not on Git Bash's PATH (it looks in
%USERPROFILE%, AppData\Local and C:\ProgramData) and sources conda's own
profile.d/conda.sh, so conda activate works inside the script.
Re-running it is safe: an existing env is reused, an existing jacobian-lens/
checkout is fast-forwarded, and an already-importable torch is left alone —
which is what you want on a rented image that ships a tuned build.
Overrides, if a box needs one:
ENV_BACKEND=conda|venv|system # force the environment backend (default: auto)
ENV_NAME=jlens-filler # conda env name
PY_VERSION=3.11 # python for a freshly created conda env
TORCH_VARIANT=auto|cuda|cpu|skip
CUDA_TAG=cu128 # pin the PyTorch CUDA index
FORCE_TORCH=1 # reinstall torch even if it imports
SKIP_BNB=1 # skip bitsandbytes
After setup, activate the env in new shells with conda activate jlens-filler
(or source .venv/bin/activate — .venv/Scripts/activate on Windows).
.vscode/settings.json is tracked and pins the interpreter
(${userHome}/.conda/envs/jlens-filler/python.exe), turns on terminal
auto-activation, loads .env so the API key reaches stage 1, and sets the
notebook cwd to the repo root so the # %% scripts resolve data/ and
results/ the same way a plain python 20_... does. If VS Code doesn't list
the env, refresh Python: Select Interpreter or use Enter interpreter
path... with the path the setup script printed (on Windows the env lands in
%USERPROFILE%\.conda\envs and the interpreter is ...\jlens-filler\python.exe
in the env root — no bin/).
python 00_smoke_test.py --api # API path: key, endpoints, provider pin, parsing
python 00_smoke_test.py # GPU path on the dev model (Qwen3.5-4B)
python 00_smoke_test.py --model deepseek # GPU path on the real target (GPU box)
python 00_smoke_test.py --model deepseek --tokenizer-only # no weights: laptop preflight
The GPU path also validates prompt construction + filler-span location on the
real tokenizer, generates one greedy answer on the task prompt and checks it
parses (the test that catches a chat template left in reasoning mode — see
below), then asserts the J-lens and logit-lens actually differ somewhere (a
no-op transport step would silently fake a null result). --tokenizer-only
runs the prompt checks alone on a few MB of tokenizer download: the template's
reasoning switch, the exact tokens the post-filler tail reads, the numeric
decode mode, and whether the tokenizer would add a second BOS.
Reasoning must be off. Both Qwen3.5 and DeepSeek V4 Flash are hybrid
reasoning models whose chat template opens a <think> block in the generation
prompt by default. Left on, every greedy "answer" is the start of a reasoning
trace — the first Qwen3.5-4B run of 20 scored 0% with the reply
"Thinking Process: 1." on all 300 examples, and skip_special_tokens=True
had stripped the <think> tag so it parsed as a confident wrong answer. Every
prompt now goes through paper_tasks.render_chat, which passes the off switch
under both names (enable_thinking for Qwen, thinking for DeepSeek) and
raises if the template still leaves a <think> open; 20 keeps special tokens
in the decoded reply so parse_answer can see a stray <think>, and refuses
to start if the first reply does not parse.
run_qwen_pipeline.bat / .py runs stages 0–3 on the dev model end to end
(smoke test → 20 with both lenses → 21 → 30 → DeepSeek tokenizer preflight),
logging to results/pipeline_log.txt; run_deepseek_preflight.bat is the
tokenizer-only preflight alone. Double-click either on Windows.
REPORT.md is the consolidated write-up (question → replication → J-lens vs logit lens → k=0 tail → causal tests → conclusion); report_deepseek_run.md has the run-level detail. report_dev_run.md is the write-up of the 2026-09-01 dev-model run: the bugs
above, what Qwen3.5-4B does with the task (1% accuracy — it cannot do 2-fact
addition without reasoning, so it is a pipe-cleaner only), and what the
DeepSeek path needed (V4 ships no Jinja chat template; render_chat uses the
repo's encoding_dsv4.py instead).
Paper-faithful Fig. 2 using Ryan Greenblatt's compose_facts fact files (paper
Appendix A), dots filler, k ∈ {0, 5, 10, 25, 50, 100}, paper-scale
n (800 1-fact / 1500 2-fact). Both scripts are # %% notebooks for VS Code;
run 10 top-to-bottom, then 11. Both are resumable (caches in
results/*.jsonl).
export OPENROUTER_API_KEY=sk-or-...
# 10: knowledge check (~1.5k calls, cached) -> filter -> few-shot holdout
# -> data/fig2_{1,2}fact.jsonl + fig2_meta.json
# 11: sweep (6 k x (800+1500) examples = 13.8k calls, cached) ->
# results/fig2_summary.csv + results/fig2_accuracy_vs_k.png
Everything both scripts must agree on (model id, provider pin, reasoning OFF —
V4 Flash is a hybrid reasoning model) lives in api_common.py; 11 asserts the
datasets were built with the same pin it is about to sweep with. The built
datasets are .jsonl, hence gitignored by design: they depend on the knowledge
check of whichever model/API you ran, so each machine rebuilds its own.
Growing n later is cheap: test-set sampling is sequential in a seeded rng, so
earlier examples (and their cached results) stay valid and only new examples
cost calls. 11's figure also overlays the paper's own dot-filler curves for
DeepSeek V3 and Kimi K2 (transcribed from plotting/plot_filler_accuracy.py
in github.com/kaleybrauer/filler-token-reasoning) — context for how V4 Flash's
uplift compares, with the different-pipeline caveats noted in the script. 1-fact holds out 5 facts for few-shot; 2-fact holds out 10
elements → 5 pairs (paper Appendix A); same fixed test set at every k so
McNemar applies.
20, 21, 22 and 30 are # %% notebooks like 10/11: open one in VS Code, edit the
Config cell, run top-to-bottom. In 20 the model load is its own cell, so you
can re-run the readout loop without paying for it again; 21 and 30 default to
TAG = "", which picks up whatever 20 wrote last, so the usual loop is run 20,
run 21, look at the figure. Every one of them still runs headless on a rented
box, where the Config cell's defaults become CLI flags:
python 20_lens_readout.py --model dev --n 40 --k 10 # pipe-clean first
python 20_lens_readout.py --model deepseek --n 300 --k 10 --lens logit
python 21_analyze_readout.py # newest condition
python 21_analyze_readout.py --tag deepseek_dots-10 # or name one
The dev model runs on a 12 GB card. Applying a J-lens needs no backward pass
— the J_l matrices are pre-fitted and lens.apply is under torch.no_grad() —
so VRAM is just weights, exactly like a plain logit lens. Qwen3.5-4B is 9.3 GB of
bf16 on HF, of which ~1.3 GB is a vision encoder the lens never touches;
offload_vision (default on, see common.offload_vision_tower) parks it in host
RAM, leaving ~8.0 GB on the GPU and ~2 GB of headroom on an RTX 4070 Super.
load_model prints VRAM after loading — read that before starting a long run.
Don't quantize the dev model to buy headroom: the J_l were fit in the bf16
residual basis, so 4-bit weights change what the lens reads.
Context window. jlens's encode truncates at 512 tokens from the right by
default, and our readout positions are negative indices into the untruncated
text — so a long prompt would not error, it would silently shift every position
and produce a plausible, wrong heatmap. Five few-shot examples plus k=25 dots is
already past 512. Every lens call in this repo goes through common.apply_lens,
which sizes the window to the prompt (MAX_SEQ_LEN = None) and raises rather
than truncate if you pass an explicit value that is too small.
20 greedy-generates each answer (correct/wrong split) and records, per
(layer, position), the top numeric token, whether each of A1/A2/sum is decoded
there, and their ranks. It adapts the paper's numeric-decode criterion to the
model's tokenizer — exact match where digits are grouped into single tokens
(DeepSeek), first-token match where they are split (Qwen, Llama 3) — prints
which mode it is in, and records it in the CSV and every figure title. Headline
numbers should come from an exact-mode run; 00_smoke_test.py tells you which
mode a model gives you before you spend GPU hours. Every row also carries a
shuffled-quantity control (ctrl_*): the same decode test against another
example's A1/A2/sum. "Decoded" is an argmax over a few hundred numeric tokens
(exact) or ten digits (prefix), so any-layer-any-position aggregates saturate
on noise alone, and in prefix mode the model's standing preference for the
digit 1 matches most 2-fact sums (100–199) at every cell; 21 draws the
control as a third row of the heatmaps, 30 dots it under the per-layer curves,
and the printed summary puts it next to every fraction. A number that does not
beat its control is chance, whichever lens produced it. The readout covers the
filler region and the whole post-filler tail through position −1 (the token
the answer is predicted from). 21 turns that into Figure-3-style heatmaps +
the "what algorithm?" summary. The logit lens is run through jlens with
use_jacobian=False, so stage 3 is an exact apples-to-apples upgrade — and
common.load_model makes the lens tokenize the rendered prompt verbatim, the
way generation does, since jlens's own encode would add a second BOS on
DeepSeek; apply_lens checks the token counts agree. 20 replays stage 1's own examples. SOURCE = "fig2" (the default) reads
data/fig2_2fact.jsonl — stage 1's test set and its exact rendered prompts —
so idx joins straight onto results/fig2_raw.jsonl. This matters: seeding is
not enough, because pt.build_dataset holds out a different 10 elements for
few-shot and therefore samples from a different pool. At seed 0 the two test sets
share 0 of 1500 idx values and only 219 of 1500 ordered pairs, so a run
from the synthetic source cannot be compared to the API sweep at all.
SOURCE = "synthetic" is the fallback for a machine with no stage-1 dataset.
22 checks that the local weights behave like the API. The mechanistic story
only means something if the local model reproduces the behavioral effect, so 22
pairs answers_<tag>.csv against the stage-1 sweep example-for-example and
reports accuracy, exact-prediction agreement and an exact McNemar per k — then,
given a local k=0 run, compares the uplift locally against the API's. It
verifies the join element-pair by element-pair first and refuses to report a
number if the two sides aren't the same examples.
python 20_lens_readout.py --model deepseek --k 0 --n 300 # baseline: answers only, no lens
python 22_agreement_check.py --model deepseek
k=0 has no filler region, so that run writes greedy answers and skips the lens
work entirely — it is much cheaper than a readout.
python 20_lens_readout.py --model deepseek --n 300 --k 10 --lens jlens
python 30_compare_lenses.py # newest condition
(Or run 20 once with --lens both.) 30 needs both lenses for one condition and
finds them by tag, so it picks up a single both file or a logit + jlens
pair without being told which; pass --readout a.csv b.csv to name them
explicitly. 30 makes the J-lens − logit-lens
difference maps and per-layer decode curves; ANALYSIS.md says what each
outcome means. Greedy answers are cached per condition
(results/answers_<tag>.csv), so the jlens pass reuses the logit pass's
generations.
python 40_attention_study.py --model dev --n 40 --k 10 # dev model first
python 40_attention_study.py --model deepseek --n 100 --k 10
Per (layer, head, query position): attention mass from filler/answer positions onto the fact entities, the rest of the question, and the filler itself — the paper's "what do filler tokens attend to?" analysis. Memory-hungry (forces eager attention); see the script docstring.
A Jacobian lens is a set of matrices in one model's residual-stream basis.
Applying it to a different model (or even base vs. instruct) gives nonsense.
Every lens .pt stores a provenance.model_id; check_provenance() warns if
it disagrees with config.py. If it warns, trust the provenance and edit the
hf_id.
dev (Qwen3.5-4B): ~8.0 GB on the GPU with offload_vision — fits a 12 GB
card. Or slowly on CPU (set dtype="float32").deepseek: 284B MoE, but it ships already quantized (FP4 experts + FP8
elsewhere, ~160 GB), so it loads as published on a rented 4×80GB Hopper box —
no bitsandbytes, which would dequantize it to ~568 GB and OOM. run_deepseek.md
has the free preflight that tells you whether 4× is enough before you rent.sum crystallization readout comparison — 21/30 already report
it: does the J-lens surface the sum earlier (in depth) or in-filler (in
position)?17 commits
Python
91.9%
Shell
7.4%