Jacobian Sparse Autoencoders: Sparsify Computations, Not Just Activations

TL;DR

This paper introduces Jacobian Sparse Autoencoders (JSAEs) that sparsify not only activations but also the computations in large language models, enabling better understanding of their internal processes while maintaining performance.

Contribution

JSAEs are a novel method that efficiently sparsify the Jacobian matrices in LLMs, revealing computational sparsity and improving interpretability over traditional SAEs.

Findings

JSAEs extract significant computational sparsity while preserving model performance.

Jacobian matrices serve as effective proxies for computational sparsity.

Pre-trained LLMs exhibit greater computational sparsity than randomized models.

Abstract

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\"ive 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.