Kronecker Factorization Improves Efficiency and Interpretability of Sparse Autoencoders
Vadim Kurochkin ⋅ Yaroslav Aksenov ⋅ Daniil Laptev ⋅ Nikita Balagansky ⋅ Daniil Gavrilov
Abstract
Sparse Autoencoders (SAEs) decompose language-model activations into sparse, interpretable features, but standard encoders usually treat the latent dictionary as a flat set of independent coordinates, leaving hierarchy and feature interactions to emerge only implicitly. We propose **KronSAE**, an encoder-side hierarchical SAE that factorizes the latent space into heads and forms post-latent features as pairwise compositions of lower-dimensional pre-latents using $\mathrm{mAND}$, a differentiable AND-like interaction. This imposes a compositional co-activation prior while remaining compatible with standard SAE objectives and variants such as TopK, Matryoshka, and Switch SAEs. Under equal-token-budget training, KronSAE achieves reconstruction quality comparable to strong baselines, improves several interpretability and absorption metrics, better reflects correlated feature structure, and yields encoder parameter and FLOP reductions as an additional benefit.
Chat is not available.
Successful Page Load