Identifying Sub-networks in Neural Networks via Functionally Similar Representations

## Abstract

Abstract Mechanistic interpretability aims to provide human-understandable insights into the inner workings of neural network models by examining their internals. Existing approaches typically require significant manual effort and prior knowledge, 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, reducing the need for human intervention.
Our method identifies similar and dissimilar layers in the network, revealing potential sub-components. 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.
Through experiments on algebraic and language tasks, we observe the emergence of sub-groups within neural network layers corresponding to functional abstractions. Additionally, we find that different training strategies influence the positioning of these sub-groups. Our approach offers meaningful insights into the behavior of neural networks with minimal human and computational cost.

## 1 Introduction

Rapid progress in transformer language models has directed attention towards understanding the underlying causes of new capabilities. Like many other neural methods, large language models (LLMs) are mostly black-box models and explainable artificial intelligence aims to offer insights and improve human understanding of these LLMs. Recently, mechanistic interpretability research has gained popularity, focusing on reverse-engineering
models into human-understandable algorithms, using methods such as computational graphs and circuits [, ]. Examples include recovering internal mechanisms of LLMs to solve typical mathematical problems such as modular sum [].

Current approaches primarily rely on extensive manual inspection and trial-and-error to reverse engineer networks by discovering a sequence of learned functions that produce a desired output. This process requires significant human prior knowledge with strategies tailored to specific tasks such as algebraic problems []. Automatic discovery of these functions is a difficult task due to the large search space involved. We investigate a simpler task that can be considered as an initial step towards automating the discovery of these functions: is it possible to detect how many distinct (complex) functions exist in a learned network, and which layers correspond to each such function?
In this paper, we take such a step towards a more automatic and task-agnostic approach to mechanistic interpretability by identifying sub-components in neural networks that are functionally different from each other. Specifically, we treat the intermediate layers of neural networks as computing different *functions* of the input and base our approach on *functional similarity*. These distinctive subnetworks would represent one function in a sequence of different ones, and knowing such a partition would facilitate further decoding these functions leading a better understanding of the network.

The nature of functional similarities differ depending on which of the two cases is true: either we have a hypothesized target function that part of the network may (approximately) compute, or we lack such a hypothesis.
In the case of a given target function, functional similarity relates to the output of the target function, and the problem is one of search to find a network layer that best approximates it. This is akin to comparing two functions based on the values they take for the same input samples. The use of probes attached to representations is a popular method for detecting such similarities.

Figure: Figure 1: Overview of our approach where we use representations to identify functionally distinct sub-networks (darker blocks) leveraging GW distance.
Refer to caption: x1.png

The more realistic case with neural networks is the absence of a target function, which is the primary focus of this paper. In this case, we propose to measure the functional similarity *between* layers of the network, and the problem is to identify distinctive sub-components within the network by finding which layers are functionally less similar to previous layers as shown in Figure 1. The idea is that less similar layers (brighter colors) may indicate boundaries of a sub-component, while more similar layers (darker colors) likely belong to the same sub-component. While it may be tempting to also use probes here for layer comparisons, the lack of targets makes it unclear how the representations in these layers should be transformed by the probes for comparison. We need a more direct way to compare these representations.
Hence, we propose to measure similarity using Gromov-Wasserstein (GW) distance [] between representations from different layers of the network.
As elaborated further in Section 4.1, GW allows distance computation between distributions supported on two different metric spaces with different supports and potentially different dimensions, which is common across different layers in neural networks. GW is also invariant to permutation of the representation within a layer, a crucial property since neural networks are known to have permutation symmetries []. As such, GW can effectively identify genuinely distinct behaviors across (groups of) layers.

We validate our approach on algebraic and NLP sentiment analysis tasks, showing that GW distance provides a systematic way to analyze and identify subnetworks. Additionally, our findings provides a holistic view on differences in representations of models trained with different strategies.
We observe clear patterns in the form of block structures among different layers, suggesting there exist sub-networks that have different functions, particularly at the transition layers where major functional changes may occur. Moreover, the GW distance can also be used to observe the emergence of subnetworks during training process. Overall, we hope that our method can improve the efficiency of mechanistic interpretation by finding subnetworks in larger models, reducing the need for extensive human effort and potentially contributing to a further understanding of neural network behaviors.

## 2 Background and Related Work

With its popularity, mechanistic interpretability has become a disparate area and we survey various directions within it.

Algorithm discovery in algebraic problems Modular sum algorithms [] are studied in the context of progress measure in emerging abilities of transformers. The authors in [] show, via various inspection such as gradient symmetry, logit patterns, and input pair relations, that one class of algorithm is preferred than another. Other math problems such as the greatest common divisor [] also show insights into the inner mechanisms of neural networks.

Other types of discovery Automated circuit discovery [, ] intends to find a computational graph that is far sparser without sacrificing performance, using KL measures. Heuristic Core [] hypothesizes that there exists a set of attention heads that encompasses all subnetworks, learning shallow and non-generalizable features. Generalization is achieved by incorporating additional attention heads with higher-level features. Average gradient outer product [] uses eigen directions of weight matrices to identify task-relevant gradient directions. Moreover, various work have explored neuron semantics and possible disentanglement of semantics [, , ], as well as concept representations [].

Real applications An intervention study [] comparing clean and corrupted images via causal tracing shows that vision and language components are processed independently in the intermediate layers, interacting only in the final layers. [] studies factual and counterfactual statements, investigating how they compete with each other through their differences in attention logits and weights. [] finds that relation decoding can be represented by a linear relation embedding with causal semantics. [] studies the hallucination problem in factual completion tasks.

Methods of Studying Mechanisms Weight inspection and manipulation are commonly employed techniques to gain insights into the inner work of networks, including studying periodicity [] and weight gradients []. Modifying or ablating the activation of a specific model components [, ], including attention knockout [], and even direct modification of attention matrix [, ] are prevalent. Another popular approach involves representation and output inspection [], including logit patterns [], residual stream [], and periodicity []. Causal mediation analysis is used to compute the indirect treatment effect [, ] with perturbed embedding, often with Gaussian noises.
Rather than just inspection of outputs, many works have proposed to map the output to some target and is a popular technique for analyzing how neural activations correlate with high-level concepts []. Linear probes are generally used. [] uses a nonparametric probe ($k-$nearest-neighbor) to classify outputs for reasoning tasks.

Similarity Measure between Neural Network Layers Similarity measure between layers have been studied. Cross product between each layer output and final output [] is used to approximate each layer’s contribution to the final prediction. There are studies that quantify the similarity between different groups of neurons [], typically layers [], to understand and compare different neural networks. Generally a normalized representation is used to compare different transformer blocks, with different desired invariance properties, such as invariance to invertible linear transformation in canonical correlation analysis [], orthogonal transformation, isotropic scaling, and different initializations in centered kernel alignment []. Other measures include representational similarity
analysis [], which studies all pairwise distances across different inputs. Wasserstein distance has been explored in measuring similarities in the context of neural networks [, , ], but they assume that different layer representations belong to the same metric space, which is very unlikely even if they have the same dimensionality as the semantics captured by each layer are likely to differ significantly.

## 3 Functional Similarity Within Neural Networks

We aim to identify sub-components of a neural network based on their mechanistic functions. When we have a hypothesized target function that the sub-component may compute, this can be formulated as a similarity search problem
In this context, we search for candidate representations within neural networks’ outputs that represents changes in function computation. The search problem consists of three key elements: the search space, the search target, and the similarity measures used to evaluate how closely the candidates in the search space match the target.

### 3.1 Similarity Measures

