A JumpReLU loader fix, so SAEs run as the SAEs they are
A merged contribution to SAELens: loading JumpReLU sparse autoencoders trained with dictionary_learning either failed outright or, before v6.23, silently ran them as plain ReLU SAEs. The fix, the evidence, and the tests that pin it.
SAELens is an open-source library for training and analysing sparse autoencoders (SAEs) on language models. Researchers publish trained SAEs in many formats, and SAELens converts them on load. A converter has to get the architecture, the learned weights, and any extra parameters exactly right, or the SAE still loads but computes something different from what was trained.
This contribution is about one such converter, the one for SAEs trained with dictionary_learning, including the JumpReLU SAE suites in SAEBench.
Two failure modes, one cause
The converter mapped JumpReluAutoEncoder to architecture="standard". Since #592 (v6.23.0) the loader also copies the per-latent threshold into the state dict, so those SAEs could no longer be loaded:
RuntimeError: Error(s) in loading state_dict for StandardSAE:
Unexpected key(s) in state_dict: "threshold".
Before v6.23 the same SAEs loaded without error but ran as plain ReLU SAEs with the threshold dropped, giving silently wrong latents and L0. That is the same failure class that #608 fixed for BatchTopKSAE. The loud failure is the better of the two; the silent one is the dangerous one, because nothing tells the user their feature activations are off.
Map the architecture, and fix b_dec handling
Two changes, both needed:
1. Map JumpReluAutoEncoder to architecture="jumprelu". This lets the threshold load into a JumpReLU SAE instead of a standard one.
2. Set apply_b_dec_to_input=False for it. dictionary_learning's JumpReluAutoEncoder sets apply_b_dec_to_input = False in its __init__, so its encode computes x @ W_enc + b_enc without subtracting b_dec. Mapping the architecture alone loads cleanly but gives wrong latents whenever b_dec is nonzero. A load that succeeds is not evidence that the SAE is correct.
The other dictionary_learning classes are unchanged. No release registered in pretrained_saes.yaml with dictionary_learning_1 is a JumpReLU SAE, so the bug affected loading by repo id with converter=dictionary_learning_sae_huggingface_loader_1.
What the second change is worth
Measured on a real SAEBench SAE, adamkarvonen/saebench_pythia-160m-deduped_width-2pow14_date-0108, with pythia-160m-deduped layer-8 residual activations over 152 tokens:
| Loader | L0 | Explained variance |
|---|---|---|
| This PR | 20.8 | 0.976 |
Architecture mapping only (apply_b_dec_) | 206.8 | 0.221 |
Figures are from the PR description. One SAE and 152 tokens: a sanity check on this loader, not a benchmark. L0 is the average number of latents active per token; explained variance is the share of the activation's variance the SAE reconstructs. Loaded with this PR the SAE is a JumpReLUSAE whose encode and decode match dictionary_learning's reference implementation exactly.
Two tests, both failing on main
test_dictionary_learning_sae_huggingface_loader_1_jumprelu checks that the converted config is jumprelu with apply_b_dec_to_input=False and that the per-latent threshold is kept.
test_from_pretrained_dictionary_learning_jumprelu_matches_reference_encoding checks that SAE.from_pretrained(..., converter=dictionary_learning_sae_huggingface_loader_1) reproduces dictionary_learning's encode and decode exactly, with inputs chosen so the threshold actually zeroes some positive pre-activations. That choice matters: without it a wrong threshold would pass unnoticed. The encoding test also fails with only the architecture mapping fixed.
Both tests mock hf_hub_download and need no network. The change is +109 / −1 across sae_lens/loading/pretrained_sae_loaders.py and tests/loading/test_pretrained_sae_loaders.py.
Merged to date: 1 pull request
- #747 fix: load dictionary_learning JumpReluAutoEncoder SAEs in dictionary_learning_1 converter Merged 2026·10·01 · bug fix · +109 / −1 · 2 files