sae-vis 0.3.3

Open-source SAE visualizer, based on Anthropic's published visualizer.

Note - I'm still open to accepting PRs on this library, and am very happy for other people to build on it, but I won't be actively maintaining it going forwards since I'll be focusing on my job. The SAELens library will continue to have more development and iteration, and it uses a fork of this repo as well as containing a much larger suite of tools for working with SAEs, so depending on your use case you might find that library preferable!

---

This codebase was designed to replicate Anthropic's sparse autoencoder visualisations, which you can see here. The codebase provides 2 different views: a feature-centric view (which is like the one in the link, i.e. we look at one particular feature and see things like which tokens fire strongest on that feature) and a prompt-centric view (where we look at once particular prompt and see which features fire strongest on that prompt according to a variety of different metrics).

Install with pip install sae-vis. Link to PyPI page here.

See here for a demo Colab notebook (all the code to produce it is also in this repo, in the file sae_vis/demos/demo.py, as well as the files containing the created visualizations).

The library supports two types of visualizations:

1. Feature-centric vis, where you look at a single feature and see e.g. which sequences in a large dataset this feature fires strongest on.

2. Prompt-centric vis, where you input a custom prompt and see which features score highest on that prompt, according to a variety of possible metrics.

Citing this work

To cite this work, you can use this bibtex citation:

@misc{sae_vis,
    title  = {{SAE Visualizer}},
    author = {Callum McDougall},
    howpublished    = {\url{https://github.com/callummcdougall/sae_vis}},
    year   = {2024}
}

Contributing

This project is uses Poetry for dependency management. After cloning the repo, install dependencies with poetry install.

This project uses Ruff for formatting and linting, Pyright for type-checking, and Pytest for tests. If you submit a PR, make sure that your code passes all checks. You can run all checks with make check-all.

Version history (recording started at 0.2.9)

• 0.2.9 - added table for pairwise feature correlations (not just encoder-B correlations)
• 0.2.10 - fix some anomalous characters
• 0.2.11 - update PyPI with longer description
• 0.2.12 - fix height parameter of config, add videos to PyPI description
• 0.2.13 - add to dependencies, and fix SAELens section
• 0.2.14 - fix mistake in dependencies
• 0.2.15 - refactor to support eventual scatterplot-based feature browser, fix ’ HTML
• 0.2.16 - allow disabling buffer in feature generation, fix demo notebook, fix sae-lens compatibility & type checking
• 0.2.17 - use main branch of sae-lens
• 0.2.18 - remove circular dependency with sae-lens
• 0.2.19 - formatting, error-checking
• 0.2.20 - fix bugs, remove use of batch_size in config
• 0.2.21 - formatting
• 0.3.0 - major refactor which makes several improvements, removing complexity and adding new features:
  • OthelloGPT SAEs with linear probes (input / output space)
  • Attention output SAEs with max DFA visualized
  • Tokens labelled with their (batch, seq) indices as well as the change in correct-token probability on feature ablation, when hovered over
• 0.3.1 - fix transformerlens dependency
• 0.3.2 - adjust pyright type-checking
• 0.3.3 - remove pyright type-checking