Identifying Sub-networks in Neural Networks via Functionally Similar Representations

We thank the reviewer for their valuable feedback. In response, we have clarified notations and improved the presentation in the revision.

Metric in a quantitative manner: The proposed GW distance can absolutely be used in a quantitative manner. However, it would generally require some kind of ground truth for comparison, which is challenging to obtain in the context of sub-network discovery. Our synthetic experiment of modular sum is designed to provide ground truth subnetworks that compute different functions, but this does not guarantee the absence of further subnetworks within them (in fact, modular sum experiments have been shown to contain further sub-functions, as noted by Nanda et al., 2023).

Nevertheless, we design downstream tasks of model fine-tuning (Appendix G) and pruning/compression (Appendix J) to evaluate how well the discovered sub-networks would perform in prediction. Shown in Appendix J, we take the original pre-trained BERT and only use the first $n = \{12, 8, 4, 2, 1, 0 \}$ transformer blocks while discarding the rest. Here $n = 12$ corresponds to using all the transformer blocks, resulting the same BERT model. $n=0$, on the other hand, means that we only use a (linear) classifier layer after the embedding layer to predict the class label. The results are shown in Table 7 in Appendix J. As a reminder, GW distance suggest the last 4 blocks in YELP (see Figure 6) and the last 2 blocks in SST (see Figure 17) are mostly different, which is marked by star $(*)$ in the table. It shows that by using a limited number of layers, we can achieve similar performance to the full 12 block model, with $0.01\%$ and $0.54\%$ differences in YELP and SST, respectively. In contrast, using one fewer transformer block can risk much larger performance drop, with $0.10\%$ in YELP and $8.60\%$ differences in SST2 (approximately a 10-fold worse reduction). These accuracy results further justify the quality of the subnetworks identified.

Missing comparison with state-of-the-art metrics: In the original draft, we tested cosine and Wasserstein distances, used by some of the papers suggested. In the revision, we further added 6 additional baselines to provide a more comprehensive comparison. The results demonstrate that our method consistently identifies subnetworks with greater clarity and reliability compared to these other approaches.

Observations reported in Sec 5.2 and 5.3: We added an experiment, as discussed above, to show that a small subset of layers can fit the function well. Regarding the exact positions and gradient values, we checked the gradients of all layers and found that there is no significant differences among them, except for the gradients in the linear classification layer (i.e., the very last layer), where the gradient values were more than 10 times higher. As a result, we have removed the statement in the revision.

Vision task: As requested, We test our approach on the lightweight ResNet9 model [5], for ease of visualization and comparison. We compare a randomly initialized ResNet9 and a trained model that achieves 91.63% accuracy. We show the distance measures computed using all methods in Figure 18, with GW distance give 4 subnetworks. To further examine how the sub-network structures align with learned representation visually, we visualize the computed distances along with the learned representations of a "ship" image across all layers in Figure 9. The top row shows the representations of a ship at each layer. To observe the gradual changes over layers, we visualize the distance between each layer and its previous layer, using various methods that can handle different dimensions between compared spaces.

As shown in Figure 9, RSM, RSA, MSID, and CKA show indicate significant changes across many layers, without providing clear evidence of sub-network structures. AGW highlights the changes in the final few layers only. In comparison, GW distance demonstrates the most consistency with the image representations visually. Specifically, the 3rd convolution layer (Layer ID 2.ReLU) introduces the first notable differences, where the ship's shape becomes less distinct, signaling the learning of mid-level features. The shapes become increasingly blurred in the 5th convolution layers (Layer ID 4.Conv2d ) and by Layer 4.ReLU the ship's shape is nearly absent. The final convolutional layer (Layer ID 7.Conv2d) shows significant changes from its preceding layer (Layer ID 6.ReLU), marking the point where class-specific information is consolidated. These results suggest that GW distance aligns most effectively with the learned image representations, providing strong evidence that it can reveal meaningful subnetwork structures in vision models.

[5] Park et al. "Trak: Attributing model behavior at scale." ICML 2023.

We thank the review for the valuable comments, and have fixed the typo per suggestion.

