Adaptive Sparse Allocation with Mutual Choice & Feature Choice Sparse Autoencoders

要約

スパース オートエンコーダ (SAE) は、ニューラル ネットワークから特徴を抽出するための有望なアプローチであり、モデルの解釈可能性とモデル内部への因果的介入を可能にします。
SAE は、一連のトークンと特徴の一致を暗黙的に定義するスパース化活性化関数を使用して、スパースな特徴表現を生成します。
トークンと特徴のマッチングを、総スパース性の上限によって制約されるリソース割り当て問題として組み立てます。
たとえば、TopK SAE は、各トークンが最大 $k$ の特徴と一致するという追加の制約を使用して、この割り当ての問題を解決します。
TopK SAE では、一部のトークンは他のトークンよりも再構築が難しいにもかかわらず、トークンあたりの $k$ アクティブな機能の制約はどのトークンでも同じです。
この制限に対処するために、私たちは 2 つの新しい SAE バリアント、Feature Choice SAE と Mutual Choice SAE を提案します。これらはそれぞれ、トークンごとに可変数のアクティブな機能を許可します。
機能選択 SAE は、各機能が最大 $m$ トークンと一致するという追加の制約の下で、スパース性割り当ての問題を解決します。
Mutual Choice SAE は、総スパース性バジェットをトークンと機能間で自由に割り当てることができる、無制限の割り当て問題を解決します。
さらに、新しい補助損失関数 $\mathtt{aux\_zipf\_loss}$ を導入します。これは、$\mathtt{aux\_k\_loss}$ を一般化し、無効な機能や十分に活用されていない機能を軽減します。
私たちの方法では、固有の適応計算の結果として、デッド特徴が少なく、等価スパース性レベルでの再構築損失が改善された SAE が得られます。
より正確でスケーラブルな特徴抽出方法は、基礎モデルのより深い理解とより正確な制御への道を提供します。

要約(オリジナル)

Sparse autoencoders (SAEs) are a promising approach to extracting features from neural networks, enabling model interpretability as well as causal interventions on model internals. SAEs generate sparse feature representations using a sparsifying activation function that implicitly defines a set of token-feature matches. We frame the token-feature matching as a resource allocation problem constrained by a total sparsity upper bound. For example, TopK SAEs solve this allocation problem with the additional constraint that each token matches with at most $k$ features. In TopK SAEs, the $k$ active features per token constraint is the same across tokens, despite some tokens being more difficult to reconstruct than others. To address this limitation, we propose two novel SAE variants, Feature Choice SAEs and Mutual Choice SAEs, which each allow for a variable number of active features per token. Feature Choice SAEs solve the sparsity allocation problem under the additional constraint that each feature matches with at most $m$ tokens. Mutual Choice SAEs solve the unrestricted allocation problem where the total sparsity budget can be allocated freely between tokens and features. Additionally, we introduce a new auxiliary loss function, $\mathtt{aux\_zipf\_loss}$, which generalises the $\mathtt{aux\_k\_loss}$ to mitigate dead and underutilised features. Our methods result in SAEs with fewer dead features and improved reconstruction loss at equivalent sparsity levels as a result of the inherent adaptive computation. More accurate and scalable feature extraction methods provide a path towards better understanding and more precise control of foundation models.

arxiv情報

提供元, 利用サービス

arxiv.jp, Google