[Feature] Add support for label-supervised SAEs, and evaluate hypothesis generation results on Congress dataset

[rmovva]

Feature Request

Recall that in HypotheSAEs, our goal is to hypothesize interpretable concepts present in input texts (e.g. Yelp reviews) that predict a target variable (e.g. ratings).

Description

Currently:

• In Step 1, we train the SAE solely to reconstruct the input text embeddings: this step doesn't use our labels at all.
• In Step 2, we use the labels to perform feature selection.

Instead:

• In Step 1, we could use the labels by training the SAE to use its sparse encoding to perform both reconstruction and target prediction.
• Step 2 would remain the same. As long as we have the reconstruction loss, many of the neurons will not be relevant to the target variable, so we'll still need to select a subset.
• The relative weighting of the two losses would be a hyperparameter.

Motivation

The hope is that this would encourage more of the SAE's latent neurons to be relevant to the target (e.g., more neurons about different reasons why the service was bad, and fewer neurons about how many times the review used words starting with 'D').

Proposed Solution

One possible implementation:

• Add a SupervisedSparseAutoencoder class that extends SparseAutoencoder from [sae.py].
• This will require re-writing the following functions: forward, compute_loss, and fit.
• The forward pass should compute the latent SAE encoding, and then use the encoding for two tasks: first, to reconstruct the input embedding, and second, to predict y.
• User will need to pass in labels y and a weight alpha which specifies weight on the supervised loss (set a reasonable default for this based on experiments, ideally).
• Add a unit test to make sure the implementation works in [test_workflow.py].

Benchmarking

• Add a notebook running this implementation on the Congress dataset. (Headlines is pairwise, so it may be difficult to adapt the implementation for it; Yelp is relatively "easy", so not as useful.)
• Follow the structure of [experiment_reproduction.ipynb].
• If you don't have access to an OpenAI key and would need one to benchmark your implementation, feel free to reply to the issue with a link to your fork showing your progress so far, and I can give you a key.

Notes

I am happy to advise on this. If you have an implementation in progress on a fork, feel free to reply asking for a review or any questions.

[Nuna7]

Hey @rmovva , I try working on this feature enhancement. I create a new class SupervisedSparseAutoencoder and the corresponding train_supervised_sae function. The main addition is using the output of the SAE (the activation) performing regression (technically, I guess as no softmax or sigmoid is used as of now) using single linear layer and use MSE as loss (not BCE or CE).

I have download the [hein-daily.zip](https://stacks.stanford.edu/file/druid:md374tz9962/hein-daily.zip) (guess this is the congress dataset) and try train the SupervisedSAE on binary classification task of very small subset. I use MSE for binary classification because I think this is what you use for feature selection as per the paper. I cannot go forward much as of now, as I don't have access to gpu (using vllm in mac is messy) and could only use small 1-2 Billion model.

If you are interested, you could review and add anything to fix. I could make PR as well.

[rmovva]

@Nuna7 thanks for working on this! I didn't see the changes on your fork, can you push?

I think we'd want to look at some results first before adding to the repo, but totally understand if you can't run it. I had a chance to try implementing this too recently and didn't see a clear difference on Yelp, but would be interested to see more results on other datasets, which I can look into in the future.

BTW, the exact data we used are available on HF. Also, we do use BCE for binary classification as per footnote 3 on page 4, sorry, that was easy to miss!

[Nuna7]

@rmovva Hey, I actually already push the codes to my forked repo on different branch feature/supervised_sae and also add the BCE and CE support. You can mention anything you would like to perform.