Could the authors provide more explanation about the relationship between model interpretation and subnetwork?: As discussed in the introduction, we believe that discovering sub-functions in the form of distinct subnetworks can provide valuable insights into how many functions a neural network utilizes to achieve specific outcomes, such as making predictions. This represents an important first step toward the automated identification of these functions. To address the broader question of "What kind of functions does the neural network learn?", it is both important and beneficial first to explore "How many distinct (sub)-functions does the neural network learn?". We believe understanding neural networks through the identification of subnetworks is essential due to the complexity and opacity of modern deep learning models. Neural networks, especially those with many layers and parameters, often exhibit behaviors that are difficult to interpret holistically. Identifying subnetworks allows us to break down the model into smaller, more interpretable units, providing insights into how individual components contribute to the model’s overall performance.

The key idea to identify subnetworks is the use of Gromov-Wasserstein (GW) distance, which has been proposed before and might diminish the novelty of the work. Our work proposes to use GW for understanding similarity among neural network representations. Although we do not technical deviate from existing GW measure, we justify its usage by invariance properties and empirical evaluations, demonstrating its promising application in detecting substructures in neural networks.

Additionally, the field of model pruning involves identifying and retaining stronger subnetworks. What are the differences and connections between this approach and model pruning?: In our approach, the primary focus is on subnetwork discovery, which aims to identify and ultimately understand the distinct components or sub-functions within a neural network that contribute to its overall task, such as prediction. This process provides insight into the network’s internal structure and the roles of different components, answering questions like 'How many different sub-functions does the neural network learn?". In contrast, model pruning focuses on reducing the size of the model by removing less important or redundant parts, such as neurons or layers, to improve computational efficiency while retaining performance. The goal of pruning is not necessarily to understand the model's internal structures, but to streamline it without significantly impacting its functionality. While these two approaches have different objectives, they do share connections. Insights from subnetwork discovery can inform pruning decisions by identifying which parts of the model are essential for maintaining performance (as we have tested in Appendix L). Additionally, pruning can simplify the model, making it easier to identify and interpret the important subnetworks.

Could the author provide more explanation about Figure 3, particularly the relationship between the three identified groups and the model architecture?: The 3 identified groups correspond to different components within transformer blocks. Group i (layers roughly from 22 to 44, and 52 and 74) and group iii (the initial and final layers) mainly involve MLP and/or residual computations. Figure 3 shows these groups share strong similarities among themselves, indicating that these layers computes relatively simple transformations. In contrast, Group ii (layer roughly from 12 to 20, 43 to 51, and around 74 to 82) consists of attention-computations, suggesting more complex transformations compared to MLP layers. We provide a more explanation in Figure 5 and the accompanying paragraph to the left of the figure.

Thank you for your time to review our rebuttal and for your thoughtful consideration.

Motivation: Regarding the motivation, we have revised the introduction and related work sections to provide a clearer context for our study. We would like to further emphasize that our work addresses significant manual effort typically required to understand the underlying functions, which remains a challenge and previous methods have not adequately addressed. For example, Nanda et al, 2023 and Zhong et al, 2024 both rely on manual inspections to reverse engineer these functions, which are limited in smaller networks (e.g., one-layer transformer) and simpler problems (e.g., modular sum between two numbers). These limitations underscore the relevance of our work in automating this effort for larger networks and real datasets.

Novelty:Furthermore, while it is true that the Gromov-Wasserstein distance has been explored in prior research, our work makes a novel contribution by demonstrating its applicability for automated functional discovery. This not only advances its practical usage in the field but also uncovers new and interesting insights for real datasets.

We hope these clarifications address your concerns. We would greatly appreciate any additional suggestions or guidance to further improve the clarity and impact of our work.

We thank the reviewer for their constructive feedback.

