Subspace-Aware Sparse Autoencoders for Effective Mechanistic Interpretability

원문 초록

Sparse Autoencoders (SAEs) are widely used for mechanistic interpretability in large language models, yet their standard formulation assigns each latent feature a single decoder direction, implicitly assuming features to be one-dimensional. We show that this assumption can induce feature splitting when model features are multi-dimensional. Geometrically, accurately reconstructing a feature under a fixed co-activation budget can require exponentially many decoder directions as its dimension grows. From an optimization perspective, adding decoder directions can strictly lower the $\ell_1$-regularized population risk, even when a basis can already reconstruct the feature exactly. Motivated by these results, we introduce Subspace-Aware Sparse Autoencoders (SASA), which replace single-vector decoders with learned subspaces, enforce sparsity through Top-$s$ group gating, and adapt group rank with nuclear-norm regularization. Under an idealized single-feature model, we prove that the SASA objective is minimized by a single group whose reconstruction map recovers the feature subspace. We also derive polynomial sample complexity bounds for subspace estimation within a group and a lower bound on the samples needed to observe every atom in a fixed fragmented dictionary. Experiments on GPT-2 and Mistral-7B show reduced absorption compared with SAEs. Matched-budget GPT-2 comparisons show lower absorption for concept families and more accurate prediction of their members using a small number of latent coordinates, at a reconstruction cost. The code is available at https://github.com/arshandalili/sasa.