SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Shreyas Malakarjun Patil<sup>1</sup> Loizos Michael<sup>2,3</sup> Constantine Dovrolis<sup>4,1</sup>
<sup>1</sup>Georgia Institute of Technology, Atlanta, USA <sup>2</sup>Open University of Cyprus, Nicosia, Cyprus
$^{3}$ CYENS Center of Excellence, Nicosia, Cyprus  $^{4}$ The Cyprus Institute, Nicosia, Cyprus
sm_patial@gatech.edu, loizos@ouc.ac.cy, c.dovrolis@cyi.ac.cy
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.
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 engineering 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]. Translating this hierarchical modularity 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 sufficiently deep NN, how can we uncover the underlying hierarchical organization of sub-functions in that task?
Recent studies through NN unit clustering have demonstrated that certain modular structures can emerge during the training of NNs [12, 13, 14, 15, 16]. However, it is unclear whether the structures extracted reflect the underlying hierarchy of sub-functions in a task. Csordas et al. [17] proposed a method that identifies sub-networks in NNs that learn specific sub-functions. This method however,
requires knowledge of the exact sub-functions within a task's hierarchy. 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 hierarchical modularity 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. Within a hierarchically modular task, we consider sub-functions that are input-separable and reused. We first show that conventionally trained NNs do not acquire structural properties that reflect the previous sub-function properties. To address this, we introduce a sequential unit and edge pruning method to train NNs. Unit pruning first conditions the NN to learn reused sub-functions, while edge pruning subsequently reveals the sparse connectivity between the learned sub-functions. Finally, we propose a network analysis tool to uncover modules and their hierarchical organization within the sparse NNs.
We demonstrate the capability of Neural Sculpting to reveal the structure of diverse hierarchically modular Boolean tasks and two MNIST digits based 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 methods to uncover and harness structural properties in NNs.
1.1 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 \to \{0,1\}^m$ , maps  $n$  input bits to  $m$  output bits. We define the set of gates  $G$  as  $\{\land, \lor, \mathbb{I}\}$ , where  $\land$  represents logical conjunction (AND),  $\lor$  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\to \{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. The sub-function is characterized by having strong internal connectivity within 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  $(\mathcal{X}_t,\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  $(\mathcal{X}_v,\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$ .
$^{1}$ Distinction between function reuse and operation reuse. Consider the AND gate which is an operation that is independent of its input variables. On the other hand, a sub-function is a combination of such operations along with specific inputs. Specific operations have to be relearned if applied to different inputs due to the fixed data flow in NNs [17].
2 Standard training of NNs is not enough
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.
2.1 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 1a). 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
Figure 1: a. Function graph with input separable sub-function, b. Edge-weight product of paths from input units to output units in a trained NN.
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  $\pmb{\theta} \in \mathbb{R}^a$  and the set of all paths,  $\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 \mathcal{P}$ . We define the edge-weight product of a path as  $\pi_p = \prod_{i=1}^a |\theta_i^{p_i}|$ , and  $\pi(i,j) = \sum_{p \in \mathcal{P}_{i \rightarrow j}} \pi_p$  represents the sum of  $\pi_p$  for all paths from input unit  $i$  to output unit  $j$ . We evaluate  $\pi(i,j)$  for  $i = 1, 2$  and  $j = 3, 4$ , and  $\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.
Width	24	36	48
Hidden Layers
1	✓	✓	✓
2	X	X	X
3	X	X	X
Table 1: NN architectures with varying widths and depths for which both null hypotheses are rejected.
We perform a two-sample mean test with unknown standard deviations, with  $(\mu_1,s_1)$  representing the mean and sample standard deviation of  $\pi (i,j)$  for  $i = 1,2$  and  $j = 3,4,$  and  $(\mu_{2},s_{2})$  representing the mean and sample standard deviation
of  $\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 1b shows an example heat map for  $\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.
2.2 Reused
Consider the function graph shown in Figure 2a, 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. We analyze the edge-weight product of the paths from hidden units to output units. Let  $\pi_p^l$  denote the sum of  $\pi_p$  for all paths originating from hidden layer  $l$ , which contains  $N_l$  hidden units. We compute
Figure 2: a. Function graph with reused subfunction, b. the number of units covering  $90\%$  of the total edge-weight product of paths.
the minimum number of units,  $N_P^l$ , that are necessary to achieve  $P\%$  of the total  $\pi_p$ . By comparing  $N_P^l$  to that of the function graph, we can conclude whether the NN has learned the reused states.
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.
3.1 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 subfunctions. 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.
Figure 3: a. Edge-weight product of paths from input units to output units in edge-pruned NN; b. the number of units covering  $90\%$  of the total edge-weight product of paths in edge-pruned NNs.
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.
3.2 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.
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:
$$
S _ {i} ^ {l} = \left| \frac {\partial \mathcal {L} (\boldsymbol {\mathscr {D}} _ {\boldsymbol {v}}, \boldsymbol {\theta})}{\partial a _ {i} ^ {l}} \times a _ {i} ^ {l} \right| \tag {1}
$$
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. For additional details please refer to appendix section H.
We conducted the previous 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.
Figure 4: a. Edge-weight product of paths from input units to output units in unit-edge pruned NN; b. the number of units covering  $90\%$  of the total edge-weight product of paths in unit-edge pruned NNs.
Figure 5: Proposed module detection pipeline.
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 _ {i i} - \left[ \sum_ {j = 1} ^ {k} A _ {i j} \right] ^ {2}\right) \tag {2}
$$
The first term in this equation measures the total intra-cluster distance between the unit features while the second term is the expected intra-cluster distance under the null hypothesis that unit distances were randomly assigned based on joint probabilities in  $\tilde{D}$ . A negative value of  $M$  indicates that pair-wise distance within a cluster is lower than the random baseline. We iterate through values of  $k$  ranging from 2 to  $N_{l} - 1$ , which are obtained through Agglomerative clustering, and select the value of  $k$  that minimizes the modularity metric.
The modularity metric can accurately detect the presence of multiple unit clusters ( $K_{l} = 2, \dots, 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,\dots,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.