Could you elaborate on the necessity and benefits of understanding neural networks through the identification of sub-networks? Understanding neural networks through the identification of subnetworks is essential due to the complexity and opacity of modern deep learning models. Neural networks, especially those with many layers and parameters, often exhibit behaviors that are difficult to interpret holistically. Identifying subnetworks allows us to break down the model into smaller, more interpretable units, providing insights into how individual components contribute to the model’s overall performance. It has also many benefits. First, it enhances interpretability by enabling us to understand how specific subnetworks process information. Second, it aids in model optimization by identifying which components are critical for performance and which can be removed or simplified. Third, understanding subnetworks helps in debugging and improving the model. By identifying and analyzing subnetworks, we can simplify models, enhance their efficiency, and improve their transparency, leading to more reliable and interpretable neural networks.

Could you detail the technical innovations presented in this paper that extend beyond the Gromov-Wasserstein method? While we do not introduce technical modifications to the existing GW measure, we adapt its application to analyze the similarity among neural network representations and facilitate subnetwork discovery. We justify this approach through its invariance properties and empirical evaluations. Notably, our experiments demonstrate that the GW distance effectively identifies distinct substructures and aligns visually with the learned representations, as shown in Section 5.4.

Could you present additional experimental results across diverse tasks, datasets, and learning methodologies? Per suggestion, we have included experiment results for a convolutional neural network (ResNet 9) on image classification dataset CIFAR-10 in Section 5.4. Our method has now been tested in algebraic, NLP, and computer vision datasets with different neural network architectures. Additionally, we have incorporated the suggested baseline [6] along with several others. Overall, GW demonstrated the clearest subnetwork structures and the strongest visual consistency on image datasets. We believe this further supports the validity of our approach for subnetwork identification.

[6] Demetci, Pinar, et al. "Revisiting invariance and introducing priors in Gromov-Wasserstein distances." arXiv preprint arXiv:2307.10093 (2023).

Could you discuss a real-world application of sub-network identification, such as improving the inference speed of neural networks? While our paper focuses on a different objective -- finding the number of components within a network, rather than identifying the most critical components for a specific task -- with downstream applications aimed at enabling a more automated process for mechanistic interpretability by facilitating the study of each component independently, the reviewer likely suggests an interesting direction on using a subnetwork to speed up inference. We have tested only fine-tuning specific layers in a BERT while freezing the other layers, in Appendix G, which can speed up training. In our revision, we have added an additional experiment involving network pruning, where only certain layers are retained for training. As shown in Appendix L, we identified a subset of transformer blocks (4 for YELP and 2 for SST2, compared to all 12) that can achieve comparable performance. This approach not only aids in understanding but can also accelerates both training and testing.

We thank the reviewer for their constructive feedback and insightful suggestions.

In particular on real data there is no analysis on correlation on accuracy on downstream tasks and GW similarity of layers ... In order to understand this one way could be just to train leaner readout on the downstream from representations of the two blocks separately and see how the performance changes and correlates with the GW distance. Even when termediate layers learn different functions, their representations may not be effective at prediction. Per suggestion, we have trained a logistic regression and a 2-layer MLP probe on every transformer block to predict the final label on YELP datasets. The accuracies of logistic regression readout across transformer blocks are all around 50%, until the very last block, while the MLP readouts gradually increase from 50% to 59%, until the very last block. Hence leaner readouts may not predict well due to its lack of predictive power. To further check whether layers can identify with a given complex function, we use a vision dataset CIFAR-10 to visually inspect the subnetworks. Results are discussed below.

GW distance vs other baselines: Per suggestion, we have added results of MSID and CKA, along with several others, with publicly available implementation and default parameters in experiments. In comparison to CKA and other kernel distance measures, GW shares some of the invariance properties (including orthogonal transformation and isotropic scaling) but additionally considers the geometric property of the metric space. Moreover, compared to MSID, we find that directly optimizing the GW distance yields better results in identifying subnetwork structures.

