sparsify trains sparse autoencoders without caching activations on disk

Activations are recomputed, not stored

sparsify trains k-sparse autoencoders and transcoders on the activations of HuggingFace language models, roughly following the recipe detailed in Scaling and evaluating sparse autoencoders by Gao et al. The library is described as lean with few configuration options, and the contrast drawn is with other SAE libraries such as SAELens. The difference is that those libraries cache activations on disk, while this one computes them on the fly. The stated benefit is that it allows scaling to very large models and datasets with zero storage overhead. The cost is stated just as plainly. Trying different hyperparameters for the same model and dataset will be slower than if activations were cached, because activations are re-computed each time. Caching may be added as an option in the future. For a researcher, that is a clean exchange. Storing the residual stream activations of a large model over a large dataset is a storage problem, and a compute problem, and this design avoids the first while making the second worse.

TopK is the objective, and L1 is not coming

Following Gao et al., the library uses a TopK activation function which directly enforces a desired level of sparsity in the activations, in contrast to other libraries which use an L1 penalty in the loss function. The position is stated as a belief rather than a preference, that TopK is a Pareto improvement over the L1 approach, and hence no plan to support it. The distinction is not academic for anyone comparing libraries. An L1 penalty is a soft pressure toward sparsity, and the amount of sparsity you get depends on how hard the penalty pulls. A TopK activation hard-codes it, exactly k latents fire, and the objective becomes reconstruction with a fixed budget. That is a materially different training dynamic, and it is why a library that picks one for you is making a decision you may want to make differently. The dependency list in pyproject.toml includes schedulefree, which is consistent with a modern optimiser choice for this kind of objective, along with torch, transformers, datasets, accelerate, einops, safetensors and simple-parsing.

Loading pretrained SAEs and hooking them into a forward pass

Pretrained sparse autoencoders come from the HuggingFace Hub and load with a single method:

from sparsify import Sae

sae = Sae.load_from_hub("EleutherAI/sae-llama-3-8b-32x", hookpoint="layers.10")

That loads the SAE for residual stream layer 10 of Llama 3 8B, trained with an expansion factor of 32. To get every layer at once there is a load_many, which returns a dictionary guaranteed to be naturally sorted by hook point name, so for the common case of embed_tokens, layers.0 through layers.n the SAEs come back in layer order. The library then shows how to gather activations for a model forward pass, running the model with output_hidden_states enabled under inference mode, flattening each hidden state to the shape the encoder expects, and calling sae.encode on it. The documentation is candid that this pattern is for residual stream SAEs specifically, and for anything beyond that it recommends PyTorch hooks, with a link to a gist showing the approach. The difference matters because a general hook can target any submodule, while the built-in path is tuned for the residual stream.

Training from the command line and from Python

The command line entry point is a module invocation with a model, an optional dataset and a flag to switch to transcoders:

python -m sparsify EleutherAI/pythia-160m [optional dataset] [--transcode]

By default the training dataset is EleutherAI/SmolLM2-135M-10B, and you can substitute any dataset from the HuggingFace Hub or any local dataset in Hub format, since the string is passed to load_dataset from the datasets library. Extra dataset loading arguments are passed with --data_args in an arg1=x,arg2=y form. The CLI exposes all of the config options provided by the TrainConfig class, and you can list them with python -m sparsify --help. The programmatic path is the same work in a script, with the pieces being SaeConfig for the autoencoder shape, TrainConfig for the training settings, a Trainer that takes the config, the tokenized data and the model, and a fit call. The data preparation step is a function called chunk_and_tokenize, which takes the dataset and the tokenizer.

Finetuning and custom hookpoints

Finetuning an existing SAE is a path argument rather than a new workflow, where the finetune argument takes the path of a pretrained SAE. Hookpoints are the other important knob. By default the SAEs are trained on the residual stream activations, but you can train on the activations of any other submodule by specifying custom hookpoint patterns. The patterns are like standard PyTorch module names, for example h.0.ln_1, and they also allow Unix pattern matching syntax with wildcards and character sets. To train on the output of every attention module and the inner activations of every MLP in GPT-2, the patterns are given as separate arguments:

python -m sparsify gpt2 --hookpoints "h.*.attn" "h.*.mlp.act"

Restricting to the first three layers uses a character set instead of a wildcard. There is one limit worth knowing. The library does not support fine-grained manual control over the learning rate, number of latents or other hyperparameters on a hookpoint by hookpoint basis. By default the expansion_factor option selects the number of latents for each hookpoint based on the width of that hookpoint's output, and the default learning rate for each hookpoint is set using an inverse square root scaling law based on the number of latents. If you set the number of latents or the learning rate manually, it applies to all hookpoints.

Distributed training has a memory efficient mode

Distributed training is supported through PyTorch's torchrun command, and the default method is Distributed Data Parallel, which means the weights of each SAE are replicated on every GPU. The documentation calls this simple but very memory inefficient, and for training SAEs for many layers it recommends the distribute_modules flag, which allocates the SAEs for different layers to different GPUs. The constraint on that mode is that the number of GPUs must evenly divide the number of layers you are training SAEs for. The example given trains an SAE for every even layer of Llama 3 8B using all available GPUs, accumulating gradients over 8 minibatches and splitting each minibatch into 2 microbatches before feeding them into the SAE encoder, which is where the memory saving comes from. The command also shows the other levers in practice, with a layer stride of 2, an 8-bit load, a context length of 2048 and a k of 192. The project is released under MIT, the package is published as eai-sparsify with a sparsify console script, and the version at the time of writing is 1.3.3 with a changelog at the repository root.

Editorial conclusion

Adopt sparsify when you want to train sparse autoencoders on a model and dataset large enough that caching activations would need storage you do not have, and when you are following the Gao et al. recipe with a TopK objective. Do not adopt it if you need to sweep hyperparameters quickly on a fixed model, because activations are recomputed each time, and do not expect L1 penalty support. Verify first that your model and dataset are on the HuggingFace Hub or loadable in Hub format, that the hookpoint patterns match the submodule names of your architecture, and that the number of GPUs divides the number of layers evenly if you use the module distribution mode.

Frequently asked questions

Official sources

1. EleutherAI/sparsify on GitHub
2. Issues
3. License: MIT
4. README
5. Releases