Beyond Black Boxes: Enhancing Interpretability of Transformers Trained on Neural Data

Abstract

Transformer models have become state-of-the-art in decoding stimuli and behavior from neural activity, significantly advancing neuroscience research. Yet greater transparency in their decision-making processes would substantially enhance their utility in scientific and clinical contexts. Sparse autoencoders offer a promising solution by producing hidden units that respond selectively to specific variables, enhancing interpretability. Here, we introduce SAEs into a neural decoding framework by augmenting a transformer trained to predict visual stimuli from calcium imaging in the mouse visual cortex. The enhancement of the transformer model with an SAE preserved its original performance while yielding hidden units that selectively responded to interpretable features, such as stimulus orientation and genetic background. Furthermore, ablating units associated with a given variable impaired the model's ability to process that variable, revealing how specific internal representations support downstream computations. Together, these results demonstrate that integrating SAEs with transformers combines the power of modern deep learning with the interpretability essential for scientific understanding and clinical translation.