Entrena Autoencoders Dispersos de Manera Eficiente Aprovechando la Correlación de Características

Train Sparse Autoencoders Efficiently by Utilizing Features Correlation

May 28, 2025

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

cs.AI

Resumen

Los Autoencoders Escasos (SAEs, por sus siglas en inglés) han demostrado un potencial significativo para interpretar los estados ocultos de los modelos de lenguaje al descomponerlos en direcciones latentes interpretables. Sin embargo, entrenar SAEs a gran escala sigue siendo un desafío, especialmente cuando se utilizan tamaños de diccionario grandes. Si bien los decodificadores pueden aprovechar núcleos optimizados para la escasez, los codificadores aún requieren operaciones lineales computacionalmente intensivas con dimensiones de salida grandes. Para abordar esto, proponemos KronSAE, una arquitectura novedosa que factoriza la representación latente mediante la descomposición del producto de Kronecker, reduciendo drásticamente la sobrecarga de memoria y computación. Además, introducimos mAND, una función de activación diferenciable que aproxima la operación binaria AND, lo que mejora la interpretabilidad y el rendimiento en nuestro marco factorizado.

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.