Avançando: Melhorando a Fidelidade da Reconstrução com Autoencoders Esparsos JumpReLU

Jumping Ahead: Improving Reconstruction Fidelity with JumpReLU Sparse Autoencoders

July 19, 2024

Autores: Senthooran Rajamanoharan, Tom Lieberum, Nicolas Sonnerat, Arthur Conmy, Vikrant Varma, János Kramár, Neel Nanda

cs.AI

Resumo

Os autoencoders esparsos (SAEs) são uma abordagem não supervisionada promissora para identificar características lineares causalmente relevantes e interpretáveis em ativações de um modelo de linguagem (LM). Para serem úteis para tarefas subsequentes, os SAEs precisam decompor as ativações do LM de forma fiel; no entanto, para serem interpretáveis, a decomposição deve ser esparsa - dois objetivos que estão em tensão. Neste artigo, apresentamos os SAEs JumpReLU, que alcançam fidelidade de reconstrução de última geração em um nível de esparsidade dado nas ativações do Gemma 2 9B, em comparação com outros avanços recentes, como os SAEs Gated e TopK. Também demonstramos que essa melhoria não compromete a interpretabilidade por meio de estudos de interpretabilidade manuais e automatizados. Os SAEs JumpReLU são uma modificação simples dos SAEs convencionais (ReLU) - onde substituímos o ReLU por uma função de ativação JumpReLU descontínua - e são igualmente eficientes para treinar e executar. Ao utilizar estimadores de passagem direta (STEs) de maneira fundamentada, mostramos como é possível treinar os SAEs JumpReLU de forma eficaz, apesar da função JumpReLU descontínua introduzida na passagem direta do SAE. Da mesma forma, usamos STEs para treinar diretamente L0 de forma esparsa, em vez de treinar em proxies como L1, evitando problemas como encolhimento.

English

Sparse autoencoders (SAEs) are a promising unsupervised approach for identifying causally relevant and interpretable linear features in a language model's (LM) activations. To be useful for downstream tasks, SAEs need to decompose LM activations faithfully; yet to be interpretable the decomposition must be sparse -- two objectives that are in tension. In this paper, we introduce JumpReLU SAEs, which achieve state-of-the-art reconstruction fidelity at a given sparsity level on Gemma 2 9B activations, compared to other recent advances such as Gated and TopK SAEs. We also show that this improvement does not come at the cost of interpretability through manual and automated interpretability studies. JumpReLU SAEs are a simple modification of vanilla (ReLU) SAEs -- where we replace the ReLU with a discontinuous JumpReLU activation function -- and are similarly efficient to train and run. By utilising straight-through-estimators (STEs) in a principled manner, we show how it is possible to train JumpReLU SAEs effectively despite the discontinuous JumpReLU function introduced in the SAE's forward pass. Similarly, we use STEs to directly train L0 to be sparse, instead of training on proxies such as L1, avoiding problems like shrinkage.