[论文解读] 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 自编码器框架，用于分析重建稀疏性权衡，提出新的质量指标，并在 GPT-4 激活上训练了一个 16M-latent 自编码器以研究缩放定律。

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）自编码器，直接控制稀疏性，避免 L1 惩罚，使用纯 L2 损失评估重建性能。
• 使用初始化和辅助损失来防止死潜在变量，即便在大规模下也能实现较低的死潜在率（例如，16M-latent 模型的死潜在率为 7%）。
• 系统地在 GPT-2 小型和 GPT-4 系列激活上，按 n 潜在变量和稀疏度 k 对自编码器大小进行系统性放大，以推导 MSE 和稀疏性的缩放定律。
• 将 TopK 与基线 ReLU 及其他激活进行比较，显示在稀疏性-重建前沿上的改进以及激活收缩的减少。
• 定义并计算超越 MSE 的评估指标，包括下游损失、探针损失、可解释性，以及消融稀疏性，以评估特征质量。
• 在 GPT-4 残差流上对 40B tokens 训练一个 16M-latent 自编码器，以展示可扩展性。

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 自编码器能否在大规模下缓解死潜在变量？与其他激活相比，这种选择如何？
• RQ5 在固定稀疏性下，随着目标模型大小（GPT-2 与 GPT-4）的变化，缩放定律如何改变？

主要发现

• TopK 自编码器在稀疏性-重建前沿上优于 ReLU 和 ProLu 基线，在激活示例上具有更好的单义性。
• Dead latents are substantially reduced using encoder-decoder initialization and an auxiliary loss; in the largest 16M-latent model, only 7% latents remained dead.
• Clean scaling laws emerge for MSE with respect to autoencoder size n and sparsity k, including a joint L(n,k) relationship for GPT-4.
• Larger autoencoders generally improve downstream loss, probe loss, and explainability metrics across GPT-2 small and GPT-4 activations.
• A 16M-latent GPT-4 autoencoder trained on 40B tokens achieves downstream performance comparable to 10% of GPT-4 pretraining compute in a specific comparison.
• New evaluation metrics (downstream loss, probe loss, explainability, ablation sparsity) show improvements with larger and sparser models.

(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.