Resource	Info
Paper	http://arxiv.org/abs/2406.04093
Code & Data	https://github.com/openai/sparse_autoencoder
Public	arXiv
Date	2024.09.04

Summary Overview

Openai 关于 SAE 的一篇工作。提出了 TopK SAE。Sparse autoencoder (SAE) 提供了一种具有前景的无监督方法，通过从神经网络层中重建激活来提取具有可解释性特征。同时还介绍了一些新指标来评估特征质量。

Main Content

三个问题：

• whether certain hypothesized features were recovered
• whether downstream effects are sparse
• whether features can be explained with both high precision and recall

Methods:

Setup

Wec choose a layer near the end of the network, which should contain many features without being specialized for next-token predictions.

Evaluation: sparsity $L_0$, reconstruction meansquared error (MSE)

Baseline: ReLU autoencoder

$\begin{aligned}&z=\text{ReLU}(W_\text{enc}(x-b_\text{pre})+b_\text{enc})\\&\hat{x}=W_\text{dec}z+b_\text{pre}\end{aligned}$

The training loss is defined by $\mathcal{L}=\left|\left|x-\hat{x}\right|\right|^2_2+\lambda\left|\left|z\right|\right|_1$, where $\left|\left|x-\hat{x}\right|\right|_2^2$ is reconstruction MSE. $\left|\left|z\right|\right|_1$ is an $L_1$ penalty promoting sparsity in latent activations $z$, and $\lambda$ is a hyperarameter that need to be tuned.

TopK activation function

Using an activation function (TopK) that only keeps the $k$ largest latents, zeroning the rest.

$z=\mathrm{TopK}(W_{\mathrm{enc}}(x-b_{\mathrm{pre}}))$

The training loss is simply $\mathcal{L}=\left|\left|x-\hat{x}\right|\right|_2^2$. No penalty anymore!

Using $k$-sparse autoencoder has a number of benefits:

• In remove the need for the $L_1$ penalty. $L_1$ is an imperfect approximation of $L_0$, and it introduces a bias of shrinking all positive activations towward zero.
• It enables setting the $L_0$ directly, as opposed to tuning an $L_1$ coefficient $\lambda$, enabling simpler model comparison and rapid iteration. It can also be used in combination with arbitrary activation functions
• It empirically outperforms baseline ReLU autoencoders on the sparsity-reconstruction frontier, and this gap increases with scale.
• It increases monosemanticity of random activating examples by effectively clamping small activations to zero.

Preventing dead latents

We find two important ingredients for preventing dead latents:

• intialize the encoder to the transpose of the decoder
• use an auxiliary loss that models reconstruction errorusing the top-$k_\mathrm{aux}$ dead latents.

Using thesetechniques, even in our largest (16 million latent autoencoder only 7% of latents are dead).

Evaluation

Following metrics:

• Downstream loss: How good is the language model loss if the residual stream latent is replaced with the autoencoder reconstruction of that latent?
• Probe loss: Do autoencoders recover features that we believe they might have?
• Explainability: Are there simple explanations that are both necessary and sufficient for the activation of the autoencoder latent?
• Ablation sparsity: Does ablating individual latents have a sparse effect on downstream logits?

Finding simple explaination for features

Neuron to Graph (N2G), a substantially less expressive but much cheaper method that outputs explanations in the form of collections of n-grams with wildcards.

Limitations and Future Directions

• TopK forces every token to use exactly $k$ latents, which is likely suboptimal. Ideally we would constrain $\mathbb{E}\left[L_0\right]$ rather than $L_0$.
• The optimization can likely be greatly improved, for example with learning rate scheduling, better optimizers, and better aux losses for preventing dead latents.
• Much more could be done to understand what metrics best track relevance to downstream applications, and to study those applications themselves. Applications include: finding vectors for steering behavior, doing anomaly detection, identifying circuits, and more.
• We're excited about work in the direction of combining MoE and autoencoders, which would substantially improve the asymptotic cost of autoencoder training, and enable much larger autoencoders.
• A large fraction of the random activations of features we find, especially in GPT-4, are not yet adequately monosemantic. We believe that with improved techniques and greater scale this is potentially surmountable.
• Our probe based metric is quite noisy, which could be improved by having a greater breadth of tasks and higher quality tasks.
• While we use N2G for its computational efficiency, it is not only able to capture very simple patterns. We believe there is a lot of room for improvement in terms of more expressive explanation methods that are also cheap enough to simulate to estimate explanation precision.
• A context length of 64 tokens is potentially too few tokens to exhibit the most interesting behaviors of GPT-4.

Related Work

Malllat 和 Zhang 引入了在超完整字典上的稀疏编码。Olshausen 和 Field 通过无监督的方法从数据中学习字典改善了此方法。这种方法在图像处理中影响极大。之后 Hinton 和 Salakhutdinov 提出了自动编码器结构以降低维度。结合这些概念，人们开发了 SAE 以训练具有稀疏先验的自动编码器，例如使用$L_1$惩罚来提取稀疏特征。Makhzani 和 Frey 通过引入 K-Sparse autoencoder，使用 TopK 激活功能而不是 $L_1$惩罚 来完善这一概念。 Makelov 等人使用度量标准来自动评估编码器，该指标可以测量从先前发现的回路中恢复特征。

最近，将 SAE 用于语言模型，并且有多种 SAE 在开源的小模型上进行训练。Mark 等人发现SAE 可以从语言模型中发现回路。Wright 和 Sharkey 指出 SAE 会受到 $L_1$惩罚激活缩小。Taggart 和 Rajamanoharan 建议使用不同的激活来解决 SAE 中的激活收缩问题。Braun 等人建议在下游 KL 上训练 SAE，而不是重建 MSE。

Kaplan 等人研究了语言模型的缩放定律，这些定律发掘了损失如何随各种超参数而变化。Clark 等人使用双线性你和探索与稀疏有关的缩放定律。Lindsey 等人研究了专门针对自动编码器的缩放定律，将损失定义为重建和稀疏性的特定平衡（而不是简单地重建）。

🤖ChatGPT

Others