SlowGuess/NeurIPS_2023_OCR huggingface dataset:
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.
Figure 6: Function graphs used to validate the proposed pipeline
5.1 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 6. 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 functions, we set the modularity metric threshold to -0.2 and the
cluster merging threshold to 0.9. Figure 7 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)
Figure 7: Success rates for the validation function graphs
5.2 Sub-functions that are reused many times are uncovered more accurately
We consider the function graph shown in Figure 6b, 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.
Figure 8: Success rate when the number of times a sub-function is used increases
Figure 9: Increasing the input overlap (reuse) between two input separable sub-functions
5.3 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 9a. 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
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).
Figure 11: The same Boolean function represented by different function graphs depending on the number of hierarchical levels
5.4 Hierarchical structures uncovered vary depending on NN depth
The function graph in Figure 11a 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 12a 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.
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
This result 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 12b). The findings demonstrate an increasing trend in the success rate, which is consistent with our previous results.
5.5 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 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 33 out of 36 trials, indicating the high effectiveness of our approach. (see appendix F for NN visualizations)
Hierarchical and Modular MNIST Task: We consider next a hierarchical and modular task that uses the MNIST digits dataset (Figure 13a). 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
Figure 13: a. Hierarchically modular task constructed with MNIST digits; b. Success rates for various modules recovered by NNs with varying depths.
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 hidden layers are often 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.
Emerging modular structures in trained NNs has been a subject of prior investigation, primarily through unit clustering algorithms. Previous methods for extracting such structures can be categorized
as either structural or functional. Structural methods organize units based on their structural characteristics, including connectivity and edge-weights [12, 13, 14], while functional methods consider attributes based on unit activation patterns [15, 16]. However, it remains unclear whether these extracted structures represent the underlying hierarchy of sub-functions in a given task.
To the best of our knowledge, this work is the first attempt to introduce an approach that combines training (pruning) and network analysis to unveil the hierarchical modularity of tasks. Our proposed method for clustering units and discovering modules aligns with the structural methods. A series of prior studies leveraged normalized spectral clustering to globally extract unit clusters and analyze their characteristics [14, 15]. Spectral clustering optimizes for N-cuts, quantifying the internal connectivity relative to the external connectivity of unit clusters. Our method is most similar to previously introduced layer-wise unit clustering techniques [12, 13]. These methods consider incoming and outgoing edge-weights to group units within a layer and consolidate edges into single incoming and outgoing connections. Importantly, all previous methods were tailored for conventionally trained NNs. In contrast, our approach is simpler and customized for the pruned NNs we obtain.
Pruning of NN units and edges has gained significant traction as a method to enhance NN computational efficiency and reduce memory demands while maintaining performance integrity [21, 24, 23, 34, 35, 30, 36, 37, 38]. Recent research has also focused on investigating how pruning impacts NN generalization [39, 40, 41]. However, prior studies have primarily concentrated on either unit pruning or edge pruning in isolation, without employing both in a sequential manner as proposed in our work. In our approach, we employ unit pruning as a means to condition or compel NNs to learn reused sub-functions effectively. Furthermore, edge pruning is employed to uncover the sparse connectivity between various units or modules.
Mechanistic interpretability seeks to reverse-engineer trained NNs to understand their internal mechanisms. One avenue of research within mechanistic interpretability involves identifying the circuits formed within NNs. Recent discoveries include curve-detecting circuits in vision models [42, 43], transformer circuits [44], induction heads [45], and indirect object identification [46], among others. Our proposed approach for uncovering the hierarchical modular task structure dissects the NN into modules responsible for learning specific sub-functions. Hence, simplifying the complexity of reverse engineering entire NNs, by focusing instead on smaller sub-networks.
Continual learning aims at learning multiple tasks presented sequentially while preserving performance on previously learned tasks [47, 48, 49, 50]. Modular NNs have been demonstrated to mitigate catastrophic forgetting by freezing and reusing certain modules while introducing and updating others to learn new tasks [6, 7, 8]. However, those methods entail manually designing modular NNs with fixed module sizes at initialization. NNs that naturally acquire a hierarchically modular structure mirroring that of the task may offer additional advantages compared to generic modular NNs. These potential benefits encompass enhanced generalization, efficient transfer learning through the reuse of modules corresponding to frequently utilized sub-functions, and informed module additions to prevent sub-linear increases in NN capacity.
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.
This work is supported by the National Science Foundation (Award: 2039741), by the EU's Horizon 2020 Research and Innovation Program (grant agreement no. 739578), and by the Government of the Republic of Cyprus through the Deputy Ministry of Research, Innovation, and Digital Policy. The authors acknowledge the constructive comments given by the NeurIPS 2023 reviewers, Cameron Taylor, Qihang Yao and Mustafa Burak Gurbuz.
[1] Herbert A Simon. The architecture of complexity. In *Facets of systems science*, pages 457–476. Springer, 1991.
[2] Günter P Wagner, Mihaela Pavlicev, and James M Cheverud. The road to modularity. Nature Reviews Genetics, 8(12):921-931, 2007.
[3] David Meunier, Renaud Lambiotte, Alex Fornito, Karen Ersche, and Edward T Bullmore. Hierarchical modularity in human brain functional networks. Frontiers in neuroinformatics, 3:37, 2009.
[4] Miguel Valencia, MA Pastor, MA Fernández-Seara, J Artieda, J Martinerie, and M Chavez. Complex modular structure of large-scale brain networks. Chaos: An Interdisciplinary Journal of Nonlinear Science, 19(2):023119, 2009.
[5] David Meunier, Renaud Lambiotte, and Edward T Bullmore. Modular and hierarchically modular organization of brain networks. Frontiers in neuroscience, 4:200, 2010.
[6] Tom Veniat, Ludovic Denoyer, and Marc'Aurelio Ranzato. Efficient continual learning with modular networks and task-driven priors. arXiv preprint arXiv:2012.12631, 2020.
[7] Jorge A Mendez and Eric Eaton. Lifelong learning of compositional structures. arXiv preprint arXiv:2007.07732, 2020.
[8] Oleksiy Ostapenko, Pau Rodriguez, Massimo Caccia, and Laurent Charlin. Continual learning via local module composition. Advances in Neural Information Processing Systems, 34:30298-30312, 2021.
[9] Dzmitry Bahdanau, Shikhar Murty, Michael Noukhovitch, Thien Huu Nguyen, Harm de Vries, and Aaron Courville. Systematic generalization: what is required and can it be learned? arXiv preprint arXiv:1811.12889, 2018.
[10] Brenden Lake and Marco Baroni. Generalization without systematicity: On the compositional skills of sequence-to-sequence recurrent networks. In International conference on machine learning, pages 2873-2882. PMLR, 2018.
[11] Dieuwke Hupkes, Verna Dankers, Mathijs Mul, and Elia Bruni. Compositionality decomposed: how do neural networks generalise? Journal of Artificial Intelligence Research, 67:757-795, 2020.
[12] Chihiro Watanabe, Kaoru Hiramatsu, and Kunio Kashino. Modular representation of layered neural networks. Neural Networks, 97:62-73, 2018.
[13] Chihiro Watanabe. Interpreting layered neural networks via hierarchical modular representation. In International Conference on Neural Information Processing, pages 376-388. Springer, 2019.
[14] Daniel Filan, Stephen Casper, Shlomi Hod, Cody Wild, Andrew Critch, and Stuart Russell. Clusterability in neural networks. arXiv preprint arXiv:2103.03386, 2021.
[15] Shlomi Hod, Stephen Casper, Daniel Filan, Cody Wild, Andrew Critch, and Stuart Russell. Detecting modularity in deep neural networks. arXiv preprint arXiv:2110.08058, 2021.
[16] Richard D Lange, David S Rolnick, and Konrad P Kording. Clustering units in neural networks: upstream vs downstream information. arXiv preprint arXiv:2203.11815, 2022.
[53] Yuanzhi Li, Colin Wei, and Tengyu Ma. Towards explaining the regularization effect of initial large learning rate in training neural networks. Advances in Neural Information Processing Systems, 32, 2019.
[54] Stanislav Fort, Paweł Krzysztof Nowak, Stanislaw Jastrzebski, and Srini Narayanan. Stiffness: A new perspective on generalization in neural networks. arXiv preprint arXiv:1901.09491, 2019.
[55] Peter J Rousseeuw. Silhouettes: a graphical aid to the interpretation and validation of cluster analysis. Journal of computational and applied mathematics, 20:53-65, 1987.
[56] Tadeusz Calinski and Jerzy Harabasz. A dendrite method for cluster analysis. Communications in Statistics-theory and Methods, 3(1):1-27, 1974.
[57] David L Davies and Donald W Bouldin. A cluster separation measure. IEEE transactions on pattern analysis and machine intelligence, 1(2):224-227, 1979.
A Appendix structure and details
In this section, we discuss the structure of the appendix. The code and datasets used can be found here: https://github.com/ShreyasMalakarjunPatil/Neural-Sculpting.
Appendix sections (B - F) accompany section 5 in the main paper and provide visualizations of neural networks (NNs) and the hierarchical modularity uncovered by our proposed pipeline. Each section includes NN visualizations corresponding to the five experiments summarized in the main paper. The main figures in these sections are organized such that the first row represents different function graphs that the NNs are trained to learn, with each column representing a distinct function graph relevant to the experiment. The subsequent rows show visualizations of the uncovered hierarchical modular structures for NNs with increasing depths.
Visualization figure details: In the appendix sections, each function graph is associated with visualizations from a specific NN trial at a particular depth value (with one width value and one seed value). This visualization represents the structure that is commonly recovered for the given function graph and NN depth. While the NN width value and seed value remain consistent in the visualizations, they may vary if the recovered structure does not represent the majority. We also present and discuss other structures that are frequently recovered, using other function graphs representing the same Boolean function. For success rate computations, such trials are typically considered failures in recovering overall structures. However, for individual sub-function detection success rates, they may be considered successful if they uncover a module corresponding to that sub-function. Additional visualizations and the hyperparameters used for training and pruning the NNs are detailed in the code base.
Function graph nomenclature: In the visualizations, the term "modularity" indicates the function graph being used. The modularity is represented as a list that indicates the number of sub-functions in the hierarchy. For example, if the modularity is [1, 2], it means there are two sub-functions in the upper hierarchical level that take as input the output of a single sub-function lower in the hierarchy. In the code base, we presuppose the number of input nodes and append the number of output nodes to this list.
Appendix section G accompanies section 2 in the main paper. Section 2 presents two tests developed to demonstrate that NNs, through standard training, do not acquire structural properties reflecting input separable and reused sub-functions. In appendix section G, we provide detailed results for these tests conducted on nine different NN architectures to further support our claims. We also report the results of the same tests conducted on edge-pruned and unit-edge-pruned NNs.
Appendix section H accompanies section 3 in the main paper, which presents the method for pruning NNs. This appendix section provides a detailed explanation of the unit and edge pruning algorithms, along with the exploration strategy and results for hyperparameter tuning.
Appendix section I accompanies section 4 in the main paper, which introduces the proposed method for detecting modules corresponding to various sub-functions. In this appendix section, we present a detailed algorithm description and the methodology used for hyperparameter exploration.
Note that the pruning and module detection hyperparameters were explored concurrently. However, in the appendix section, we first present experiments with the pruning hyperparameters while keeping the module detection method fixed, and vice versa.
Resources: The unit and edge pruning processes for the Boolean functions were executed on a CPU, while the pruning for the MNIST experiment was performed using RTX6000 GPUs.
B NN visualizations: Function graphs used for validation
Input separable sub-function
Reused sub-function
Modularity[1, 2]; Arch[4, 24, 16]; Density = 8.75%
Input separable and reused
Modularity[3, 3]; Arch[6, 24, 3]; Density = 5.56%
Single sub-function
Modularity[1]; Arch[4, 24, 4]; Density = 8.33%
Modularity[2]: Arch[4, 24, 4]; Density = 8.33%
$lr = 0.05$  , batch size  $= 8$  epochs  $= 15$  seed  $= 0$  , arch  $= \{4,4,4\}$ $P_{\mathrm{v}} = 70.0$ $P_{\mathrm{v}} = 2.5$
Modularity[2]: Arch[4, 24, 24, 4]: Density = 2.08%
$ir = 0.05$  batch size  $= 8$  epochs  $= 10$  seed  $= 0$  arch  $\equiv [4,4,2,4]$
P=65.0,Pc=1.5
Modularity[2]: Arch[4, 24, 24, 24, 4]: Density = 1.26%
ir = 0.05, batch size =8, epochs =20, seed =0, arch =[4,3,3,2,4]
P=70.0,P=1.5
Figure 14: The visualizations provide examples of NNs successfully uncovering the exact hierarchical structure as the validation function graphs. Each column represents a distinct function graph, characterized by: 1. input separable sub-functions, 2. reused sub-function, 3. input separable and reused sub-functions, or 4. a single sub-function. Each row showcases NNs with different depths: 1. one hidden layer, 2. two hidden layers, and 3. three hidden layers.
$I = 0.1$  , batch size  $= 8$  , epochs  $= 40$  , seed  $= 0$  , arch  $= [4,3,16]$ $P_{\mathrm{c}} = 50.5$ $P_{\mathrm{c}} = 20$
Modularity[1, 2]: Arch[4, 24, 24, 16]: Density = 2.75%
$I = 0.1$  batch size  $= 8$  epochs  $= 40$  seed  $= 0$  arch  $\equiv [4,3,2,16]$
$P_{u} = 70.0, P_{e} = 2.5$