Studying Cross-cluster Modularity in Neural Networks
Abstract
An approach to improve neural network interpretability is via clusterability, i.e., splitting a model into disjoint clusters that can be studied independently. We define a measure for clusterability and show that pre-trained models form highly enmeshed clusters via spectral graph clustering. We thus train models to be more modular using a "clusterability loss" function that encourages the formation of non-interacting clusters. We then investigate the emerging properties of these highly clustered models. We find our trained clustered models do not exhibit more task specialization, but do form smaller circuits. We investigate CNNs trained on MNIST and CIFAR, small transformers trained on modular addition, and GPT-2 and Pythia on the Wiki dataset, and Gemma on a Chemistry dataset. This investigation shows what to expect from clustered models.
Cite
@article{arxiv.2502.02470,
title = {Studying Cross-cluster Modularity in Neural Networks},
author = {Satvik Golechha and Maheep Chaudhary and Joan Velja and Alessandro Abate and Nandi Schoots},
journal= {arXiv preprint arXiv:2502.02470},
year = {2025}
}
Comments
8 pages, under review. arXiv admin note: text overlap with arXiv:2409.15747 (author note: this is an extension of that paper but has different authors)