AI ALIGNMENT FORUM
AF

Noa Nabeshima — AI Alignment Forum

This is a special post for quick takes (aka "shortform"). Only the owner can create top-level comments.

Replying toMatryoshka Sparse Autoencoders

You know I was thinking ab this-- say that there are two children and they're orthogonal to the parent and each have probability 0.4 given the parent. If you imagine the space it looks like three clusters, two with probability 0.4, norm 1.4 and one with probability 0.2 and norm 1. They all have high cosine similarity with each other. From this frame, having the parent 'include' the children directions a bit doesn't seem that inappropriate. One SAE latent setup that seems pretty reasonable is to actually have one parent latent that's like "one of these three clusters is active" and three child latents pointing to each of the three clusters. The parent latent decoder in that setup would also include a bit of the child feature directions.

This is all sketchy though. It doesn't feel like we have a good answer to the question "How exactly do we want the SAEs to behave in various scenarios?"

Replying toMatryoshka Sparse Autoencoders

Noa Nabeshima1y

Matryoshka Sparse Autoencoders

This is cool! I wonder if it can be fixed. I imagine it could be improved some amount by nudging the prefix distribution, but it doesn't seem like that will solve it properly. Curious if this is a large issue in real LMs. It's frustrating that there aren't ground-truth features we have access to in language models.

I think how large of a problem this is can probably be inferred from a description of the feature distribution. It'd be nice to have a better sense of what that distribution is (assuming the paradigm is correct enough).

Replying toMatryoshka Sparse Autoencoders

Noa Nabeshima1y*

Matryoshka Sparse Autoencoders

Tree Methodology

To generate the trees in sparselatents.com/tree_view, I use a variant of Masked Cosine Similarity (MCS), a metric introduced in Towards Monosemanticity. The original MCS is calculated like this: For any two latents A and B, first compute the cosine similarity between their activations, but only considering tokens where latent A is active. Then compute the same similarity, but only for tokens where latent B is active. The final MCS value is the larger of these two similarities.

Instead of taking the max, I do a directed MCS where I just consider the cosine similarity between A and B's activations on tokens where B is active. Then, I multiply this directed MCS score... (read more)

Replying toMatryoshka Sparse Autoencoders

Noa Nabeshima1y

Matryoshka Sparse Autoencoders

The Matryoshka SAE trained on the toy model learn the true features on most runs, not all of them. Sometimes a small number of latents, modally one, seem to get stuck in a bad state.

Replying toMatryoshka Sparse Autoencoders

Noa Nabeshima1y

Matryoshka Sparse Autoencoders

I thought I had some informal evidence that permuting the latents was good and after double checking some evidence I don't feel confident that it is good.

Training without permutation seems to attain slightly better FVU/L0, has reasonable looking features at a quick glance, seems to solve the toy model at comparable rates to permuted, and is simpler to code.

Replying toLearning Multi-Level Features with Matryoshka SAEs

Noa Nabeshima1y*

Learning Multi-Level Features with Matryoshka SAEs

This is great work! I like that you tested on large models and your very comprehensive benchmarking. I also like the BatchTopK architecture.

It's interesting to me that MSE has a smaller hit than cross-entropy.

Here are some notes I made:

We suspect that using a fixed group size leads to more stable training and faster convergence.

This seems plausible to me!

Should the smallest sub-SAE get gradients from all losses, or should the losses from larger sub-SAEs be stopped?

When I tried stopping the gradient from flowing from large sub-SAE losses to small it made later latents much less interpretable. I tried an approach where early latents got less gradient from larger sub-SAE losses and it seemed... (read more)

Replying toMatryoshka Sparse Autoencoders

Noa Nabeshima1y*

Matryoshka Sparse Autoencoders

