Source-linked AI summary
Are Neural Nets Modular? Inspecting Functional Modularity Through Differentiable Weight Masks
Róbert Csordás, Sjoerd van Steenkiste, Jürgen Schmidhuber
TL;DR
Neural networks are expected to benefit from reusable functional modules, but prior inspection has not established whether discovered structures correspond to functions. This paper learns binary weight masks to identify function-specific subnetworks and studies modularity across standard networks. It finds that typical networks specialize without naturally reusing modules, while algorithmic-task solutions can rely on combination-specific weights rather than compositional structure.
Problem
Whether neural networks develop function-specific modules matters because modularity is associated with compositionality, which is relevant to systematic generalization.
Source: §1 · p. 1Method
Results
Takeaways & Limitations
Functional weight masks reveal that weight sharing often tracks shared I/O rather than task similarity, with redundant re-learning of the same functionality.
Source: §7 · p. 9
Abstract
from arXiv · showhide
Neural networks (NNs) whose subnetworks implement reusable functions are expected to offer numerous advantages, including compositionality through efficient recombination of functional building blocks, interpretability, preventing catastrophic interference, etc. Understanding if and how NNs are modular could provide insights into how to improve them. Current inspection methods, however, fail to link modules to their functionality. In this paper, we present a novel method based on learning binary weight masks to identify individual weights and subnets responsible for specific functions. Using this powerful tool, we contribute an extensive study of emerging modularity in NNs that covers several standard architectures and datasets. We demonstrate how common NNs fail to reuse submodules and offer new insights into the related issue of systematic generalization on language tasks.
1 INTRODUCTION
The introduction motivates a functional account of modularity and presents weight masks as a way to identify subnetworks responsible for specific functions. It reports that typical networks often specialize without reusing modules, and that learned solutions on algorithmic tasks can be non-compositional.
1 INTRODUCTION
- Modularity offers a route to compositionality, which is considered important for systematic generalization, an area where typical neural networks perform poorly.The introduction frames this as a motivation for understanding whether networks develop functional modules.Source: §1 · p. 1
- Prior methods found modular structures through activation statistics or connectivity clustering, but did not establish whether those structures represented functional decompositions.Explicitly designed modules have shown generalization benefits, while whether standard networks acquire functional modules themselves remained an open question.Source: §1 · p. 1
- The paper defines functional modules as subsets of weights responsible for a target function, such as a subtask, making modules interpretable at different levels of granularity.This functional definition addresses whether identified modules correspond to a functional decomposition.Source: §1 · p. 1
- The proposed analysis trains differentiable binary masks over frozen weights using an auxiliary task, identifying the subnetwork needed for the target function across several standard architectures.The paper applies the tool to recurrent, Transformer, feedforward, and convolutional networks.Source: §1 · p. 1
- Typical networks exhibit specialization but not reuse, and their weight sharing is driven more by shared I/O than task similarity.The resulting redundancy means the same functionality may be re-learned; the authors also suggest standard networks lack a bias for separating data transformation from information routing.Source: §1 · p. 2
- On SCAN, encoder-decoder LSTMs and Transformers learn combination-specific weights even for command combinations governed by the same rules, indicating non-compositional solutions.The paper reports identical behavior on the Mathematics Dataset.Source: §1 · p. 2
- For image-classification CNNs, some weight subsets are solely responsible for individual classes, and removing them significantly reduces performance on those classes.Confusion matrices identify classes that rely on similar features.Source: §1 · p. 2
2 DISCOVERING MODULES VIA WEIGHT-LEVEL INTROSPECTION
The section introduces weight-level introspection: train masks for a target task while keeping a pretrained network's weights fixed, then use the resulting subnetwork to inspect functional modularity. It also describes binary mask training, sharing metrics, and performance-based tests, including a limitation of mask inversion.
2 DISCOVERING MODULES VIA WEIGHT-LEVEL INTROSPECTION
- Training masks on a target task while freezing network weights identifies the subnetwork responsible for that function; targets can be subtasks or dataset splits.Weight-level analysis is used because units can act as information-transmitting wires even when different weights perform the computation.Source: §2 · p. 2
- Learned logits parameterize each weight’s mask probability, and Gumbel-Sigmoid sampling with a straight-through estimator produces binary keep-or-remove masks.The logits are initialized to favor keeping weights; regularization discourages retaining weights unless they are needed for the target task.Source: §2 · pp. 2–3
- The paper distinguishes reuse of a module from composition: a compositional solution need not reuse a module, and reused weights need not be composed into a compositional solution.Specialization is treated as modularization in which modules implement semantically meaningful functions.Source: §2 · p. 2
- Performance drops after applying or inverting task masks provide an interpretable test of which tasks need a subnetwork, without assuming a desirable sharing level.Mask inversion is limited to cases with entirely disjoint weights.Source: §2 · p. 3
- Intersection over Union measures overlap relative to the union of task weights, while Intersection over Minimum measures overlap relative to the smaller task mask.IoMin equals one when one task’s weights are a strict subset of the other’s, whereas IoU remains below one.Source: §2 · p. 3
3 ANALYZING FUNDAMENTAL PROPERTIES OF MODULES
These synthetic experiments test whether standard networks specialize weights for different functions and reuse weights for identical functions, under shared or separate input/output arrangements. Across the experiments, networks tend to specialize but not reuse; the authors suggest that routing inputs between functions may help explain this pattern.
3 ANALYZING FUNDAMENTAL PROPERTIES OF MODULES
- The experiments isolate specialization and reuse by contrasting shared I/O for different functions with separate I/O for repeated instances of the same function.The setup uses shared I/O to test specialization and separate I/O to test reuse.Source: §3 · p. 4
- Typical networks tend to specialize weights for different functions but not reuse them, with sharing driven more by shared I/O than task similarity.The authors report resulting redundancy and reduced data efficiency.Source: §3 · p. 4
3.1 ADDITION/MULTIPLICATION EXPERIMENTS
- With shared I/O, FNNs and LSTMs show more weight sharing in input/output layers than hidden layers, suggesting some specialization but potentially insufficient separation to prevent interference.Multiplication uses 3.8 times more weights than addition, partly explaining the low IoU; FNNs tend to perform the masked operation, while LSTMs tend to produce invalid outputs on the other operation.Source: §3.1 · p. 4
3.2 DOUBLE-ADDITION EXPERIMENTS
- With separate I/O, weight sharing is low in double addition, and removing weights needed for one pair only slightly reduces performance on the other, suggesting independent modules.The authors conclude that Preuse is violated even in this simple case.Source: §3.2 · p. 5, Figure 2 · p. 5
3.3 TRANSFER LEARNING EXPERIMENTS
- On permuted MNIST, little sharing occurs while capacity remains; sharing emerges only after capacity is saturated, especially in the output layer.The experiment resets the first layer, making the observed transfer an upper bound when first-layer conflicts do not occur; explicitly biasing masks toward occupied weights can force later-layer sharing, but requires knowing which weights to reuse.Source: §3.3 · p. 6, Figure 3 · p. 6
3.4 A POTENTIAL EXPLANATION FOR LACK OF WEIGHT SHARING
- The authors propose that standard networks struggle to route data between inputs and outputs while preserving a representation suitable for shared processing.They argue that suitable inductive biases need further research.Source: §3.4 · p. 6
4 ANALYZING SYSTEMATIC GENERALIZATION ON ALGORITHMIC TASKS
The section tests whether poor systematic generalization reflects missed symmetries or solutions that do not use reusable rules. Masked-network evaluations provide evidence for combination-specific weights on SCAN and difficulty-specific weights on the Mathematics Dataset.
4 ANALYZING SYSTEMATIC GENERALIZATION ON ALGORITHMIC TASKS
- Systematic generalization requires applying learned operations to unseen combinations governed by the same rules, while the two proposed explanations differ in whether Preuse holds.The section distinguishes missed symmetries despite a learned algorithm from failure to learn the correct algorithm.Source: §4 · p. 7
- Masking SCAN training splits leaves a large generalization gap on systematically different test sets, supporting the view that the learned solutions rely on patterns rather than reusable rules.Pretraining on IID data is intended to rule out insufficient evidence about concept symmetries as the only explanation.Source: §4 · p. 7
- Masks trained on the Mathematics Dataset’s easy split show a performance drop on its hard test split, indicating reliance on difficulty-specific weights despite shared underlying rules.The figure compares the unmasked model with masks trained on IID data and on the easy set; it reports five seeds per task.Source: §4 · p. 8, Figure 5 · p. 8
- On SCAN’s add-jump split, nearly half of the decoder weights for “I_JUMP” are removed, suggesting pattern detection in the output layer rather than reusable variable manipulation.The proposed variable-manipulation account is presented as a hypothesis about a generalizing algorithm.Source: §4 · p. 8
5 ANALYZING CONVOLUTIONAL NEURAL NETWORKS
The CNN analysis finds that individual classes depend substantially on class-exclusive feature-detector weights, while other classes show only small performance changes when a target class’s weights are removed.
5 ANALYZING CONVOLUTIONAL NEURAL NETWORKS
- The analysis trains a full-dataset control mask and then masks with each class removed, keeping the target output-layer mask fixed to inspect feature detectors rather than classifier connections.The procedure yields 11 masks across the classes.Source: §5 · p. 8
- Removing class-exclusive weights significantly reduces target-class performance, while non-target classes show only a small drop, indicating class-exclusive feature-detector weights.The interpretation assumes the network has sufficient capacity relative to dataset size.Source: §5 · p. 8
- Removing airplane-related weights can shift predictions toward birds and ships, which the authors note share a blue-background feature.The authors report further class-specific misclassification patterns in Appendix C.8.Source: §5 · p. 8
6 RELATED WORK
Prior modularity analyses often group neurons by connectivity or activation statistics, whereas this paper identifies weight-level modules by function. The related work also distinguishes its compositionality analysis from representation-focused measures and from masking approaches with different objectives.
6 RELATED WORK
- Connectivity and activation-based clustering identify neuron groups, but the authors argue that weight contributions are needed to reason about functional modularity.Mutual-information pathways are also not tied to particular functions or analyzed for compositionality.Source: §6 · p. 9
- Andreas (2019) measures compositionality by reconstructing representations from primitives and a learned composition function, rather than by analyzing modularity of computation.Source: §6 · p. 9
- Related masking work targets continual learning, pruning, attribution, or conditional execution; the paper instead aims to let functional modularity emerge naturally.Mallya et al. (2018) explores differentiable binary weight masks in multi-task learning, but deterministically rather than with the Gumbel-Sigmoid.Source: §6 · p. 9
7 CONCLUSION
The conclusion reports that typical networks’ weight sharing often reflects shared input-output interfaces rather than task similarity, and that repeated functions may not reuse weights. The appendices describe the Gumbel-Sigmoid sampling case and a straight-through estimator for hard samples with gradient flow.
7 CONCLUSION
- Across diverse network types, weight sharing often reflects trivial shared input-output interfaces rather than task similarity, while repeated uses of the same function lack weight sharing.The authors state that this lack of reuse makes learning data inefficient because functions must be relearned repeatedly.Source: §7 · p. 9
A.1 FROM GUMBEL-SOFTMAX TO GUMBEL-SIGMOID
- Gumbel-Sigmoid is the special case of Gumbel-Softmax with k = 2 and l2 = 0, with its sample corresponding to y1.The appendix renames s = y1 and l = l1 to express the sampling formula.Source: §A.1 · p. 14
A.2 STRAIGHT-THROUGH ESTIMATOR
- The straight-through estimator converts the Gumbel-Sigmoid sample into a hard binary sample while permitting gradient flow.For the two-category case, the threshold test is yi > 0.5.Source: §A.2 · p. 14
A.3 THE EXPECTED VALUE OF THE SAMPLES
The appendix derives the expected value of the binary mask samples, explains the temperature choice, and checks mask stability and gradient approximation. The reported masks are highly consistent across repeated runs.
A.3 THE EXPECTED VALUE OF THE SAMPLES
- The mask samples are Bernoulli distributed, with mean μ = σ(l), as derived from the sampling process.The derivation uses independent uniform samples and the sigmoid-transformed logit.Source: §A.3 · pp. 15–16
A.4 CHOOSING THE TEMPERATURE
- Changing temperature τ leaves the binarized sample’s output and its gradients unchanged, while scaling the logit gradient by 1/τ.The authors set τ = 1 for all experiments and note that Adam normalization can mitigate the gradient scaling.Source: §A.4 · p. 16
B.1 STABILITY OF THE MASKS
- 93.26 ± 0.96% mean IoU across 100 mask pairs confirms that the discovered masks are stable, though redundant weight configurations can yield different subsets.The authors note that dropout would encourage such redundant configurations.Source: §B.1 · p. 16
B.1.1 POTENTIAL ERRORS INTRODUCED BY THE STRAIGHT-THROUGH ESTIMATOR
- The straight-through estimator approximates gradients, but each gradient path contains at most one such estimator, so inaccuracies do not accumulate across repeated estimation steps.The masking and estimator are applied directly to the network parameters.Source: §B.1.1 · p. 16
B.2 DOES MASKING CHANGE THE PERFORMED OPERATION?
The appendix tests whether learned masks preserve a network’s operation and discusses how target functions, routing, and regularization shape mask analysis. It also reports a sanity check and practical choices for tuning the method.
B.2 DOES MASKING CHANGE THE PERFORMED OPERATION?
- For almost all tested configurations, masking half the weights causes only a low performance drop; the big FNN on double addition is the exception when early layers are masked.The exceptional network remains above chance, and the authors speculate that masking may cut alternative pathways; masking early layers is generally more difficult than masking later layers.Source: §B.2 · p. 17, Figure 7 · p. 17
B.3 CHOOSING TARGET FUNCTIONALITY
- A target functionality can be any operation the network can perform, but experiments use training-set subsets so the mask identifies weights for an already learned subproblem without requiring generalization.The authors recommend evaluating the mask or its inverse on a test subset where a measurable performance difference is expected.Source: §B.3 · p. 17
B.4 IS ATTENTION THE SOLUTION?
- Attention-based routing is limited when functional modules need arbitrary activation-vector parts or disjoint inputs, and the appropriate granularity of a routable concept remains unclear.Attention in recurrent models routes data to monolithic transformation blocks rather than routing between internal functional modules.Source: §B.4 · p. 18
B.5 EXPLICITLY MODULAR NETWORKS
- Hardcoded blocks exchange a fixed-size activation vector, which can require operation-specific access to stored variables and may not suit independent processing of different state subsets.The authors contrast this with functional modules and argue that a general bias toward function-based specialization may be preferable to explicit modularity.Source: §B.5 · p. 18
C.1 SANITY CHECKING THE MASK DISCOVERY PROCESS
- When input and output weights are manually copied between double-addition pairs, the method discovers nearly perfect sharing in the hidden layers, while the first and last layers retain separate copies.The copying makes the hidden layers unable to distinguish the pairs.Source: §C.1 · p. 19, Figure 8 · p. 19
C.2 COMMON HYPERPARAMETER CHOICES
- The standard setup uses Adam, batch size 128, learning rate 10^-3, gradient clipping 1, four mask samples per batch, and temperature τ = 1.Unless otherwise noted, figures report means and standard deviations over 10 seeds.Source: §C.2 · p. 19
C.3 CHOOSING THE REGULARIZATION HYPERPARAMETER
- Tune α upward until full-task performance drops, then lower it slightly; this heuristic transfers across network sizes, while insufficient capacity can confound conclusions about sharing.The authors suggest retaining adequate performance, such as above 95% of the original, and use Figure 10 to show transfer across sizes.Source: §C.3 · p. 19, Figure 10 · p. 20
C.4 ADDITION/MULTIPLICATION EXPERIMENTS
The addition/multiplication and double-addition experiments examine sharing and inverted-mask behavior across networks. They find limited separation in addition versus multiplication and report a distinct hidden-layer result when double-addition masks are inverted.
C.4 ADDITION/MULTIPLICATION EXPERIMENTS
- The proportion of shared weights changes with network size, even though smaller networks can solve the task.Figure 11 compares small, medium, big, and huge networks.Source: §C.4 · p. 21, Figure 11 · p. 21
- Despite some natural separation between addition and multiplication modules, substantial weight sharing—especially in input and output layers—means inverted masks perform poorly.Table 2 reports accuracy for both operations with FNNs and LSTMs.Source: §C.4 · p. 21, Table 2 · p. 21
C.5 DOUBLE ADDITION EXPERIMENTS
- The double-addition setup varies whether inputs are presented together or one tuple at a time, with the LSTM processing three steps per operation.The FNN uses β = 4 × 10^-4; the LSTM reads its output after the final step.Source: §C.5 · p. 22
C.5.1 ADDITIONAL INVERTED MASK EXPERIMENTS
- With input and output layers unmasked, inverted hidden-layer masks perform around 90% on both pairs, but nonzero performance on the original task requires including both original and inverted weights.The authors leave the cause of this behavior unclear.Source: §C.5.1 · p. 22
C.6 TRANSFER LEARNING EXPERIMENTS
The transfer-learning experiment tests whether masks reuse weights across sequential MNIST permutations. Biasing mask initialization toward old weights increases sharing compared with the standard setup.
C.6 TRANSFER LEARNING EXPERIMENTS
- The experiment trains on 11 MNIST permutations using the same network, with separate 30,000-step phases for weight and mask training.It uses eight mask samples per batch and a mask loss with α = 10^-5.Source: §C.6 · p. 22
- Initializing old-weight mask logits to 2 and new-weight logits to 0 or −1 significantly increases sharing compared with the earlier setup.These initializations correspond to old-weight probability ≈0.88 and new-weight probabilities 0.5 or ≈0.27.Source: §C.6 · pp. 22–23
- Figure 14 reports the proportion of weights shared per layer after every second task on permuted MNIST, using masks initialized to prefer old weights sampled with P ≈0.88.Each task corresponds to a permutation; decreasing the probability of new weights forces increased sharing.Source: Figure 14 · p. 23
C.7.1 SCAN EXPERIMENTS
On SCAN, training weights on the IID split modestly improves performance, but masks still reveal a need for task-specific weights.
C.7.1 SCAN EXPERIMENTS
- Full-size word embeddings produced redundant input-to-hidden configurations likely to be removed by thresholding, so embeddings were reduced to 16 and excluded from masking.Teacher forcing was used with 50% probability for each batch.Source: §C.7.1 · pp. 23–24
- The Transformer used 100-dimensional model and 200-dimensional inner layers, four heads, and three encoder and decoder layers; mask training used a 10^-2 learning rate.Weights were trained for 25k steps before freezing, followed by 25k steps for each mask training phase.Source: §C.7.1 · pp. 23–24
- IID-trained SCAN weights modestly improve split performance, but masking away split-specific weights significantly degrades it, showing the solution still requires task-specific weights.Figure 15 compares direct training and testing on the indicated splits with IID weight training followed by mask training on other splits.Source: §C.7.1 · p. 24, Figure 15 · p. 24
C.7.2 EXPERIMENTS ON THE MATHEMATICS DATASET
Across these experiments, masks expose difficulty-specific and class-specific weight dependence in neural networks, with the CIFAR-10 confusion patterns similar across CNN architectures.
C.7.2 EXPERIMENTS ON THE MATHEMATICS DATASET
- Training on all Mathematics Dataset difficulty levels and then masking on easy examples produces a hard-split performance drop, indicating separate weights are needed across difficulty levels and violating Preuse.The experiments use five selected tasks and train each separately; hard-split accuracy is difficult to interpret because hard sets may include easier samples.Source: §C.7.2 · pp. 24–25
- Hard-split accuracy includes samples from easier splits, so a nonzero score does not establish that hard examples were solved correctly.The dataset’s split structure makes the size of the easy-to-hard performance drop nontrivial to interpret.Source: §C.7.2 · p. 25
C.8.1 SIMPLE CNN
- On the simple CNN, removing class-exclusive feature detectors decreases each CIFAR-10 class’s performance by 40% to 60%.Airplane is confused with bird and ship, cat with dog, and truck with car; the authors relate these patterns to similar backgrounds or shapes.Source: §C.8.1 · p. 26
C.8.2 SIMPLE CNN WITHOUT DROPOUT
- With dropout, removing class-exclusive feature detectors causes roughly 30%-40% higher per-class performance drops than in the same architecture without dropout.The authors interpret this as greater dependence on class-specific modules.Source: §C.8.2 · p. 26, Figure 19 · p. 29