From the repository README
Goodfire’s Variational Parameter Decomposition (VPD) rewrites every weight matrix of a model as a sum of rank-one components and trains a causal-importance (CI) network that switches components on per token. This repository ships one such decomposition of facebook/opt-125m (55,296 components, 0.6% on per position) and asks a narrow question: does the decomposition, run at its own operating point, still recall what the model knows? It does not. The CI-masked forward loses factual recall (the two-shot “capital of France is” drops from 0.76 to 0.06 and answers “London”) and it loses rare-word recall the same way (Hong Kong from 0.99 to 0.52, Kuala Lumpur from 0.99 to 0.04, Vincent van Gogh from 0.96 to 0.001), while frequent continuations survive or improve (In order to 0.81 to 0.99). Damage is governed by the answer token’s corpus frequency, not by whether the prediction needs a fact.
This repository holds the model, the data, the training recipe, and the scripts that re-derive every number below in about a minute on one GPU.
Results
All numbers are P(answer | prompt) at the last prompt position, original OPT-125m versus the same weights run through the decomposition with the CI mask applied (delta = 0). “All-on” (every component on, no mask) reproduces the original to within 0.002 nats of cross-entropy, so everything the mask loses is inside the decomposition. Full distributions, ranks, and per-position traces: reference/probes.json.
| prompt (answer) | original | CI-masked | rank | CI-masked top-1 |
|---|---|---|---|---|
The capital of Italy is Rome. The capital of Japan is Tokyo. The capital of France is ( Paris) | 0.761 | 0.057 | 1 → 2 | London 0.250 |
Rome is the capital of Italy. Tokyo is the capital of Japan. Paris is the capital of ( France) | 0.780 | 0.090 | 1 → 2 | England 0.140 |
Paris is the capital of ( France), bare | 0.323 | 0.119 | 2 → 2 | the |
The capital of France is ( Paris), bare | 0.007 | 0.00005 | 22 → 1543 | a |
...The capital of Greece is ( Athens) | 0.853 | 0.023 | 1 → 5 | Wales |
...The capital of Hungary is ( Budapest) | 0.859 | 0.003 | 1 → 42 | London |
...The capital of Kuwait is ( Kuwait) | 0.466 | 0.001 | 1 → 122 | London |
...The capital of England is ( London) | 0.608 | 0.517 | 1 → 1 | London |
The Chinese city of Hong ( Kong) | 0.988 | 0.518 | 1 → 1 | Kong |
The flight landed in Kuala ( Lumpur) | 0.990 | 0.040 | 1 → 1 | Lumpur |
A painting by Vincent van ( Gogh) | 0.956 | 0.001 | 1 → 102 | de |
He ordered a plate of macaroni and ( cheese) | 0.867 | 0.025 | 1 → 1 | cheese |
In the museum there was a xyl(ophone) | 0.903 | 0.0005 | 1 → 244 | |
The city of Los ( Angeles) | 0.950 | 0.692 | 1 → 1 | Angeles |
A patient with Alzheimer('s) | 0.697 | 0.647 | 1 → 1 | 's |
In order ( to) | 0.813 | 0.991 | 1 → 1 | to |
The pattern over the whole battery (399 probe/condition rows, medians; reference/battery_summary.json):
| family | n | P all-on | P CI-masked | damage (nats) | random mask of the same size |
|---|---|---|---|---|---|
| world facts, two-shot (inventory) | 59 | 0.628 | 0.032 | −2.70 | −11.9 |
| held-out world facts, two-shot | 46 | 0.631 | 0.014 | −3.86 | −12.8 |
multi-token entity names (Rio de Janeiro) | 13 | 0.993 | 0.231 | −1.44 | −9.7 |
collocations (Pros and cons) | 14 | 0.898 | 0.166 | −1.55 | −12.0 |
| word-piece and rigid-name completions | 82 | 0.861 | 0.026 | −2.93 | −10.3 |
| function-word slots (syntax controls) | 11 | 0.573 | 0.828 | +0.21 | −8.0 |
| nonce subjects, same template | 46 | 0.003 | 0.003 | +0.27 | −8.4 |
What decides the damage is the answer token’s log unigram frequency in the training corpus: r = 0.75 over the 125 known items (slope 0.68 nats per log-unit), r = 0.83 inside the relational facts alone and 0.66 inside the locally determined completions, and r = 0.66 over 2,423 confident positions of held-out natural text, where the median damage runs −5.5, −2.1, −0.95, −0.06, +0.06, +0.05 nats across log-frequency bins from rarest to most frequent. Facts with frequent answers (England → London, −0.13 nats) barely move; word-piece continuations with rare answers (xyl → ophone, −7.5 nats) die.
The decomposition itself (TRAINING.md): 72 matrices (q, k, v, out, fc1, fc2 of all 12 layers), 768 components each, 56,981 training steps on 128-token Pile windows; exact reconstruction to 1.7e-4 max logit difference; all-on cross-entropy 3.319 against the original’s 3.317 on held-out text; CI-masked 3.959; 339 to 344 components on per position (0.62%); 683 of 55,296 alive on Pile.
Quick start
Tested on Linux with Python 3.10 and 3.13, one RTX 4090 (about 3 GB of GPU memory for the analyses; the decomposition needs ~14 GB to retrain). Everything after setup.sh runs in under a minute.
git clone https://github.com/czhs/vpd-factual-recall && cd vpd-factual-recall
bash setup.sh # .venv + deps, downloads the checkpoint (818 MB) and corpus (139 MB) from the release, caches facebook/opt-125m, sanity check
bash reproduce.sh # probes -> battery -> natural text -> figures -> check against reference/
setup.sh honours TORCH_SPEC / TORCH_INDEX for a specific CUDA build (the reference machine used TORCH_SPEC=torch==2.8.0 TORCH_INDEX=https://download.pytorch.org/whl/cu128) and VFR_CKPT / VFR_PILE / VFR_RESULTS to relocate the assets and outputs (vfr/paths.py).
One prompt, four forwards, no scripts:
source .venv/bin/activate && python -m vfr.sanity
prints P(' Paris') after the two-shot prompt under the original model (0.7605), the exact reconstruction (0.7605), all components on (0.7736) and the CI mask (0.0567, rank 2 behind ` London`).
What the pipeline does
| step | command | notes |
|---|---|---|
| showcase probes | python -m vfr.probe | 38 prompts ([data/probes/showcase.jsonl](https://github.com/czhs/vpd-factual-recall/blob/main/data/probes/showcase.jsonl)): top-10 next-token distributions under original / exact / all-on / CI-masked / random-matched forwards, per-position traces, gate counts → results/probes.json |
| frequency battery | python -m vfr.battery | 194-item 2 x 2 battery (rare/frequent answer × relational/locally determined) + 113 inventory facts + 46 held-out facts with nonce twins → results/battery.jsonl, battery_summary.json (families, frequency-law fit, capital ladder) |
| natural text | python -m vfr.wild --n-windows 64 | every position of held-out 128-token Pile windows (the CI net's training regime) → results/wild_positions.jsonl, wild_summary.json, wild_passages.json |
| figures | python -m vfr.figures / python -m vfr.figures_svg | PNGs in figures/ (matplotlib) and theme-aware SVG includes in figures/svg/ |
| check | python -m vfr.check | compares results/ with reference/: exact-reconstruction identity, per-probe log-probs within 0.05 nats, ranks, the frequency-law fit, per-family medians, natural-text cross-entropy |
| retrain | bash scripts/train_pileopt_long.sh | the exact command that produced the checkpoint (4.5 h on a 4090); python -m vfr.tokenize_pile rebuilds the corpus from NeelNanda/pile-10k |
Reproduction log. 2026-09-02, fresh clone on the reference machine (RTX 4090, driver 535, system Python 3.10.12, TORCH_SPEC=torch==2.8.0 TORCH_INDEX=https://download.pytorch.org/whl/cu128): setup.sh built the venv, downloaded and checksummed both release assets (41 s), and passed the sanity check; reproduce.sh re-derived every number above and vfr.check passed 9/9 against reference/. The 2026-08-14 campaign’s numbers for the same probes agree to the third decimal (reference/README.md). On other GPUs expect log-probabilities to agree within ~0.01 nats, not bit-for-bit.
Conventions worth knowing before writing new probes
- The CI transformer is bidirectional over the 128-token window, so prompts are never padded:
vfr.batterybatches only prompts of identical token length, and the answer is never inside the window (it is scored at the last prompt position). - OPT’s tokenizer auto-prepends BOS (id 2) and distinguishes
" Rome"(one token) from"Rome"(two); answers carry their leading space, and multi-token answers are scored on their first piece and flagged (answer_single_token). Every probe is checked for BPE seam consistency (tok(prompt + answer)must equaltok(prompt)followed by the answer’s first piece). - “Damage” is
log P_cim(answer) − log P_allon(answer)in nats; the random control is a mask with exactly CI’s per-position, per-matrix active count (seeded), so “sparse” is never the explanation of anything. - Zero-shot capitals sit at P ≈ 0.007 under this model (rank correct, buried under ` the
/a`); the two-shot template lifts the same facts to P ≈ 0.66, which is why the factual probes are two-shot. The demonstrations never contain the scored answer.
Layout
setup.sh, reproduce.sh entry points (see above)
scripts/train_pileopt_long.sh retrain the decomposition from scratch
TRAINING.md the CI network, losses, schedule, every hyper-parameter, what the run converged to
vfr/ the package; run modules with `python -m vfr.<name>`
paths.py all locations; env-overridable
decomp.py loads the checkpoint; original / exact / all-on / CI-masked / random forwards
probe.py battery.py wild.py the three analyses
figures.py figures_svg.py check.py figures and the reference check
train_full_vpd.py common.py the training driver (VPD via vfr/nano)
tokenize_pile.py unigram.py corpus and unigram table
nano/run.py nano/compat_2d.py Goodfire's nano_param_decomp reference implementation (MIT, commit 53965b5e, unmodified) + the OPT 2-D shim
data/probes/ showcase.jsonl, frequency_battery.jsonl, localization_set.jsonl, heldout_probes.jsonl, inventory_battery.jsonl
data/unigram_pile_opt.pt unigram log-frequency table over the training corpus (50,272 tokens)
reference/ the numbers in this README, the training log and config, the 2026-08-14 campaign's artifacts
figures/ PNGs; figures/svg/ the site includes
artifacts/ (downloaded by setup.sh) pileopt_long_ckpt.pt, pile_opt.pt
Release assets
| file | size | sha256 | what |
|---|---|---|---|
pileopt_long_ckpt.pt | 818 MB | 00e608dd…482f49c | the decomposition: V, U for 72 matrices (fp32), the CI network's state dict, cfg, args, module_paths (torch.load(..., weights_only=True)) |
pile_opt.pt | 139 MB | fd56a8a3…166012 | NeelNanda/pile-10k tokenized with OPT's tokenizer, 17,361,935 int64 tokens; the first 90% trained the decomposition |
Full hashes in setup.sh. The target model is not stored; it is facebook/opt-125m from the Hub, loaded in fp32 with eager attention.
Provenance and license
The decomposition and the 2026-08-14 knowledge campaign (inventory, localization, two adversarial verification passes) were produced in the VPD-Copying project; the campaign’s reports and raw files are in reference/campaign/. Code and data here are MIT (LICENSE); vfr/nano/run.py is Goodfire’s, MIT, vendored from https://github.com/goodfire-ai/param-decomp.
Figures
