Jumping Ahead: Improving Reconstruction Fidelity with JumpReLU Sparse Autoencoders

要約

スパース オートエンコーダ (SAE) は、言語モデル (LM) のアクティベーションにおいて因果関係があり、解釈可能な線形特徴を特定するための有望な教師なしアプローチです。
下流のタスクに役立つように、SAE は LM アクティベーションを忠実に分解する必要があります。
まだ解釈可能であるためには、分解はまばらである必要があります。つまり、緊張している 2 つの目的です。
この論文では、Gemma 2 9B アクティベーションで特定のスパース性レベルで最先端の再構成忠実度を達成する JumpReLU SAE を、Gated SAE や TopK SAE などの最近の進歩と比較して紹介します。
また、手動および自動の解釈可能性調査を通じて、この改善が解釈可能性を犠牲にするものではないことも示します。
JumpReLU SAE は、バニラ (ReLU) SAE を単純に変更したもの (ReLU を不連続な JumpReLU アクティベーション関数に置き換えたもの) であり、トレーニングと実行が同様に効率的です。
原理的な方法でストレートスルー推定器 (STE) を利用することにより、SAE のフォワード パスに導入された不連続な JumpReLU 関数にもかかわらず、JumpReLU SAE を効果的に訓練することがどのように可能であるかを示します。
同様に、L1 などのプロキシでトレーニングするのではなく、STE を使用して L0 をスパースになるように直接トレーニングし、縮小などの問題を回避します。

要約(オリジナル)

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.

arxiv情報

提供元, 利用サービス

arxiv.jp, Google