Let $f_{0}$ be a function that produces a set of output $Y^{0}$ given a set of input $X$, where $Y^{0}\in{\mathbb{R}}^{n\times d_{y}}$ and $X\in{\mathbb{R}}^{n\times d_{x}}$ are assumed to have dimensions $d_{y}$ and $d_{x}$, respectively, with $n$ being the set size, without loss of generality. Let $f$ be another function that produces $Y$ given $X$.

Functional Similarity. Similarity between two functions, $f_{0}$ and $f$, over a set of input $X$, can be measured by the similarity between their outputs $Y_{0}=f_{0}(X)$ and $Y=f(X)$.

We can use a scoring or distance function $D(Y^{0},Y)$ as a measure between output similarity and hence the functional similarity between $f_{0}$ and $f$, where we regard $Y^{0}$ as the function/search target. If they are close according to $D$, then the function values should be similar to each other locally at a set of points $X$. Otherwise, these functions should be different at $X$. Popular measures such as Euclidean distance have been used for this purpose []. Each intermediate representations of a neural network can be naturally treated as function outputs, given inputs $X$.

The Need for Complex Functional Similarity Measure. Since we cannot exactly control the behavior of a trained neural network, the layer-wise functions $f$ that it learns can be complex and thus the learned representation $Y$ from each layer may be a complex function of the target $Y^{0}$ rather than a simpler transformation. For example, let $Y^{0}=\sin(X)$ and a candidate $Y=\sin^{2}(X)=(Y^{0})^{2}$. They share strong similarity, but a linear transformation will not be able to capture their functional similarity. If we want to truly understand where function $f_{0}$ might be approximately computed, we should consider functions of target $Y^{0}$, but naively listing out all possibilities can be prohibitive. As a consequence, one may need to use more complex measures to deal with such a space.

### 3.2 Search Space

We consider multiple candidate $Y$’s to form the search space for target $Y^{0}$. In the context of MLP neural networks for example, where $\sigma(.)$ denotes the non-linearity and $W$s are the parameter matrices, we have $Y^{*}=W_{n}(\sigma(W_{n-1}\dots\sigma(W_{1}X)))$ for the whole network. We can extract many $Y$’s from intermediate functions of the model, for instance $Y_{1}=W_{1}X$, $Y_{2}=\sigma(W_{1}X)$, and so on. These $Y$’s are often called representations, activations, or sometimes even “outputs” from each layer. We use these terms interchangeably here. For attention modules in transformer neural networks [], we can similarly extract $Y$’s from attention key, query, and value functions as well as MLP functions. More specifically, a deep transformer architecture of depth $l$ is formed by sequentially stacking $l$ transformer blocks. Each transformer block takes the representations of a sequence ${\bm{X}}_{\text{in}}\in\mathbb{R}^{T\times d}$, where ${\bm{X}}_{\text{in}}$ = Emb(${\bm{X}}$) with embedding layer Emb and input ${\bm{X}}$, $T$ is the number of tokens and $d$ is the embedding dimension, and outputs ${\bm{X}}_{\text{out}}$, where:

$$ $\begin{split}{\bm{X}}_{\text{out}}&=\alpha_{\text{FF}}\hat{{\bm{X}}}+\beta_{ \text{FF}}\text{MLP}(\text{Norm}(\hat{{\bm{X}}}))\\ \text{where,}\quad\text{MLP}({\bm{X}}_{m})&=\sigma({\bm{X}}_{m}{\bm{W}}^{1}){ \bm{W}}^{2}\\ \hat{{\bm{X}}}&=\alpha_{\text{SA}}{\bm{X}}_{\text{in}}+\beta_{\text{SA}}\text{ MHA}(\text{Norm}({\bm{X}}_{\text{in}})),\\ \text{MHA}({\bm{X}})&=[\text{Attn}_{1}({\bm{X}}),\dots,\text{Attn}_{H}({\bm{X} }))]{\bm{W}}^{P},\\ \text{Attn}({\bm{X}})&={\bm{A}}({\bm{X}}){\bm{X}}{\bm{W}}^{V},\\ {\bm{A}}({\bm{X}})&=\text{softmax}\left(\frac{1}{\sqrt{d_{k}}}{\bm{X}}{\bm{W}} ^{Q}{{\bm{W}}^{K}}^{\top}{\bm{X}}^{\top}+{\bm{M}}\right),\end{split}$ (1) $$

with scalar weights $\alpha_{\text{FF}}$, $\beta_{\text{FF}}$ , $\alpha_{\text{SA}}$, and $\beta_{\text{SA}}$ usually set to $1$ by default. Here FF stands for feedforward network, SA stands for self-attention, MHA is Multi-Head Attention, and Norm is a normalization layer. MLP usually has a single hidden layer with dimension $d$ and ReLU activation. The MHA sub-block shares information among tokens by using self-attention with ${\bm{W}}^{Q},~{}{\bm{W}}^{K}$ and ${\bm{W}}^{V}$ indicating query, key and value matrices. We list the exact locations of representations considered in the transformer models in Table 1, which serves as the focus of this paper.

**Table 1: All representations $Y$ considered in experiments as per equation 1. Omitting $Y$ in most names for readability.**
| (Across Blocks) |  |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- |
| Name | Resid-Pre^l | $Y^{l}$, at each block |  |  |  |  |
| Value | $={\bm{X}}_{\text{in}}^{l}$ | $={\bm{X}}_{\text{out}}^{l}$ |  |  |  |  |
| (Within Each Block $l$) |  |  |  |  |  |  |
| Name | Attn-Out^l | Resid-Mid^l | Pre | Post | MLP-out^l | Resid-Post^l |
| Value | $=$MHA(${\bm{X}}$)^l | $=\hat{{\bm{X}}}$ | $=\hat{{\bm{X}}}{\bm{W}}^{1}$ | $=$MLP($\hat{{\bm{X}}}$) | $=\text{MLP}(\hat{{\bm{X}}})$ | $={\bm{X}}_{\text{out}}$ |
| (Within Each Attention Head $h$) |  |  |  |  |  |  |
| Name | $k_{h}$ | $q_{h}$ | Attn-Preh | Attnh | vh | zh |
| Value | $={\bm{X}}{\bm{W}}^{K}$ | $={\bm{X}}{\bm{W}}^{Q}$ | $=q_{h}k_{h}^{T}$ | $={\bm{A}}({\bm{X}})$ | $={\bm{X}}{\bm{W}}^{V}$ | $=$Attn(${\bm{X}}$) |

### 3.3 Search Targets

#### Known Targets

When the search target, denoted as $T$, is a value from a known function, we can directly compare outputs between representations from each layer and known function output $T$. Representations from each layer can be directly compared with the target via a probe. Popular linear probes can be used to assess the similarity between a target and any layer’s representation. For instance, linear regression can be used to model each target $T$ from each representation candidate $Y$, and the residual error is used as the search criteria between $Y$ and $T$. As discussed previously, to deal with the potentially large search space of functions of the target, a more powerful probe (such as a nonlinear MLP function) may have to be used so that it can detect more complex similarities to $T$.

One challenge with using predictive probes to compute the distance measure $D$ is that the target function must be known. In practice, however, we often lack knowledge of specific targets. While it’s possible to experiment with various target functions with power probes,
the vast number of potential targets makes this approach inefficient. This calls for an alternative strategy to distinguish sub-components in a network through representation similarity.

## 4 Gromov-Wasserstein Distance as a Similarity Measure

#### Unknown Targets

When the search target is unknown, functionally similar parts by comparing them cannot be identified by comparison to a predefined set of target functions. Instead, we propose to identify the similarities and subnetworks among the representations at each intermediate layer. Each layer, however, posits a representation that potentially has a different distribution, not to mention even different dimensionality depending on the architectures and layers one considers (viz. mlp_in layer and attention layer in transformer blocks). Consequently, representations across layers may be incomparable using standard distance metrics, such as the $\ell_{p}$ norm amongst others.

