Sparsity-driven emergence of functional modularity in multitask neural networks

Bowen Shi; Qi Ouyang; Long Qian

cs.LGOct 23, 2026

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

Abstract

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+\mathrm{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 $\gamma=1$ (upper plots), and their retained-weight ratios (lower plot). (b) The corresponding quantities for regularizer exponents $\gamma=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 $\tau=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 $\gamma=1$ (c) and across regularizer exponents $\gamma=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 $\epsilon_{S}=10^{-9}$ , $\epsilon_{\mathrm{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 $\gamma=1$ . Masks are recomputed at each checkpoint using Taylor output-sensitivity profiling with four probes and the same extraction quantile, $\tau=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 $\epsilon_{S}=10^{-9}$ , $\epsilon_{\mathrm{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

\gamma=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 $a^{2}+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 $E_{w}=A(C,5)/\rho_{\mathrm{nz},5}(l)$ , where $A(C,5)$ is the relative accuracy for configuration $C$ , and $\rho_{\mathrm{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).

Modular Neural Networks

Modular neural networks organize computation into specialized, reusable components that can be routed, composed, or updated independently.

See all 74 papers

Sparse Competition during Training For the Emergence of Specialized Modules
Aug 31, 2026Baptiste Rossigneux, Karim HarounSparse Neural Networks Neural Network Training Dynamics
Univ. Rennes · Inria, IRISA · University of Paris 8
Neural networks are more modular than single neurons suggest
Oct 9, 2026Shreya Saha, Benjamin Bergen, Meenakshi KhoslaModular Neural Networks Neural Network Interpretability
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
Unifying Dynamical Systems and Graph Theory to Mechanistically Understand Computation in Neural Networks
May 5, 2026Jatin Sharma, Dan F. M Goodman, Danyal AkarcaRecurrent Neural Networks
Department of Electrical and Electronic Engineering Imperial College London

Aug 31, 2026J/K move · Enter open · S save

Sparse Competition during Training For the Emergence of Specialized Modules

Baptiste Rossigneux, Karim Haroun

Univ. Rennes · Inria, IRISA · University of Paris 8

Sparse Neural Networks Neural Network Training Dynamics Modular Neural Networks

Open paper AskPDF arXiv

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.

Abstract

Modular Neural Networks

Explore similar work

Sparse Competition during Training For the Emergence of Specialized Modules

Neural networks are more modular than single neurons suggest

Unifying Dynamical Systems and Graph Theory to Mechanistically Understand Computation in Neural Networks