Эффективное обучение разреженных автокодировщиков с использованием корреляции признаков

Train Sparse Autoencoders Efficiently by Utilizing Features Correlation

May 28, 2025

Авторы: Vadim Kurochkin, Yaroslav Aksenov, Daniil Laptev, Daniil Gavrilov, Nikita Balagansky

cs.AI

Аннотация

Разреженные автокодировщики (SAE) продемонстрировали значительный потенциал в интерпретации скрытых состояний языковых моделей путем их декомпозиции на интерпретируемые латентные направления. Однако обучение SAE в масштабе остается сложной задачей, особенно при использовании больших размеров словарей. Хотя декодеры могут использовать специализированные ядра для повышения эффективности, кодировщики по-прежнему требуют вычислительно затратных линейных операций с большими выходными размерностями. Для решения этой проблемы мы предлагаем KronSAE — новую архитектуру, которая факторизует латентное представление с помощью декомпозиции Кронекера, значительно снижая затраты памяти и вычислений. Кроме того, мы представляем mAND — дифференцируемую функцию активации, аппроксимирующую бинарную операцию AND, что улучшает интерпретируемость и производительность в нашей факторизованной структуре.

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.