SAELens: A Library for Training and Analyzing Sparse Autoencoders on Language Model Activations
SAELens is an open-source Python library developed by Decode Research to facilitate the training and analysis of sparse autoencoders (SAEs) on language model activations. Its primary purpose is to advance mechanistic interpretability, a subfield of AI safety focused on understanding how neural networks process information. The library integrates with existing tools such as TransformerLens and Neuronpedia, expanding its utility for researchers working on model transparency and alignment.
SAELens was introduced in 2023 and is maintained by Joseph Bloom, Curt Tigges, Anthony Duong, and David Chanin. It is licensed under the permissive MIT license, ensuring broad accessibility for the research community. The project reflects a collaborative effort among contributors aiming to improve humanity’s understanding of neural networks, particularly in mitigating risks associated with artificial intelligence.
The library supports both training custom sparse autoencoders and analyzing pre-trained models. Researchers can leverage SAELens to extract insights from language model activations, which can then inform the development of safer and more aligned AI systems. Its design accommodates a wide range of frameworks beyond TransformerLens, including Hugging Face Transformers and NNight, by allowing users to extract activations and pass them to the SAE’s encoding and decoding methods.
SAELens is structured to be modular and extensible. Users can download and analyze pre-trained SAEs for various models, train their own SAEs, and generate feature dashboards using the companion SAE-Vis library. The library’s documentation provides step-by-step guidance on these processes, ensuring accessibility for both beginners and experienced practitioners.
How It Works
SAELens operates by enabling researchers to train sparse autoencoders on the internal activations of language models. In the context of language models, SAEs help identify and isolate specific features or patterns within the model’s activations, such as syntactic structures, semantic roles, or attention mechanisms.
The library provides tools for extracting activations from PyTorch-based models, regardless of the framework used. For instance, users can integrate SAELens with TransformerLens via the HookedSAETransformer class, which simplifies the process of applying SAEs to transformer architectures. Alternatively, activations can be manually extracted and fed into the SAE’s encode and decode functions, making SAELens compatible with Hugging Face Transformers, custom neural networks, or other frameworks.
A key feature of SAELens is its support for pre-trained SAEs. Researchers can import these models directly, avoiding the computational cost of training from scratch. The library maintains a curated list of pre-trained SAEs, which are regularly updated to include new architectures and models. This functionality accelerates research by providing ready-to-use tools for analyzing model behavior without extensive prior setup.
SAELens also includes mechanisms for visualizing SAE features through integration with the SAE-Vis library. This allows researchers to generate interactive dashboards that map activations to specific features, such as tokens, attention heads, or positional encodings.
Use Cases
SAELens is primarily used in the field of mechanistic interpretability, where researchers seek to dissect the inner workings of language models.
One common application is the analysis of pre-trained language models, such as those developed by Hugging Face or proprietary architectures.
SAELens also supports the evaluation of SAE architectures through synthetic benchmarks. The SynthSAEBench suite, integrated with SAELens, allows researchers to test the performance of different SAE designs on controlled datasets. This capability is essential for comparing the effectiveness of various training methods, sparsity constraints, or architectural choices, ensuring that the resulting SAEs are both accurate and generalizable.
The library’s integration with Neuronpedia further extends its utility.
Tutorials provided by the library guide users through the process of training SAEs, analyzing features, and generating visualizations. These resources are particularly valuable for newcomers to mechanistic interpretability, offering hands-on experience with state-of-the-art tools and techniques.
Key Features and Capabilities
SAELens includes several features that distinguish it from other SAE training libraries. One of its standout capabilities is the ability to work with any PyTorch-based model, not just those compatible with TransformerLens. This flexibility ensures that researchers can apply SAEs to a broad range of architectures, from small feedforward networks to large-scale transformers.
The library provides deep integration with TransformerLens through the HookedSAETransformer class. Users can also employ SAELens with other frameworks, such as Hugging Face Transformers, by manually extracting activations and passing them to the SAE’s encode and decode functions.
The library’s pre-trained SAEs are a significant asset, as they eliminate the need for users to train models from scratch. SAELens maintains a repository of pre-trained SAEs for various language models, which can be imported and analyzed directly. This feature accelerates research by providing ready-to-use tools for interpretability studies, allowing practitioners to focus on analysis rather than training infrastructure.
Visualization is a core component of SAELens, facilitated by its integration with the SAE-Vis library. Users can generate interactive dashboards that map SAE features to model activations, providing intuitive representations of how models process information. These visualizations are essential for mechanistic interpretability, as they enable researchers to identify and study specific patterns in model behavior.
SAELens also includes tools for evaluating SAE performance through synthetic benchmarks. The SynthSAEBench suite allows researchers to test different SAE architectures on controlled datasets, ensuring that the resulting models are both accurate and generalizable.
The library’s documentation is comprehensive, offering tutorials on topics such as loading and analyzing pre-trained SAEs, training custom models, and generating feature visualizations.
The migration guide for version 6 outlines changes to the training code, ensuring that users can transition smoothly to the latest features.
SAELens is part of a broader ecosystem of tools and libraries dedicated to mechanistic interpretability. It collaborates with projects such as dictionary-learning, Sparsify, and Overcomplete, each offering unique approaches to SAE training. By integrating with these tools, SAELens provides researchers with a comprehensive suite of resources for exploring model internals and advancing AI safety.