Neural Sculpting: Uncovering hierarchically modular task structure in
  neural networks through pruning and network analysis
Natural target functions and tasks typically exhibit hierarchical modularity
-- they can be broken down into simpler sub-functions that are organized in a
hierarchy. Such sub-functions have two important features: they have a distinct
set of inputs (input-separability) and they are reused as inputs higher in the
hierarchy (reusability). Previous studies have established that hierarchically
modular neural networks, which are inherently sparse, offer benefits such as
learning efficiency, generalization, multi-task learning, and transfer.
However, identifying the underlying sub-functions and their hierarchical
structure for a given task can be challenging. The high-level question in this
work is: if we learn a task using a sufficiently deep neural network, how can
we uncover the underlying hierarchy of sub-functions in that task? As a
starting point, we examine the domain of Boolean functions, where it is easier
to determine whether a task is hierarchically modular. We propose an approach
based on iterative unit and edge pruning (during training), combined with
network analysis for module detection and hierarchy inference. Finally, we
demonstrate that this method can uncover the hierarchical modularity of a wide
range of Boolean functions and two vision tasks based on the MNIST digits
dataset.
Shreyas Malakarjun Patil, Loizos Michael, Constantine Dovrolis
Neural Sculpting: Uncovering hierarchically modular task structure through pruning and network analysis

###### Abstract

Natural target functions and tasks typically exhibit hierarchical modularity -- they can be broken down into simpler sub-functions that are organized in a hierarchy. Such sub-functions have two important features: they have a distinct set of inputs (_input-separability_) and they are reused as inputs higher in the hierarchy (_reusability_). Previous studies have established that hierarchically modular neural networks, which are inherently sparse, offer benefits such as learning efficiency, generalization, multi-task learning, and transferability. However, identifying the underlying sub-functions and their hierarchical structure for a given task can be challenging. The high-level question in this work is: if we learn a task using a sufficiently deep neural network, how can we uncover the underlying hierarchy of sub-functions in that task? As a starting point, we examine the domain of Boolean functions, where it is easier to determine whether a task is hierarchically modular. We propose an approach based on iterative unit and edge pruning (during training), combined with network analysis for module detection and hierarchy inference. Finally, we demonstrate that this method can uncover the hierarchical modularity of a wide range of Boolean functions and two vision tasks based on the MNIST digits dataset.

## 1 Introduction

Modular tasks typically consist of smaller sub-functions that operate on distinct input modalities, such as visual, auditory, or haptic inputs. Additionally, modular tasks are often hierarchical, with simpler sub-functions embedded in, or reused by, more complex functions [1]. Consequently, hierarchical modularity is a key organizing principle studied in both artificial and biological systems. In neuroscience, the hierarchical modularity of the brain's neural circuits is believed to play a crucial role in its ability to process information efficiently and adaptively [2; 3; 4; 5]. This hierarchical organization allows the brain to break down complex tasks into simpler sub-tasks, which can be processed in a distributed manner. Translating the hierarchical modularity of the brain's neural circuits to artificial neural networks (NNs) can potentially lead to more efficient, adaptable, and interpretable learning systems. Prior works have already shown that modular NNs efficiently adapt to new tasks [6; 7; 8] and display superior generalization over standard NNs [9; 10; 11]. However, those studies assume knowledge of the task's hierarchy, or hand-design, modular NNs at initialization. Given an arbitrary task, however, we normally do not know its underlying sub-functions or their hierarchical organization. The high-level question in this work is: if we learn a task using a NN, how can we uncover the underlying hierarchical organization of sub-functions in that task?

