SAFR: Neuron Redistribution for Interpretability

TL;DR

SAFR is a regularization method that enhances neural network interpretability by promoting clearer neuron feature representations, while maintaining prediction accuracy, through targeted loss function modifications.

Contribution

This paper introduces SAFR, a novel regularization technique that improves interpretability of transformer models by controlling neuron superposition, a feature not previously exploited for this purpose.

Findings

SAFR improves interpretability without reducing accuracy

Neuron representations become more monosemantic with SAFR

SAFR enables visualization of neuron allocations in models

Abstract

Superposition refers to encoding representations of multiple features within a single neuron, which is common in deep neural networks. This property allows neurons to combine and represent multiple features, enabling the model to capture intricate information and handle complex tasks. Despite promising performance, the model's interpretability has been diminished. This paper presents a novel approach to enhance model interpretability by regularizing feature superposition. We introduce SAFR, which simply applies regularizations to the loss function to promote monosemantic representations for important tokens while encouraging polysemanticity for correlated token pairs, where important tokens and correlated token pairs are identified via VMASK and attention weights respectively. We evaluate SAFR with a transformer model on two classification tasks. Experiments demonstrate the effectiveness of SAFR in improving model interpretability without compromising prediction performance. Besides, SAFR provides explanations by visualizing the neuron allocation within the intermediate layers.