SHIFTとTPPを使ったスパースオートエンコーダの評価

スパースオートエンコーダー（SAE）は、ニューラルネットワークの動作を理解する手助けをし、活性化を分解してわかりやすい部分にします。この分野の大きな問題は、SAEのパフォーマンスを測定するための良い方法がないことです。過去の研究のほとんどは、あまり信頼性のない方法を使用していました。この研究では、SHIFTという手法を使ってSAEを評価する新しい方法を提案します。SHIFTは、ニューラルネットワークのどの部分がタスクに役立っていないかを特定するのに役立ちます。また、ターゲティッドプローブパータービュレーション（TPP）手法を導入し、SAEが似た概念をどれだけ区別できるかを測定します。

#課題

SAEはニューラルネットワークを理解するための便利なツールです。今年は、TopKやゲーテッドSAEなど、多くの新しいSAEの種類が開発されました。しかし、進捗を測定するための信頼できるメトリクスがないという大きな問題があります。他の機械学習タスクは明確な目標がありますが、SAEの解釈可能性を評価するのには明確な基準がありません。

通常のメトリクスであるスパース性や忠実度は、モデルをよりよく理解するという点で私たちが求めるものと必ずしも一致しません。この乖離により、SAEの改善が本当に解釈可能性を向上させているのか、単に代理メトリクスが改善されているのかがわかりにくくなります。

#解決策

これに対処するために、SAEが訓練以外のタスクでどれだけ機能するかに基づいて評価することを提案します。SHIFT手法は、SAEが偏った予測に寄与するモデルの部分を特定し、排除できるかを評価します。SHIFTを使用することで、研究者はニューラルネットワークの出力に影響を与える特徴とそうでない特徴を確認できます。私たちは、異なる概念を分離するSAEの有効性を評価するために、SHIFTに基づいたスパリオスコリレーションリムーバル（SCR）という新しい評価を作成しました。

しかし、SCRはさまざまなデータタイプに対してスケールする際に限界があります。これを克服するために、特定のクラスを識別して変更し、他のクラスには影響を与えないようなSAEを探るTPP手法を開発しました。SCRとTPPの両方において、分類タスクにどれだけ影響を与えるかを反映したスコアを使用してSAEの特徴を選択します。

#方法と貢献

私たちの主な貢献は：

1. SHIFTの適応：SHIFTのスパリオスコリレーションリムーバルタスクをSAEの評価ツールとして機能するように調整しました。
2. TPPの導入：様々なデータセット上でSAEを評価するために、ターゲティッドプローブパータービュレーションメトリクスを開発しました。
3. オープンソーススイート：SAEのコレクションを訓練し、さまざまな言語モデルやデータセットを使ってメトリクスをテストしました。

SAEはニューラルネットワークの内部動作から理解可能な特徴のセットを見つけることを目的としています。良いSAEはモデルのプロセスに忠実であり、人間が理解できる概念を分離できる必要があります。

従来、人々はSAEを評価するために主に2つの無監督メトリクスを使用してきました：

1. クロスエントロピ損失の回復：これは、SAEの予測を使って元のモデルのパフォーマンスをどれだけ模倣できるかを確認します。
2. 特徴の活性化のL0ノルム：これは、特定の入力に対してどれだけの特徴が活性化されるかを測定します。

最近の研究では、ボードゲーム、回路、特定の言語概念を使用してSAEを評価することに注目しています。概念除去の目的は、モデルから不要なアイデアを見つけて排除し、その全体的なパフォーマンスを維持することです。私たちの目標は、概念を取り除くための現在の方法を改善することではなく、これらのタスクをSAEの進捗を評価するメトリクスに変えることです。

#概念の分離評価

この研究では、SAEが異なる概念をどれだけうまく分離できるかを、その品質の主な指標として焦点を当てます。私たちの方法をテストするために、体系的なアプローチに従います：

1. 特定の概念のために分類器を訓練します。
2. その概念に関連するSAEの特徴を特定します。
3. 概念に関連する特徴を除去することが、分類器にどのように影響するかを確認します。

良いSAEは、関連する特徴を削除したときに分類器の精度に大きな影響を与えるはずです。私たちのSHIFTとTPPメトリクスはこのアイデアを具体化します。

#妥当性とサニティチェック

私たちのメトリクスが有効であることを確認するために、いくつかのテストを実施して、期待されるSAEの特性と一致するかどうかを確認します。以下の各サブセクションでは評価手順を詳述し、詳細は付録で提供します。

#SAE潜在選択

どのSAEの特徴を評価するかを選ぶには、特定の概念にとってどの特徴が最も関連性が高いかを見極める必要があります。私たちはこれを、分類器に対する影響をランク付けすることによって行い、理解可能性のためにこれらの特徴をフィルタリングすることがあります。

