Compute Optimal Inference and Provable Amortisation Gap in Sparse Autoencoders

TL;DR

This paper demonstrates that simple sparse autoencoders are inherently limited in accurate sparse inference, and shows that more expressive inference methods significantly improve interpretability and performance, especially in large language models.

Contribution

The paper proves the limitations of SAE encoders for sparse inference and empirically shows that advanced inference methods enhance interpretability and accuracy.

Findings

SAE encoders are insufficient for accurate sparse inference.

More expressive inference methods outperform traditional SAE encoders.

Enhanced inference improves interpretability in large language models.

Abstract

A recent line of work has shown promise in using sparse autoencoders (SAEs) to uncover interpretable features in neural network representations. However, the simple linear-nonlinear encoding mechanism in SAEs limits their ability to perform accurate sparse inference. Using compressed sensing theory, we prove that an SAE encoder is inherently insufficient for accurate sparse inference, even in solvable cases. We then decouple encoding and decoding processes to empirically explore conditions where more sophisticated sparse inference methods outperform traditional SAE encoders. Our results reveal substantial performance gains with minimal compute increases in correct inference of sparse codes. We demonstrate this generalises to SAEs applied to large language models, where more expressive encoders achieve greater interpretability. This work opens new avenues for understanding neural network representations and analysing large language model activations.