Modularity in Transformers: Investigating Neuron Separability & Specialization

TL;DR

This paper investigates the modularity and task specialization of neurons within transformer architectures, focusing on both vision and language models, and finds that neuron clusters identified through MoEfication correspond more strongly to task-specific neurons in earlier and later layers of the models.

Abstract

Transformer models are increasingly prevalent in various applications, yet our understanding of their internal workings remains limited. This paper investigates the modularity and task specialization of neurons within transformer architectures, focusing on both vision (ViT) and language (Mistral 7B) models. Using a combination of selective pruning and MoEfication clustering techniques, we analyze the overlap and specialization of neurons across different tasks and data subsets. Our findings reveal evidence of task-specific neuron clusters, with varying degrees of overlap between related tasks. We observe that neuron importance patterns persist to some extent even in randomly initialized models, suggesting an inherent structure that training refines. Additionally, we find that neuron clusters identified through MoEfication correspond more strongly to task-specific neurons in earlier and later layers of the models. This work contributes to a more nuanced understanding of transformer internals and offers insights into potential avenues for improving model interpretability and efficiency.

Figures (6)

• Figure 1: Relative Top1 % image classification accuracy compared to baseline for ViT in Cifar100, after selectively pruning MLP neurons to get a target drop in performance of 50% (left) and 80% (right) for the 'unlearned' class (shown as the vertical strips).
• Figure 2: Drop in Top1 accuracy in Mistral 7B for next-token-prediction datasets as well as zero-shot subject-split MMLU question answering. classes within each super-class of Cifar100. We target a relative drop in performance of 20% for the class being 'unlearned', and each vertical strip shows on the other data subsets.
• Figure 3: Intersection analysis between selected neurons ranked in the top $12.5\pm6.0\%$ for trained ViT model (left) and random model (right) in relative importance to Cifar20 classes. Each square shows the percentage overlap between neurons identified as being important for each class pair.
• Figure 4: Intersection analysis between selected neurons ranked top 1% for Mistral 7b model that is trained (left) and randomly initialised (right) in relative importance to different text classes. Each square shows percentage overlap between neurons identified as being important for each class pair.
• Figure 5: Plot of the overall overlaps between classes of Vit (left) and subtexts of Mistral (right) for both random and pre-trained selected neurons. We also show the difference between the random and pre-trained overlaps at each point.