最も関連性のある特徴を見つけるために、線形分類器を用いてモデル出力からの関連を探ります。各特徴がどれだけ貢献しているかを反映したスコアを集め、上位のものを選択します。さらに、特徴が活性化される文脈に基づいて理解可能かどうかを評価するためにLLMジャッジを使用します。

#SHIFTとTPPの適用

SHIFT手法には、テキストを2つのバイナリラベルに結びつけるデータセットが必要です。私たちは、職業と性別の分類にバイアスインバイオスデータセット、製品カテゴリと評価にアマゾンレビューのデータセットを使用します。

両方のデータセットを2つのラベルでフィルタリングし、バイアスのかかったデータセットで分類器を訓練します。前述のプロセスを使って分類器から特徴を削除し、バイアスなしで分類器がどれだけ機能するかを確認します。

TPPアプローチはSHIFTを一般化し、任意のテキスト分類データセットに対応します。ここでは、クラスを区別するのに役立つSAEの特徴を見つけ、削除がモデル精度にどれだけ影響を与えるかを確認します。

#実験結果

Pythia-70MとGemma-2-2Bの2つのモデルでSAEを訓練し、私たちのメトリクスをテストしました。両方のメトリクスは、SAEがバイアスを効果的に除去し、分類器の精度を向上させることを示しています。SHIFT評価は、さまざまなSAEの種類やアーキテクチャを区別します。

#発見

結果は一貫して、TopKとJumpReLUのアーキテクチャがスタンダードSAEを上回っていることを示しています。また、SAEのパフォーマンスは訓練中に改善され、訓練の最初の段階が全体的なスコア向上に大きく寄与しています。

私たちの発見は、選択方法に関わらず、ほとんどの上位SAE特徴がLLMジャッジによって解釈可能と見なされることを示しています。ノイズ情報に基づく方法はLLMを必要とせず、より速く、良好な評価を提供します。

#討論と限界

私たちの実験は、SHIFTとTPPが異なるSAEアーキテクチャを成功裏に区別することを確認しました。しかし、各メトリクスに対して最適なスパース性レベルは異なります。TPPメトリクスをスパース性測定に関連付けるためには、さらに作業が必要です。

使用したLLMジャッジは、他の実装と比べて解釈可能性に対する基準が低いです。私たちの簡単な方法はより速く、コストがかかりませんが、いくつかの解釈を見落とす可能性があります。したがって、LLMジャッジを使用するかどうかを決定するときには、品質と効率のバランスが必要です。

SHIFTとTPPは、SAEが学ぶべき理想を人間が設定することに依存しており、これはモデルが実際に表すものと一致しない場合があります。この依存関係は重要な特徴を見落とす可能性があります。

強みがあるにもかかわらず、両方のメトリクスには複雑さや未定義のパラメーターに関する限界があります。これらは他の評価方法を補完するものであり、単独の測定手段として機能するべきではありません。

#結論

SHIFTおよびTPP手法は、スパースオートエンコーダーを評価するための貴重なツールを提供します。これらはさまざまなデータセットに簡単に適用でき、訓練中の改善を示し、迅速に計算できます。研究者は、私たちのメトリクスを利用して自分のSAEを評価し、訓練の進捗を追跡することをお勧めします。

#謝辞

この研究は、MLアライメント理論奨学金プログラムの支援を受けました。このプロジェクトにおいて貢献してくださった皆さんの洞察と専門知識に感謝します。さらに、さまざまな研究所から提供された計算資源にも感謝します。

#今後の方向性

今後は、因果的な分離だけでなく、SAEの他の重要な特性をカバーする評価を改善することを目指します。SAEの品質のすべての側面を検討する包括的なフレームワークの開発が、重要な課題であることを認識しています。

#プローブ訓練の洞察

バイアスのあるデータセット上でプローブを訓練する際には、検出された信号のバランスを取ることが重要です。プローブが一つのラベルに偏りすぎると、望ましくない特徴を除去する効果が制限されます。バッチサイズや学習率を調整することで、プローブの精度に大きな影響を与えることができました。

データセットのラベルへの依存を最小限に抑えるために、複数のクラスペアでスコアを平均化しました。両クラスの正確さが60%以上のペアを選択することで、評価の信頼性を向上させることができました。

#スパースオートエンコーダーの訓練手順

Pythia-70MおよびGemma-2-2Bモデルに基づいて、さまざまなSAEを訓練し、利用可能にしています。私たちの訓練パラメータは、異なるデータセットで良好な特徴の識別を確保することを目指しています。

私たちの発見をもとに、SAE評価方法のさらなる研究を奨励し、これらのモデルがどのように機能し改善されていくのかの理解を深めることを望んでいます。