nielsr's picture
nielsr HF Staff
Update metadata: add pipeline tag, remove arXiv tag
8b4b1f4 verified
|
Raw History Blame
16.5 kB
---
datasets:
- hematteo/wu-crosscoder-snapshots
license: mit
pipeline_tag: other
tags:
- crosscoder
- sparse-autoencoder
- mech-interp
- parameter-trajectory
- pythia
- olmo
---
# Parameter-trajectory crosscoders for the LM readout
Trained dictionaries from **Learning to Read Out: Unembedding Dynamics in Language
Model Pretraining** (NeurIPS 2026). A parameter-trajectory crosscoder fits one
sparse dictionary to every token's readout (unembedding, $W_U$) row at every
checkpoint, using only checkpoint weights. Each feature keeps one index across
checkpoints and has its own decoder direction at each, so you can follow how its
strength and the tokens it promotes change during pretraining.
[Paper (arXiv:2610.21025)](https://arxiv.org/abs/2610.21025) 路
[Project page](https://matteohe.com/learning-to-read-out/) 路
[Feature atlas](https://huggingface.co/spaces/hematteo/readout-feature-atlas) 路
[Code](https://github.com/hematteo/learning-to-read-out) 路
[Collection](https://huggingface.co/collections/hematteo/learning-to-read-out-neurips-2026-6ad25219e26a69daa47880e0)
## Quick start
Follow one feature of the Pythia-160M dictionary across training (downloads
1.6 GB for the dictionary and 155 MB for one readout snapshot):
```python
import json
import torch
from huggingface_hub import hf_hub_download
from safetensors import safe_open
from transformers import AutoTokenizer
repo = "hematteo/parameter-trajectory-crosscoders"
name = "pythia-160m/W_U/cross-snapshot-32/d8192/seed0"
weights = hf_hub_download(repo, f"{name}.safetensors")
steps = json.load(open(hf_hub_download(repo, f"{name}.config.json")))["steps"] # the 32 checkpoint steps
feature = 6046
with safe_open(weights, framework="pt") as f:
D = f.get_slice("W_D")[:, feature] # (32, d_model): this feature's decoder row at each checkpoint
rho = D.norm(dim=-1) / D.norm(dim=-1).max() # decoder norm relative to its peak
print({s: round(rho[i].item(), 2) for i, s in enumerate(steps) if s in (0, 512, 2000, 8000, 143000)})
# {0: 0.05, 512: 0.07, 2000: 0.44, 8000: 0.78, 143000: 0.93}
snap = hf_hub_download("hematteo/wu-crosscoder-snapshots",
"pythia-160m/EleutherAI_pythia-160m_step143000_wu.pt", repo_type="dataset")
W_U = torch.load(snap, weights_only=True) # (vocab, d_model) readout at step 143000
tok = AutoTokenizer.from_pretrained("EleutherAI/pythia-160m")
top = (W_U @ D[steps.index(143000)]).topk(10).indices
print([tok.decode([t]) for t in top])
# [' tumour', ' realise', ' characterised', ' realised', ' colour', ' recognise', ' tumours', ' whilst', ' recognised', ' favour']
```
Feature 6046 promotes British spellings at the end of training, and its decoder
norm grows from 7% of its peak at step 512 to 93% at the final checkpoint. The
[feature atlas](https://huggingface.co/spaces/hematteo/readout-feature-atlas)
has a page like this for every feature of the Pythia-160M, 1B and 6.9B
dictionaries.
To download whole dictionaries:
```bash
hf download hematteo/parameter-trajectory-crosscoders --include "pythia-1b/**" --local-dir ./parameter-trajectory-crosscoders
hf download hematteo/parameter-trajectory-crosscoders --local-dir ./parameter-trajectory-crosscoders # everything, about 180 GB
```
## Files
Each artifact is `<name>.safetensors` + `<name>.config.json` (training
hyperparameters and recomputed quality metrics) + `<name>.md` (card).
`index.json` is the machine-readable inventory of everything below.
| Tensor | Shape | Meaning |
|---|---|---|
| `W_D` | (checkpoints, d_sae, d_model) | Decoder: one direction per feature and checkpoint |
| `b_D` | (checkpoints, d_model) | Decoder bias per checkpoint |
| `W_E` | (checkpoints, d_model, d_sae) | Encoder (unrelated to the input embedding $W_E$ below) |
| `b_E` | (checkpoints, d_sae) | Encoder bias |
| `activation_function.log_jumprelu_threshold` | (d_sae,) | JumpReLU threshold (log scale) |
The checkpoint steps are listed under `steps` in each `.config.json`. Inputs were
centered and scaled per checkpoint (`preprocess_mode` in the config); the code
release reproduces that preprocessing.
## What you probably want
| What | Path |
|---|---|
| Headline 5-seed Pythia-160M $W_U$ crosscoder | `pythia-160m/W_U/cross-snapshot-32/d8192/seed{0..4}.safetensors` |
| High-resolution 160M instrument (atlas) | `pythia-160m/W_U/cross-snapshot-32/d24576/seed0.safetensors` |
| Cross-scale (Pythia-1B) | `pythia-1b/W_U/cross-snapshot-32/d24576/seed0.safetensors` |
| Large-scale, selected sparse run | `pythia-6.9b/W_U/cross-snapshot-32/d32768/seed0-sparse.safetensors` |
| Cross-family (OLMo-2-7B) | `olmo-2-7b/W_U/cross-snapshot-32/d32768/seed0.safetensors` |
| Input embedding $W_E$ (released; the paper's analyses use only $W_U$) | `pythia-160m/W_E/cross-snapshot-32/...` |
| Activation-rate aggregates (lifecycle figures) | `derived/aggregates/`, `derived/rates/` |
| Attribution-patching artifacts | `attribution/pythia-160m/` |
| Held-out eval token corpus | `evaluation/eval-corpus/eval_tokens.pt` |
## Full inventory
| Path | Model | Matrix | Kind | d_sae | Seed | Quality |
|---|---|---|---|---|---|---|
| `olmo-2-7b/W_U/cross-snapshot-32/d32768/seed0.safetensors` | allenai/OLMo-2-1124-7B | W_U | cross-snapshot-32 | 32768 | 0 | EV 0.853 / L0 557 |
| `pythia-160m/W_E/cross-snapshot-32/d24576/seed0.safetensors` | EleutherAI/pythia-160m | W_E | cross-snapshot-32 | 24576 | 0 | EV 0.831 / L0 118 |
| `pythia-160m/W_E/cross-snapshot-32/d8192/seed0.safetensors` | EleutherAI/pythia-160m | W_E | cross-snapshot-32 | 8192 | 0 | EV 0.581 / L0 82 |
| `pythia-160m/W_E/cross-snapshot-32/d8192/seed1.safetensors` | EleutherAI/pythia-160m | W_E | cross-snapshot-32 | 8192 | 1 | EV 0.580 / L0 82 |
| `pythia-160m/W_E/cross-snapshot-32/d8192/seed2.safetensors` | EleutherAI/pythia-160m | W_E | cross-snapshot-32 | 8192 | 2 | EV 0.582 / L0 82 |
| `pythia-160m/W_E/cross-snapshot-32/d8192/seed3.safetensors` | EleutherAI/pythia-160m | W_E | cross-snapshot-32 | 8192 | 3 | EV 0.581 / L0 82 |
| `pythia-160m/W_E/cross-snapshot-32/d8192/seed4.safetensors` | EleutherAI/pythia-160m | W_E | cross-snapshot-32 | 8192 | 4 | EV 0.583 / L0 83 |
| `pythia-160m/W_U/architecture-comparison/d8192/batchtopk/seed0.safetensors` | EleutherAI/pythia-160m | W_U | architecture-comparison/d8192 | 8192 | 0 | EV 0.725 / L0 203 |
| `pythia-160m/W_U/architecture-comparison/d8192/gated/seed0.safetensors` | EleutherAI/pythia-160m | W_U | architecture-comparison/d8192 | 8192 | 0 | EV 0.214 / L0 12 |
| `pythia-160m/W_U/architecture-comparison/d8192/gated-retuned/seed0.safetensors` | EleutherAI/pythia-160m | W_U | architecture-comparison/d8192 | 8192 | 0 | EV 0.827 / L0 654 |
| `pythia-160m/W_U/cross-snapshot-16/d8192/seed0.safetensors` | EleutherAI/pythia-160m | W_U | cross-snapshot-16 | 8192 | 0 | EV 0.773 / L0 216 |
| `pythia-160m/W_U/cross-snapshot-32/d16384/seed0.safetensors` | EleutherAI/pythia-160m | W_U | cross-snapshot-32 | 16384 | 0 | EV 0.780 / L0 103 |
| `pythia-160m/W_U/cross-snapshot-32/d24576/seed0.safetensors` | EleutherAI/pythia-160m | W_U | cross-snapshot-32 | 24576 | 0 | EV 0.920 / L0 286 |
| `pythia-160m/W_U/cross-snapshot-32/d24576/seed1.safetensors` | EleutherAI/pythia-160m | W_U | cross-snapshot-32 | 24576 | 1 | EV 0.920 / L0 286 |
| `pythia-160m/W_U/cross-snapshot-32/d24576/seed2.safetensors` | EleutherAI/pythia-160m | W_U | cross-snapshot-32 | 24576 | 2 | EV 0.920 / L0 286 |
| `pythia-160m/W_U/cross-snapshot-32/d8192/seed0.safetensors` | EleutherAI/pythia-160m | W_U | cross-snapshot-32 | 8192 | 0 | EV 0.776 / L0 203 |
| `pythia-160m/W_U/cross-snapshot-32/d8192/seed1.safetensors` | EleutherAI/pythia-160m | W_U | cross-snapshot-32 | 8192 | 1 | EV 0.776 / L0 203 |
| `pythia-160m/W_U/cross-snapshot-32/d8192/seed2.safetensors` | EleutherAI/pythia-160m | W_U | cross-snapshot-32 | 8192 | 2 | EV 0.776 / L0 203 |
| `pythia-160m/W_U/cross-snapshot-32/d8192/seed3.safetensors` | EleutherAI/pythia-160m | W_U | cross-snapshot-32 | 8192 | 3 | EV 0.776 / L0 203 |
| `pythia-160m/W_U/cross-snapshot-32/d8192/seed4.safetensors` | EleutherAI/pythia-160m | W_U | cross-snapshot-32 | 8192 | 4 | EV 0.777 / L0 203 |
| `pythia-160m/W_U/final-snapshot-saes/d16384.safetensors` | EleutherAI/pythia-160m | W_U | final-snapshot-saes | 16384 | 0 | EV 0.870 / L0 1913 |
| `pythia-160m/W_U/final-snapshot-saes/d32768.safetensors` | EleutherAI/pythia-160m | W_U | final-snapshot-saes | 32768 | 0 | EV 0.926 / L0 3410 |
| `pythia-160m/W_U/final-snapshot-saes/d6144.safetensors` | EleutherAI/pythia-160m | W_U | final-snapshot-saes | 6144 | 0 | EV 0.765 / L0 862 |
| `pythia-160m/W_U/final-snapshot-saes/d65536.safetensors` | EleutherAI/pythia-160m | W_U | final-snapshot-saes | 65536 | 0 | EV 0.964 / L0 5943 |
| `pythia-160m/W_U/final-snapshot-saes/d8192.safetensors` | EleutherAI/pythia-160m | W_U | final-snapshot-saes | 8192 | 0 | EV 0.799 / L0 1084 |
| `pythia-160m/W_U/lambda-sweep/d8192/lam0p40_seed0.safetensors` | EleutherAI/pythia-160m | W_U | lambda-sweep | 8192 | 0 | EV 0.748 / L0 160 |
| `pythia-160m/W_U/lambda-sweep/d8192/lam1p00_seed0.safetensors` | EleutherAI/pythia-160m | W_U | lambda-sweep | 8192 | 0 | EV 0.632 / L0 58 |
| `pythia-160m/W_U/lambda-sweep/d8192/lam1p20_seed0.safetensors` | EleutherAI/pythia-160m | W_U | lambda-sweep | 8192 | 0 | EV 0.603 / L0 45 |
| `pythia-160m/W_U/lambda-sweep/d8192/lam1p35_seed0.safetensors` | EleutherAI/pythia-160m | W_U | lambda-sweep | 8192 | 0 | EV 0.582 / L0 38 |
| `pythia-160m/W_U/lambda-sweep/d8192/lam1p35_seed1.safetensors` | EleutherAI/pythia-160m | W_U | lambda-sweep | 8192 | 1 | EV 0.582 / L0 38 |
| `pythia-160m/W_U/lambda-sweep/d8192/lam1p35_seed2.safetensors` | EleutherAI/pythia-160m | W_U | lambda-sweep | 8192 | 2 | EV 0.582 / L0 38 |
| `pythia-160m/W_U/lambda-sweep/d8192/lam1p80_seed0.safetensors` | EleutherAI/pythia-160m | W_U | lambda-sweep | 8192 | 0 | EV 0.528 / L0 23 |
| `pythia-160m/W_U/per-snapshot-saes/d8192/step0.safetensors` | EleutherAI/pythia-160m | W_U | per-snapshot-saes | 8192 | 0 | EV 0.729 / L0 1150 |
| `pythia-160m/W_U/per-snapshot-saes/d8192/step1.safetensors` | EleutherAI/pythia-160m | W_U | per-snapshot-saes | 8192 | 0 | EV 0.729 / L0 1150 |
| `pythia-160m/W_U/per-snapshot-saes/d8192/step1000.safetensors` | EleutherAI/pythia-160m | W_U | per-snapshot-saes | 8192 | 0 | EV 0.786 / L0 997 |
| `pythia-160m/W_U/per-snapshot-saes/d8192/step102000.safetensors` | EleutherAI/pythia-160m | W_U | per-snapshot-saes | 8192 | 0 | EV 0.800 / L0 983 |
| `pythia-160m/W_U/per-snapshot-saes/d8192/step116000.safetensors` | EleutherAI/pythia-160m | W_U | per-snapshot-saes | 8192 | 0 | EV 0.812 / L0 958 |
| `pythia-160m/W_U/per-snapshot-saes/d8192/step128.safetensors` | EleutherAI/pythia-160m | W_U | per-snapshot-saes | 8192 | 0 | EV 0.729 / L0 1150 |
| `pythia-160m/W_U/per-snapshot-saes/d8192/step130000.safetensors` | EleutherAI/pythia-160m | W_U | per-snapshot-saes | 8192 | 0 | EV 0.820 / L0 940 |
| `pythia-160m/W_U/per-snapshot-saes/d8192/step14000.safetensors` | EleutherAI/pythia-160m | W_U | per-snapshot-saes | 8192 | 0 | EV 0.791 / L0 996 |
| `pythia-160m/W_U/per-snapshot-saes/d8192/step143000.safetensors` | EleutherAI/pythia-160m | W_U | per-snapshot-saes | 8192 | 0 | EV 0.824 / L0 924 |
| `pythia-160m/W_U/per-snapshot-saes/d8192/step16.safetensors` | EleutherAI/pythia-160m | W_U | per-snapshot-saes | 8192 | 0 | EV 0.729 / L0 1150 |
| `pythia-160m/W_U/per-snapshot-saes/d8192/step2.safetensors` | EleutherAI/pythia-160m | W_U | per-snapshot-saes | 8192 | 0 | EV 0.729 / L0 1150 |
| `pythia-160m/W_U/per-snapshot-saes/d8192/step2000.safetensors` | EleutherAI/pythia-160m | W_U | per-snapshot-saes | 8192 | 0 | EV 0.799 / L0 969 |
| `pythia-160m/W_U/per-snapshot-saes/d8192/step21000.safetensors` | EleutherAI/pythia-160m | W_U | per-snapshot-saes | 8192 | 0 | EV 0.791 / L0 998 |
| `pythia-160m/W_U/per-snapshot-saes/d8192/step256.safetensors` | EleutherAI/pythia-160m | W_U | per-snapshot-saes | 8192 | 0 | EV 0.732 / L0 1142 |
| `pythia-160m/W_U/per-snapshot-saes/d8192/step27000.safetensors` | EleutherAI/pythia-160m | W_U | per-snapshot-saes | 8192 | 0 | EV 0.791 / L0 999 |
| `pythia-160m/W_U/per-snapshot-saes/d8192/step3000.safetensors` | EleutherAI/pythia-160m | W_U | per-snapshot-saes | 8192 | 0 | EV 0.799 / L0 972 |
| `pythia-160m/W_U/per-snapshot-saes/d8192/step32.safetensors` | EleutherAI/pythia-160m | W_U | per-snapshot-saes | 8192 | 0 | EV 0.729 / L0 1150 |
| `pythia-160m/W_U/per-snapshot-saes/d8192/step34000.safetensors` | EleutherAI/pythia-160m | W_U | per-snapshot-saes | 8192 | 0 | EV 0.791 / L0 1000 |
| `pythia-160m/W_U/per-snapshot-saes/d8192/step4.safetensors` | EleutherAI/pythia-160m | W_U | per-snapshot-saes | 8192 | 0 | EV 0.729 / L0 1150 |
| `pythia-160m/W_U/per-snapshot-saes/d8192/step4000.safetensors` | EleutherAI/pythia-160m | W_U | per-snapshot-saes | 8192 | 0 | EV 0.798 / L0 977 |
| `pythia-160m/W_U/per-snapshot-saes/d8192/step47000.safetensors` | EleutherAI/pythia-160m | W_U | per-snapshot-saes | 8192 | 0 | EV 0.791 / L0 1000 |
| `pythia-160m/W_U/per-snapshot-saes/d8192/step5000.safetensors` | EleutherAI/pythia-160m | W_U | per-snapshot-saes | 8192 | 0 | EV 0.796 / L0 982 |
| `pythia-160m/W_U/per-snapshot-saes/d8192/step512.safetensors` | EleutherAI/pythia-160m | W_U | per-snapshot-saes | 8192 | 0 | EV 0.754 / L0 1087 |
| `pythia-160m/W_U/per-snapshot-saes/d8192/step6000.safetensors` | EleutherAI/pythia-160m | W_U | per-snapshot-saes | 8192 | 0 | EV 0.795 / L0 985 |
| `pythia-160m/W_U/per-snapshot-saes/d8192/step61000.safetensors` | EleutherAI/pythia-160m | W_U | per-snapshot-saes | 8192 | 0 | EV 0.790 / L0 1002 |
| `pythia-160m/W_U/per-snapshot-saes/d8192/step64.safetensors` | EleutherAI/pythia-160m | W_U | per-snapshot-saes | 8192 | 0 | EV 0.729 / L0 1150 |
| `pythia-160m/W_U/per-snapshot-saes/d8192/step7000.safetensors` | EleutherAI/pythia-160m | W_U | per-snapshot-saes | 8192 | 0 | EV 0.794 / L0 988 |
| `pythia-160m/W_U/per-snapshot-saes/d8192/step75000.safetensors` | EleutherAI/pythia-160m | W_U | per-snapshot-saes | 8192 | 0 | EV 0.790 / L0 1004 |
| `pythia-160m/W_U/per-snapshot-saes/d8192/step8.safetensors` | EleutherAI/pythia-160m | W_U | per-snapshot-saes | 8192 | 0 | EV 0.729 / L0 1150 |
| `pythia-160m/W_U/per-snapshot-saes/d8192/step8000.safetensors` | EleutherAI/pythia-160m | W_U | per-snapshot-saes | 8192 | 0 | EV 0.793 / L0 990 |
| `pythia-160m/W_U/per-snapshot-saes/d8192/step89000.safetensors` | EleutherAI/pythia-160m | W_U | per-snapshot-saes | 8192 | 0 | EV 0.792 / L0 1002 |
| `pythia-160m/W_U/per-snapshot-saes/d8192/step9000.safetensors` | EleutherAI/pythia-160m | W_U | per-snapshot-saes | 8192 | 0 | EV 0.792 / L0 992 |
| `pythia-1b/W_U/cross-snapshot-32/d16384/seed0.safetensors` | EleutherAI/pythia-1b | W_U | cross-snapshot-32 | 16384 | 0 | EV 0.781 / L0 499 |
| `pythia-1b/W_U/cross-snapshot-32/d24576/seed0.safetensors` | EleutherAI/pythia-1b | W_U | cross-snapshot-32 | 24576 | 0 | EV 0.861 / L0 517 |
| `pythia-1b/W_U/cross-snapshot-32/d8192/seed0.safetensors` | EleutherAI/pythia-1b | W_U | cross-snapshot-32 | 8192 | 0 | EV 0.628 / L0 374 |
| `pythia-1b/W_U/cross-snapshot-32-matched-window/d24576/seed0.safetensors` | EleutherAI/pythia-1b | W_U | cross-snapshot-32-matched-window | 24576 | 0 | EV 0.884 / L0 264 |
| `pythia-6.9b/W_U/cross-snapshot-32/d32768/seed0-sparse.safetensors` | EleutherAI/pythia-6.9b | W_U | cross-snapshot-32 | 32768 | 0 | EV 0.808 / L0 742 |
| `pythia-6.9b/W_U/cross-snapshot-32/d32768/seed0.safetensors` | EleutherAI/pythia-6.9b | W_U | cross-snapshot-32 | 32768 | 0 | EV 0.833 / L0 1957 |
Quality metrics are recomputed from the released weights on the released
snapshot schedule (see the code repo's `reproduce/validation/recompute_quality.py`).
The `gated` architecture-comparison run intentionally documents
位-transfer failure (default 位=0.3 moved across architectures); see
`gated-retuned` (位=0.05) for the tuned comparison point.
## Citation
```bibtex
@inproceedings{he2026learningtoreadout,
title = {Learning to Read Out: Unembedding Dynamics in Language Model Pretraining},
author = {He, Matteo and Shen, William F. and Iacob, Alex and Jovanovic, Andrej
and Qiu, Xinchi and Lane, Nicholas D.},
booktitle = {Advances in Neural Information Processing Systems},
year = {2026},
eprint = {2610.21025},
archivePrefix = {arXiv},
primaryClass = {cs.CL},
}
```
MIT. W_U/W_E source tensors derive from public Apache-2.0 checkpoints
(EleutherAI Pythia, AllenAI OLMo-2). The eval corpus derives from Wikipedia
(CC-BY-SA 4.0).