SAEs are early enough that there’s tons of low hanging fruit and ideas to try. They also require relatively little compute (often around $1 for a training run), so AI agents could afford to test many ideas. I wouldn’t be surprised if SAE improvements were a good early target for automated AI research, especially if the feedback loop is just “Come up with idea, modify existing loss function, train, evaluate, get a quantitative result”.
Adam Karvonen
Karma: 369
If you’re looking for a hackable SAE training repo for experiments, I’d recommend our dictionary_learning repo. It’s been around for a few months, but we’ve recently spent some time cleaning it up and adding additional trainer types.
It’s designed to be simple and hackable—you can add a new SAE type in a single file (~350 lines). We have 8 tested implementations, including JumpReLU, TopK, BatchTopK, Matryoshka, Gated, and others, with BatchTopK recommended as a good default. Training is quick and cheap—training 6 16K width SAEs on Gemma-2-2B for 200M tokens takes ~6 3090 hours, or ~$1.20.
The repo integrates with SAE Bench and includes reproducible baselines trained on Pythia-160M and Gemma-2-2B. While it’s not optimized for large models like Eleuther’s (no Cuda kernels/multi-GPU support) and has fewer features than SAE Lens, it’s great for experiments and trying new architectures.
Here is a link to the repo: https://github.com/saprmarks/dictionary_learning
Adam Karvonen’s Shortform
The forward hook for our best performing approach is here. As Sam mentioned, this hasn’t been deployed to production. We left it as a case study because Benchify is currently prioritizing other parts of their stack unrelated to ML.
For this demonstration, we added a forward hook to a HuggingFace Transformers model for simplicity, rather than incorporating it into a production inference stack.
Rejection sampling is a strong baseline that we hadn’t considered, and it’s definitely worth trying out—I suspect it will perform well here. Currently, our focus is on identifying additional in-the-wild tasks, particularly from other companies, as many of Benchify’s challenges involve sensitive details about their internal tooling that they prefer to keep private. We’re especially interested in tasks where it’s not possible to automatically measure success or failure via string matching, as this is where techniques like model steering are most likely to be the most practical.
I also agree with Sam that rejection sampling would likely need to operate on entire blocks rather than individual lines. By the time an LLM generates a line containing a regular expression, it’s often already committed to that path—for example, it might have skipped importing required modules or creating the necessary variables to pursue an alternative solution.
SAEBench: A Comprehensive Benchmark for Sparse Autoencoders
Evaluating Sparse Autoencoders with Board Game Models
I agree. In particular, there’s a huge advantage to the defender due to the scale of the model weights. All the defense has to do is raise the bar high enough that an adversary can’t exfiltrate model weights during the lifetime of the model.
If the adversary gains access to the weak model, it still reduces the possible encoded information density, as I discuss here. I haven’t done the experiments, but I would guess this reduction improves significantly if the adversary doesn’t have access to the weak model.
Various ways of changing the weak model weights or retraining the weak model with new data could also increase the size of the defender’s advantage.
I think it does stack if the perplexity filter is applied before the compression step, which relies on an underlying assumption that it’s easier to secure the internet connection choke point rather than the entire data center.
I will try to continue the discussion here for simplicity.
The perplexity filter works best if the attacker doesn’t have complete control of the entire data center, but instead limited access to some machines. An underlying assumption is that it’s easier to secure the internet connection choke point, rather than everything going on inside the data center.
This can be stacked with compression by applying the perplexity filter before the compression step.
Assuming that the attacker has access to the perplexity filter model, it would still reduce the possible information density if they have to perform some sort of arithmetic encoding scheme that has low perplexity. I had coworkers working on exactly this problem of encoding information into natural language using arithmetic encoding and GPT-2 small, and there’s a major trade-off between the perplexity of the language generated and the density of the encoded information, although I don’t have numbers available to quantify this trade-off.
Another potential option (that could stack with text compression) is to filter outgoing data. In particular, if we only allow outgoing data to be language tokens, and we filter the outgoing data to have low perplexity as measured by a small, secret language model, then it could become very difficult to encode terabytes of model weights in a way that has low perplexity.
I discuss this idea more here.
Thanks for this comment, by the way! I added a paragraph to the beginning to make this post more clear.
The purpose of this proposal is to limit anyone from transferring model weights out of a data center. If someone wants to steal the weights and give them to China or another adversary, the model weights have to leave physically (hard drive out of the front door) or through the internet connection. If the facility has good physical security, then the weights have to leave through the internet connection.
If we also take steps to secure the internet connection, such as treating all outgoing data as language tokens and using a perplexity filter, then the model weights can be reasonably secure.
We don’t even have to filter all outgoing data. If there was 1 Gigabyte of unfiltered bandwidth per day, it would take 2,000 days to transfer GPT-4′s 2 terabytes of weights out (although this could be reduced by compression schemes).
Using an LLM perplexity filter to detect weight exfiltration
I would guess that it would learn an exact algorithm rather than heuristics. The challenging part for OthelloGPT is that the naive algorithm to calculate board state from input tokens requires up to 60 sequential steps, and it only has 8 layers to calculate the board state and convert this to a probability distribution over legal moves.
I think it’s pretty plausible that this is true, and that OthelloGPT is already doing something that’s somewhat close to optimal within the constraints of its architecture. I have also spent time thinking about the optimal algorithm for next move prediction within the constraints of the OthelloGPT architecture, and “a bag of heuristics that promote / suppress information with attention to aggregate information across moves” seems like a very reasonable approach.
In Othello, pieces must be played next to existing pieces, and the game is initialized with 4 pieces in the center. Thus, it’s impossible for the top left corner to be played within the first 5 moves, and extremely unlikely in the early portion of a randomly generated game.
A $1 training run would be training 6 SAEs across 6 sparsities at 16K width on Gemma-2-2B for 200M tokens. This includes generating the activations, and it would be cheaper if the activations are precomputed. In practice this seems like large enough scale to validate ideas such as the Matryoshka SAE or the BatchTopK SAE.