提出新型稀疏自编码器,让大模型的计算过程也变稀疏。
Jacobian Sparse Autoencoders: Sparsify Computations, Not Just Activations
- 通过雅可比矩阵实现输入输出与计算路径的双重稀疏化
- 在保持模型性能前提下,计算稀疏度显著提升
- 适合研究大模型内部学习到的计算规律
稀疏自编码器(SAEs)已被成功用于发现大语言模型(LLMs)隐层激活的稀疏且可解释的表示。然而,我们最终希望理解的是模型执行的计算过程,而不仅仅是其表示。现有SAEs仅能稀疏化激活值,无法反映计算本身的稀疏性。为此,本文提出雅可比稀疏自编码器(JSAEs),不仅能稀疏化模型组件的输入/输出激活,还能稀疏化连接它们的计算(即雅可比矩阵)。针对大模型中雅可比矩阵计算量过大的问题,本文提出一种高效计算方法。实验表明,JSAEs在保持下游任务性能的同时,实现了较高的计算稀疏度;且在重写后的基底下,前馈网络近似线性,验证了雅可比作为计算稀疏代理的合理性。更重要的是,预训练模型的计算稀疏度高于随机模型,说明计算图稀疏性是模型训练所习得的特性,因此JSAEs更适用于解析变压器模型的深层计算机制。
原文摘要 · Abstract (English)
Sparse autoencoders (SAEs) have been successfully used to discover sparse and human-interpretable representations of the latent activations of LLMs. However, we would ultimately like to understand the computations performed by LLMs and not just their representations. The extent to which SAEs can help us understand computations is unclear because they are not designed to "sparsify" computations in any sense, only latent activations. To solve this, we propose Jacobian SAEs (JSAEs), which yield not only sparsity in the input and output activations of a given model component but also sparsity in the computation (formally, the Jacobian) connecting them. With a naïve implementation, the Jacobians in LLMs would be computationally intractable due to their size. One key technical contribution is thus finding an efficient way of computing Jacobians in this setup. We find that JSAEs extract a relatively large degree of computational sparsity while preserving downstream LLM performance approximately as well as traditional SAEs. We also show that Jacobians are a reasonable proxy for computational sparsity because MLPs are approximately linear when rewritten in the JSAE basis. Lastly, we show that JSAEs achieve a greater degree of computational sparsity on pre-trained LLMs than on the equivalent randomized LLM. This shows that the sparsity of the computational graph appears to be a property that LLMs learn through training, and suggests that JSAEs might be more suitable for understanding learned transformer computations than standard SAEs.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。