simonbernier/jlens-filler

0

stars

17

commits

Python

primary language

Sep 2, 2026

updated

Browse cluster: Python AI/ML Research and Analysis Tools

README

Filler tokens × Jacobian lens

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.

stagescriptswhatwhere
000_smoke_test.pyprove the machine works (API path or GPU/lens path)anywhere
110_build_fig2_dataset.py11_run_fig2_sweep.pyFig. 2: accuracy vs filler length k, paper-scale nAPI (no GPU)
220_lens_readout.py (LENS="logit") → 21_analyze_readout.py, 22_agreement_check.pyFig. 3: the paper's logit-lens picture on V4 Flash, plus a local-vs-API accuracy checkGPU box
320_lens_readout.py (LENS="jlens") → 30_compare_lenses.pythe new result: does the J-lens see more than the logit lens?GPU box
440_attention_study.pyattention study as in the paper (optional, code ready)GPU box
550_filler_patching.pycausal tests: donor residual patches on the dots (single layer / band / all layers) and J-lens coordinate swaps & ablations, with pre-registered predictionsGPU 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.

Get the code on a new machine

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)

Install

setup_env.sh detects the machine it is on and needs no flags:

it checksGPU box (no conda)local box with miniconda
environmentpython -m venv .venvconda create -n jlens-filler python=3.11
torchCUDA 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
bitsandbytesinstalled (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).

VS Code

.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/).

Stage 0 — smoke test

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).

Stage 1 — Figure 2 (accuracy vs k, API only)

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.

Stage 2 — the paper's logit-lens picture (GPU box)

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.

Stage 3 — J-lens vs logit-lens (the new result)

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.

Stage 4 (optional) — attention study

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.

The one thing to get right: model ↔ lens matching

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.

Compute

  • 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.
  • Stage 1: API only — no GPU. The paper found uplift robust to API vs local 4-bit.
  • An inference API cannot be used for the lens work: the lens needs the residual stream, which APIs don't expose. White-box weights are required.

Suggested next steps (beyond this scaffold)

  1. Confirm the Fig. 2 shape on V4 Flash, then check the lens picture tracks the behavioral uplift across k (run 20 at k=10/25/50).
  2. Add the sum crystallization readout comparison — 21/30 already report it: does the J-lens surface the sum earlier (in depth) or in-filler (in position)?
  3. Stage 4's attention analysis, plus a KV-cache transplant at filler positions (causal check), for triangulation.
  4. If V4 Flash shows weak 2-fact uplift (the paper saw only ~3 points on V3), the 1-fact task (54%→72% in the paper) is already in the stage-1 data.

Contributors

simonbernier

17 commits

simonbernier/jlens-filler

0

stars

17

commits

Python

primary language

Sep 2, 2026

updated

Browse cluster: Python AI/ML Research and Analysis Tools

README

Filler tokens × Jacobian lens

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.

stagescriptswhatwhere
000_smoke_test.pyprove the machine works (API path or GPU/lens path)anywhere
110_build_fig2_dataset.py11_run_fig2_sweep.pyFig. 2: accuracy vs filler length k, paper-scale nAPI (no GPU)
220_lens_readout.py (LENS="logit") → 21_analyze_readout.py, 22_agreement_check.pyFig. 3: the paper's logit-lens picture on V4 Flash, plus a local-vs-API accuracy checkGPU box
320_lens_readout.py (LENS="jlens") → 30_compare_lenses.pythe new result: does the J-lens see more than the logit lens?GPU box
440_attention_study.pyattention study as in the paper (optional, code ready)GPU box
550_filler_patching.pycausal tests: donor residual patches on the dots (single layer / band / all layers) and J-lens coordinate swaps & ablations, with pre-registered predictionsGPU 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.

Get the code on a new machine

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)

Install

setup_env.sh detects the machine it is on and needs no flags:

it checksGPU box (no conda)local box with miniconda
environmentpython -m venv .venvconda create -n jlens-filler python=3.11
torchCUDA 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
bitsandbytesinstalled (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).

VS Code

.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/).

Stage 0 — smoke test

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).

Stage 1 — Figure 2 (accuracy vs k, API only)

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.

Stage 2 — the paper's logit-lens picture (GPU box)

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.

Stage 3 — J-lens vs logit-lens (the new result)

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.

Stage 4 (optional) — attention study

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.

The one thing to get right: model ↔ lens matching

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.

Compute

  • 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.
  • Stage 1: API only — no GPU. The paper found uplift robust to API vs local 4-bit.
  • An inference API cannot be used for the lens work: the lens needs the residual stream, which APIs don't expose. White-box weights are required.

Suggested next steps (beyond this scaffold)

  1. Confirm the Fig. 2 shape on V4 Flash, then check the lens picture tracks the behavioral uplift across k (run 20 at k=10/25/50).
  2. Add the sum crystallization readout comparison — 21/30 already report it: does the J-lens surface the sum earlier (in depth) or in-filler (in position)?
  3. Stage 4's attention analysis, plus a KV-cache transplant at filler positions (causal check), for triangulation.
  4. If V4 Flash shows weak 2-fact uplift (the paper saw only ~3 points on V3), the 1-fact task (54%→72% in the paper) is already in the stage-1 data.

Contributors

simonbernier

17 commits

Languages

Python

91.9%

Shell

7.4%