Treine Autoencoders Esparsos de Forma Eficiente Utilizando a Correlação 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

Resumo

Os Autoencoders Esparsos (SAEs) têm demonstrado um potencial significativo na interpretação dos estados ocultos de modelos de linguagem, decompondo-os em direções latentes interpretáveis. No entanto, o treinamento de SAEs em grande escala continua sendo um desafio, especialmente quando são usados tamanhos de dicionário grandes. Embora os decodificadores possam aproveitar kernels esparsos para eficiência, os codificadores ainda exigem operações lineares computacionalmente intensivas com grandes dimensões de saída. Para resolver isso, propomos o KronSAE, uma nova arquitetura que fatoriza a representação latente por meio da decomposição do produto de Kronecker, reduzindo drasticamente a sobrecarga de memória e computação. Além disso, introduzimos o mAND, uma função de ativação diferenciável que aproxima a operação binária AND, melhorando a interpretabilidade e o desempenho em nossa estrutura fatorizada.

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.