arXiv preprint ·

Shaped SAEs: Expanding the Toolkit for Mechanistic Interpretability

Oveys Delafrooz Noroozi, Bhagyashree Puranik and Upamanyu Madhow

Published
Source
arXiv preprint, posted before peer review
Original
Read it at arXiv
Topics
Machine Learning

Abstract

Sparse autoencoders (SAEs) seek to decompose the dense activations of large language models (LLMs) into sparse, potentially interpretable features using a learned overcomplete dictionary. In this work, we ask whether it is possible to enhance the semantic selectivity of the learned dictionary atoms by linearly \emph{shaping} the dense activations prior to learning an SAE. We explore rank-reducing transformations based on the recently introduced J-space corresponding to the averaged Jacobian from the current layer to the LLM output, projecting activations onto the ``live'' subspace spanned by eigenmodes of the Jacobian with larger singular values (which we expect to be more semantically relevant). We also introduce a new approach for training SAEs in which (a)~an affine shaping transformation is learned along with the dictionary atoms, (b)~the atom selection mechanism is based on the TopK projections along the one-dimensional subspace spanned by the atoms, and (c)~reconstruction is carried out using the same atoms. We call this approach a Matched Filter SAE (MF-SAE) based on analogies with communication theory. The promise of shaping is demonstrated via evaluations against dense activations and state-of-the-art SAEs on benchmarks including SAEBench and SAE Probes.

The authors’ abstract, as arXiv lists it. arXiv’s listings, abstracts included, are free for anyone to reuse.

Read it on ReadSearch

Each opens after you create an account, and brings you back to this paper.

Already a reader? Sign in

All new papers in Machine Learning

Get the few papers worth your time in Machine Learning every morning.

ReadSearch reads the new papers in your field each day and brings you the few worth reading, each with its key finding in a sentence.

Papers belong to their authors and publishers; every page here links to the original. All research topics