Effizientes Training von Sparse Autoencodern durch Nutzung von Merkmalskorrelationen

Train Sparse Autoencoders Efficiently by Utilizing Features Correlation

May 28, 2025

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

cs.AI

Zusammenfassung

Sparse Autoencoders (SAEs) haben ein erhebliches Potenzial bei der Interpretation der verborgenen Zustände von Sprachmodellen gezeigt, indem sie diese in interpretierbare latente Richtungen zerlegen. Das Training von SAEs in großem Maßstab bleibt jedoch eine Herausforderung, insbesondere wenn große Wörterbuchgrößen verwendet werden. Während Decoder sparsitätsbewusste Kernel zur Effizienzsteigerung nutzen können, erfordern Encoder weiterhin rechenintensive lineare Operationen mit großen Ausgabedimensionen. Um dies zu adressieren, schlagen wir KronSAE vor, eine neuartige Architektur, die die latente Darstellung durch Kronecker-Produkt-Zerlegung faktorisiert und dadurch den Speicher- und Rechenaufwand drastisch reduziert. Darüber hinaus führen wir mAND ein, eine differenzierbare Aktivierungsfunktion, die die binäre AND-Operation approximiert und die Interpretierbarkeit und Leistung in unserem faktorisierten Framework verbessert.

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.