[論文レビュー] Scaling and evaluating sparse autoencoders

Leo Gao, Tom Dupré la Tour+7 more|arXiv (Cornell University)|Jun 6, 2024

Generative Adversarial Networks and Image Synthesis被引用数 10

ひとこと要約

この OpenAI の研究は、再構成‑スパース性のトレードオフを分析するスケーラブルな TopK-sparse autoencoder フレームワークを開発し、新しい品質指標を導入し、GPT-4 活性化に対して 16M-latent autoencoder を訓練してスケーリング法則を研究する。

ABSTRACT

Sparse autoencoders provide a promising unsupervised approach for extracting interpretable features from a language model by reconstructing activations from a sparse bottleneck layer. Since language models learn many concepts, autoencoders need to be very large to recover all relevant features. However, studying the properties of autoencoder scaling is difficult due to the need to balance reconstruction and sparsity objectives and the presence of dead latents. We propose using k-sparse autoencoders [Makhzani and Frey, 2013] to directly control sparsity, simplifying tuning and improving the reconstruction-sparsity frontier. Additionally, we find modifications that result in few dead latents, even at the largest scales we tried. Using these techniques, we find clean scaling laws with respect to autoencoder size and sparsity. We also introduce several new metrics for evaluating feature quality based on the recovery of hypothesized features, the explainability of activation patterns, and the sparsity of downstream effects. These metrics all generally improve with autoencoder size. To demonstrate the scalability of our approach, we train a 16 million latent autoencoder on GPT-4 activations for 40 billion tokens. We release training code and autoencoders for open-source models, as well as a visualizer.

研究の動機と目的

• 言語モデルの活性化に対して、極めて広くてスパースな自動エンコーダを信頼性高く訓練する最先端の手法を開発する。
• GPT-2 と GPT-4 にまたがるスパース性、オートエンコーダのサイズ、対象モデルサイズのスケーリング法則を特徴づける。
• 潜在特征の品質を評価する新しい指標を提案・検証する。特徴回復、説明可能性、下流への影響を含む。

提案手法

• TopK (k-sparse) autoencoders を採用して、L1 ペナルティなしに直接スパース性を制御し、純粋な L2 ロスを用いて再構成性能を評価する。
• 初期化と補助損失を用いてデッド latent を防ぎ、大規模なスケールでもデッド latent レートを低く保つ（例：16M-latent モデルでデッド latent が 7%）
• GPT-2 small および GPT-4 ファミリーの活性化を横断して、autoencoder サイズ（n latents）とスパース性（k）をSystematically 拡大し、MSE とスパース性のスケーリング法則を導出する。
• TopK をベースラインの ReLU や他の活性化と比較し、スパース性-再構成前線の改善と活性化の収縮の低減を示す。
• MSE を超える評価指標を定義・計算する。下流損失、プローブ損失、説明可能性、およびアブレーション・スパース性を含み、特徴品質を評価する。
• GPT-4 残差ストリームで 40B tokens を対象に 16M-latent autoencoder を訓練してスケーラビリティを示す。

Figure 1 : Scaling laws for TopK autoencoders trained on GPT-4 activations. (Left) Optimal loss for a fixed compute budget. (Right) Joint scaling law of loss at convergence with fixed number of total latents $n$ and fixed sparsity (number of active latents) $k$ . Details in Section 3 .

実験結果

リサーチクエスチョン

• RQ1 スパース性とオートエンコーダのサイズは、大規模言語モデル全体で再構成品質と潜在的有用性にどう相互作用するか。
• RQ2 より大きくよりスパースなオートエンコーダは、下流の予測性能と解釈可能な特徴回復を改善するか。
• RQ3 言語モデルで学習されるスパースなオートエンコーダの特徴の品質と解釈可能性を最もよく捉える指標は何か。
• RQ4 TopK-sparse autoencoders は大規模化におけるデッド latent を低減できるか、そしてこの選択は代替活性化とどう比較されるか。
• RQ5 対象モデルサイズ（GPT-2 対 GPT-4）と固定スパース性で、スケーリング法則はどう変化するか。

主な発見

• TopK autoencoders は sparsity-reconstruction frontier で ReLU および ProLU ベースラインを上回り、活性化例の単義性（monosemanticity）が向上する。
• デッド latent はエンコーダ-デコーダ初期化と補助損失を用いて大幅に削減。最大の 16M-latent モデルではデッド latent は 7% のみだった。
• MSE と autoencoder のサイズ n およびスパース性 k に対する明確なスケーリング法則が現れ、GPT-4 には joint L(n,k) の関係が含まれる。
• 大きなオートエンコーダは、GPT-2 small および GPT-4 の活性化全体で下流損失、プローブ損失、説明可能性の指標を一般に改善する。
• 40B tokens で訓練された 16M-latent GPT-4 autoencoder は、特定の比較で GPT-4 の事前学習計算の 10% に相当する下流性能を達成する。
• 新しい評価指標（下流損失、プローブ損失、説明可能性、アブレーション・スパース性）は、より大きく、よりスパースなモデルで改善を示す。

(a) At a fixed number of latents ( $n=32768$ ), TopK has a better reconstruction-sparsity trade off than ReLU and ProLU, and is comparable to Gated.