arXivDaily arXiv每日学术速递 周一至周五更新
arXiv 2407.14435cs.LG

跳跃前进:利用JumpReLU稀疏自编码器提升重构保真度

Jumping Ahead: Improving Reconstruction Fidelity with JumpReLU Sparse Autoencoders

Google DeepMind(谷歌DeepMind)

Senthooran Rajamanoharan, Tom Lieberum, Nicolas Sonnerat, Arthur Conmy, Vikrant Varma, János Kramár, Neel Nanda

AI总结:

本文提出JumpReLU稀疏自编码器,通过用不连续JumpReLU激活函数替换ReLU并利用直通估计器直接训练L0稀疏性,在Gemma 2 9B激活上以给定稀疏度实现最先进的重构保真度,且不损失可解释性。

摘要:

稀疏自编码器(SAEs)是一种很有前景的无监督方法,用于识别语言模型(LM)激活中具有因果相关性和可解释性的线性特征。为了对下游任务有用,SAEs需要忠实地分解LM激活;然而,为了可解释,分解必须是稀疏的——这两个目标之间存在张力。在本文中,我们引入了JumpReLU SAEs,在Gemma 2 9B激活上,与Gated和TopK SAEs等其他近期进展相比,它在给定稀疏度水平下实现了最先进的重构保真度。我们还通过人工和自动化可解释性研究证明,这种改进并未以牺牲可解释性为代价。JumpReLU SAEs是对普通(ReLU)SAEs的简单修改——我们将ReLU替换为不连续的JumpReLU激活函数——并且在训练和运行方面同样高效。通过以有原则的方式利用直通估计器(STEs),我们展示了尽管在SAE的前向传播中引入了不连续的JumpReLU函数,仍然可以有效地训练JumpReLU SAEs。类似地,我们使用STEs直接训练L0使其稀疏,而不是在L1等代理目标上训练,从而避免了收缩等问题。

英文摘要:

Sparse autoencoders (SAEs) are a promising unsupervised approach for identifying causally relevant and interpretable linear features in a language model's (LM) activations. To be useful for downstream tasks, SAEs need to decompose LM activations faithfully; yet to be interpretable the decomposition must be sparse -- two objectives that are in tension. In this paper, we introduce JumpReLU SAEs, which achieve state-of-the-art reconstruction fidelity at a given sparsity level on Gemma 2 9B activations, compared to other recent advances such as Gated and TopK SAEs. We also show that this improvement does not come at the cost of interpretability through manual and automated interpretability studies. JumpReLU SAEs are a simple modification of vanilla (ReLU) SAEs -- where we replace the ReLU with a discontinuous JumpReLU activation function -- and are similarly efficient to train and run. By utilising straight-through-estimators (STEs) in a principled manner, we show how it is possible to train JumpReLU SAEs effectively despite the discontinuous JumpReLU function introduced in the SAE's forward pass. Similarly, we use STEs to directly train L0 to be sparse, instead of training on proxies such as L1, avoiding problems like shrinkage.

机构、总结、中文摘要均由AI分析而来,可能存在错误,请以论文正文为准。

↑