Vooruit Springen: Verbetering van Reconstructiekwaliteit met JumpReLU Sparse Autoencoders

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

Samenvatting

Sparse autoencoders (SAE's) vormen een veelbelovende onbewaakte aanpak voor het identificeren van causaal relevante en interpreteerbare lineaire kenmerken in de activaties van een taalmodel (LM). Om nuttig te zijn voor downstream taken, moeten SAE's de LM-activaties nauwkeurig ontbinden; maar om interpreteerbaar te zijn, moet de ontbinding spaarzaam zijn – twee doelstellingen die met elkaar in spanning staan. In dit artikel introduceren we JumpReLU SAE's, die state-of-the-art reconstructiegetrouwheid bereiken bij een bepaald spaarzaamheidsniveau op Gemma 2 9B-activaties, vergeleken met andere recente ontwikkelingen zoals Gated en TopK SAE's. We laten ook zien dat deze verbetering niet ten koste gaat van de interpreteerbaarheid, aan de hand van zowel handmatige als geautomatiseerde interpreteerbaarheidsstudies. JumpReLU SAE's zijn een eenvoudige aanpassing van standaard (ReLU) SAE's – waarbij we de ReLU vervangen door een discontinue JumpReLU-activatiefunctie – en zijn even efficiënt om te trainen en uit te voeren. Door straight-through-estimators (STE's) op een principiële manier te gebruiken, laten we zien hoe het mogelijk is om JumpReLU SAE's effectief te trainen ondanks de discontinue JumpReLU-functie die in de forward pass van de SAE wordt geïntroduceerd. Evenzo gebruiken we STE's om L0 direct spaarzaam te trainen, in plaats van te trainen op proxies zoals L1, waardoor problemen zoals krimp worden vermeden.

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.