Even with all possible prefixes included in every batch the toy model learns the same small mixing between parent and children (this was best out of 2, for the first run the matryoshka didn't represent one of the features): https://sparselatents.com/matryoshka_toy_all_prefixes.png

Here's a hypothesis that could explain most of this mixing. If the hypothesis is true, then even if every possible prefix is included in every batch, there will still be mixing.

Hypothesis:

Regardless of the number of prefixes, there will be some prefix loss terms where
1. a parent and child feature are active
2. the parent latent is included in the prefix
3. the child latent isn't included in the prefix.

The

Noa Nabeshima1y

Matryoshka Sparse Autoencoders

That's very cool, I'm looking forward to seeing those results! The Top-K extension is particularly interesting, as that was something I wasn't sure how to approach.

I imagine you've explored important directions I haven't touched like better benchmarking, top-k implementation, and testing on larger models. Having multiple independent validations of an approach also seems valuable.

I'd be interested in continuing this line of research, especially circuits with Matryoshka SAEs. I'd love to hear about what directions you're thinking of. Would you want to have a call sometime about collaboration or coordination? (I'll DM you!)

Really looking forward to reading your post!

Replying toMatryoshka Sparse Autoencoders

Noa Nabeshima1y

Matryoshka Sparse Autoencoders

Yes, follow up work with bigger LMs seems good!

I use number of prefix-losses per batch = 10 here; I tried 100 prefixes per batch and the learned latents looked similar at a quick glance, so I wonder if naively training with block size = 1 might not be qualitatively different. I'm not that sure and training faster with kernels on its own seems good also!

Maybe if you had a kernel for training with block size = 1 it would create surface area for figuring out how to work on absorption when latents are right next to each other in the matryoshka latent ordering.

Matryoshka Sparse Autoencoders

Noa Nabeshima

View trees here
Search through latents with a token-regex language
View individual latents here
See code here (github.com/noanabeshima/matryoshka-saes)
Alternate version of this document with appropriate-height interactives.

Abstract

Sparse autoencoders (SAEs)^[1]^[2] break down neural network internals into components called latents. Smaller SAE latents seem to correspond to more abstract concepts while larger SAE latents seem to represent finer, more specific concepts.

While increasing SAE size allows for finer-grained representations, it also introduces two key problems: feature absorption introduced in Chanin et al. ^[3], where latents develop unintuitive "holes" as other latents in the SAE take over specific cases, and what I term fragmentation, where meaningful abstract concepts in the small SAE (e.g. 'female names' or 'words in quotes') shatter (via... (read 3078 more words →)

Replying toSparse Autoencoders Work on Attention Layer Outputs

Noa Nabeshima2y

Sparse Autoencoders Work on Attention Layer Outputs

I wonder if multiple heads having the same activation pattern in a context is related to the limited rank per head; once the VO subspace of each head is saturated with meaningful directions/features maybe the model uses multiple heads to write out features that can't be placed in the subspace of any one head.

Poll: Which variables are most strategically relevant?

Daniel Kokotajlo

Daniel Kokotajlo, Noa Nabeshima, Noa Nabeshima

Which variables are most important for predicting and influencing how AI goes?

Here are some examples:

Timelines: “When will crazy AI stuff start to happen?”
Alignment tax: “How much more difficult will it be to create an aligned AI vs an unaligned AI when it becomes possible to create powerful AI?
Homogeneity: "Will transformative AI systems be trained/created all in the same way?"
Unipolar / Multipolar: "Will transformative AI systems be controlled by one organization or many?"
Takeoff speeds: "Will takeoff be fast or slow (or hard or soft, etc.)?"

We made this question to crowd-source more entries for our list, along with operationalizations and judgments of relative importance. This is the first step of a... (read 155 more words →)

AI ALIGNMENT FORUM
AF

AI ALIGNMENT FORUM
AF

Noa Nabeshima

Noa Nabeshima

Noa Nabeshima's Shortform

Matryoshka Sparse Autoencoders

Poll: Which variables are most strategically relevant?

Noa Nabeshima

Noa Nabeshima

Noa Nabeshima's Shortform

Matryoshka Sparse Autoencoders

Poll: Which variables are most strategically relevant?

Tree Methodology

Abstract