Ускорение: улучшение точности восстановления с помощью разреженных автокодировщиков с функцией активации JumpReLU.

Jumping Ahead: Improving Reconstruction Fidelity with JumpReLU Sparse Autoencoders

July 19, 2024

Авторы: Senthooran Rajamanoharan, Tom Lieberum, Nicolas Sonnerat, Arthur Conmy, Vikrant Varma, János Kramár, Neel Nanda

cs.AI

Аннотация

Разреженные автокодировщики (SAE) представляют собой многообещающий метод без учителя для выявления причинно значимых и интерпретируемых линейных признаков в активациях модели языка (LM). Для того чтобы быть полезными для последующих задач, SAE должны верно декомпозировать активации LM; однако для интерпретируемости декомпозиция должна быть разреженной - две цели, которые находятся в напряжении. В данной статье мы представляем JumpReLU SAE, которые достигают передовой точности восстановления на определенном уровне разреженности на активациях Gemma 2 9B по сравнению с другими недавними достижениями, такими как Gated и TopK SAE. Мы также показываем, что это улучшение не происходит за счет интерпретируемости через ручные и автоматизированные исследования интерпретируемости. JumpReLU SAE представляют собой простое изменение обычных (ReLU) SAE - где мы заменяем ReLU на разрывную активационную функцию JumpReLU - и аналогично эффективны в обучении и выполнении. Используя прямые оценщики (STEs) в принципиальной манере, мы показываем, как можно эффективно обучать JumpReLU SAE, несмотря на разрывную функцию JumpReLU, введенную в прямом проходе SAE. Аналогично мы используем STE для прямого обучения L0 на разреженность, вместо обучения на прокси, такие как L1, избегая проблем, таких как уменьшение.

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.