[I'm writing this quickly because the results are really strong. I still need to do due diligence & compare to baselines, but it's really exciting!]
Last post I found 600+ features in MLP layer-2 of Pythia-70M, but ablating the direction didn't always make sense nor the logit lens. Now I've found 3k+ in the residual stream, and the ablation effect is intuitive & strong.
The majority of features found were single word or combined words (e.g. not & 't ) which, when ablated, effected bigram statistics, but I also wanted to showcase the more interesting ones:
German Feature?
Uniform examples:
Logit lens:
Stack Exchange
Ablated Text (w/ red showing removing that token drops activation of that feature)
In the Pile dataset, stackexchange text is formatted as "Q:/n/n....", and removing the "Q" greatly decreasing the feature activation is evidence for this indeed detecting stackexchange format.
Title Case
Ablating Direction (in text):
Here, the log-prob of predicting " Of" went down by 4 when I ablate this feature direction, along w/ other Title Case texts.
Logit lens
One for the Kids
Last Name:
Ablated Text:
As expected, removing the first name affects the feature activation
Logit Lens:
More specifically, this looks like the first token of the last name, which makes since to follow w/ Jr or Sr or completing last names w/ "hoff" & "worth". Additionally, the "(\@" makes sense for someone's social media reference (e.g. instagram handle).
[word] and [word]
Ablated Text:
Logit Lens:
alike & respectively make a lot of since in context, but not positive about the rest.
Beginning & End of First Sentence?
On
logit lens:
"on behalf", "oncology", "onyx" (lol), "onshore", "onlook", "onstage", "onblur"(js?), "oneday" ("oneday" is separated as on eday?!?), "onloaded" (js again?), "oncomes"(?)
High Level Picture
I could spend 30 minutes (or hours) developing very rigorous hypotheses on these features to really nail down what distribution of inputs activate it & what effect it really has on the output, but that's not important.
What's important is that we have a decomposition of the data that:
1. Faithfully represents the original data (in this case, low reconstruction loss w/ minimal effect on perplexity)
2. Can simply specify desired traits in the model (e.g. honesty, non-deception)
There is still plenty of work to do to ensure the above two work (including normal hyperparameter search), but it'd be great to have more people working on this (I don't have time to mentor, but I can answer questions & help get you set up) if that's helping our work or integrating it in your work.
This work is legit though, and there's a recent paper that finds similar features in a BERT model.
Setup:
Model: Pythia-70m (actually named 160M!)
Transformer lens: "blocks.2.hook_resid_post" (so layer 2)
Data: Neel Nanda's Pile-10k (slice of pile, restricted to have only 25 tokens, same as last post)
Dictionary_feature sizes: 4x residual stream ie 2k (though I have 1x, 2x, 4x, & 8x, which learned progressively more features according to the MCS metric)
Uniform Examples: separate feature activations into bins & sample from each bin (eg one from [0,1], another from [1,2])
Logit Lens: The decoder here had 2k feature directions. Each direction is size d_model, so you can directly unembed the feature direction (e.g. the German Feature) you're looking at. Additionally I subtract out several high norm tokens from the unembed, which may be an artifact of the pythia tokenizer never using those tokens (thanks Wes for mentioning this!)
Ablated Text: Say the default feature (or neuron in your words) activation of Token_pos 10 is 5, so you can remove all tokens from 0 to 10 one at a time and see the effect on the feature activation. I select the token pos by finding the max feature activating position or the uniform one described above. This at least shows some attention head dependencies, but not more complicated ones like (A or B... C) where removing A or B doesn't effect C, but removing both would.
[Note: in the examples, I switch between showing the full text for context & showing the partial text that ends on the uniformly-selected token]