Recent studies have demonstrated that certain hierarchically modular structures can emerge during the training of NNs. Previous methods to extract those structures in NNs can be categorized into _structural_ and _functional_. Structural methods organize trained networks based on the structural properties of units, such as their connectivity and edge-weights [12; 13; 14], while functional methods
consider characteristics based on unit activation [15; 16]. However, it is unclear whether the structures extracted reflect the underlying hierarchy of sub-functions in a task. In a recent study, Csordas et al. ([17]) proposed a method that identifies sub-networks or modules in NNs that learn specific sub-functions. However, this method requires knowledge of the exact sub-functions involved, and cannot be applied when such knowledge is unavailable. Ideally, modules corresponding to specific sub-functions should emerge through a training strategy, and a method should be available to detect these modules without explicit knowledge of the corresponding sub-functions.

Biological networks exhibit hierarchically modular structures where clusters of nodes with relatively dense internal connectivity and sparse external connectivity learn specific sub-functions [4; 5; 18; 19]. Drawing inspiration from this, we propose _Neural Sculpting_, a novel approach to train and structurally organize NN units to reveal the underlying hierarchy of sub-functions in a task. We begin by considering sub-functions that are input-separable and reused within a hierarchically modular task. We show that standard training of NNs does not result in the acquisition of structural properties that reflect these sub-function properties. To address this, we propose an iterative unit and edge pruning method to train NNs, which results in sparse networks that do acquire those previous structural properties. Further, we propose a method based on network analysis to uncover modules within these sparse NNs. Finally, we showcase the capability of the proposed methodology to reveal the structure of diverse hierarchically modular tasks. To the best of our knowledge, this paper is the first to analyze specific sub-function properties within a hierarchically modular task and propose an end-to-end methodology to uncover its structure. This work also sheds light on the potential of pruning and network analysis methods to uncover and harness structural properties in NNs.

### Preliminary

To represent the decomposition of a task, we visualize it as a graph. Initially, we consider Boolean functions and their corresponding function graphs. A Boolean function \(f:\{0,1\}^{n}\rightarrow\{0,1\}^{m}\), maps \(n\) input bits to \(m\) output bits. We define the set of gates \(G\) as \(\{\wedge,\vee,\mathbb{I}\}\), where \(\wedge\) represents logical conjunction (AND), \(\vee\) represents logical disjunction (OR), and \(\mathbb{I}\) represents the identity function (ID). Additionally, we define the set of edge-types \(E\) as \(\{\rightarrow,\neg\}\), where \(\rightarrow\) represents a transfer edge and \(\neg\) represents a negation edge.

**Function Graph:** A \((G,E)\)**-graph** representing a Boolean function \(f:\{0,1\}^{n}\rightarrow\{0,1\}^{m}\) is a directed acyclic graph comprising: \(n\) sequentially ordered vertices with zero in-degree, designated as the _input nodes_; \(k\) vertices with non-zero in-degree, designated as the _gate nodes_, with each vertex associated with a gate in \(G\) and each of its in-edges associated with an edge-type in \(E\); and \(m\) sequentially-ordered vertices with zero out-degree, designated as the _output nodes_.

We use the \(G\) and \(E\) defined above, as NN units have been demonstrated to learn these universal gates [20]. Function graphs, which break down complex computations into simpler ones, are typically sparse. Modularity in sparse graphs refers to the structural organization of subsets of nodes that exhibit strong internal and weak external connectivity, which are grouped into modules.

**Sub-function:** A _sub-function_ is a subset of nodes within the function graph that collectively perform a specific task or computation. The sub-function is characterized by having strong internal connectivity between its nodes, meaning that they are highly interdependent and work together to achieve the desired output. At the same time, nodes in the sub-function have weak external connectivity, indicating that they are relatively independent from the rest of the graph.

**Input Separable:** Two sub-functions are considered _input separable_ if their in-edges originate from distinct and non-overlapping subsets of nodes in the function graph.

**Reused:** A sub-function is _reused_ if it has two or more outgoing edges in the function graph.