ResNet and Vision Experiments: We focus on attention-based architectures because it is the state-of-the-art architecture for NLP tasks and are widely used in existing works for mechanistic interpretability studies (e.g., Nanda et al, 2023). As the reviewer suggested, our methods can indeed generalize to other models. We have tested our method on ResNet-9 trained on the CIFAR 10 dataset. The results, including pairwise distances and visualizations, are provided in Section 5.4 and Appendix N. As shown in Figure 9, GW distance demonstrates the most consistency with the image representations visually. Specifically, the 3rd convolution layer (Layer ID 2.ReLU) introduces the first notable differences in GW distance, where the ship's shape becomes less distinct, signaling the learning of mid-level features. The shapes become increasingly blurred in the 5th convolution layers (Layer ID 4.Conv2d ) and by Layer 4.ReLU the ship's shape is nearly absent. The final convolutional layer (Layer ID 7.Conv2d) shows significant changes from its preceding layer (Layer ID 6.ReLU), marking the point where class-specific information is consolidated. These results confirm that that early layers sub-components focus on local statistics of the input such as shapes, while later layers learn higher levels of features and rely on global information. These results also suggest that GW distance aligns most effectively with the learned image representations, providing strong evidence that it can find meaningful subnetwork structures in vision models.

Model merging related works: Thank you for the suggestions. [1] uses GW distance as a regularization to fuse models. We have cited the works in the related work section.

[1] Singh, Sidak Pal, and Martin Jaggi. "Model fusion via optimal transport." Advances in Neural Information Processing Systems 33 (2020)

[2] Stoica, George, et al. "Zipit! merging models from different tasks without training." ICLR

Implementation details: "Why the choice of the Fused Gromov Wasserstein method in particular?" We do not use Fused GW method. "What is the exact difference between the implemented version of Wasserstein and GW distance in the paper?" We have added a discussion in Appendix H on the exact implementation details of W and GW distances. Specifically, we solve W distance based on [3] and GW distance based on conditional gradient [4].

[3] Bonneel et al. Displacement interpolation using Lagrangian mass transport. In ACM Transactions on Graphics (TOG), 2011

[4] Titouan et al. “Optimal Transport for structured data with application on graphs”. ICML, 2019.

Random Seeds and Similarity Results: To study the impact of initialization seeds on the learned representations, we train the same BERT model on YELP datasets with different seeds, with identical hyperparameters for a total of 27,000 iterations. As shown in Figure 19 in Appendix M, while the learned representations vary across seeds, but the general block structures remain consistent when analyzed using GW distances.

We thank the reviewer for your time to review our rebuttal and for your further consideration.

Validation: Thank you for highlighting this important concern. We understand the significance of validating that the subspaces of the network correspond to specialized tasks, as it directly relates to the robustness and interpretability of our approach. We acknowledge that direct validation is necessary, which is why we designed the synthetic experiments. Real datasets generally lack such ground truth labels on subspaces, so it is difficult and impractical to compare methods quantitatively. In our revision, besides the modular sum experiment, we provide indirect validation through a set of downstream tasks of model fine-tuning (Appendix G) and pruning/compression (Appendix J) to evaluate how well the discovered sub-networks would perform in prediction. Specifically, shown in Appendix J, we take the original pre-trained BERT and only use the first $n = \{12, 8, 4, 2, 1, 0 \}$ transformer blocks while discarding the rest. Here $n = 12$ corresponds to using all the transformer blocks, resulting the same BERT model. $n=0$, on the other hand, means that we only use a (linear) classifier layer after the embedding layer to predict the class label. The results are shown in Table 7 in Appendix J. As a reminder, GW distance suggest the last 4 blocks in YELP (see Figure 6) and the last 2 blocks in SST (see Figure 17) are mostly different, which is marked by star $(*)$ in the table. It shows that by using a limited number of layers, we can achieve similar performance to the full 12 block model, with $0.01\%$ and $0.54\%$ differences in YELP and SST, respectively. In contrast, using one fewer transformer block can risk much larger performance drop, with $0.10\%$ in YELP and $8.60\%$ differences in SST2 (approximately a 10-fold worse reduction). These accuracy results further justify the quality of the subnetworks identified. We appreciate your feedback and welcome any additional suggestions or methodologies you believe would further strengthen this aspect of our work.

