SAEs, Part I: What the Heck is an SAE?

Have you ever tried to read a blogpost/paper/article on ML interpretability and wondered, what the heck is a sparse autoencoder?

If you have, this article is for you!

First off, what is a neural network?

A neural network (NN) is a computational model.

They’re called “neural” because they are loosely inspired by real neurons in mammalian brains.

So, can we give computers human intelligence with these “neural networks”?

Sort of!

ChatGPT/Claude et al are very large neural networks.

Why are they so big? Because people noticed that the bigger you made these networks, all else being equal, the smarter the models become. This is known as a “scaling law”.

It’s cool that we’ve figured out how to make these neural networks act smarter than humans in some domains!

But we also don’t understand how the heck they work so well.

Some researchers who study these models think it’s good for us to figure out what makes these models tick. This is known as model interpretability.

Large language models (LLMs) have only been popular for some years now, so interpretability is also a new field of study. Broadly speaking, we can study models behaviorally or mechanistically.

Mechanistic means we poke around the insides of the model.

Here’s one question we don’t understand about models,

Q: “Where does a model store its facts?” A: “Not in the neurons.”

For example, Anthropic found that a single example neuron corresponded to English academic citations, HTTP requests, and Korean text.

Should that be surprising? I’m not sure! On the one hand, human neurons don’t seem to store facts per neuron either. But also, coming back to our goal of understanding models better, this means we can’t just inspect neurons and understand their behavior. And unlike human brains, we can actually engineer neural networks, so it makes sense to aspire for structures that are more interpretable.

But before we go out and try to design new models based on what feels maybe promising, consider this: we already have really big and expensive models! And while understanding them better is sometimes a side goal, the main purpose for the institutions training these is to get them to be good. So we can’t just out and propose new designs just because they’re interpretable. (although, if you’re interested in this, check out Eleuther AI’s Pythia and AI2’s OLMo).

So let’s try and understand existing LLMs!

Anthropic researchers call this phenomenon of the same neuron representing Korean text and HTTP polysemanticity.

But what if our models were monosemantic instead? If the model is cramming lots of concepts into some number of dimensions, what if we expanded into a higher dimensional space, so each direction can represent a single concept cleanly?

Enter the sparse auto-encoder (SAE).

We also want to limit how all-over-the-place our autoencoder is.

If it’s always “active” for every concept, it’s not much easier to interpret than our original model. Ideally, only a small handful of hidden units should activate for any given input. This property is called sparsity. In matrix form, this looks like a lot of entries that are zero or nearly zero.

How do we get sparsity? Let’s look at our loss equation:

This measures how well the autoencoder rebuilds the original signal. We want to discourage larger values (which is the same as encouraging values to go to zero). So let’s add a sparsity penalty, an L1 norm on the activations that penalizes units for being active at all.

This is commonly expressed as:

And that’s how you train an SAE.

Now we understand how SAEs work! But how good are they for understanding large models? Stay tuned for my next post to find out!

p.s. if you liked this article, you might enjoy:

• This is a more in-depth / technical yet still intuitive explanation of how SAEs work: SAE Intuitions

• This paper by Anthropic goes over a theory of how polysemanticity occurs: Toy Models of Superposition