**Training NNs on target Boolean functions:** We start by obtaining the truth table of each function graph, which serves as our data source. The training set \((\mathbf{\mathcal{X}}_{t},\mathbf{\mathcal{Y}}_{t})\) includes the complete truth table with added random noise (\(\mathcal{N}(0,0.1)\)) during each iteration to increase the number of training samples. The validation set \((\mathbf{\mathcal{X}}_{v},\mathbf{\mathcal{Y}}_{v})\) consists of the noise-free rows of the truth table. We use multi-layered perceptrons (MLPs) with ReLU activation for the hidden units and Kaiming weight initialization to learn the Boolean functions. The loss function is bitwise cross-entropy with Sigmoid activation. The NNs are trained using Adam optimizer with L2 regularization of \(1e-4\).
## 2 Standard training of NNs is not enough

In this section, we show that NNs through standard training, do not acquire structural properties that reflect the properties of sub-functions. Specifically, we consider the two properties in isolation by constructing two graphs: one with input separable sub-functions, and the other with a reused sub-function.

**Input Separable:** The first function graph we consider has 4 input nodes and 4 output nodes. The output nodes \(\{y_{1},y_{2}\}\) depend only on \(\{x_{1},x_{2}\}\), while the output nodes \(\{y_{3},y_{4}\}\) depend only on \(\{x_{3},x_{4}\}\) (Figure 0(a)). We train various NN architectures to learn the target function with perfect accuracy on the validation set. The property of the function graph that reflects input-separable sub-functions is that there are no paths from input nodes \(\{x_{1},x_{2}\}\) to output nodes \(\{y_{3},y_{4}\}\) and from input nodes \(\{x_{3},x_{4}\}\) to output nodes \(\{y_{1},y_{2}\}\). However, in a neural network, all input units are connected to all output units through the same number of paths. Therefore, we analyze the strength of their learned relationship by considering the product of weight magnitudes along those paths.

Consider a trained NN parametrized by \(\mathbf{\theta}\in\mathbb{R}^{a}\) and the set of all paths, \(\mathbf{\mathcal{P}}\). Each edge weight \(\theta_{i}\) is assigned a binary variable \(p_{i}\) to indicate whether it belongs to a given path \(p\in\mathbf{\mathcal{P}}\). We define the edge-weight product of a path as \(\mathbf{\pi}_{p}=\prod_{i=1}^{a}|\theta_{i}^{p_{i}}|\), and \(\mathbf{\pi}(i,j)=\sum_{p\in\mathbf{\mathcal{P}}_{i\to j}}\mathbf{\pi}_{p}\) represents the sum of \(\mathbf{\pi}_{p}\) for all paths from input unit \(i\) to output unit \(j\). We evaluate \(\mathbf{\pi}(i,j)\) for \(i=1,2\) and \(j=3,4\), and \(\mathbf{\pi}(i,j)\) for \(i=1,2\) and \(j=1,2\). If the former is significantly smaller than the latter, then the input units \(1\) and \(2\) are not used by output units \(3\) and \(4\).

We perform a two-sample mean test with unknown standard deviations, with \((\mu_{1},s_{1})\) representing the mean and sample standard deviation of \(\mathbf{\pi}(i,j)\) for \(i=1,2\) and \(j=3,4\), and \((\mu_{2},s_{2})\) representing the mean and sample standard deviation of \(\mathbf{\pi}(i,j)\) for \(i=1,2\) and \(j=1,2\). The null hypothesis \(H_{0}\) is \(\mu_{1}=\mu_{2}\), and the alternative hypothesis is \(\mu_{1}<\mu_{2}\). A similar test is performed for input units \(3\) and \(4\). Figure 0(b) shows an example heat map for \(\mathbf{\pi}(i,j)\), where \(i,j\in\{1,2,3,4\}\). For the example shown we cannot reject the null hypotheses. Similarly, we conducted statistical tests on nine different NNs with varying architectures and seed values, and the results are summarized in Table 1. For NNs with a single hidden layer, we can reject the null hypotheses, but for deeper NNs, both the null hypotheses cannot be rejected.

