Are Neural Nets Modular? Inspecting Functional Modularity Through Differentiable Weight Masks

1 INTRODUCTION . Modularity is an important organization principle in both artificial ( Ballard , 1987 ; Baldwin & Clark , 2000 ) and biological ( von Dassow & Munro , 1999 ; Lorenz et al. , 2011 ; Clune et al. , 2013 ) systems . It provides a natural way of achieving compositionality , which appears essential for systematic generalization , one of the areas where typical artificial neural networks ( NNs ) do not yet perform well ( Fodor et al. , 1988 ; Marcus , 1998 ; Lake & Baroni , 2018 ; Hupkes et al. , 2020 ) . Recently , NNs with explicitly designed modules have demonstrated superior generalization capabilities ( Clune et al. , 2013 ; Andreas et al. , 2016 ; Kirsch et al. , 2018 ; Chang et al. , 2019 ; Bahdanau et al. , 2019 ; Goyal et al. , 2021b ) , which support this intuition . An implicit assumption behind such models is that NNs without hand-designed modularity do not learn to become modular by themselves . In contrast , it was recently shown that certain types of modular structures do emerge in standard NNs ( Watanabe , 2019 ; Filan et al. , 2020 ) . However , due to defining modules in terms of activation statistics or clustering connectivity , it remains unclear whether these correspond to a functional decomposition . This paper contributes new insights into the generalization capabilities of popular neural networks by investigating whether modules implementing specific functionality emerge and to what extent they enable compositionality . This calls for a functional definition of modules , which has not previously been studied in prior work . In particular , we consider functional modules given by subsets of weights ( i.e . subnetworks ) responsible for performing a specific ‘ target functionality ’ , such as solving a subtask of the original task . By associating modules with performing a specific function they become easier to interpret . Moreover , depending on the chosen target functionality , modules at multiple different levels of granularity can be considered . To unveil whether a NN has learned to acquire functional modules we propose a novel analysis tool that works on pre-trained NNs . Given an auxiliary task corresponding to a particular target function of interest ( e.g. , train only on a specific subset of the samples from the original dataset ) , we train probabilistic , binary , but differentiable masks for all weights ( while the NN ’ s weights remain frozen ) . The result is a binary mask exhibiting the module necessary to perform the target function . Our approach is simple yet general , which readily enables us to analyze several popular NN architectures on a variety of tasks in this way , including recurrent NNs ( RNNs ) , Transformers ( Vaswani et al. , 2017 ) , feedforward NNs ( FNNs ) and convolutional NNs ( CNNs ) . To investigate whether the discovered functional modules are part of a compositional solution , we analyze whether the NN has the following two desirable properties : ( Pspecialize ) it uses different modules for very different functions , and ( Preuse ) it uses the same module for identical functions that may have to be performed multiple times1 . Here we treat Pspecialize and Preuse as continuous quantities , which lets us focus on the degree to which functional modularity emerges . Further , since for many tasks it is unclear what precise amount of sharing is desirable , we will measure Pspecialize and Preuse by considering the change in performance as a result of applying different masks corresponding to a target function . This yields an easy to interpret metric that does not assume precise knowledge about the desired level of weight sharing . We experimentally show that many typical NNs exhibit Pspecialize but not Preuse . By additionally analyzing the capacity for transfer learning , we provide further insight into this issue . We offer a possible explanation : while simple data routing between modules in standard NNs is often highly desirable , it is hard to learn since the weights must also implement the data transformation . Indeed , our findings suggest that standard NNs have no bias towards separating these conceptually different goals of data transformation and information routing . We also demonstrate how the functional modules discovered by typical NNs do not tend to encourage compositional solutions . For example , we analyze encoder-decoder LSTMs ( Hochreiter & Schmidhuber , 1997 ) and Transformers ( Vaswani et al. , 2017 ) on the SCAN dataset ( Lake & Baroni , 2018 ) designed to test systematic generalization based on textual commands . We show that combinationspecific weights are learned to deal with certain command combinations , even when they are governed by the same rules as the other combinations . The existence of such weights indicates that the learned solution is non-compositional and fails at performing the more symbolic manipulation required for systematic generalization on SCAN . To demonstrate that this issue is present even in more real-world scenarios , we highlight identical behavior on the challenging Mathematics Dataset ( Saxton et al. , 2019 ) . Finally , we study whether functional modules emerge in CNNs trained for image classification , which are thought to rely heavily on shared features . Surprisingly , we can identify subsets of weights solely responsible for single classes : when removing these weights the performance on its class drops significantly . By analyzing the resulting confusion matrices , we identify classes relying on similar features . 2 DISCOVERING MODULES VIA WEIGHT-LEVEL INTROSPECTION . To investigate whether functional modules emerge in neural networks one must perform a weight-level analysis . This precludes the use of existing methods , which discover modular structure in NNs based on clustering individual units according to their similarity ( Watanabe , 2019 ; Filan et al. , 2020 ) and that may not always be enough to draw meaningful conclusions . Units can be shared even when their weights , which perform the actual computation , are not . Indeed , units can be viewed as mere “ wires ” for transmitting information . Consider for example a gated RNN , such as an LSTM , where gates can be controlled either by the inputs or the state , yet make use of different weights to project to the same gating units . To overcome this limitation , we propose a novel method to inspect pre-trained NNs at the level of individual weights . It works as follows . First , we formulate a target task corresponding to the specific function for which we want to investigate if a module has been learned . For example , this can be a subset of the original problem ( i.e . a subtask ) , or based on a particular dataset split , e.g . to test generalization . Next , we train a weight mask on this target task while keeping the weights themselves frozen . The resulting mask then reveals the module ( subnetwork ) responsible for the target task . To train the mask , we treat all N weights seperately of each other . Let i ∈ [ 1 , N ] to denote the weight index . The mask ’ s probabilities are represented as learned logits li ∈ R , which are initialized to keep the weights with high probability ( 0.9 ) . If one were to apply continuous masks to the weights it would be possible to scale them arbitrarily , thereby potentially modifying the function the network performs . To prevent this , we binarize masks , which only provides for keeping or removing individual weights . The binarization is achieved using a Gumbel-Sigmoid with a straight-through estimator , which we derive from the Gumbel-Softmax ( Jang et al. , 2017 ; Maddison et al. , 2017 ) in Appendix A.1 . A 1We emphasize the distinction between the ability to reuse modules and the ability to compose them : a compositional solution may fail to reuse a module to implement the same behavior multiple times . Similarly , weights can be reused without them being composed to yield a compositional solution . Further , we consider specialization of modules a special case of modularization where modules are specialized to implement a particular functionality that is semantically meaningful . sample si ∈ [ 0 , 1 ] from the mask can be drawn as follows : si = σ ( ( li − log ( logU1/ logU2 ) ) /τ ) with U1 , U2 ∼ U ( 0 , 1 ) , ( 1 ) where τ ∈ ( 0 , ∞ ) is the temperature and σ ( x ) = 11+e−x is the sigmoid function . Next , we can use a straight-through estimator ( Hinton , 2012 ; Bengio et al. , 2013 ) to obtain a binarized sample bi ∈ { 0 , 1 } : bi = [ 1si > 0.5 − si ] stop + si , ( 2 ) where 1x is the indicator function and [ · ] stop is an operator for preventing backward gradient flow . In this case the bi are samples from a Bernoulli ( σ ( li ) ) random variable ( proof in Appendix A.3 ) . The masks are applied elementwise : w′i = wi∗bi . Training is done by applying the loss function defined by the target task and backpropagating into the logits li . Typically multiple ( between 4–8 ) binary masks are sampled and applied to different parts of a batch to improve the quality of the estimated gradient . The goal of the masking is to remove weights that are not necessary to perform the target function . Thus , the logits li should be regularized , such that the probability for weight wi to be active is small unlesswi is necessary for the task . This is achieved by adding a regularization term r = α ∑ i li to the loss , where α ∈ [ 0 , ∞ ) is a hyper-parameter responsible for the strength of the regularization . How to best choose α is described in detail in Appendix C.3 . At the end of the training process , deterministic binary masks Mi ∈ { 0 , 1 } for weights i are obtained via thresholding Mi = 1σ ( li ) > 0.52 . Applying the full mask M then uncovers the module responsible for the target task . A preliminary study confirmed that the mask training process is stable and thereby suitable for inspection ( Appendix B.1 ) . In the following sections we will analyze several standard NNs using this technique of mask-training34 . Throughout our experiments we will avoid drawing conclusions based on the measured amount of sharing alone as much as possible , since it is unclear what degree of sharing can be expected or is desirable . Rather , we will analyze the performance drop caused by removing weights corresponding to certain functionality , which offers a more consistent and easier to interpret metric5 . For example , to show that a module is responsible for a particular subtask ( A ) but not for another ( B ) , we train masks on A and test on both . A performance drop is expected on task B only . In contrast , to show that this module is exclusively needed for a particular subtask , we can invert the masks and test on both tasks . The inverted masks are expected to perform well on the complementary task , but not on the original one . However , we note that this mask inversion method is limited to analyzing entirely disjoint weights . We analyze weight sharing between two tasks using two different metrics : one is Intersection over Union ( IoU ) , which measures how much the weights used for solving the tasks overlap . We call the other Intersection over Minimum ( IoMin ) , which measures the number of overlapping weights ( intersection ) divided by the minimum of the total number of weights used for each task . In that sense IoMin is a measure of “ subsetness ” . Intuitively , if no weights are shared , both IoU and IoMin are zero . If all weights are shared , both IoU and IoMin are one . However , when the weights needed for one task are a strict subset of the other , then IoMin is one , while IoU < 1 . 2In general we find that li concentrates at either 0 or 1 and so thresholding is safe ( see also Fig . 9 ) . 3A complete overview of all experimental details is available in Appendix C. Mean and standard deviations shown in the figures are calculated over 10 runs unless otherwise noted . 4Code for all experiments is available at https : //github.com/RobertCsordas/modules . 5Exceptions only include cases where the observed amount of weight sharing can be clearly interpreted . However , even in these cases , our analysis will focus on general trends rather than the precise amounts observed .

Summary: The paper presents an empirical study of whether modularity can emerge within neural networks. It starts by proposing a novel definition of modularity that identifies modules by their functionality. To discover the module that implements a specific target functionality, the paper proposes to first pretrain the full network on the original task, then freeze the pretrained weights, and train a binary mask for each weight using Gumbel-Sigmoid. The training objective for the masks is given by the target functionality (e.g., a subtask of the original task), plus some sparsity regularization. The paper then investigates the discovered modules in terms of specialization, reusability, and compositionality. The main findings are: (1) Neural nets tend to satisfy specialization but not reusability; (2) Weight sharing between modules tends to be affected more by whether I/O are shared than by task similarity, and there tends to be less sharing in larger networks; (3) When trained on algorithmic tasks, neural nets fail to learn compositional rules, and thus generalize poorly; (4) CNNs trained for image classification contain class-specific, non-shared weights in the feature detectors.