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. 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, 37th Conference on Neural Information Processing Systems (NeurIPS 2023).