Jacobian Sparse Autoencoders: Sparsify Computations, Not Just Activations

作者: Lucy Farnik, Tim Lawson, Conor Houghton, Laurence Aitchison

发布时间: 2025-02-27

来源: arxiv

研究方向: 机器学习,深度学习,可解释人工智能

主要内容

本文提出了 Jacobian Sparse Autoencoders (JSAEs),一种用于发现大型语言模型 (LLMs) 中稀疏计算的新方法。该方法通过优化计算图(Jacobian 矩阵)的稀疏性,旨在更深入地理解 LLMs 的内部计算过程。

主要贡献

1. 提出了 Jacobian Sparse Autoencoders (JSAEs),一种结合了稀疏激活和稀疏计算的自动编码器。

2. 开发了一种高效计算 Jacobian 矩阵的方法,使得在 LLMs 中应用 JSAEs 成为可能。

3. 证明了 JSAEs 能够在保持下游 LLM 性能的同时,显著提高计算图的稀疏性。

4. 表明了 Jacobians 是计算稀疏性的合理代理,因为 MLPs 在 JSAE 基础下近似线性。

5. 展示了 JSAEs 在预训练 LLMs 上比随机 LLMs 实现了更高的计算稀疏度,表明计算图的稀疏性是 LLMs 通过训练学习到的属性。

6. 提供了 JSAEs 的源代码,方便其他研究人员使用和扩展。

研究方法

1. 稀疏自动编码器 (SAEs)

2. Jacobian 矩阵计算

3. TopK 激活函数

4. 自动微分

5. 线性近似

实验结果

实验结果表明,JSAEs 成功地在输入和输出 SAE 潜在激活之间引入了稀疏的 Jacobian 矩阵,同时保持了重建质量和模型性能。此外,JSAEs 在预训练 LLMs 上实现了更高的计算稀疏度,表明它们可以更好地理解学习到的 Transformer 计算结构。

未来工作

未来工作将包括将 JSAEs 扩展到其他类型的神经网络,如 GLUs,以及将它们应用于整个模型以发现整个模型中的计算结构。此外,研究如何将 JSAEs 与其他可解释人工智能技术相结合,以更全面地理解 LLMs 的内部工作原理。