To overcome these challenging issues, we propose computing distances between representations produced at the same layer for different inputs and match the vertices of this weighted graph – where each dimension of the representation of the inputs are vertices and the distances indicate weights on the edges – with the vertices of a similarly constructed weighted graph from another layer. Essentially, we assume the representations in a layer are samples of the underlying distribution, and we want the best permutation of representation dimensions in one layer that aligns with vertices in another layer, thereby deriving the inter-layer distance. If this inter-layer distance is low, then the two layers we consider to be functionally similar, given the same input data.

Formally, without loss of generality, let ${\bm{Y}}_{1}\in{\mathbb{R}}^{n\times d_{1}}$ and ${\bm{Y}}_{2}\in{\mathbb{R}}^{n\times d_{2}}$ be representations of $n$ examples from two different layers, where the discrete distributions over the representations are $\mu_{1}$ and $\mu_{2}$ respectively, with dimension $d_{1}$ possibly being different from $d_{2}$. Direct distance computation between them is not reasonable. Instead, we seek to compute a coupling or matching $\pi\in\Pi(\mu_{1},\mu_{2})$ between the $n$ examples in each set such that given the pairwise distances $D_{1},D_{2}\in{\mathbb{R}}^{n\times n}$ within representations ${\bm{Y}}_{1}$ and ${\bm{Y}}_{2}$ respectively, the sum of differences between the distances of the matched examples is minimized. Loosely speaking, we aim to find a matching that preserves the pairwise distance as much as possible. In particular, we want to minimize the following:

$$ $\displaystyle\rho({\bm{Y}}_{1},{\bm{Y}}_{2},\mu_{1},\mu_{2},D_{1},D_{2}) \triangleq\min\limits_{\pi\in\Pi(\mu_{1},\mu_{2})}$ $\displaystyle\sum_{i,j,k,l}(D_{1}(i,k)-D_{2}(j,l))^{2}\pi_{i,j}\pi_{k,l}$ $\displaystyle\text{s. t.}\quad\pi{\bm{I}}=\mu_{1};\pi^{T}{\bm{I}}=\mu_{2};\pi \geq 0.$ (2) $$

It turns out that $\rho$ corresponds to the Gromov-Wasserstein (GW) distance [], used to map two sets of points in optimal transport. We thus utilize this distance as a measure of inter-layer functional similarity in the setting where the target is unknown.

### 4.1 Justification for GW Distance as a Functional Similarity Measure

Let $({\bm{Y}}_{1},D_{1},\mu_{1})$ and $({\bm{Y}}_{2},D_{2},\mu_{2})$ be two given metric measure space (mm-space), where $({\bm{Y}},D)$ is a compact metric space and $\mu$ is a Borel probability measure with full support: $\text{supp}(\mu)={\bm{Y}}$. An isomorphism between ${\bm{Y}}_{1},{\bm{Y}}_{2}$ is any isometry $\Psi:{\bm{Y}}_{1}\rightarrow{\bm{Y}}_{2}$, i.e., a distance-preserving transformation between metric spaces, such that $\Psi_{\#\mu_{1}}=\mu(\Psi^{-1})=\mu_{2}$.

###### Theorem 4.1 .

[]. The Gromov-Wasserstein distance in equation 4 defines a proper distance on the collection of isomorphism classes of the mm-spaces.

Remark. The Gromov-Wasserstein distance itself is defined on isomorphism-classes of metric measure spaces, which means that any distance preserving (isometric) transformation of a space should preserve GW distance between the points in that space and any other space []. These isometric transformations include rigid motions (translations and rotations) and reflections or compositions of them. Additionally, permutations of points in a space also preserve GW distances, as the points are unlabeled. Hence, GW distance captures much richer transformations across layers.

The computed GW distance represents the minimal distance over all possible transportation plans between two sets of points from different spaces. In our context, we can also view GW as a measure that quantifies the distance between distance-based (i.g., Euclidean-distance) graphs, with a set of points as its nodes. Hence this would be low if the graph undergoes (nearly) isomorphic transformations between layers. Conversely, a high GW distance indicates a non-distance preserving transformation across layers, potentially reflecting a highly non-linear operation. While GW distance does not reveal the exact function operation, it highlights specific layers for further investigations.

Favorable Properties of GW. Besides the above noteworthy property of GW to map between different spaces, it also has other favorable properties [, ]:
i) It is symmetric and satisfies triangle inequality. ii) It is invariant under any isometric transformation of the input, which is advantageous because we do not want rotations and reflections to affect our similarity search. This invariance also includes permutation invariance, which is a beneficial property since we want the distance between layer representations to remain unaffected by permutations within the representations in each layer.
iii) GW is scalable since it does not require estimating high-dimensional distributions, which is often the case with intermediate representations in large models; instead, it only compare them to obtain a distance measure.
iv) GW is monotonic in (positive) scaling of pairwise distances, and hence the same layers should appear to be closer than others even with scaling.

Figure: (a) Pretrained
Refer to caption: extracted/5943928/figs/YELP_bert-base-uncased_target_unknown_spectral_GW_pre_scores_histo.png

Distance Distributions. As an illustrative example, we plot the histogram on pairwise distances for a batch of samples across all transformer blocks in pretrained and fine-tuned BERT models from the YELP review dataset in Figure 2(a) and  2(b). For more details on YELP, we provide a comprehensive discussion in the experiments section
5.2. The results show the distributions on pairwise distances in the fine-tuned model begin to differ from layer 9, consistent with GW distance observed in Figure 6 suggests that significant transformations occur and can be effectively captured by GW.

Neighborhood Change. Complementary to the distribution of pairwise distances, the changing representations of samples could also alter their relative neighborhoods across layers. Using a fine-tuned BERT model, we plot a tSNE projection [] of representations from a batch of samples on YELP, and visualize it in Figure 2(c). As shown, the sample neighborhood changes across layers, which can be indicative of functional changes that are not captured by comparing distributions alone. However, GW can account for such changes as well. More tSNE figures are shown 9(e) in Appendix D.

Computation Details. We use an existing optimal transport toolbox, [pythonot](pythonot) [], for computing GW distance. Specifically, we use an approximate conditional gradient algorithm proposed in [], which has a complexity of
$O(mn^{2}+m^{2}n)$, where $m$ and $n$ are the dimensions of two spaces (here the number of data samples from two layers being compared). In comparison, the Wasserstein distance may require $O(n^{3}log(n))$ for exact computation, and the Sinkhorn solver [] operates at nearly $O(n^{2})$. When the dataset is large, we can also sub-sample the dataset to improve the computational efficiency.

###### Theorem 4.1 .

## 5 Empirical Study and Findings

### 5.1 Synthetic Modular Sum Tasks

We begin by validating the Gromov-Wasserstein distance by comparing it against known partitions of the networks to determine whether it can successfully identify sub-networks. We first introduce the setup for the experiment, including data generation and models to be investigated.

Setup As a test case, we focus on a modular sum problem, following existing works []. We consider two datasets: the first generated by a single modular sum function with $c=f_{\text{mod}}(a+b)=(a+b)\text{mod p}$, where $a,b,c=0,1,\dots,p-1$, with $p=59$. The second dataset is more complex, with $c=f_{\text{mod3}}(a,b)$ of three levels of modular sums, namely: $c_{1}=(a+b)\text{mod }p_{1},c_{2}=(c_{1}+b)\text{mod }p_{2},c=(c_{2}+b)\text{
mod }p_{3}$, where $p=[59,31,17]$.