**Reused:** Consider the function graph shown in Figure 1(a), consisting of 4 input nodes and 16 output nodes. First, we construct an intermediate sub-function \(g(X)\) such that the three gate nodes depend on all the inputs. This sub-function is utilized by two separate gates, \(f_{1}(g)\) and \(f_{2}(g)\), which are then used 8 times by different output nodes. All paths to the output nodes pass through three gate nodes in the first hierarchical level and two gate nodes in the second level. We analyze the edge-weight product of the paths from hidden units to output units. Let \(\mathbf{\pi}_{p}^{l}\) denote the sum of \(\mathbf{\pi}_{p}\) for all paths originating from hidden layer \(l\), which contains \(N_{l}\) hidden units. We compute the minimum number of units, \(N_{P}^{l}\), that are necessary to achieve \(P\%\) of the total \(\mathbf{\pi}_{p}\). By comparing \(N_{P}^{l}\) to that of the function graph, we can conclude whether the NN has learned to reuse intermediate states.

\begin{table}
\begin{tabular}{c|c|c|c} \hline Width & 24 & 36 & 48 \\ \hline Hidden Layers & & & \\ \hline
1 & & & \\ \hline
2 & & & \\ \hline
3 & & & \\ \hline \end{tabular}
\end{table}
Table 1: NN architectures with varying widths and depths for which both null hypotheses are rejected.

Figure 1: a. Function graph with input separable sub-function, b. Edge-weight product of paths from input units to output units in trained NNs.

Figure 2: a. Function graph with reused sub-function, b. the number of units covering 90% of the total edge-weight product of paths.
We independently trained NNs with two hidden layers and different widths on the target Boolean function. The resulting \(N_{90}^{l}\) values for different layers are shown in Figure 2b. Our analysis indicates that, as the width of the NN increases, so does \(N_{90}^{l}\) in both the hidden layers, and these values are consistently close to the actual width of the NN. We trained an NN with hidden layers of width 3 and 2, respectively, to confirm that NNs with those widths can learn the function well.

## 3 Iterative pruning of NNs

In the previous section, we demonstrated that standard NN training is inadequate for learning structural properties that reveal input separable or reused sub-functions. Nevertheless, we observed that NNs with a relatively low number of parameters could still acquire these properties. Since hierarchical modularity is a property of sparse networks, we explored NN pruning algorithms as a means of achieving this sparsity. Prior research has shown that pruning NNs can reduce their number of parameters without compromising their performance [21; 22; 23]. Edge pruning methods typically require specifying a target NN density or pruning ratio [21; 24], but determining a density value that preserves the NN's performance remains an open question. Recently, iterative magnitude pruning [25; 26] has emerged as a natural solution to this problem. The algorithm prunes edges iteratively, removing some edges in each iteration and then retraining the NN. This process can be repeated until the NN achieves the same validation performance as the dense NN while being as sparse as possible.

**Iterative edge pruning:** Consider the initial edge pruning step \(p_{e}\). We train the dense NN and prune \(p_{e}\%\) of the edges with the lowest weight magnitude. The sparse NN resulting from this pruning is then trained with the same number of epochs and learning rate schedule as the original NN. We repeat this process until the sparse NN can no longer achieve the same validation accuracy as the dense NN. At this point, we rewind the NN to the previous sparse NN that achieved the same validation accuracy as the dense NN and update the value of \(p_{e}\) as \(p_{e}=p_{e}/2\). We repeat this process until \(p_{e}\) becomes lower than the required step to prune a single edge, thereby ensuring that the lowest possible density is achieved.

We next apply the tests developed in the previous section to analyze whether edge-pruned NNs acquire structural properties resembling those of the sub-functions. Specifically, we train and prune different NN architectures to learn the function graph with input-separable sub-functions. For the example shown in Figure 3a, as well as for NNs with different architectures, we can reject the null hypotheses in favor of the alternate hypotheses. This implies that edge-pruned NNs acquire structural properties that enable them to recognize input-separable sub-functions.

