Revisiting End To End Sparse Autoencoder Training — A Short Finetune is All You Need

要約

スパース自動エンコーダー（SAE）は、言語モデルのアクティベーションの解釈に広く使用されています。
主要な評価メトリックは、モデルの活性化をSAEの再構成に置き換える際のエントロピー損失の増加です。
通常、SAEは、事前に計算されたシャッフルされたアクティベーションを使用して、平均二乗エラー（MSE）のみで訓練されています。
最近の研究では、KL発散とMSE（「エンドツーエンド」SAE）の組み合わせでトレーニングSAEを直接導入し、大幅に増加した計算のコストで再建の精度を大幅に改善し、広範囲にわたる採用が制限されています。
同等の改善を達成する最終的な25Mトレーニングトークン（典型的なトレーニング予算のほんの数パーセント）にのみ適用される短いKL+MSE微調整ステップを提案し、エントロピーの損失ギャップを20-50％削減し、最小限の追加の計算コストを帯びます。
さらに、複数の微調整方法（KL微調整、LORAアダプター、線形アダプター）が同様の非加法交差エントロピーの改善をもたらし、MSE訓練を受けたSAEで一般的で簡単に容易なエラーソースを示唆していることがわかります。
KLとMSEの損失の規模の違いにもかかわらず、ハイパーパラメーターとスパースのペナルティを効果的に転送するための簡単な方法を実証します。
ReluとTopk SAEの両方が大幅にエントロピーの喪失の改善を見ていますが、監視されたSaebenchメトリックの評価は混合結果をもたらし、実際の利点がSAEアーキテクチャと特定のダウンストリームタスクの両方に依存することを示唆しています。
それにもかかわらず、私たちの方法は、わずかな追加コストで回路分析などの解釈可能性アプリケーションの有意義な改善を提供します。

要約(オリジナル)

Sparse autoencoders (SAEs) are widely used for interpreting language model activations. A key evaluation metric is the increase in cross-entropy loss when replacing model activations with SAE reconstructions. Typically, SAEs are trained solely on mean squared error (MSE) using precomputed, shuffled activations. Recent work introduced training SAEs directly with a combination of KL divergence and MSE (‘end-to-end’ SAEs), significantly improving reconstruction accuracy at the cost of substantially increased computation, which has limited their widespread adoption. We propose a brief KL+MSE fine-tuning step applied only to the final 25M training tokens (just a few percent of typical training budgets) that achieves comparable improvements, reducing the cross-entropy loss gap by 20-50%, while incurring minimal additional computational cost. We further find that multiple fine-tuning methods (KL fine-tuning, LoRA adapters, linear adapters) yield similar, non-additive cross-entropy improvements, suggesting a common, easily correctable error source in MSE-trained SAEs. We demonstrate a straightforward method for effectively transferring hyperparameters and sparsity penalties despite scale differences between KL and MSE losses. While both ReLU and TopK SAEs see significant cross-entropy loss improvements, evaluations on supervised SAEBench metrics yield mixed results, suggesting practical benefits depend on both SAE architecture and the specific downstream task. Nonetheless, our method offers meaningful improvements in interpretability applications such as circuit analysis with minor additional cost.

arxiv情報

提供元, 利用サービス

arxiv.jp, Google