Training procedure We train 3 different neural networks with transformer blocks to predict $c$ given $(a,b)$. These networks contain input embeddings for $a$ and $b$, each of size $d$, i.e., $[{\bm{E}}_{a},{\bm{E}}_{b}]\in{\mathbb{R}}^{2d}$, and predict a categorical output $c$ via an unembedding/decoding layer. All parameters in the network are learned. For the first simpler $f_{\text{mod}}$ dataset, we train a neural network consisting of a one-layer ReLU transformer [], following the same protocol and hyperparameter choices as previous works [, ]. We call this Model 0. For the more complex $f_{\text{mod3}}$ dataset, we train two neural networks consisting of three-layer ReLU transformers, with 3 transformer blocks corresponding to the three levels of modular sum functions, and 4 attention heads within each block. The first network, which we call Model E, employs an end-to-end training procedure to directly learn output $c$ given input $(a,b)$. For the second network, which we call Model L, we use the same architecture as Model E but with a layer-wise training approach instead of end-to-end training. Specifically, we use the following procedure:

1. We train the first transformer block layer of Model L to predict $(c_{1},b)$ using an additional linear layer on top, given inputs $(a,b)$.

2. Once block 1 is fully trained, we discard the linear layer, freeze everything before the linear layer, and use its representations of $(c_{1},b)$ to train the second block to predict $(c_{2},b)$, again incorporating an additional layer on top.

3. Finally, we repeat the above step by freezing the first and second block and training the last block to predict $c$, using representations of $(c_{2},b)$.

In all these models, we are able to achieve $100\%$ prediction accuracy. Note that extra linear layers can also be considered as probes but used in training. More details can be found in appendix A. We use the GW measure on Model L with layer-wise training to verify if there is consistency between GW distances and known output $c$’s.
To evaluate the capability of handling different dimensions, we directly measure GW distance between 93 intermediate representation $Y$ (see appendix A for search space details) and $c$’s. To speed up computation of GW distance, we randomly sub-sample $1000$ data from a total of $3600$ samples, reducing time from $2$ min to $5$ seconds for each computation.

**Table 2: Gromov-Wasserstein Distance Results for Various Targets, for $f_{\text{mod3}}$ dataset.**
| Model L | GW-D for | Top Similar Layers | $D_{\text{min}}=$ |
| --- | --- | --- | --- |
|  | $c_{1}$ | Resid-Post^1 | $0.02$ |
|  | $c_{2}$ | Resid-Post^2 | $0.03$ |
|  | $c$ | Resid-Post^2, Resid-Pre^2, Resid-Post^3, and 6 others | $0.04$ |

Results The results are shown in the Table 2. We see that in the Model L, the GW distance correctly identifies the most similar layers in accordance with different intermediate $c$’s. The final target $c$ contains 9 similar layers all with distance around $0.04$. In Appendix B, we also test probes since the targets are known. Results shows GW distance can be a reliable alternative to the probes.

Figure: Figure 3: Model L (layer-wise training) pairwise GW distance, on $f_{\text{mod3}}$ dataset.
Refer to caption: extracted/5943928/figs/gw_model2.png

Moreover, as previously mentioned GW distance can naturally compare representations across layers with different dimensions. In Fig 4 and Fig 4, we visualize the pairwise GW distance between layer representations without a target for Model L and Model E. Looking at Model L we see predominantly 3 groupings of layers: i) layers roughly from $20$ to $44$ are similar to each other and to layers $52$ to $72$, ii) layers roughly $12$ to $19$ are similar to each other and layers $45$ to $51$ and iii) the initial and last few layers are mainly similar to themselves. Interestingly, the number of groupings corresponds to the 3 functions trained layer-wise in Model L. We also observe differences in patterns across Model L and Model E, suggesting layer-wise and end-to-end training return different networks. Compared to the fixed layer-wise training, end-to-end training in Model E may learn faster in the earlier layers and may not have much to learn in later layers, as the function may not be particularly challenging for it. This could explain why, starting from layer 64, all layers in Model E exhibit similar representations.
Moreover, magnitudes of the distances are also different, with Model L showing larger distances, indicating that learning the targets $c_{1},c_{2}$ result in more functional differences. One possible explanation could be that Model E directly operates in the trigonometry space [], without having to predict the exact integer values until later, thereby suppressing the distances.

Figure: Figure 5: Model 0 Pairwise GW distance, , on $f_{\text{mod}}$ dataset.
Refer to caption: extracted/5943928/figs/gw_model0.png

To gain a deeper understanding of the operations within each transformer block, we visualize pairwise GW distances among layers for Model 0 for dataset $f_{\text{mod}}$ in Figure 5. In this case, we have a total of $31$ representations since only one transformer block is used. We notice the first major difference occurs between layers $13$ and $16$, which are $4$ Attn-Pre (computing key and value product). The second difference occurs between layers $17$ and $20$, which are the first $3$ Attn (computing ${\bm{A}}({\bm{X}})$). This suggests that major computation seems to be done by the attention mechanism. Note that distances are not monotonically increasing across layers, which is expected as the representation spaces can change significantly given the heterogeneity of the operations such as those performed by residual connections and attention within a transformer block.

### 5.2 NLP Tasks

