Identifying Sub-networks in Neural Networks via Functionally Similar Representations

TL;DR

Identifies functionally distinct subnetworks by comparing intermediate representations with the Gromov-Wasserstein distance. The method is task-agnostic and does not rely on predefined targets, enabling automated discovery of layer groupings that correspond to different abstractions. Empirical results across algebraic, NLP, and vision tasks show clear block structures and demonstrate utility for model compression and targeted fine-tuning, while requiring minimal human effort. The approach is permutation-invariant and scalable to differing layer shapes, distributions, and dimensions, addressing fundamental representational mismatches across layers.

Abstract

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.

Figures (18)

• Figure 1: Overview of our approach, where we use representations from different neural network layers to identify functionally distinct sub-networks (darker blocks) leveraging GW distance. The figure is an illustration based on a single transformer block, and the proposed technique can be applied to other types of network structures.
• Figure 2: Histogram on pairwise distances for outputs from all transformer blocks in a fine-tuned BERT model trained on YELP dataset.
• Figure 3: Pairwise GW distance on the synthetic Modular Sum dataset.
• Figure 4: Pairwise (layer) distances on Yelp, across different BERT models, using the proposed GW distance, from top to bottom. Different columns: first column is the pre-trained BERT and the rest are fine tuned BERT models with increasing sparsity (dense, $25\%$, $70\%$ and $95\%$ sparsity). As can be seen GW clearly demarcates the (functional) sub-network blocks. Due to page limit, we show baseline results in Appendix \ref{['appendix:yelp_morebaselines']}.
• Figure 5: Pairwise GW distance in YELP datasets, over training iterations.

Theorems & Definitions (1)

• Theorem 4.1