先を見据えて：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の活性化を忠実に分解する必要がありますが、解釈可能性を確保するためにはその分解がスパースでなければなりません。これら2つの目的はしばしば相反します。本論文では、JumpReLU SAEを紹介します。これは、Gemma 2 9Bの活性化において、Gated SAEやTopK SAEなどの最近の進歩と比較して、特定のスパース性レベルで最先端の再構成忠実度を達成します。また、手動および自動化された解釈可能性研究を通じて、この改善が解釈可能性を犠牲にしないことを示します。JumpReLU SAEは、バニラ（ReLU）SAEのシンプルな修正版であり、ReLUを不連続なJumpReLU活性化関数に置き換えたもので、同様に効率的に訓練および実行できます。ストレートスルーエスティメーター（STE）を原理的に活用することで、SAEの順伝播に導入された不連続なJumpReLU関数にもかかわらず、JumpReLU SAEを効果的に訓練する方法を示します。同様に、L1などの代理指標を訓練する代わりに、STEを使用して直接L0をスパースに訓練し、収縮などの問題を回避します。

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.