[論文レビュー] BatchTopK Sparse Autoencoders

Bart Bussmann, Patrick Leask+1 more|arXiv (Cornell University)|Dec 9, 2024

Generative Adversarial Networks and Image Synthesis被引用数 4

ひとこと要約

この論文は、バッチ全体の活性化の上位k個を選択することで、従来のTopK SAEsの1サンプルあたりの上位k個の活性化制約を緩和するBatchTopK Sparse Autoencodersを提案する。これにより、入力の複雑さに応じてサンプルごとの潜在変数活性化数を可変化でき、平均スパarsityを犠牲にせずに再構成性能を向上させることができる。また、直接的にスパarsityレベルを指定可能となり、ハイパーパramータス윕のコストを回避する。性能はTopKを上回り、JumpReLU SAEsと同等の再構成性能を達成する。

ABSTRACT

Sparse autoencoders (SAEs) have emerged as a powerful tool for interpreting language model activations by decomposing them into sparse, interpretable features. A popular approach is the TopK SAE, that uses a fixed number of the most active latents per sample to reconstruct the model activations. We introduce BatchTopK SAEs, a training method that improves upon TopK SAEs by relaxing the top-k constraint to the batch-level, allowing for a variable number of latents to be active per sample. As a result, BatchTopK adaptively allocates more or fewer latents depending on the sample, improving reconstruction without sacrificing average sparsity. We show that BatchTopK SAEs consistently outperform TopK SAEs in reconstructing activations from GPT-2 Small and Gemma 2 2B, and achieve comparable performance to state-of-the-art JumpReLU SAEs. However, an advantage of BatchTopK is that the average number of latents can be directly specified, rather than approximately tuned through a costly hyperparameter sweep. We provide code for training and evaluating BatchTopK SAEs at https://github.com/bartbussmann/BatchTopK

研究の動機と目的

• TopK SAEsにおける固定された1サンプルあたりのスパarsity制限を是正すること。これは、入力の複雑さに関係なく、すべての入力が同じ数の活性化潜在変数を必要とするという仮定に基づく。
• 入力の複雑さに応じてサンプルごとの活性化数を可変化させることで、平均スパarsityを増加させることなく再構成性能を向上させること。
• 平均スパarsityレベルの直接指定を可能とし、他のSAEバリアントで一般的なコストの高いハイパーパramータース윕を回避すること。
• バッチレベルの活性化選択を活用することで、解釈可能性を維持しつつ再構成性能を向上させること。
• GPT-2 SmallおよびGemma 2 2Bを用いて、標準的な指標を用いてTopKおよびJumpReLU SAEsと性能を比較すること。

提案手法

• BatchTopKは、1サンプルごとの上位k個の選択を、バッチ全体の上位k個の活性化を選択するバッチレベルの上位k個選択に置き換える。
• 勾配伝搬を可能とするため、上位k個の操作の微分可能緩和を用いることで、潜在辞書のエンドツーエンド最適化を可能にする。
• 推論時、複数のバッチにおける最小の正の活性化の平均として、グローバルなしきい値θを推定し、バッチ依存の選択をしきい値ベースの活性化関数に置き換える。
• しきい値θは、複数のバッチにおける最小の正の活性化の期待値として計算され、一貫した推論動作を保証する。
• 標準的なSAEエンコーダ・デコーダ構造にReLU活性化関数とL2再構成損失を用い、死活潜在変数の防止を目的としたスパarsityペナルティと補助損失を追加する。
• BatchTopKの補助損失は、TopK SAEsと同一であり、死活潜在変数の小規模な再構成を用いてその有用性を維持する。

Figure 1: On GPT-2 Small activations, BatchTopK largely achieves better NMSE and CE than standard TopK across different dictionary sizes, for a fixed number of active latents of 32 (Left). JumpReLU SAEs are omitted from this comparison as their L0 cannot be fixed to a value. For fixed dictionary siz

実験結果

リサーチクエスチョン

• RQ1 1サンプルあたりの上位k個の制約をバッチレベルの制約に緩和することで、スパースオートエンコーダーの再構成性能が向上するか？
• RQ2 1サンプルあたりの活性化潜在変数の数を可変化させることで、平均スパarsityを増加させずに再構成性能が向上するか？
• RQ3 BatchTopKは、スパarsityレベルの直接指定が可能でありながら、最先端のJumpReLU SAEsと同等の性能を達成できるか？
• RQ4 BatchTopKと固定上位k個の方法との間で、1サンプルあたりの活性化潜在変数の分布にはどのような違いがあるか？
• RQ5 BatchTopKの向上した柔軟性が、交差エントロピー劣化という観点で、より優れた下流性能をもたらすか？

主な発見

• GPT-2 SmallおよびGemma 2 2Bの両モデルにおいて、テストされたすべての辞書サイズおよびスパarsityレベルで、BatchTopK SAEsはTopK SAEsよりも顕著に低い正規化平均二乗誤差（NMSE）を達成する。
• 活性化潜在変数数を固定（L0=32）した場合、両モデルにおいてBatchTopKはNMSEおよび交差エントロピー（CE）劣化の両面でTopKを上回る。
• GPT-2 Smallでは、高活性化領域における下流のロバスト性に優れるJumpReLU SAEsでさえ、すべてのスパarsityレベルでBatchTopKがNMSEおよびCE劣化の両面で上回られる。
• Gemma 2 2Bでは、BatchTopKは再構成性能（NMSE）においてJumpReLUと同等またはそれを上回るが、CE劣化では唯一最も低いスパarsityレベル（k=16）でのみ上回る。
• BatchTopKにおける1サンプルあたりの活性化潜在変数の分布は広範なばらつきを示しており、一部のサンプルでは1つの潜在変数しか使用せず、他のサンプルでは80以上を用いるものも存在し、この方法の適応的活性化戦略を確認する。
• BatchTopKは、JumpReLU SAEsとは異なり、ハイパーパramータース윕を経由してターゲットスパarsityを達成する必要がなく、平均スパarsityレベルを直接指定可能である。

Figure 2: On Gemma 2 2B activations, BatchTopK outperforms TopK SAEs across different values of k. Although BatchTopK has a better reconstruction performance (left), it only outperforms JumpReLU in terms of downstream CE degradation in the setting where k=16 (right).