Consider the previous function graph that has reused sub-functions and NNs with increasing widths. We independently train and prune each NN architecture and determine the number of units \(N_{90}^{l}\) (Figure 3b). We find that in a majority of the trials, \(N_{90}^{2}=2\), which are reused by the output units. Although there is a significant reduction in \(N_{90}^{1}\), the edge-pruned NNs fail to identify 3 units that are being reused. We observe that edge-pruned NNs identify sparse connectivity between units and reuse those units. However, they do not identify units corresponding to sub-functions that are not sparsely connected yet reused. We hypothesize that this may be due to the initial pruning iterations where some edges from all units are removed, leaving no hidden units with dense connectivity to learn densely connected and reused sub-functions. Additionally, this could also be due to very large NN widths and the absence of an objective conditioning NNs to utilize as few hidden units as possible. Therefore, next we introduce an iterative unit pruning method to limit the number of units used.

**Iterative hidden unit pruning:** To prune hidden units, we first train the neural network and then assign each hidden unit a score. We eliminate \(p_{u}\%\) of the hidden units with the lowest scores and train the resulting pruned network for the same number of epochs and learning rate schedule as the original network. We repeat this two-step pruning process iteratively.

Figure 3: a. Edge-weight product of paths from input units to output units in edge-pruned NNs; b. the number of units covering 90% of the total edge-weight product of paths in edge-pruned NNs.
We use loss sensitivity as the scoring metric for the units [27; 28; 29; 30], which approximates the change in the network's loss when the activation of a particular hidden unit is set to zero. Specifically, for a unit \(i\) in hidden layer \(l\), with activation \(a_{i}^{l}\), the score is computed as :

The iterative unit pruning process continues until the NN can no longer achieve the same validation accuracy as the original NN. To ensure that we have pruned as many units as possible, we revert to the latest NN that achieved the same validation accuracy as the original NN and halve the value of \(p_{u}\). This process is repeated until the unit pruning step becomes lower than the step size needed to prune a single unit. We perform unit pruning before the edge pruning process to identify the minimum widths required in each hidden layer. Once the minimum widths are determined, edges are pruned to reveal the sparse connectivity between those units.

We conducted tests to analyze whether the unit-edge pruned NNs acquire structural properties resembling those of the sub-functions, as presented in the previous section. The results of these tests are shown in Figure 4, and show that the unit-edge pruned NNs do acquire structural properties resembling both input separable sub-functions and reused sub-functions.

## 4 Detecting modules within sparse NNs

In this section, we propose a method to uncover sub-networks or modules that approximate the sub-functions within the target function. We approach this by projecting the problem as a two-step partitioning problem due to the layered structure of NNs. First, we cluster the units belonging to the same layer. We assume that given a layer, there exist various subsets of units that participate in learning the same sub-function. Next, we merge the previous unit clusters across layers such that they exhibit strong connectivity. The overview of our proposed method is illustrated in Figure 5.

**Layer-wise unit clustering:** Let us consider a single layer \(l\) with \(N_{l}\) units. For each unit, we construct a feature vector based on its outgoing connectivity. The feature vector for a unit \(i\) is a binary vector \(f_{i}^{l}\in\{0,1\}^{g}\), where \(g\) is the total number of units in all the later layers. If unit \(j\) is connected to unit \(i\) through at least one directed path, then \(f_{i}^{l}(j)=1\), otherwise \(f_{i}^{l}(j)=0\). Our hypothesis is that the units that participate in learning the same sub-function are reused by similar units in the later layers. To partition the units into clusters such that their feature vectors have low intra-cluster and high inter-cluster distances, we use the Agglomerative clustering method with cosine distance and average linkage. To identify the optimal number of clusters \(K_{l}\), we use the modularity metric, which we have modified for our problem [31; 32; 33].

