The Nonlinear Library AF - Attention SAEs Scale to GPT-2 Small by Connor Kissane
Feb 3, 2024
13:13
Welcome to The Nonlinear Library, where we use Text-to-Speech software to convert the best writing from the Rationalist and EA communities into audio. This is: Attention SAEs Scale to GPT-2 Small, published by Connor Kissane on February 3, 2024 on The AI Alignment Forum.
This is an interim report that we are currently building on. We hope this update + open sourcing our SAEs will be useful to related research occurring in parallel. Produced as part of the ML Alignment & Theory Scholars Program - Winter 2023-24 Cohort
Executive Summary
In a previous post, we showed that
sparse autoencoders (SAEs) work on the attention layer outputs of a two layer transformer. We scale our attention SAEs to GPT-2 Small, and continue to find sparse interpretable features in every layer. This makes us optimistic about our ongoing efforts scaling further, especially since we didn't have to do much iterating
We open source our SAEs. Load them from
Hugging Face or this
colab notebook
The SAEs seem good, often recovering more than 80% of the loss relative to zero ablation, and are sparse with less than 20 features firing on average. The majority of the live features are interpretable
We continue to find the same three feature families that we found in the
two layer model:
induction features,
local context features, and
high level context features. This suggests that some of our lessons interpreting features in smaller models may generalize
We also find new, interesting feature families that we didn't find in the two layer model, providing hints about fundamentally different capabilities in GPT-2 Small
See our
feature interface to browse the first 30 features for each layer
Introduction
In
Sparse Autoencoders Work on Attention Layer Outputs we showed that we can apply SAEs to extract sparse interpretable features from the last attention layer of a two layer transformer. We have since applied the same technique to a 12-layer model, GPT-2 Small, and continue to find sparse, interpretable features in every layer. Our SAEs often recover more than 80% of the loss[1], and are sparse with less than 20 features firing on average. We perform shallow investigations of the first 30 features from each layer, and we find that the majority (often 80%+) of non-dead SAEs features are interpretable.
interactive visualizations for each layer.
We open source our SAEs in hope that they will be useful to other researchers currently working on dictionary learning. We are particularly excited about using these SAEs to better understand attention circuits at the feature level. See the SAEs on
Hugging Face or load them using this
colab notebook. Below we provide the key metrics for each SAE:
L0 norm
loss recovered
dead features
% alive features interpretable
L0
3
99%
13%
97%
L1
20
78%
49%
87%
L2
16
90%
20%
95%
L3
15
84%
8%
75%
L4
15
88%
5%
100%
L5
20
85%
40%
82%
L6
19
82%
28%
75%
L7
19
83%
58%
70%
L8
20
76%
37%
64%
L9
21
83%
48%
85%
L10
16
85%
41%
81%
L11
8
89%
84%
66%
It's worth noting that we didn't do much differently to train these,[2] leaving us optimistic about the tractability of scaling attention SAEs to even bigger models.
Excitingly, we also continue to identify feature families. We find features from all three of the families that we identified in the two layer model: induction features, local context features, and high level context features. This provides us hope that some of our lessons from interpreting features in smaller models will continue to generalize.
We also find new, interesting feature families in GPT-2 Small, suggesting that attention SAEs can provide valuable hints about new[3] capabilities that larger models have learned. Some new features include:
Successor features, which activate when predicting the next item in a sequence such as "15, 16" -> "17" (which are partly coming from
Successor Heads in the model), and boost the logits of the next item.
Name mover features, which predict a name in the context, such as in the
IOI task
Duplicate token f...
This is an interim report that we are currently building on. We hope this update + open sourcing our SAEs will be useful to related research occurring in parallel. Produced as part of the ML Alignment & Theory Scholars Program - Winter 2023-24 Cohort
Executive Summary
In a previous post, we showed that
sparse autoencoders (SAEs) work on the attention layer outputs of a two layer transformer. We scale our attention SAEs to GPT-2 Small, and continue to find sparse interpretable features in every layer. This makes us optimistic about our ongoing efforts scaling further, especially since we didn't have to do much iterating
We open source our SAEs. Load them from
Hugging Face or this
colab notebook
The SAEs seem good, often recovering more than 80% of the loss relative to zero ablation, and are sparse with less than 20 features firing on average. The majority of the live features are interpretable
We continue to find the same three feature families that we found in the
two layer model:
induction features,
local context features, and
high level context features. This suggests that some of our lessons interpreting features in smaller models may generalize
We also find new, interesting feature families that we didn't find in the two layer model, providing hints about fundamentally different capabilities in GPT-2 Small
See our
feature interface to browse the first 30 features for each layer
Introduction
In
Sparse Autoencoders Work on Attention Layer Outputs we showed that we can apply SAEs to extract sparse interpretable features from the last attention layer of a two layer transformer. We have since applied the same technique to a 12-layer model, GPT-2 Small, and continue to find sparse, interpretable features in every layer. Our SAEs often recover more than 80% of the loss[1], and are sparse with less than 20 features firing on average. We perform shallow investigations of the first 30 features from each layer, and we find that the majority (often 80%+) of non-dead SAEs features are interpretable.
interactive visualizations for each layer.
We open source our SAEs in hope that they will be useful to other researchers currently working on dictionary learning. We are particularly excited about using these SAEs to better understand attention circuits at the feature level. See the SAEs on
Hugging Face or load them using this
colab notebook. Below we provide the key metrics for each SAE:
L0 norm
loss recovered
dead features
% alive features interpretable
L0
3
99%
13%
97%
L1
20
78%
49%
87%
L2
16
90%
20%
95%
L3
15
84%
8%
75%
L4
15
88%
5%
100%
L5
20
85%
40%
82%
L6
19
82%
28%
75%
L7
19
83%
58%
70%
L8
20
76%
37%
64%
L9
21
83%
48%
85%
L10
16
85%
41%
81%
L11
8
89%
84%
66%
It's worth noting that we didn't do much differently to train these,[2] leaving us optimistic about the tractability of scaling attention SAEs to even bigger models.
Excitingly, we also continue to identify feature families. We find features from all three of the families that we identified in the two layer model: induction features, local context features, and high level context features. This provides us hope that some of our lessons from interpreting features in smaller models will continue to generalize.
We also find new, interesting feature families in GPT-2 Small, suggesting that attention SAEs can provide valuable hints about new[3] capabilities that larger models have learned. Some new features include:
Successor features, which activate when predicting the next item in a sequence such as "15, 16" -> "17" (which are partly coming from
Successor Heads in the model), and boost the logits of the next item.
Name mover features, which predict a name in the context, such as in the
IOI task
Duplicate token f...