Train Sparse Auto-encoders Efficiënt door Gebruik te Maken van Kenmerken Correlatie

Train Sparse Autoencoders Efficiently by Utilizing Features Correlation

May 28, 2025

Auteurs: Vadim Kurochkin, Yaroslav Aksenov, Daniil Laptev, Daniil Gavrilov, Nikita Balagansky

cs.AI

Samenvatting

Sparse Autoencoders (SAEs) hebben aanzienlijke belofte getoond in het interpreteren van de verborgen toestanden van taalmodel(len) door deze te ontbinden in interpreteerbare latente richtingen. Het trainen van SAEs op grote schaal blijft echter een uitdaging, vooral wanneer grote woordenboekgroottes worden gebruikt. Hoewel decoders gebruik kunnen maken van sparse-aware kernels voor efficiëntie, vereisen encoders nog steeds rekenintensieve lineaire operaties met grote uitvoerdimensies. Om dit aan te pakken, stellen we KronSAE voor, een nieuwe architectuur die de latente representatie factoriseert via Kronecker-productontbinding, waardoor het geheugen- en rekenoverhead aanzienlijk wordt verminderd. Bovendien introduceren we mAND, een differentieerbare activatiefunctie die de binaire AND-bewerking benadert, wat de interpreteerbaarheid en prestaties in ons gefactoriseerde raamwerk verbetert.

English

Sparse Autoencoders (SAEs) have demonstrated significant promise in interpreting the hidden states of language models by decomposing them into interpretable latent directions. However, training SAEs at scale remains challenging, especially when large dictionary sizes are used. While decoders can leverage sparse-aware kernels for efficiency, encoders still require computationally intensive linear operations with large output dimensions. To address this, we propose KronSAE, a novel architecture that factorizes the latent representation via Kronecker product decomposition, drastically reducing memory and computational overhead. Furthermore, we introduce mAND, a differentiable activation function approximating the binary AND operation, which improves interpretability and performance in our factorized framework.