Setup We now apply GW distance to real natural language processing tasks. We experiment on benchmark sentiment analysis datasets, Yelp reviews (^1^11 [https://www.kaggle.com/code/suzanaiacob/sentiment-analysis-of-the-yelp-reviews-data](https://www.kaggle.com/code/suzanaiacob/sentiment-analysis-of-the-yelp-reviews-data)) and Stanford Sentiment Treebank-v2 (SST2) from the GLUE NLP benchmark [], with the goal to predict of the text has positive or negative sentiment, and analyze how different layers from fine-tuning BERT(-base) [] models perform on these datasets. We use the pretrained BERT to generate 4 fine-tuned models, corresponding to a dense model and 3 sparse models with sparsity levels of $25\%$, $70\%$ and $95\%$ using a state-of-the-art structured pruning algorithm []. Training details are in Appendix C. In these BERT models, there are $12$ transformer blocks and we compare the ${\bm{X}}_{\text{in}^{1}}$ and $Y^{l}={\bm{X}}_{\text{out}}^{l}$, for $l=\{1,\dots,12\}$ layers.

Results In the last row of Figure 6(a), we see that the pre-trained BERT does not have major differences among layers, which is not surprising given its accuracy on YELP is only 49.3% (roughly equivalent to random guessing). In the last row of Figures 6(b) to 6(e), we see an interesting pattern emerge, revealing two-to-three major block structures in the (sparse) fine-tuned BERT models identified by our approach. The first major differences occur at block $9$ and then the last three layers ($10$, $11$, $12$) seem to form a distinct block. This seems to indicate that most of the function/task fitting occurs at these later layers. This may be explained by the fact that gradient values are potentially larger at these later layers so they change faster from the pretrained models. The presence of block structures in the GW-distance matrices indicates major functional changes may concentrate at these transition layers. This finding may suggest that for other downstream tasks, we may consider using the representation from layer $8$ and only fine-tuning the later layers. We validate this observation in appendix E.

Figure: (a) Pre-trained
Refer to caption: extracted/5943928/figs/euc_Yelp_pre.png

When sparsifying these models, we observe the more sparse models have lesser differences among the layers (with 95% sparse model in 6(e) having the least differences). This is expected as fewer parameters contribute to the final function output besides others being suppressed. Nonetheless, a similar pattern persists, indicating that later layers differ significantly from earlier ones. This observation is consistent with fine-tuning and sparsification literature [, ], where it has been observed that later layers typically undergo substantial changes during fine-tuning as they focus on task-specific solutions, while the earlier middle layers remain stable as they capture syntactic and semantic patterns of the language necessary for various tasks. In appendix F, we further investigate the GW distance between layers from different models, providing insights into how representations vary across architectures.

We also compare the proposed GW distance with Euclidean, Cosine, and Wasserstein distance as baselines in Figure 6, on the same YELP dataset and with the same settings. Euclidean distance between two layers’ outputs, shown in the first row of Figure 6, can be seen as the GW distance with a fixed identity-mapping transportation plan for each sample. This validates the low-valued diagonal elements. Off-diagonal elements show greater variation, and it is less obvious there are two distinct sub-groups within layers. The similar pattern is also observed with Cosine and Wasserstein distances, with similar strong diagonal pattern but more pronounced block structures than Euclidean distance. On SST2 dataset, we also observe very similar patterns with the GW distance and 3 baselines, for which we refer the readers to appendix G for detailed results. In both datasets, low distance measure are consistently observed in the diagonal elements, but
the overall block structures are not as obvious in the baselines as they are with GW distance, highlighting the effectiveness of the GW distance.

Clustering Besides visualization above, one can also utilize clustering methods to automatically identify the subnetworks from the GW distance. We tested spectral clustering [] on a similarity matrix computed as the reverse pairwise GW distance matrix. This method successfully identified 2 groups with layers $1\sim 8$ and layers $9\sim 12$.

### 5.3 Emergence of Subnetworks During Training

We also visualize the GW distance between layers while fine-tuning the pretrained BERT model on YELP datasets in the entire training process, in order to observe when these subnetwork structures begin to emerge. Figure 7 show a few visualization on GW distances at selected training iterations. Layers distances are low in the beginning (observed in Figure 6(a)), but by iteration 300 the last layer begin to differ from other layers. As training progresses, layers 9, 10, and 11 begin to show at iteration 3k and 15k. These growing differences in GW distance reflect the model’s increasing F1 score on the test data. Overall it show the gradual specialization of layers into distinct sub-networks, with each sub-network potentially focusing on different aspects of the task.

Figure: (a) Iteration = $300$
Refer to caption: extracted/5943928/figs/score_GW_YELP_YELP_bert-base-uncased_batch_300_target_unknown_spectral_GW_scores.json.png

We also plot the mean GW distance of all layer pairs in Figure 8. Figure 8a show the mean GW distance over training iterations, and show it grows over time. Figure 8b shows that mean GW distance versus two different accuracy metric on the test dataset. GW distance grow slowly at first, followed by a rapid increase as the model achieves better accuracy and F1 scores. Such observation is consistent with existing “grokking" behavior, where validation accuracy can suddenly increases well after achieving near perfect training accuracy []. Similarly, Figure 8c shows a rapid increase in mean GW distance in order to achieve a lower training loss.

Figure: Figure 8: Pairwise GW distance in YELP datasets, over training iterations.
Refer to caption: extracted/5943928/figs/iter_vs_meanGW.png

## 6 Discussion

We proposed a novel approach to model interpretation based on functional similarity within intermediate layers of neural networks, using Gromov-Wasserstein (GW) distance to compute such similarities. To the best of our knowledge, our application of GW distance in this context is novel. On algebraic and real NLP tasks, we identified the existence of major sub-components amongst layers, corresponding to functionally meaningful abstractions.
Overall, our method provides an automatic low-cost approach to find sub-components within neural networks, facilitating human understanding. Future work could investigate additional models to observe general trends and applications of functional similarity to additional models and tasks. Theoretical study of other properties of GW distances within the context of neural network interpretability is also an interesting future direction.

#### Limitations and Broader Impact

Our approach first assumes we have access to the intermediate layer representations, which may not be available for some black-box models. Our approach is general, but assumes the proposed distances correctly represent the functional similarities. Our findings are also limited to the datasets and models studied and are not guaranteed to be observed in other scenarios. In terms of broader impact, our approach could be applied widely given its simplicity for identifying functional units in neural networks. However, more investigations on inner mechanisms will have to be done, perhaps building on our approach, in order to fully understand the behavior of neural models.

## Acknowledgments and Disclosure of Funding

We thank colleagues and anonymous reviewers for their helpful feedback in many versions of this paper.

## Appendix A Modular Sum Experiment Details

We use the same architecture and protocols in training, as previous modular papers [, ], based on their available Github repos. Specifically, we use transformer width $d=128$, and each attention head has $32$ dimensions. As a result, MLP has $512$ hidden neurons. ReLU is used as the activation throughout the models,

#### Data

Among all data points ($59^{2}=3481$ of them), we randomly select 80% as training
samples and 20% as validation samples.

#### Hyperparameters

We used AdamW optimizer
[] with learning rate $\gamma=0.001$ and weight decay factor
$\beta=2$. We use the shuffled data as one batch in every epoch. We train models from scratch and train for 26,000 epoches.

#### Search Space

For the $f_{\text{mod3}}$ dataset, we consider all layers in the network, including all representations within transformer blocks. As shown in Table 1, each attention head has 6 intermediate layers, for a total of 24. Each block has an additional 7 layers (1 input layer, Resid-Pre, and 6 intermediate layers). Hence, for three blocks each with four attention heads, we have a total of $93$ representations to evaluate, as each block has $31=24+7$ representations.

## Appendix B Probes on Modular Sum Dataset: When Target is Known

When the target is a value from a known function, we can directly compare outputs between representations from each layer and the known function output. Representations from each layer can be directly compared with the target via a probe. We first consider Model E and then Model L.

#### Linear Probe

Popular linear probes can be used to assess the similarity between a target and any layer’s representation. We perform linear regression of each target $(c_{1},c_{2},c)$ on each of the 93 representations $Y$, and report the residual error as the scoring distance function between $Y$ and $c$’s.

**Table 3: Linear and Nonlinear Probe Results, for $f_{\text{mod3}}$ dataset.**
| Model L | Linear Probe for | Perfect Match? | Top Similar Layers | $D_{\text{min}}=$ |
| --- | --- | --- | --- | --- |
|  | $c_{1}$ | $\checkmark$ | Resid-Post^1 and 21 others | 0 0 |
|  | $c_{2}$ | $\checkmark$ | Resid-Post^2 and 21 others | 0 0 |
|  | $c$ | $\checkmark$ | Resid-Post^3 and Post^2 | 0 0 |
| Model E | Linear Probe for | Perfect Match? | Top Similar Layers | $D_{\text{min}}=$ |
|  | $c_{1}$ | $\times$ | Post^2 | $0.522$ |
|  | $c_{2}$ | $\times$ | Post^1 | $0.93$ |
|  | $c$ | $\checkmark$ | Resid-Post^3 and 5 others | 0 0 |
| Model E | Nonlinear Probe for | Perfect Match? | Top Similar Layers | $D_{\text{min}}=$ |
|  | $c_{1}$ | $\checkmark$ | Resid-Post^1 and 15 others | 0 0 |
|  | $c_{2}$ | $\checkmark$ | Resid-Post^1 and 4 others | 0 0 |
|  | $c$ | $\checkmark$ | Resid-Post^3 and 9 others | 0 0 |

#### Results

Since we perform layer-wise training with Model L, we know the true locations of $c_{1}$ and $c_{2}$, which sit at ${\bm{X}}_{\text{out}}^{1}$ and ${\bm{X}}_{\text{out}}^{2}$ with names Resid-Post^1 and Resid-Post^2, respectively. As shown in the top part of Table 3, a linear regression probe can predict targets perfectly with these two layers. In fact, there are 21 other layers which also show perfect accuracy. For $c_{1}$, these consist of Post^0 and MLP-out^0 from the same block and some layers from the next block, including linear operations with all $k$’s, $q$’s, $v$’s. The final prediction $c$ can be linearly predicted as expected, due to the model’s perfect prediction accuracy.

Naturally we would like to confirm if the same happens with Model E: if we use the same linear probe, does each block in Model E learn the corresponding $c$ at the output of the transformer block? As shown in the mid part of Table 3, we are not able to find any layer that produces a representation that is linearly predictive of $c_{1}$ and $c_{2}$, with the lowest prediction errors at $52\%$ and $93\%$, respectively. Moreover, the most similar layers to $c_{1}$ and $c_{2}$ are in the 2nd block and 1st block respectively, instead of the expected 1st and 2nd blocks. This seems to suggest that Model E does not actually learn any function of $c_{1}$ and $c_{2}$.

#### Non-linear Probe

As discussed previously, to deal with the potentially large search space of functions of the target, a more powerful probe (such as a nonlinear MLP function) may have to be used so that it can detect more complex similarities to $c$. Therefore, we train a two-layer MLP(^2^22We use the neural network classifier from the scikit-learn package, with default parameters.) to predict $c$’s. As shown at the bottom of Table 3, these two-layer MLPs have more predictive power and can perfectly predict the targets, while still showing differences among various layers indicating that the matched layers do capture the intended target functions while other layers do not. Many layers in the 3rd block, for example, have only $1\%$ accuracy relative to $c_{1}$. This indicates that non-linear probes can be used to find subgroups of layers in neural networks. Unlike existing work that primarily focuses on linear probes, we show that non-linear probes, still with limited capacity, are useful.

One issue with using predictive probes to compute the distance measure $D$ is that the target function has to be known. In practice, however, we may not know any intermediate targets, as suggested in the end-to-end training of Model E. While we still can try different target functions and use non-linear probes, the infinite number of possible targets makes such an approach inefficient. This calls for a different strategy to differentiate sub-components in a network through representation similarity.

## Appendix C Real NLP Experiment Details

We analyze a BERT-base-uncased [] model based on our optimal matching inspired mechanistic interpretability approach. We fine tune it on two well known datasets in NLP; i) Yelp reviews ([https://www.kaggle.com/code/suzanaiacob/sentiment-analysis-of-the-yelp-reviews-data](https://www.kaggle.com/code/suzanaiacob/sentiment-analysis-of-the-yelp-reviews-data)) and ii) Stanford Sentiment Treebank-v2 (SST2), which is part of the GLUE NLP benchmark []. Both of these are sentiment analysis tasks, where the goal is to predict if a piece of text has positive or negative sentiment. The Yelp dataset has hundreds of thousands of reviews, while the SST2 dataset has tens of thousands of sentences. The training details are as follows: i) Hardware: $1$ A100 Nvidia GPU and $1$ intel CPU, ii) Max. Sequence Length : $256$, iii) Epochs: $1$, iv) Batch Size: $16$ and v) Learning Rate: $2e^{-5}$ with no weight decay. The accuracy on Yelp was $97.87\%$, while that on SST2 was $92.4\%$. Without fine tuning the pre-trained BERT models accuracy on Yelp and SST2 was $49.29\%$ and $50.34\%$ respectively indicative of random chance performance.

We also fine tuned a series of sparse models on these datasets. The method we used to sparsify was a state-of-the-art dynamic sparse training approach NeuroPrune [], which leads to high performing structured sparse models. Using this approach and the same training settings as above we created BERT models with $25\%$, $70\%$ and $95\%$ sparsity which had accuracies of $96.31\%$, $97.53\%$ and $96.22\%$ respectively for the Yelp dataset and accuracies of $90.25\%$, $88.5\%$ and $84.4\%$ respectively for the SST2 dataset. We then used the resultant models for our analysis.

## Appendix D Alignment from GW distance

We plot a tSNE projection [] down to 2 dimensions, on a batch of 16 samples (color indicative of sample) on YELP, and visualize it in Figure 9(e). As one can see, the sample neighborhood changes across layers, which can be indicative of functional changes but something that is not captured by comparing distributions. However, GW can also account for such changes.

Figure: (a) Pretrained
Refer to caption: extracted/5943928/figs/YELP_bert-base-uncased_target_unknown_spectral_GW_pre_scores_tsne.png

To show the exact transportation plan from GW distances, we choose plot one batch of data with size 16, and show the transportation plan over 5 random layer pairs in Figure 10. As one can see, the transportation plan does not conform to identity-mapping. Both Wasserstein and Euclidean distance will likely have trouble handle in this case. We also note that the transportation plan shown Figure 10 is a permutation of the original data, rather than a distributed transportation plan. This behavior is consistent with existing Wasserstein optimal transport plan under certain conditions [].

Figure: (a)
Refer to caption: extracted/5943928/figs/alignment_GW_1.png

To complement Figure 2 on other fine-tuned BERT models on YELP, we also plot all the histograms across 5 models in Figure 11.

Figure: (a) Pretrained
Refer to caption: extracted/5943928/figs/YELP_bert-base-uncased_target_unknown_spectral_GW_pre_scores_histo.png

## Appendix E Fine-tuning with Different Layers

Since the GW distance indicates significant changes occurs only at later layers in YELPS, we investigate performance of fine-tuning only partial layers from pretrained models, by freezing early layers during training and training only later layers alongside a classification layer (denoted as C) at the end. In Table 4, we can see that there is no significant performance differences between fine-tuning layer 8 to 12 and fine-tuning layer 9 to 12 (0.04% drop). On the other hand, the accuracy drops 6 times more by freezing layer 1 to 9, with 0.25%. Freezing layer 1 to 10 results 0.49% drop, and finally fine-tuning only 12 results $3.59\%$ drop. These findings validate that the later layers are crucial for significant functional changes.

**Table 4: Accuracy of fine-tuning partial layers in various BERT models. C denotes the classification layer on top of BERT models.**
| Fine-tune | All | 8$\sim$12 + C | 9$\sim$12 + C | 10$\sim$12 + C | 11$\sim$12 + C | Only 12 + C |
| --- | --- | --- | --- | --- | --- | --- |
| Accuracy (%) | $97.87$ | 97.47 | 97.43 | 97.19 | 96.7 | 93.11 |

## Appendix F Cross Model Comparison

We can also use GW distance to compare layers from different BERT models. Shown in Figure 12, pretrained and densely fine-tuned BERT models exhibit different similarity measures when compared to fine-tuned BERT models with different levels of sparsity.

Figure: Figure 12: Pairwise distances on YELP dataset, of layers across two different BERT models. TOP: Densely fine-tuned BERT model vs fine-tuned BERT models with different sparsity levels. Bottom: Pretrained BERT model vs fine-tuned BERT models with different sparsity levels.
Refer to caption: extracted/5943928/figs/gw_Yelp_dense_vs_sparse25.png

## Appendix G SST2 Datasets

Besides YELP Datasets, we also tested the GW distance on SST2 dataset. Results on SST2 dataset are shown in Figure 13 again confirm there exist two-three different groups in terms of functional similarity. The first major difference is seen at layers $10$ and $11$, while layer $12$ forms its own block. When sparsifying these models, lesser differences are observed in general as also seen on the YELP dataset. Other baselines provide less clarity on the division of sub-components.

Figure: (a) Pretrained
Refer to caption: extracted/5943928/figs/euc_SST_pre.png

## Peer Reviews

confidence: 3
contribution: 3
initial_rating: 6
presentation: 2
questions: - Can we use this metric in a quantitive manner?
- Can observations reported in Sec 5.2 and 5.3 be shown experimentally? For example, the paper states, "The first major differences occur at block 9 and then the last three layers (10, 11, 12) seem to form a distinct block. This seems to indicate that most of the function/task fitting occurs at these later layers. This may be explained by the fact that gradient values are potentially larger at these later layers so they change faster from the pretrained models." This can be shown experimentally.
- The algebraic task is simple which is suitable for a controlled study, can a computer vision task be included to compare with the observations reported for the NLP task?
soundness: 3
strengths: - A new metric for measuring the functional similarity of layers which has a number of desirable properties.
- The metric behaves experimentally as expected which may aid in improving our understanding of the inner workings of neural networks.
- Experiments on algebraic and language tasks under a number of settings, e.g., fine-tuning, sparsity, etc., reveal some sensible patterns.
summary: This paper proposes a new metric based on the Gromov-Wasserstein distance to measure the similarity between layers in neural networks, which can be used for mechanistic interpretability and understanding of the inner workings of neural networks. Experiments are conducted on algebraic and language tasks and a number of interesting patterns and observations are reported.
weaknesses: - The metric relies on qualitative assessment of the numerical results, which may be complicated for large neural networks, or may be subjective or ambiguous.
- Missing comparison with state-of-the-art metrics, e.g., Yu et al., 2024, Klabunde et al., 2023, Lohit & Jones, 2022 (listed in Sec 2)
- The mathematical notation is inconsistent and in some minor instances incorrect, e.g., "Let f0 be a function that produces a set of output Y0 given a set of input X, where Y0 ∈ Rn×dy", i.e., so is Y0 a set or a matrix?
- Overall, the paper was hard to follow, and I had to read it multiple times to appreciate the details.

---

confidence: 3
contribution: 1
initial_rating: 3
presentation: 2
questions: - *Unconnected Related works*. I think that the related work section could benefit from some restructuring. For example, it is not clear to me why the "Algorithm discovery in algebraic problems" and "Real applications" (e.g., here "vision and language components" are mentioned, why?) paragraphs are connected to this work. On the contrary, I found the last one (Similarity Measure between Neural Network Layers) to be clear and well-connected to the rest of the paper. I think a couple of sentences for each paragraph explaining the connection with the rest of the work could help.
- *Section 5.2, Setup paragraph*: when the three sparse models are introduced, I would suggest explaining why they are used/useful in this experiment. Even a single sentence like "The sparse models are used to force their respective models to condense more the information in the few weights they have remaining, so we can analyze any link between this constraint and their structural similarity".
- *Section 5.1, Training procedure paragraph, line 356:* "In all these models, we are able to achieve 100% prediction accuracy". Is it on a validation set hinting at perfect modeling of the target function or on a training one useful to understand if the models could overfit or not?
- There's a typo in line 146/147 (intro to Section 3). A full stop is missing after "as a similarity search problem".
- *Experiment 5.2, Results paragraph*: on line 439, it is written: "from layer 8 and only fine-tuning the later layers". As a minor comment, I would suggest changing this sentence to “freezing the model up to layer 8 and only fine-tuning the layers after that” to make the sentence clearer.
- The figures are dense and have really small ticks/labels. For example, Figure 6 and Figure 2 require a lot of zooming in. My suggestion to improve the readability would be to start by sharing the x and y axes labels and avoid adding details in the titles (i.e., the KL divergence, especially because it appears only there and is not commented on in the text).
- Section 4, start of the “Unknown Targets” paragraph (line 230/231): I think there are missing words, the opening sentence is not grammatically correct. A rewrite could be “When the search target is unknown, functionally similar parts cannot be identified by comparison to a predefined set of target functions.”
soundness: 2
strengths: - The research question is extremely fitting for the venue, as it directly studies neural representations and is not one of the many deep learning topics that only indirectly involve them.
- The approach is well-founded and investigated even for training dynamics, with an insightful analysis in Figure 8.
- The high-level structure of the paper is robust: overall idea -> synthetic investigation -> real-world case -> in-depth analysis for the real-world one. I particularly appreciated the experiment structuring with their modularity in "setup" and "results".
summary: The paper studies the functional similarity of sub-modules by applying Gromov-Wasserstein distance (GW) to their outputs. The core idea is that sub-modules (e.g., attention head/layer outputs at different depths along a transformer) can be seen as individual functions over their input, and these functions can be studied through the similarity of their outputs. However, existing methods in the literature involve either indirectly comparing the information content with respect to another function's output (a downstream task applied independently on each sub-module output) or directly with some (dis)similarity measure that assumes their pre-existing compatibility/comparability between representations. The approach proposed in this paper falls in the second category of methods, and by using GW models, a more generic measure (i.e., that takes into account more invariants) than the considered baselines. The approach is tested on a controlled environment (synthetic dataset and controlled training dynamic of an ad-hoc model) and on an NLP task (sentiment analysis) on a pre-trained model (BERT).
weaknesses: - I had a hard time understanding the **scope and goals** of the work: the paper is claimed to be aiming at mechanistic interpretability, but the methods, to my understanding, do not provide any human-understandable finding. 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.
- The **novelty** of the approach is also not entirely clear — in particular, in what aspect is the approach novel? Prior work has already investigated the (dis)similarity of different layers, both intra- and inter-networks [1, 2], and has also studied the block structure hinted at in this paper [2]. The similarity measure, by itself, reminds me of a topology-based measure in [3]. Some discussion on the novelty of the measure would help to understand. If it is novel, then how does it compare to widely established measures such as CKA [4] or CCA-based measures? The former is cited in the work, but I don’t see it in any experiments. CCA and CKA take into account (in different ways) the “sample-to-sample” relationships, so they are more suited as baseline methods to see the advantages of GW. In general, [1, 2] are too relevant to not be discussed in the work.
- **Clarity of the method.** While it is clear from the text why GW could be a strong measure to apply in this context, the method part is a bit confusing. From my understanding, GW in this paper is just “evaluated” on the representation pairs since the correspondence ($\pi$) is known (i.e., the representation for sample X at layer i corresponds to the representation of sample X at layer j). Is it correct? If this is the case, I would suggest specifying that $\pi$ is not optimized in this setting because equation 2 suggests the opposite. If this is not the case, it’s unclear to me why we would want to optimize a correspondence that is already known at the risk of matching different samples across layers. Can the authors kindly expand on this?
- I found the writing to be a bit confusing, often mixing low-level details with high-level takeaways. I suggest moving any detail not immediately necessary to understand the experiments to an appendix to allow the reader to focus on the core message.
- **Experiment inconsistencies.** There are a few mismatches between the tables/figures content and the text. For example:
    - Table 1 introduces "All representations considered in experiments". However, all the other experiments have only one of them considered. Table 2 is the only exception, where apparently, all of them were used to find the most similar ones across all possible representations. I would suggest to introduce all the variations in this experiment (the one related to Table 2) since they are applied only there, otherwise it comes up as confusing, especially since it's the first table in the paper.
    - Figure 2 (related to the experiment for section 4.1), contains the “Pretrained” (a) and “Fine-Tuned” (b) sub-figures, but I didn’t find in the text any description of what they are. The associated paragraph “Distance Distributions” just refers to the generic Figure 2 without describing its content or commenting on the different results according to the model.
    - The paragraph “Neighborhood Change” refers to a T-SNE plot showing the results that are commented on, but the figure is in the appendix (Figure 9e, Appendix D). Since the message of this paragraph is to study neighborhood consistency, I would use a quantitative metric for measuring neighborhood overlap (even a Jaccard similarity between the top-k neighbors, I think, could suffice), adding it to the main text rather than (only) a qualitative plot in the appendix.
    - Table 2 for experiment 5.1 shows results that I think are not well-commented in the text. What's missing is the message from them: what's the meaning of the top similar layers being the ones of those specific kinds and not others?
    - In the synthetic data experiment (section 5.1), 3 models are used: Model 0, Model E, and Model L. According to the *Training procedure/setup* paragraphs, Model E and L have three layers each. However, in the Results paragraph of the same section and in the related figures (3 and 4), they appear to have more than 90 layers each.
    
    [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.

---

confidence: 4
contribution: 2
initial_rating: 5
presentation: 3
questions: - Modular arithmetic experiment: could the authors justify why the choice of the GW distance represent a good choice for extracting compared to measures as CKA [b], canonical correlations analysis and other kernel distance measures mentioned in (Klabunde,2023)  

- Networks initialized with different feed converge to different representations (see e.g. [a,b,c]), although they might perform similarly. In the modular arithmetic experiment how is this taken into account? It would be interesting to measure the GW distance between networks trained on the same data and task but with different initializations.  Alternatively the experiments could be repeated averaging results for multiple seeds.  


    - [a] Moschella, L. "Relative representations enable zero-shot latent space communication (2023)." ICLR 2023

    - [b] Li, Yixuan, et al. "Convergent learning: Do different neural networks learn the same representations?." PMLR

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



- What is the exact difference between the implemented version of Wasserstein and GW distance in the paper?
soundness: 3
strengths: - The direction proposed by the paper is interesting and relevant: extracting subcomponents of neural networks based on similarity measures is a core direction for fields such as representation alignment, interpretability and model selection potentially leading to new strategies for transfer learning and neural network pruning and merging, among others.


- The paper writing is clear and the proposed method is simple.
summary: The paper proposes to compute functional similarity between layers of neural networks using Gromov-Wasserstein distance in order to identify subcomponents in the architecture of trained models. In particular, the goal is compare different neural networks layers in order to identify subcomponents of the networks, which may be specific to solve particular tasks. Experiments are conducted on synthetic modular arithmetic tasks and on BERT text models fine-tuned on sentiment analysis tasks.
weaknesses: While the direction of the paper is interesting the realization seems not to be as broad as it stated in the introductory section and title. In particular: 

- As stated in the paper one core objective is "to detect how many distinct (complex) functions exist in a learned network, and which layers correspond to each such function". While the question is very relevant, I don't see direct evidence of the fact that blocks of layers with identify with a given complex function, except for the synthetic modular experiment. In particular on real data there is no analysis on correlation on accuracy on downstream tasks and GW similarity of layers, which to the bet of my understanding seems necessary in order to assess if a subcomponent learned a specialized function of the input. 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.  Related to this: 
      
    - In Figure 8 shows just the mean GW distances across all layers which doesn't allow to validate any task specialization of the block of layers with similar GW distance


- Experiments and method section are focused on transformer networks, not general neural network architectures as discussed in th intro/title. However measuring GW distances in representation space should be agnostic of the architecture. Why the focus just on this architecture? 


- The paper focuses on the synthetic task of modular arithmetic and the one of sentiment analysis on text. In order to make the statements more general more tasks, datasets and /or modalities should be taken into account. For example, it might be interested to see in Vision models based on ResNet or convolutional architectures, early layers subcomponents would be able to solve more task which depends on local statistics of the input, while layers closer to the inputs are more expect to solve task which rely on global information, such as image classification.

- The synthetic task on modular arithmetic is trained from scratch but it doesn't seem to take into account possible effects given by different random initializations. See questions section for more details on this. 

- The following papers should be included for discussion and/or comparison:

    - Tsitsulin, Anton, et al. "The shape of data: Intrinsic distance for data distributions." ICLR 2020 

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

In the former, it is proposed a similarity measures based on heat kernels of graph on latent representations which gives a lower bound on the GW distances between embeddings. 
In the latter is proposed the CKA measure to compute similarity across data representations. This shares most of the favorable properties highlighted for the GW distance and is computationally efficient. Also it shows how comparing different layers of networks reflects the functional structure of each component (e.g. see the Figure 4 in the paper). 

More broadly there are many contributions from the field of model merging which might be useful to take into account and discuss, for example: 

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

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

- More details on the method to compute GW distance could be provided in the paper in order to make it self-contained. Specifically, why the choice of the Fused Gromov Wasserstein method in particular?

---

confidence: 4
contribution: 1
initial_rating: 3
presentation: 2
questions: - Could you elaborate on the necessity and benefits of understanding neural networks through the identification of sub-networks?
- Could you detail the technical innovations presented in this paper that extend beyond the Gromov-Wasserstein method?
- Could you present additional experimental results across diverse tasks, datasets, and learning methodologies?
- Could you discuss a real-world application of sub-network identification, such as improving the inference speed of neural networks?
soundness: 1
strengths: - The proposed method discerns similar and dissimilar layers within the network, revealing potential sub-components, and addresses challenges from varying distributions and dimensionalities in intermediate representations.
- The presentation of the experimental results is precise and well-organized.
summary: This paper proposes a novel method for identifying sub-components within neural networks, addressing challenges arising from varying distributions and dimensionalities in intermediate representations.
weaknesses: - The experiments are limited due to insufficient validation across diverse tasks, datasets, and learning methods. This includes gaps in testing with image (classification) datasets, self-supervised learning methods, etc.
- The contribution of the proposed method is unclear. Can the identified sub-network effectively serve as a proxy, replacing the entire original network to accelerate inference for specific downstream tasks?
- The motivation for identifying subnetworks requires clarification. Specifically, the phrase "understanding the network by investigating the existence of distinct subnetworks" needs a more explicit explanation. It is unclear why identifying these subnetworks contributes to a deeper comprehension of the neural network as a whole.
- The proposed technical solution appears straightforward. To address challenges associated with varying distributions and dimensionalities in intermediate representations, the primary approach involves identifying subnetworks using the Gromov-Wasserstein (GW) distance. However, since GW distance has been previously proposed, it could potentially reduce the novelty of this work [1].

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

---

confidence: 3
contribution: 2
initial_rating: 3
presentation: 2
questions: 1. The motivation for why it is necessary to identify subnetworks is unclear. Could the authors provide more explanation about the relationship between model interpretation and subnetwork?
2. Additionally, the field of model pruning involves identifying and retaining stronger subnetworks. What are the differences and connections between this approach and model pruning?
3. Could the author provide more explanation about Figure 3, particularly the relationship between the three identified groups and the model architecture?

I am happy to engage further during the discussion phase.
soundness: 2
strengths: 1. This paper proposes a novel perspective for model interpretation that minimizes both human and computational expenses.
2. The proposed methodology can effectively identify subnetworks.
summary: This paper proposes a novel perspective on model interpretation by utilizing the Gromov-Wasserstein (GW) distance to compute similarities between model layers and identify sub-networks. For algebraic and real-world NLP tasks, the proposed methods effectively identify sub-components.
weaknesses: 1. The motivation for identifying subnetworks is unclear. Once identified, I'm curious about how they can enhance our understanding of networks in real-world applications. 
2. 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.
3. There are some typos, such as a missing comma on line 146.

Meta-review: This paper introduces a method to automate mechanistic interpretability by identifying subnetworks that are functionally distinct/similar. This uses the GW distance, which allows matching representations across different intermediate representations. They perform experiments on algebraic, language and vision tasks. 

Ultimately, no reviewer championed the paper. The main limitations called by the reviewers are:

Improvement on motivation and how it is reflected in the experiments (in particular validation on real data sets. This was raised by all reviewers, with the two most critical ones being dissatisfied with the authors answer (and the two more positive ones also calling this as an unaddressed liitation)

Decision: Reject