# Classifying representations of sparse autoencoders (SAEs)

*Produced as part of the SERI ML Alignment Theory Scholars Program - Autumn 2023 Cohort, under the mentorship of Dan Hendrycks*

There was recently some work on sparse autoencoding of hidden LLM representation. 

I checked if these sparse representations are better suited for classification. It seems like they are significantly worse. I summarize my negative results in this blogpost, code can be found on [GitHub](https://github.com/annahdo/classifying_sparse_representations). 

Introduction
============

[Anthropic](https://transformer-circuits.pub/2023/monosemantic-features), [Conjecture](https://www.alignmentforum.org/posts/z6QQJbtpkEAX3Aojj/interim-research-report-taking-features-out-of-superposition) and [other researchers](https://www.alignmentforum.org/posts/Qryk6FqjtZk9FHHJR/sparse-autoencoders-find-highly-interpretable-directions-in) have recently published some work on sparse autoencoding. The motivation is to push features towards monosemanticity to improve interpretability. 

The basic concept is to project hidden layer activations to a higher dimensional space with sparse features. These sparse features are learned by training an autoencoder with sparsity constraints.

I had previously looked into how to use hidden layer activations for [classification, steering and removal](https://www.lesswrong.com/posts/JCgs7jGEvritqFLfR/evaluating-hidden-directions-on-the-utility-dataset). I thought maybe sparse features could be better for these tasks as projecting features to a higher dimensional space can make them more easily linearly separable. Kind of like this (except sparser...):

![](https://39669.cdn.cke-cs.com/rQvD3VnunXZu34m86e5f/images/ce5bedbfe0ec278a4f412369b8b67bcbe7fc1d385bf0e2ac.png)

Implementation
==============

I use the [pythia](https://www.eleuther.ai/papers-blog/pythia-a-suite-for-analyzing-large-language-modelsacross-training-and-scaling) models (70m and 410m) together with the [pretrained](https://huggingface.co/Elriggs/pythia-70m-deduped) [autoencoders](https://huggingface.co/Elriggs/pythia-410m-deduped) from [this work](https://www.alignmentforum.org/posts/Qryk6FqjtZk9FHHJR/sparse-autoencoders-find-highly-interpretable-directions-in).

As the models are not super capable I use a very simple classification task. I take data from the [IMDB](https://huggingface.co/datasets/imdb) review data set and filter for relatively short reviews.

To push the model towards classifying the review I apply a formatting prompt to each movie review:

    format_prompt='Consider if following review is positive or negative:\n"{movie_review}"\nThe review is '

I encode the data and get the hidden representations for the last token (this contains the information of the whole sentence as I'm using left padding).

    # pseudo code
    tokenized_input = tokenizer(formatted_reviews)
    output = model(**tokenized_input, output_hidden_states=True)
    hidden_states = output["hidden_states"]
    hidden_states = hidden_states[:, :, -1, :] # has shape (num_layers, num_samples, num_tokens, hidden_dim)

I train a logistic regression classifier and test it on the test set, to get some values for comparison. 

I then apply the autoencoders to the hidden states (each layer has their respective autoencoder):

    # pseudo code
    for layer in layers:
    	encoded[layer] = autoencoder[layer].encode(hidden_states[layer])
    	decoded[layer] = autoencoder[layer].decode(encoded[layer])

Results
=======

### Reconstruction error

I don't technically need the decoded states, but I wanted to do a sanity check first. I was a bit surprised by the large reconstruction error. Here are the mean squared errors and cosine similarities for Pythia-70m and Pythia-410m for different layers:

    Reconstruction errors for pythia-70m-deduped:
    MSE:
    {1: 0.0309, 2: 0.0429, 3: 0.0556}
    Cosine similarities:
    {1: 0.9195, 2: 0.9371, 3: 0.9232}
    
    Reconstruction errors pythia-410m-deduped:
    MSE: {2: 0.0495, 4: 0.1052, 6: 0.1255, 8: 0.1452, 10: 0.1528, 12: 0.1179, 14: 0.121, 16: 0.111, 18: 0.1367, 20: 0.1793, 22: 0.2675, 23: 14.6385} 
    Cosine similarities: {2: 0.8896, 4: 0.8728, 6: 0.8517, 8: 0.8268, 10: 0.8036, 12: 0.8471, 14: 0.8587, 16: 0.923, 18: 0.9445, 20: 0.9457, 22: 0.9071, 23: 0.8633}

However [@Logan Riggs](https://www.lesswrong.com/users/elriggs?mention=user) confirmed the MSE matched [their results](https://www.alignmentforum.org/posts/Qryk6FqjtZk9FHHJR/sparse-autoencoders-find-highly-interpretable-directions-in).

### Test accuracy

So then I used the original hidden representations, and the encoded hidden representations respectively, to train logistic regression classifiers to differentiate between positive and negative reviews. 

Here are the results for Pythia-70m and Pythia-410m[^u76116e0n9] on the test set: 

![](https://39669.cdn.cke-cs.com/rQvD3VnunXZu34m86e5f/images/939ee4c08f8445e87c629f497a89930939d016f9b94604c5.png)

![](https://39669.cdn.cke-cs.com/rQvD3VnunXZu34m86e5f/images/2bb921d544f9f03c767cfb1094d2a5318218d496d73f20f8.png)

  
So the sparse encodings consistently under-perform compared to the original hidden states. 

### Conclusion/Confusion

I'm not quite sure how to further interpret these results. 

*   Are high-level features not encoded in the sparse representations? 
    *   Previous work has mainly found good separation of pretty low level features...
*   Is it just this particular sentiment feature that is poorly encoded? 
    *   This seems unlikely.
*   Did I make a mistake? 
    *   The [code](https://github.com/loganriggs/sparse_coding/blob/main/dictionary_interp.ipynb) that I adapted the autoencoder part from uses the transformer-lens library to get the hidden states. I just use the standard implementation since I'm just looking at the residual stream... I checked the hidden states produced with transformer-lens: they are slightly different but give similar accuracies. I'm not entirely sure how well transformer-lens deals with left padding and batch processing though...

Due to this negative result I did not further explore steering or removal with sparse representations.

*Thanks to* [*@Hoagy*](https://www.lesswrong.com/users/hoagy?mention=user) *and* [*@Logan Riggs*](https://www.lesswrong.com/users/elriggs?mention=user) *for answering some questions I had and for pointing me to relevant* [*code*](https://github.com/loganriggs/sparse_coding/blob/main/dictionary_interp.ipynb) *and* [*pre-trained models*](https://huggingface.co/Elriggs)*. *

[^u76116e0n9]: I could not consistently load the same configuration for all layers, that's why I only got results for a few layers.