プロダクトキーメモリースパースコーダー

EleutherAIの研究チームは、プロダクトキーメモリー（PKM）を用いてスパースコーダーを改善する手法についてのブログ記事を公開しました。スパースコーダーはAIモデルの解釈や特徴抽出に重要ですが、高い再構成誤差が課題です。従来のTopKスパースコーダーは活性化関数で最も活性化する特徴を選択しますが、エンコーダ部分のパラメータと計算負荷が大きいという欠点があります。PKMは入力次元を二つの部分空間に分割し、すべての半重みの組み合わせを考慮することで、探索速度を保ちつつより大きな潜在特徴空間を実現し、エンコーダのパラメータ数を削減します。

実験では、SmolLM 2 135Mモデルの複数層に対してPKMトランスコーダーと通常のスキップトランスコーダー（SST）を比較しました。その結果、PKMトランスコーダーは拡大係数（特徴数と入力次元の比率）が256倍までの範囲で、同等の再構成品質を持つTopKモデルよりも訓練が速いことが分かりました。しかし、拡大係数が512倍に達すると、PKMの訓練時間が大幅に増加し、再構成損失も同規模のSSTに劣ります。PKMのエンコーダパラメータが少ないため、訓練時に大きなバッチサイズを使用できますが、極端な拡大係数ではバッチサイズを小さくせざるを得ず、その利点が相殺されます。

解釈可能性に関しては、PKMトランスコーダーの方が優れています。Llama 3.1 70Bモデルを用いた自動解釈評価では、PKMの検出スコアとファジングスコアがベースラインを上回りました。さらに、PKMの潜在変数は自然なグループを形成し、同じグループ内の特徴方向や説明テキストの類似度が高く、特徴分析に役立ちます。研究のコードとモデルはオープンソースとして公開されており、チームはこれらの発見がスパースコーダーの発展に貢献することを期待しています。