GW vs CKA: Thank you for pointing out the apparent qualitative similarity between the GW and CKA results in Figure 15. We appreciate your careful examination of the results and the opportunity to clarify the benefits of our approach. As briefly discussed in Appendix I, while CKA performs similarly qualitatively to GW in Figure 15, it also shows greater variability within block structures. Moreover, on CIFAR-10 dataset per Figure 20 (which may present a more challenging task due to its lower accuracy), CKA struggles to provide clear subnetwork structures, often showing different block structures at each layer. These results demonstrate that geometric-based Gromov-Wasserstein distance capture more variances and better highlight the complex transformations of data.

We thank the reviewer for many detailed comments and great suggestions.

Studying whether different layers provide similar or dissimilar representations is definitely interesting (and in fact, many existing works already do this), but I find the relation to mechanistic interpretability to be far-fetched. As we discussed in the introduction, we believe that sub-function discovery in the form of different subnetworks could provide us insights into the number and type of sub-functions that exist in a network for modeling the prediction function. This serves a good first step towards identifying these functions automatically. We believe understanding neural networks through the identification of subnetworks is essential due to the complexity and opacity of modern deep learning models. Neural networks, especially those with many layers and parameters, often exhibit behaviors that are difficult to interpret holistically. Identifying subnetworks allows us to decompose the model into smaller, more interpretable units, providing insights into how individual components contribute to the model’s overall performance.

Clarity of the method: "it’s unclear to me why we would want to optimize a correspondence..." Our goal is to uncover layer similarities, with Gromov-Wasserstein (GW) serving as a metric for comparing unaligned manifolds—effectively measuring the distance between two sets of samples. If the identity map (for $\pi$) yields the minimal distance, GW would recover such an identity mapping. If we assume that the semantics of the examples do not change across layers, setting $\pi$ to identity would work. However, in neural networks, it is plausible that the semantics of the examples could change, as evidenced by the significant geometric changes that often observed in local neighborhoods (visualized in the tSNE projection in Figure 11). These shifts arise because different network layers may focus on distinct aspects of sentences or images, suggesting that using simple identity mappings would result higher distance at every layer regardless of other transformations. To account for these semantic shifts, it is desirable to optimize optimizing the correspondence, which enables us to uncover meaningful structure and relationships that would otherwise be obscured if we use the fixed identity map. Empirically, baseline RSA is similar to GW distance with the self transportation plan, based on inter-sample distances. On YELP dataset, RSA is more sensitive to layer changes and give less definitive boundaries between subnetworks structures (Figure 16). On CIFAR-10 dataset, RSA also tends to show higher discrepancies between visually similar layers and does not provide clear sub-network structures (Figure 9).

Baselines and Experiments: Per suggestion, we conducted experiments on 6 additional baselines including the suggested CCA and CKA, on modular, NLP, and computer vision datasets. The results show our methods can give clear division of subnetworks, and visually align well with learned representations. Moreover, we also give a quantitative measure on discovered network. We also added a discussion in Related works on [1, 2, 3, 4] accordingly.

[1] Raghu, Maithra, et al. "Do vision transformers see like convolutional neural networks?." Advances in neural information processing systems 34 (2021): 12116-12128.

[2] Nguyen, Thao, Maithra Raghu, and Simon Kornblith. "On the origins of the block structure phenomenon in neural network representations." TMLR, 2022.

[3] Klabunde, Max, et al. "Similarity of neural network models: A survey of functional and representational measures." arXiv preprint arXiv:2305.06329 (2023).

[4] Kornblith, Simon, et al. "Similarity of neural network representations revisited." International conference on machine learning. PMLR, 2019.

Clarifications and editing: Per suggestion, w moved details on representation candidates in the original Table 1 to the Appendix A. Note that these representations in Table 1 is also used in Figure 3, 4, and others modular sum experiments. As discussed in Appendix B, the search space is of size 93 because all intermediate layers within transformer blocks are also compared. We have distinguished transformer blocks and intermediate layers in the revision to avoid confusion. Due to space limit, we removed the pretrained figure in Figure 2 and clarified its content further in Appendix F. We have added Jaccard similarity measures on neighborhood overlap in the main text and more details in Table 5 of Appendix F. Moreover, we have clarified KL divergence in figures in Appendix F. On the modular sum experiment, as discussed in Appendix B, we use a separate validation datasets with an 80-20 train-validation split.