Aller de l'avant : Amélioration de la fidélité de reconstruction avec les autoencodeurs éparses JumpReLU

Jumping Ahead: Improving Reconstruction Fidelity with JumpReLU Sparse Autoencoders

July 19, 2024

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

cs.AI

Résumé

Les autoencodeurs parcimonieux (SAEs) constituent une approche non supervisée prometteuse pour identifier des caractéristiques linéaires causalement pertinentes et interprétables dans les activations d'un modèle de langage (LM). Pour être utiles dans des tâches en aval, les SAEs doivent décomposer fidèlement les activations du LM ; cependant, pour être interprétables, la décomposition doit être parcimonieuse — deux objectifs qui sont en tension. Dans cet article, nous introduisons les JumpReLU SAEs, qui atteignent une fidélité de reconstruction de pointe pour un niveau de parcimonie donné sur les activations de Gemma 2 9B, par rapport à d'autres avancées récentes telles que les SAEs à portes (Gated) et TopK. Nous montrons également que cette amélioration ne se fait pas au détriment de l'interprétabilité, grâce à des études manuelles et automatisées d'interprétabilité. Les JumpReLU SAEs sont une modification simple des SAEs classiques (ReLU) — où nous remplaçons la fonction d'activation ReLU par une fonction JumpReLU discontinue — et sont tout aussi efficaces à entraîner et à exécuter. En utilisant de manière raisonnée des estimateurs directs (STEs), nous montrons comment il est possible d'entraîner efficacement les JumpReLU SAEs malgré la fonction JumpReLU discontinue introduite dans la passe avant du SAE. De même, nous utilisons les STEs pour entraîner directement la norme L0 à être parcimonieuse, au lieu de s'appuyer sur des proxys comme L1, évitant ainsi des problèmes tels que le rétrécissement.

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.