Sparsity-driven emergence of functional modularity in multitask neural networks
Authors: Bowen Shi, Qi Ouyang, Long Qian
Organizations: Institute for Advanced Study in Physics, Zhejiang University, Hangzhou 310058, China · Center for Quantitative Biology, Academy for Advanced Interdisciplinary Studies, Peking University, Beijing 100871, China · ZJU–Hangzhou Global Scientific and Technological Innovation Center, Hangzhou 311200, China · Institute of Fundamental and Transdisciplinary Research, Zhejiang University, Hangzhou, China
Neural networks can learn multiple tasks, but how learning organizes shared computation into specialized and reusable components remains poorly understood. In this work, we show that sparsity inducing regularization promotes functional modularity in multitask networks without predefined module boundaries. Using functional profiling in multilayer perceptrons and a lightweight transformer, we identified compact, task-selective subnetworks, termed functional tokens. Across rule-based task families, selectively retaining these tokens preserves their associated functions while reducing performance on other tasks. Such specialization develops preferentially in deeper layers during training, while earlier layers retain shared computation. Regularization strength and form shape module compactness and specialization, while task weighting alters functional allocation. Moreover, extracted functional tokens can be recombined through alternative pathways to reconstruct a learned composite computation without retraining, preserving their connectivity and learned weights. Our results suggest that sparsity-constrained optimization shapes not only parameter use but also the organization and reuse of learned computation.
Figures & tables
Figure 1: Network architecture, functional profiling and task-selective weight supports. (a) Shared MLP with three 300-unit fully connected (FC) ReLU hidden layers, trained jointly on five rules (R1–R5 of Binary Sequence Calculation rules). (b) Functional profiling and extraction. Each scalar weight is scored individually by output ablation or Taylor output-sensitivity profiling. Rule-specific masks are constructed separately for each profiled weight matrix and then applied simultaneously across those matrices, including the output weights; biases and parameters outside the masking set are unchanged. (c) Log-transformed weight counts, log(1+count) , for the 31 nonempty exact rule-set signatures, shown separately for the three hidden-layer weight matrices. For example, “1–3” denotes the weights selected for exactly rules 1 and 3. (d–f) Relative accuracy of subnetworks retaining the union of masks for one rule (d), two rules (e), or three to five rules (f). Colours identify the evaluated rule. Relative accuracy is retained-subnetwork bitwise accuracy divided by that of the full network on the same rule. Violins summarize 30 independent runs; black and red marks indicate means and medians, respectively.
Figure 2: Functional-token selectivity across task families, profiling methods and architectures. (a,b) Relative accuracy of ablation-derived retained masks for Sparse Pattern Logic (a) and MNIST Digit-Pair Arithmetic (b), respectively. The Binary Sequence Calculation baseline is shown in Fig. 1 d–f. (c) Mean relative accuracy matrix for single-rule Taylor-derived tokens from 30 independent MLP runs, with each retained-token row normalized to its maximum. Rows identify the retained-token rule and columns the evaluated rule. (d) Relative accuracy of Taylor-derived MLP subnetworks in Binary Sequence Calculation. (e) TinyTransformer architecture with two input-sequence tokens, a shared 32-dimensional projection and positional bias, two self-attention/FFN blocks with residual connections and layer normalization, mean pooling and a linear output head. (f) Relative accuracy of Taylor-derived TinyTransformer subnetworks in Binary Sequence Calculation. In a,b,d,f, the horizontal axis identifies the retained rule set, colors identify the evaluated rule, and accuracy is normalized to the corresponding full network on that rule. Violins summarize 30 independent runs in a,b,d and 23 batch-128 runs in f; black and red marks indicate means and medians, respectively.
Figure 3: Regularization changes the compactness and selectivity of extracted weight supports. (a) Relative accuracy of single-rule subnetworks at Lasso strengths l=0,1,5 with γ=1 (upper plots), and their retained-weight ratios (lower plot). (b) The corresponding quantities for regularizer exponents γ=0.5,1,2 at l=5 . Bar colors identify the evaluated rule; colors in the lower plots identify the rule defining each retained functional token. Bars show means and error bars show sample SD across 30 independent runs. The retained-weight ratio is the number of mask-selected scalar coordinates divided by the total number of scalar coordinates in the analyzed weight matrices. The extraction quantile is fixed at τ=0.75 ; score thresholds are determined separately for each network, matrix and rule. (c,d) Layer-wise modularity metric across Lasso strengths l=0,1,5,10 with γ=1 (c) and across regularizer exponents γ=0.5,1,2 at l=5 (d). Colored guide lines connect medians for each layer at aligned conditions (x positions). All panels use ablation profiling and 30 independent runs per condition. Panels c,d show importance-weighted normalized inverse participation (NIP), defined in Eq. ( 13 ). The analysis uses ϵS=10−9 , ϵout=10−24 and all scalar-weight candidates (Methods).
Figure 4: Training dynamics and endpoint representation geometry. (a) Importance-weighted modularity M at 1,000–5,000 epochs for l=0,1,5,10 , from left to right, with γ=1 . Masks are recomputed at each checkpoint using Taylor output-sensitivity profiling with four probes and the same extraction quantile, τ=0.75 . Violins show 30 independent runs; black and red marks show means and medians. Coloured lines connect layer-wise medians at common epochs. The analysis uses ϵS=10−9 , ϵout=10−8 , seed 0 and all scalar-weight candidates (Methods). (b) Distributions of absolute pre-activation unit correlations at the same four Lasso strengths, aligned column by column to (a) and pooled from four runs per condition. (c) PCA spectral entropy of post-ReLU hidden-layer activations and output representations. (d,e) Linear centered kernel alignment (CKA) between hidden-layer pairs (d) and between hidden and output layers (e). Endpoint metrics in c–e use 30 runs per condition over l=0,1,2,3,5,10,15 . Hidden-layer colours are shared across a-c; colors in d,e identify layer pairs, with “o” denoting the output layer.
Figure 5: Rule-loss weighting changes network performance and functional-token allocation. Each intervention assigns five-fold loss weight to one rule and unit weights to others, and is compared with an equal-weight control, with 30 independent runs per condition at l=5 and γ=1 . Rows in a–c identify the upweighted rule; columns identify the evaluated rule in a and the token-defining rule in b,c. (a) Mean change in full-network accuracy for each upweighted-rule and evaluated-rule pair, calculated as intervention minus control. (b) Mean change in the retained scalar-weight fraction, in units of 10−4 . (c) Mean specialization margin: relative accuracy on the token-defining rule minus its mean relative accuracy on the other four rules. Outlined diagonal cells in a–c identify the upweighted target. (d) Association between the five target-rule mean accuracy changes and retained-weight-ratio changes; error bars combine intervention and control marginal uncertainties as defined in Supplementary Note 7. The dashed line denotes a least-squares fit. (e) Association across 150 repeat-matched target-rule responses, with 30 responses per upweighted rule. Colors in d,e denote the upweighted rule. The solid line in e is a least-squares fit and the annotation reports the pooled Pearson correlation and nominal two-sided p value.
Figure 6: Recomposition of learned functional tokens without retraining. (a) Recomposition of rule 5 computation a(a+b) . The directly extracted rule-5 token is the control; Type 1 computes a(a+b) sequentially, and Type 2 computes a2+ab through parallel primitives followed by addition. Component outputs are thresholded and encoded using the original task interface. Component masks and learned weights remain fixed. (b) Rule-5 relative accuracy normalized to that of the full mother network. (c) Reconstruction score Ew=A(C,5)/ρnz,5(l) , where A(C,5) is the relative accuracy for configuration C , and ρnz,5(l) is the nonzero scalar-weight fraction of the directly extracted rule-5 token. The denominator is shared across all three configurations. In b,c, l denotes the mother network Lasso strength; violins summarize 30 independent runs per regularization setting, with means and medians marked in black and red. (d) Heuristic training/extraction cost comparison, normalized within each run to multitask training plus extraction. Bars show means across three benchmark runs, error bars show s.e.m., and points show individual runs. This proxy excludes recomposition (Supplementary Note 8.7).
Modularity in deep neural networks has been proposed as a means of improving both interpretability and training by promoting disentangled representations and reducing redundancy. In this work, we study the emergence of modular structure through competition dynamics between groups of neurons during training. We introduce a method that (i) maintains near-baseline accuracy, (ii) induces usage-based modularity by sparsely routing inputs to neuron groups, and (iii) encourages specialization of these modules, such that their activations are correlated with input classes. We evaluate the proposed approach on ImageNet-100 and CIFAR-100 and show that with it, specialized modules emerge without module-level supervision. These modules capture a meaningful high-level structure in the data, with individual modules responding to semantic categories (e.g., dogs or vehicles). We also study the emergence of a hierarchical partition of sub-tasks depending on the number of modules. Our results suggest that competitive dynamics can serve as a simple mechanism for inducing functional modularity in standard architectures.
Baptiste Rossigneux, Karim Haroun
Univ. Rennes · Inria, IRISA · University of Paris 8
Neural networks routinely develop representations in which individual neurons contribute to multiple distinct functions, leading to the widespread view that such networks lack modular organization. We challenge this conclusion by introducing a formally nested hierarchy of modularity: unit-level, which requires distinct computations to occupy disjoint neurons; orthogonal, which asks whether a rotation reveals functionally independent subspaces; and linear, which asks whether a bounded linear map can separate task-relevant computations into independent subspaces. We develop a unified framework for learning decompositions at each level and evaluating them causally through selective ablation. Across vision models, language models and multi-task Recurrent Neural Networks (RNNs), we find that networks which lack unit-level modularity nonetheless often exhibit clear functional modularity at the subspace level. In vision and language models, content versus style and syntax versus semantics respectively dissociate under subspace ablation despite sharing overlapping populations of neurons. These learned subspaces also align differentially with brain data, with content and style subspaces predicting higher and lower visual cortex respectively, and semantic subspaces outperforming syntactic ones in predicting language cortex activity. In RNNs trained on multiple cognitive tasks and previously shown to lack modular structure, our decomposition recovers sharp double dissociations between tasks. The absence of unit-level modularity frequently coexists with clear subspace-level modularity across domains, enabling targeted intervention on specific computations even in networks where no individual neuron is specialized, with broad implications for mechanistic interpretability and for understanding functional organization in both biological and artificial neural systems.
Shreya Saha, Benjamin Bergen, Meenakshi Khosla
Electrical and Computer Engineering, University of California, San Diego · Cognitive Science Department, University of California, San Diego · Computer Science and Engineering, University of California, San Diego
Understanding how biological and artificial neural networks implement computation from connectivity is a central problem in neuroscience and machine learning. In neural systems, structural and functional connectivity are known to diverge, motivating approaches that move beyond direct connections alone. Here, we show that the spatial and temporal function of recurrent neural networks (RNNs) trained on hierarchically modular tasks can be recovered by modelling the network as a graph and analysing the multi-hop pathways between input and output units. In particular, decomposing these pathways by hop length reveals how the network temporally routes information. This perspective reframes regularisation: if function is implemented through multi-hop communication, then standard penalties such as L1 regularisation, which act only on individual weights, constrain single-hop structure rather than the multi-hop pathways that support computation. Motivated by this view, we introduce resolvent-RNNs (R-RNNs), which constrain multi-hop pathways and thereby induce temporal sparsity beyond that achieved by standard L1 regularisation. Compared with L1 regularisation, R-RNNs achieve improved performance by inducing temporal sparsity that matches the task structure, even when the task signal is sparse. Moreover, R-RNNs exhibit stronger sparsity-function alignment, reflected in their increased robustness under strong regularisation. Together, our results identify multi-hop communication as a key principle linking structure to function in recurrent networks, and suggest that sparsity should be defined over functional pathways rather than individual parameters.
Jatin Sharma, Dan F. M Goodman, Danyal Akarca
Department of Electrical and Electronic Engineering Imperial College London