Addestra Efficientmente Autoencoder Sparse Sfruttando la Correlazione delle Caratteristiche

Train Sparse Autoencoders Efficiently by Utilizing Features Correlation

May 28, 2025

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

cs.AI

Abstract

Gli Autoencoder Sparse (SAE) hanno dimostrato un potenziale significativo nell'interpretazione degli stati nascosti dei modelli linguistici scomponendoli in direzioni latenti interpretabili. Tuttavia, l'addestramento degli SAE su larga scala rimane impegnativo, specialmente quando vengono utilizzati dizionari di grandi dimensioni. Sebbene i decoder possano sfruttare kernel ottimizzati per la sparsità per migliorare l'efficienza, gli encoder richiedono ancora operazioni lineari computazionalmente intensive con dimensioni di output elevate. Per affrontare questo problema, proponiamo KronSAE, una nuova architettura che fattorizza la rappresentazione latente attraverso la decomposizione del prodotto di Kronecker, riducendo drasticamente l'overhead di memoria e computazionale. Inoltre, introduciamo mAND, una funzione di attivazione differenziabile che approssima l'operazione binaria AND, migliorando l'interpretabilità e le prestazioni nel nostro framework fattorizzato.

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.