Let \(\tilde{D}\) be the normalized distance matrix of the unit feature vectors, where the sum of all elements is equal to \(1\). Consider the units partitioned into \(k\) clusters and a matrix \(A\in\mathbb{R}^{k\times k}\), where \(A_{ij}=\sum_{a\in C_{i},b\in C_{j}}\tilde{D}_{ab}\) represents the sum of the distances between all pairs of units in clusters \(C_{i}\) and \(C_{j}\). The modularity metric, denoted by \(M\), measures the quality of the partitioning and is defined as:

\[M=\sum_{i=1}^{k}\left(A_{ii}-\left[\sum_{j=1}^{k}A_{ij}\right]^{2}\right) \tag{2}\]

Figure 4: a. Edge-weight product of paths from input units to output units in unit-edge pruned NNs; b. the number of units covering 90% of the total edge-weight product of paths in unit-edge pruned NNs.
The modularity metric can accurately detect the presence of multiple unit clusters (\(K_{l}=2,...,N_{l}-1\)). However, it fails for the edge cases where all units may belong to a single cluster (\(K_{l}=1\)) or to separate individual clusters (\(K_{l}=N_{l}\)). To address those, we conduct a separability test if the modularity metric is lowest for \(k=2\) or \(k=N_{l}-1\), or if the modularity metric values are close to zero. Modularity metric close to zero is an indicator that the intra-cluster distances are not significantly different from the random baseline. To determine this, we set a threshold on the lowest modularity value obtained.

**Separability test:** The unit separability test is designed to evaluate whether two units in a cluster can be separated into sub-clusters. Consider two units \(i\) and \(j\), with \(o_{i}=\sum f_{i}^{l}\), \(o_{j}=\sum f_{j}^{l}\) neighbors respectively, and \(o_{ij}=f_{i}^{l}\odot f_{j}^{l}\) common neighbors. We consider a random baseline that preserves \(o_{i}\) and \(o_{j}\). The number of common neighbors is modeled as a binomial random variable with \(g\) trials and probability of success \(p=\frac{o_{i}\times o_{j}}{g^{2}}\). The units are separable if the observed value of \(o_{ij}\) is less than the expected value \(\mathbb{E}(o_{ij})\) under the random model.

Consider the partition where \(N_{l}-1\) clusters are obtained. If the two units that are found in the same cluster are separable, it implies that all units belong to separate clusters. Now let us consider the partition of units into two clusters. We merge the feature vectors of the two unit groups. If the two groups of units are not separable, it implies that all units must belong to the same cluster. In some cases, both tests yield positive results. We determine the optimal number of clusters by selecting the result that is more statistically significant.

**Merging clusters across layers:** Strongly connected clusters from adjacent layers are next merged to uncover multi-layered modules. Consider, \(C_{l}^{i}\), \(i=1,2,...,K_{l}\) to be the clusters identified at layer \(l\). Let \(e_{i,j}^{l}\) be the number of edges from cluster \(C_{l}^{i}\) to cluster \(C_{l+1}^{j}\). The two clusters are merged if : \(\frac{e_{i,j}^{l}}{\sum_{j=1}^{K_{l+1}}e_{i,j}^{l}}\geq\delta_{m}\) and \(\frac{e_{i,j}^{l}}{\sum_{i=1}^{K_{l}}e_{i,j}^{l}}\geq\delta_{m}\), where \(\delta_{m}\) is the merging threshold.

The output units are merged with the previous layer's modules, ensuring that \(\delta_{m}\) fraction of incoming edges to the unit are from that module. This allows multiple output units to be matched to the same structural module.

## 5 Experiments and results

### Modular and hierarchical Boolean function graphs

In this section, we conduct experiments on Boolean function graphs with different sub-function properties to validate our pipeline. We begin by testing the pipeline on four function graphs shown in Figure 7. These graphs include: 1) input separable sub-functions, 2) a reused sub-function, 3) sub-functions that are both input separable and reused, and 4) a function graph without any such sub-function where all nodes are strongly connected to every other node.

