Identifying Sub-networks in Neural Networks via Functionally Similar Representations
This work addresses the challenge of making neural networks more explainable and trustworthy for AI researchers and practitioners, though it appears incremental as it builds on existing interpretability efforts with a novel but task-agnostic method.
The authors tackled the problem of automating mechanistic interpretability in neural networks by identifying distinct sub-networks based on functionally similar representations, using Gromov-Wasserstein distance to handle varying distributions and dimensionalities across layers. They demonstrated the approach on algebraic, language, and vision tasks, showing it reveals functional abstractions and aids in model compression and fine-tuning with minimal human and computational cost.
Providing human-understandable insights into the inner workings of neural networks is an important step toward achieving more explainable and trustworthy AI. Existing approaches to such mechanistic interpretability typically require substantial prior knowledge and manual effort, with strategies tailored to specific tasks. In this work, we take a step toward automating the understanding of the network by investigating the existence of distinct sub-networks. Specifically, we explore a novel automated and task-agnostic approach based on the notion of functionally similar representations within neural networks to identify similar and dissimilar layers, revealing potential sub-networks. We achieve this by proposing, for the first time to our knowledge, the use of Gromov-Wasserstein distance, which overcomes challenges posed by varying distributions and dimensionalities across intermediate representations, issues that complicate direct layer to layer comparisons. On algebraic, language, and vision tasks, we observe the emergence of sub-groups within neural network layers corresponding to functional abstractions. Through downstream applications of model compression and fine-tuning, we show the proposed approach offers meaningful insights into the behavior of neural networks with minimal human and computational cost.