attention-saes-paper/attention-saes-paper-weights

Attention Output SAEs Paper

This is the repository for Attention Output SAEs for GeLU-2L, GPT-2 Small and Gemma-2B. More details are available in the accompanying paper.

Important details of SAE training include:

• SAE Widths. Our GELU-2L and Gemma-2B SAEs have width 16384. All of our GPT-2 Small SAEs have width 24576, with the exception of layers 5 and 7, which have width 49152.
• Loss Function. We trained our Gemma-2B SAE with a different loss function than the SAEs from other models. For Gemma-2B we closely follow the approach from Olah et al., while for GELU-2L and GPT-2 Small, we closely follow the approach from Bricken et al.
• Training Data. We use activations from hundreds of millions to billions of activations from LM forward passes as input data to the SAE. Following Nanda, we use a shuffled buffer of these activations, so that optimization steps don’t use data from highly correlated activations. For GELU-2L we use a mixture of 80% C4 webtext and 20% code (https://huggingface.co/datasets/NeelNanda/c4-code-tokenized-2b). For GPT-2 Small we use OpenWebText (https://huggingface.co/datasets/Skylion007/openwebtext). For Gemma-2B we use https://huggingface.co/datasets/HuggingFaceFW/fineweb.