We perform 36 trials for each function graph by training neural networks with combinations of 3 width values, 3 depth values, and 4 seed values. A trial is considered successful if the proposed pipeline detects a module corresponding to an input separable or reused sub-function (Figure 5). For Boolean

Figure 5: Proposed module detection pipeline.

Figure 6: Success rates for the validation function graphs
functions, we set the modularity metric threshold to -0.2 and the cluster merging threshold to 0.9. Figure 6 shows the success rates for each function graph and NNs with different depths. We observed that the proposed pruning and module detection pipeline has a high success rate when the depth of the NN exceeds that of the function graph. (see to appendix section B for NN visualizations)

**Sub-functions that are reused many times are uncovered more accurately:** To demonstrate this, we consider the function graph shown in Figure (b)b, which contains a single reused sub-function. We vary the number of times the two intermediate gate nodes in the second hierarchical level are used by decreasing the number of output nodes and measure the success rate. The results are shown in Figure 8, where we observe an increasing trend in the success rate as the number of output nodes using the two gate nodes increases. As the number of output units using the two gate nodes decreases, learning the two gate nodes using the previous "dense" sub-function may not be efficient. We provide visualizations of the corresponding NNs in the appendix section C. Note that NNs with only one hidden layer recover a single module, as the function graphs require four hierarchical levels. If only three hierarchical levels are available, the function graph collapses to a dense graph.

**Sub-functions with higher input separability are detected more accurately:** In this experiment, we consider a function graph with two input separable sub-functions shown in Figure (a)a. We increase the overlap between the two separable input sets by decreasing the total number of input nodes and reusing them. We also vary the number of times each sub-function is replicated or used in the output nodes. Our goal is to uncover three properties of the structure: accurately separated input units into sub-function specific and reused, two output sub-functions accurately detected in later layers, and all hidden units belong to either of the two output modules. The success rate for each of these properties as a function of input overlap and sub-function use is shown in Figure 10.

Figure 10: Success rates for various properties uncovered by the pipeline when the input overlap between two sub-functions is increased

Figure 7: Function graphs used to validate the proposed pipeline

Figure 9: Increasing the input overlap (reuse) between two input separable sub-functions
We observe that our method has a high success rate for detecting these properties for sub-functions with high input separability. However, as the overlap between input node sets increases, the success rate for detecting input modules decreases. Furthermore, the success rate decreases as the number of times a sub-function is used decreases. We also observe that for sub-functions with low input separability (and high input overlap), intermediate units are often clustered into a single module due to the hidden units pruning step. Finally, we find that the same trend is observed for detecting output sub-functions, where increasing input overlap and decreasing sub-function use results in a low success rate for our method. (See appendix section D for visualization of these results).

**Hierarchical structures uncovered vary depending on NN depth:** The function graph in Figure 10(a) consists of two input separable sub-functions, the output of those sub-functions is then used by a single intermediate sub-function. This intermediate sub-function is then reused by two additional sub-functions to produce the final output. Interestingly, we observe that the proposed pipeline uncovers different hierarchical structures depending on the depth of the NN.

Figure 11(a) shows the success rate of uncovering the specific sub-functions segregated by the NN depth. We find that NNs with depth greater than or equal to the number of hierarchical levels (5) can uncover all three types of sub-functions in the function graph. However, for NNs with lower depth, only the input separable sub-functions are uncovered, while the intermediate and output sub-functions are merged into a single module. This observation highlights the importance of selecting an appropriate depth for the NN architecture to effectively uncover hierarchical structures in Boolean function graphs. (see appendix section E for NN visualizations) In addition, we report the success rates of uncovering the exact hierarchical structure that corresponds to the depth of the NN as we vary the number of times the output sub-functions were used (Figure 11(b)). The findings demonstrate an increasing trend in the success rate, which is consistent with our previous results.

### Modular and hierarchical functions with MNIST digits

In this section, we present an experimental evaluation of hierarchically modular tasks constructed using the MNIST handwritten digits dataset. We set the modularity metric threshold to -0.3 and the cluster merging threshold to 0.6 for tasks with MNIST digits.

**Classifying two MNIST digits:** We begin by considering a simple function constructed to classify two MNIST digits. The NN is presented with two images concatenated together, and the output is a vector of length 20. The first 10 indices of the output predict the class for the first image, and the remaining 10 predict the class for the second image. We construct the dataset such that each unique combination of labels has 1000 input data points generated by randomly selecting images. We split

Figure 11: The same Boolean function represented by different function graphs depending on the number of hierarchical levels

Figure 12: Success rates: a. uncovering specific sub-functions, b. uncovering the overall hierarchical structure, for NNs with varying depths trained on function graph in Figure 11
the data into training and validation sets with an 8:2 ratio, and then train nine neural networks with varying widths (392, 784, 1568) and number of hidden layers (2, 3, 4) to learn the function for four different seed values. We use the Adam optimizer with bit-wise cross-entropy as the loss function, and select \(99\%\) as the accuracy threshold for the pruning algorithm, as all dense NNs achieve at least \(99\%\) validation accuracy. Out of the 36 trials conducted, we obtain two completely separable modules for 31 out of 36 trials, indicating the high effectiveness of our approach. (see appendix F for NN visualizations)

**Hierarchical and Modular MNIST Task:** In this section, we present a hierarchical and modular task that uses the MNIST digits dataset (Figure 12(a)). The task takes two MNIST images as input, each belonging to the first 8 classes (digits 0-7). The digit values are represented as 3-bit binary vectors and used to construct three output Boolean sub-functions. To generate the dataset, we randomly select 1000 input data points for each unique combination of digits. The data is then split into training and validation sets with an 8:2 ratio. All the NNs that we experiment with reach a validation accuracy of \(98\%\), which we use as the accuracy threshold for the pruning algorithm.

To analyze the modular structure uncovered by our methodology, we divide the success rates into three categories: (1) detecting two input separable modules, (2) detecting three output modules, and (3) middle layer unit-separability into either the input separable modules or the output separable modules. We observe that the NNs uncover the three output modules with high success rates. However, for NNs with lower depths, the NNs fail to recover the two input separable modules, and all the units in the early layers are clustered into a single module. As the depth of the NN increases, we find that the input separable modules are recovered with high accuracy as well. We observe that the success rates for middle layer unit-separability are very low, even when the depth of the NN increases. The units belonging to those 1-2 hidden layers are clustered into a single module. These units may be learning representations required to approximate the two digits well. This finding indicates that the NN depths we experiment with may not be sufficient to capture the underlying function, despite the NNs learning the task with high validation accuracy. Please refer to appendix F for NN visualizations.

## 6 Conclusion

We have introduced _Neural Sculpting_, a methodology to uncover the hierarchical and modular structure of target tasks in NNs. Our findings first demonstrated that NNs trained conventionally do not naturally acquire the desired structural properties related to input separable and reused sub-functions. To address this limitation, we proposed a training strategy based on iterative pruning of units and edges, resulting in sparse NNs with those previous structural properties. Building upon this, we introduced an unsupervised method to detect modules corresponding to various sub-functions while also uncovering their hierarchical organization. Finally, we validated our proposed methodology by using it to uncover the underlying structure of a diverse set of modular and hierarchical tasks.

As a future research direction, we could investigate the efficiency and theoretical underpinnings of modularity in function graphs, which could further motivate pruning NNs. One potential approach to overcome the computational costs and dependence on initial architecture depth could be to use neural architecture search algorithms to construct modular NNs during training. Additionally, exploring the use of attention mechanisms and transformers to uncover hierarchical modularity in tasks could be an interesting direction for future work. These approaches could provide a more efficient way of obtaining modular NNs that can also better capture the underlying structure of the target task.

Figure 13: a. Hierarchically modular task constructed with MNIST digits; b. Success rates for various modules recovered by NNs with varying depths.