Is a Modular Architecture Enough?
Anonymous Author(s)
AffiliationAddress email
Abstract
Inspired from human cognition, machine learning systems are gradually revealingadvantages of sparser and more modular architectures.
Recent work demonstratesthat not only do some modular architectures generalize well, but they also lead tobetter out-of-distribution generalization, scaling properties, learning speed, andinterpretability.
A key intuition behind the success of such systems is that the data generating system for most real-world settings is considered to consist of sparsely interacting parts, and endowing models with similar inductive biases will be helpful.
However, the field has been lacking in a rigorous quantitative assessment of suchsystems because these real-world data distributions are complex and unknown.
Inthis work, we provide a thorough assessment of common modular architectures,through the lens of simple and known modular data distributions.
We highlight the benefits of modularity and sparsity and reveal insights on the challenges faced while optimizing modular systems.
In doing so, we propose evaluation metrics that highlight the benefits of modularity, the regimes in which these benefits are substantial, as well as the sub-optimality of current end-to-end learned modular systems as opposed to their claimed potential.
Introduction
Deep learning research has an established history of drawing inspiration from neuroscience andcognitive science.
From the way hidden units combine afferent inputs, to how connectivity andnetwork architectures are designed, many breakthroughs have relied on mimicking brain strategies.
It is no surprise then that modularity and attention have been leveraged, often together, in artificialnetworks in recent years (Bahdanau et al., 2015; Andreas et al., 2016;
Hu et al., 2017; Vaswani et al.
,2017; Kipf et al., 2018;
Battaglia et al., 2018; Goyal et al., 2019, 2021), with impressive results.
Indeed, work from cognitive neuroscience (Baars, 1997; Dehaene et al., 2017) suggests that cortexrepresents knowledge in a modular way, with different such modules communicating through thebottleneck of working memory (where very few items can simultaneously be represented), in whichcontent is selected by attention mechanisms.
In recent work from the AI community (Bengio, 2017;
Goyal & Bengio, 2020), it was proposed that these characteristics could correspond to meaningfulinductive biases for deep networks, i.e., statistical assumptions about the dependencies betweenconcepts manipulated at the higher levels of cognition.
Both sparsity of the dependencies betweenthese high-level variables and the decomposition of knowledge into recomposable pieces that are asindependent as possible (Peters et al., 2017; Bengio et al., 2019; Goyal & Bengio, 2020) would makelearning more efficient.
Out-of-distribution (OoD) generalization would be facilitated by making itpossible to sequentially compose the computations performed by these modules where new situationscan be explained by novel combinations of existing concepts.
Although a number of recent results hinge on such modular architectures (Graves et al., 2014; Andreaset al., 2016; Hu et al., 2017; Vaswani et al., 2017; Kipf et al., 2018;
Santoro et al., 2018; Battagliaet al., 2018;
Goyal et al., 2019, 2021;
Locatello et al., 2020; Mittal et al., 2020; Madan et al., 2021; Keet al., 2021), the abundance of tricks and proposed architectural modifications makes it challenging to parse real, usable architectural principles.
It is also unclear whether the performance gains obtainedby such Mixture-of-Experts (MoE) based modular systems are actually due to good specialization, asis often claimed, or due to other potential confounding factors like ease of optimization.
In this work, we extend the analysis from Rosenbaum et al.
; Maziarz et al.
; Cui & Jaech(2020); Csordás et al. and propose a principled approach to evaluate, quantify, and analysecommon ingredients of modular architectures, supported by either standard MLP-like connectivity,recurrent connections or attention (Bahdanau et al., 2015; Vaswani et al., 2017) operations.
To doso, we develop a series of benchmarks and metrics aimed at probing the efficacy of a wide range ofmodular networks, where computation is factorized.
This reveals valuable insights and helps identifynot only where current approaches succeed but also when and how they fail .
Whereas previous workon disentangling (Bengio, 2013; Higgins et al., 2016; Kim & Mnih, 2018) has focused on factoringout the different high-level variables that explain the data, here we focus on disentangling systemmodules from each other, the structural ingredients of network that can facilitate this factorization,and how such ingredients relate to the data generating distributions being parsed and processed.
Given the recent increased interest in sparse modular systems (Rahaman et al., 2021; Fedus et al.
,2021; Du et al., 2021; Mittal et al., 2021), we believe that this work will provide a test-bed forinvestigating the workings of such models and allow for research into inductive biases that can pushsuch models to achieve good specialization.
Through detailed experiments and evaluation metrics,we make the following observations and contributions:
• We develop benchmark tasks and metrics based on probabilistically selected rules to quantifytwo important phenomena in modular systems, the extent of collapse and specialization .
• We distill commonly used modularity inductive biases and systematically evaluate themthrough a series of models aimed at extracting commonly used architectural attributes( Monolithic , Modular , Modular-op and GT-Modular models).
• We find that specialization in modular systems leads to significant boosts in performance
when there are many underlying rules within a task, but not so much with only few rules.
• We find standard modular systems to be often sub-optimal in both their capacity on focusing
on the right information as well as in their ability to specialize, suggesting the need for additional inductive biases.
Notation / Terminology
M o d u l e s
Rules
Choose a rule
Draw
Training Setting Regression /
Classification
A r c h i t e c t u r e
M L P /
M H A /
R NN
Prediction
True Label
[Figure]
In this paper, we study how a family of modular systems performs on a common set of tasks,prescribed by a synthetic data generating process which we call rule-based data .
Below,we introduce the notation for key ingredients:(1) rules and how they form tasks , (2) modules and how they can take different modelarchitectures , (3) specialization and how weevaluate models.
We refer the reader to Figure1 for an illustration of our setup.
Rules.
To properly understand modular systems and analyze their benefits and shortcomings, we consider synthetic settings that allowfine-grained control over different aspects oftask requirements.
In particular, operationsmust be learned on the data-generating distribution illustrated in Equations 1-3, which we also refer to as rules .
Details about the exact operationsused in experiments are described in Section 3.
[Equation]
Given this distribution, we define a rule to be an expert of thisdistribution, that is, rule r is defined as p y ( · | x , c = r ) where c isa categorical variable representing context, and x is an input sequence.
For example, consider x = (1 , 2) and c to select betweenaddition and multiplication.
Then, depending on c , the correctoutput should be either y = 3 or y = 2 .
More details about the specifics of these data distributions are presented in Section 3.
Systems will be trained to infer y given c and x .
This simple setup is meantto capture context-dependent tasks on variable data distributions, e.g. reasoning according to differentfeatures (e.g. shape, color, etc.).
However, unlike such complex systems, ground truth knowledgeof required operations is known for our synthetic task, allowing for deeper quantitative analysis.
Tasks.
A task is described by the set of rules (data-generating distribution) illustrated in Equations1-3.
Different sets of { p y ( · | x , c ) } c imply different tasks.
For a given number of rules, we trainmodels on multiple tasks to remove bias towards any particular task.
Modules.
A modular system comprises a set of neural network modules, each of which cancontribute to the overall output.
One can see this through the functional form y =
(cid:80) Mm =1 p m y m ,where y m denotes the output and p m the activation of the m th module.
Details about the differentmodular systems are outlined in Section 4.
From this point onwards, we exclusively use rules to refer to the specialized components in thedata-generating process , and modules to refer to the experts that are learned by a modular system .
Further, for ease of quantitative assessment, we always set the number of modules equal to thenumber of rules, except when evaluating monolithic models (with a single module).
Modules canbe implemented in three different architectures, as described next.
Model Architectures.
Model architectures describe the choice of architecture considered foreach module of a modular system, or the single module in a monolithic system.
Here we consider
Multi-Layer Perceptron (MLP), Multi-Head Attentions (MHA), and Recurrrent Neural Network(RNN).
Importantly, the rules (or data generating distributions) are adapted to the model architecture,and we often refer to them as such (e.g. MLP based rules).
Details about the data distributions andmodels considered in this work are provided in Sections 3 and 4 respectively.
Perfect Specialization.
When training modular systems on rule-based data, we would like the modules to specialize according to the rules in the data-generating distribution.
Thus, there is an importantneed to quantify what constitutes perfect specialization of the system to the data.
To allow for easierquantification, we always consider an equal number of modules and rules.
However, future workshould evaluate the ability of modular systems to automatically infer the required number of modules.
Data Generating Process
Since we aim to study modular systems through synthetic data, here we flesh out the data-generatingprocesses operating based on the rules scheme described above (see Equations 1-3).
We use asimple Mixture-of-Experts (MoE) Yuksel et al.
(2012); Masoudnia & Ebrahimpour (2014) styleddata-generating process (Mixture Distribution), where we expect different modules to specialize to thedifferent mixture components (rules).
It is important to note that this system is slightly different fromthe traditional flat MoE since the experts are more plug-and-play and can be composed to solve a particular problem.
As an example, if we consider a mixture of recurrent systems, different tokens (timepoints) in the input sequence can undergo computations according to different rules (e.g. a switchinglinear dynamical system), as opposed to the choice of expert being governed by the whole sequence.
We now look at more specific setups of the data-generating systems in consideration, the generaltemplate of which was outlined above.
To do so, we explain the data-generating processes amenableto our three model architectures: MLP, MHA, and RNN.
Additionally, each of the following taskshave two versions: regression, and classification.
These are included to explore potential differencesthese distinct loss types may induce.
[Equation]
MLP.
Here, we define the data scheme that is amenable for learningof modular MLP-based systems.
In this synthetic data-generatingscheme, a data sample consists of two independent numbers and achoice of rule being sampled from some distribution.
Different ruleslead to different linear combinations of the two numbers to give theoutput.
That is, the choice of linear combination is dynamically instantiated based on the rule drawn.
This is mathematically formulated in Equations 4-6, where α c and β c are the data parameters, I theidentity matrix and y denotes the label for the regression tasks and sign ( y ) for the classification tasks.
Hence, the data comes from a MoE distribution where c denotes which linear combination governsthe conditional distribution p y ( · | x 1 , x 2 , c ) .
When training modular architectures on such data, oneexpects each module in the trained system to specialize according to a unique rule.
[Equation] [Equation]
MHA.
Now, we define the data scheme that is tuned forlearning in modular MHA based systems.
Essentially,a MHA module can be understood through a set ofsearches (query-key interactions), a set of corresponding retrievals (values) and then some computation ofthe retrieved values, as explained by Mittal et al.
Accordingly, we design the data-generating distributionwith the following properties: Each rule is composedof a different notion of search, retrieval and the finallinear combination of the retrieved information respectively.
We mathematically describe the processin Equations 7-11, where n = 1 , ..., N and r = 1 , ..., R with N as the sequence length and R thenumber of rules.
We denote the tuple ( q nr , q ′ nr , v nr , v ′ nr ) as x n .
Further, y n denotes the label forthe regression tasks while for classification, we consider the categorical label to be sign ( y n ) .
Thus, we can see that c n denotes the rule for the n th token.
This rule governs which two tokensare closest to the n th token, demonstrated as s n and s ′ n .
It also governs what features are retrievedfrom the searched tokens, which are v s n c n and v s ′ n c n .
These retrieved features then undergo a 162 rule-dependent linear combination (on c n ).
Here, too, when training a modular MHA architecture, wewant each MHA module in the system to be able to specialize to a unique MHA rule in the data system.
[Equation]
[Equation]
[Equation]
[Equation]
RNN.
For recurrent systems, we define a rule as a kind oflinear dynamical system, where one of multiple rules can betriggered at any time-point.
Mathematically, this process canbe defined through Equations 12-15, where n = 1 , ...N , with
N describing the sequence length.
Each rule thus describesa different procedure for the update of the state s t as well asthe effect of the input x t to the state.
Thus, we can see that c ndenotes the rule to be used at the n th time-point.
Further, y n denotes the label for the regressiontasks while for classification, we consider the labels as sign ( y n ) .
Hence, in all settings, the data comes from a MoE distribution where c denotes the rule and governsthe conditional p y
( · | x , c ) .
When training modular architectures on such data, one expects eachmodule in the trained system to specialize according to a unique rule.
Our aim is to use these syntheticrule-based data setting to study and analyse modular systems and understand whether end-to-endtrained modular systems concentrate on the right information to specialize based on, i.e. based on c ,whether they do learn perfect specialization and whether perfect specialization actually helps in thesesettings.
To properly understand this, we detail the different kinds of models considered in Section 4as
well as the different metrics proposed in Section 5 to analyse trained systems.
For this work, we limit our analysis to infinite-data regime where each training iteration operates on anew data sample Future work would perform similar analysis in the regime of limited data.
Models
Model
Functional Form
Monolithic
[Equation]
x
[Equation]
c )
Modular
[Equation]
f
[Equation]
c )
[Equation]
y =
[Equation]
Rm
[Equation]
p m
[Equation]
y m
Modular-op ˆ y m =
f m ( x , c ) p =
g ( c ) ˆ y =
(cid:80)
[Equation]
Several works claim that end-to-end trained modular systems outperform their monolithic counterparts, especiallyin out-of-distribution settings.
However, there is a lackof step-by-step analysis on the benefits of such systemsand whether they actually specialize according to the datagenerating distribution or not.
To perform an in-depthanalysis, we consider four different types of models thatallow for varying levels of specialization, which are:
Monolithic , Modular , Modular-op , and GT-Modular .
Wegive the formulations for each of these models below andthen discuss the different analysis we can perform throughthem.
We also illustrate these models in Table 1 anddepending on the data-generating procedure described in
Section 3, f and f m can be implemented as either MLP,
MHA or RNN cells in this work.
Monolithic.
A monolithic system is a big neural networkthat takes the entire data ( x , c ) as input and makes predictions ˆ y based on it.
There is no inductive bias about modularity or sparsity explicitly bakedin the system and it is completely up to back-propagation to learn whatever functional form is neededto solve the task.
An example of such a system is a traditional Multi-Head Attention (MHA) basedsystem, eg.
a Transformer.
4 Module
R u l e
Ground Truth Rules: 4 | Modules:0.
Modular.
A modular system is composed of a numberof modules, each of which is a neural network of a givenarchitectural type (MLP, MHA, or RNN).
Each module mtakes the data ( x , c ) as input and computes an output ˆ y mand a confidence score, normalized across modules intoan activation probability p m .
The activation probability reflects the contribution of each module’s output to the finaloutput
ˆ y of the system.
Thus, there is an explicit baked-in inductive bias of modularity but it is still up to system-wideback-propagation to figure out the right specialization.
An example of such a system is a mixture of MLPs or reusable
, reusable across different positions.
Modular-op.
A modular-op (for operation only ) systemis very similar to the modular system with just one smalldifference.
Instead of the activation probability p m ofmodule m being a function of ( x , c ) , we instead makesure that the activation is decided only by the rule contextc .
Hence, unlike modular systems, modular-op cannot be distracted by x in figuring out specializationof different modules.
Even though the operation required is explicitly provided, this model still needsto learn specialization through back-propagation.
GT-Modular.
A GT-modular system (for ground truth ) serves as an oracle benchmark, i.e., a modularsystem that specializes perfectly.
In particular, the activation probability p ′
m s of modules are just setaccording to c , which is the indicator present in the data ( x , c ) .
Thus, this is a perfectly specializingsystem that chooses different modules sparsely and perfectly according to the different data rules.
Given enough capacity, we can see that there is a hierarchy of models based on the functions they canimplement, with GT-Modular ⊆
Modular-op ⊆
Modular ⊆ Monolithic .
Put differently, models from
Monolithic to GT-Modular increasingly incorporate the inductive biases for modularity and sparsity.
This is proved in Appendix C by inspecting the function classes implemented by these models.
[Figure]
In what follows, we want to analysethe benefits of having simple end-toend trained modular systems as opposed to monolithic ones.
This can beunderstood through a comparison ofvarious performance based metrics between Monolithic and Modular models, explained in the next section.
Thiswill allow us to answer if a modulararchitecture is always better for various distinct rule-based data generatingsystems.
For instance, a comparisonbetween the Modular and Modularop models will show whether the standard modular systems are able to focus on the right information and ignore the distractors in drivingspecialization.
To study this, we will look at performance as well as collapse and specializationmetrics between these class of models.
A comparison between GT-Modular and Modular-op willshow the benefits of having a sparse activation pattern with proper resource allocation of modules asopposed
Finally, we note that GT-Modular is a modular system which obtains perfect specialization.
Throughthis model, we aim to analyse whether perfect specialization is in-fact important and if so, howfar are typical modular systems from obtaining similar performance and specialization throughend-to-end training.
We now describe the metrics used for these evaluations.
[Figure]
Metrics
To reliably evaluate modular systems, we propose a suite of metrics that not only gauge the performance benefits of such systems but also evaluate them across two important modalities: collapseand specialization , which we use to analyse the extent of resource allocation (in terms of parameters/modules) and specialization respectively of a modular system
[Figure]
Performance.
The first set of evaluation metrics are based on performance of the models in bothin-distribution as well as out-of-distribution (OoD) settings.
These metrics capture how well thedifferent models perform on a wide variety of different tasks.
For classification settings, we reportthe classification error while for regression settings, we report the loss.
In-Distribution.
This refers to the in-distribution performance, evaluated by looking at both the finalperformance as well as convergence speeds of the different models.
Out-of-Distribution.
This refers to the OoD performance of different models.
We consider verysimple forms of OoD generalization: either (a) change in distribution of x by increasing variance, or(b) different sequence lengths, wherever the possibility presents (eg. in MHA and RNN).
Collapse Metrics.
We propose a set of metrics Collapse-Avg and Collapse-Worst that quantify theamount of collapse suffered by a modular system.
Collapse refers to the degree of under-utilizationof the modules.
An example of this is illustrated in Figure 2, where we can see that Module 3 is neverused.
We consider the setting where all the data rules are equi-probable and the number of modulesin the model are set to be the same as the number of data rules, to R .
High collapse thus refers tounder-utilization of resource (parameters) provided to the model, illustrating that certain modules arenever being used and concurrently meaning that certain modules are being utilized for multiple rules.
[Equation]
Collapse-Avg.
Given the data-setting with R 278 equi-probable rules, and hence R modules in themodel, we let p ( m ) be the marginal probabilitydistribution of activation of module m .
Then, we define the Collapse-Avg metric C A as in Equation
M e t r i c S c
o r e
Collapse-Avg  8 160.
Collapse-Worst  8 16 32 Rules 0.
Alignment  8 160.
Adaptation  8 160.
Inverse Mutual Information
GT-Modular
Modular-op
Modular
Random
16, where RR − 1 is for normalization.
This metric captures the amount of under-utilization of all themodules of the system.
A lower number is preferable for this metric, as a lower number demonstratesthat all the modules are equally utilized.
[Equation]
) Collapse-Worst.
Given the same data and model setting as above, 285 the Collapse-Worst metric C W is defined as in Equation 17.
Thismetric captures the amount of under-utilization of the least used module of the system.
Again, a lownumber is preferable as it signifies that even the least used module is decently utilized by the model.
Specialization Metrics.
To complement collapse metrics, we also propose a set of metrics, (1)
Alignment , (2) Adaptation and (3) Inverse Mutual Information to quantify the amount of specializationobtained by the modular systems.
We again consider the setting of equi-probable rules and the samenumber of modules and rules R .
These metrics are aimed at capturing how well the modules specializeto the rules, that is, whether different modules stick to different rules (good specialization) or whetherall modules contribute almost equally to all rules (poor specialization).
[Equation]
Alignment.
Given a modular system trained on rule-based data with 295 R rules and modules, one can obtain the activation matrix A , where
A rm denotes p ( module = m | rule = r ) , that is, the probability ofactivation of module m conditioned on rule r .
Further, given a distance metric d ( · , · ) over the spaceof matrices, perfect specialization can be quantified through Equation 18, where S R denotes the spaceof permutation matrices over R objects.
We consider d ( · , · ) as a normalized L 1 distance.
The scores d demonstrates the distance between the activation matrix A and its closest permutation matrix,with distances computed according to the metric d ( · , · ) .
Note that s d → 0 implies that each modulespecializes to a unique rule, thereby signifying perfect specialization.
Since the space of permutationmatrices S R grows exponentially at the rate of Θ ( R !) , computing s d naively soon becomes intractable.
However, we use the Hungarian algorithm (Kuhn, 1955) to compute it in polynomial time.
Thismetric shows how close the learned modular system is to a perfectly specializing one, where a lowscore implies better specialization.
[Equation]
Inverse Mutual Information.
Given R as 308 the number of rules and modules and letthe joint distribution p ( m, r ) denote theactivation probability of module m on rule r , the Inverse Mutual Information metric S IMI is definedas in Equation 19.
A low inverse mutual information metric is preferable as it denotes that the modulesare more specialized to the rules as opposed to multiple modules contributing to a single rule.
[Equation] [Equation]
Adaptation.
Let R be the number of rules and modulesand P a distribution over the R -dimensional simplex.
Further, let p ( · ) be the distribution over rules (notequi-probable in this metric) and q ( · ) the correspondingdistribution obtained over the modules.
Note that the distribution q ( · ) is dependent on p ( · ) .
Giventhese distributions, we define the Adaptation metric S A in Equation 20, where ˆ r i and ˆ m i are such thatp (ˆ r 1 ) ≤ p (ˆ r 2 ) ≤ ... ≤ p (ˆ r R ) and
q ( ˆ m 1 ) ≤ q ( ˆ m 2 ) ≤ ... ≤ q ( ˆ m R ) and P is a dirichlet distribution.
This metric can be understood as the amount by which the modules adapt (signified through thedistribution q ( · ) ) to changes in the rule distributions (which are p ( · ) sampled from P ).
The matchingbetween the rule and module is obtained through a simple sort as defined above.
A low adaptationscore implies that the marginal distribution of the modules adapt well according to the distributionof the rules.
That is, when a rule is weakly present in the data, there exists a module which weaklycontributes in the corresponding output, averaged over multiple different rule distributions.
M e t r i c S c
o r e
Collapse-Avg 0.
Collapse-Worst 0.
Alignment
GT-Modular
Modular-op
Modular
Random 0.
Adaptation 0.
Inverse Mutual Information
To understand these metrics, note that uniform random activation patterns for the modules lead to lowcollapse metrics but high alignment, adaptation and inverse mutual information metrics, implying littlecollapse but poor specialization, as expected.
On the other hand, GT-Modular systems necessarily leadto low collapse metrics as well as low alignment, adaptation and inverse mutual information, denotinglittle collapse and good specialization, which is expected since specialization is given as oracle.
Experiments
We are now ready to report experiments on the models outlined in Section 4 with associated datageneration processes described in Section 3.
For each level of modularity (i.e. Monolithic, Modular,
Modular-op, GT-Modular), we analyse models learning over five different number of rules, rangingfrom few (2) to many (32), five different model capacities (number of parameters) and two differenttraining settings, i.e. regression and binary classification.
To remove any biases towards particulartask parameters (e.g. α c , β c in Equation 6), we randomly select new rules to create five differenttasks per setting and, train five seeds per task.
In essence, we train ∼ 20,000 models 1 to properlyanalyse the benefits of modularity, the level of specialization obtained by end-to-end trained systems,the impact of number of rules and the impact of model capacity
[Figure]
Performance.
We refer the readers to Figure 3 for a compressed overview on the performance of various models.
Wesee that GT-Modular system wins most of the times ( left ),indicating the benefits of perfect specialization.
We alsosee that between standard end-to-end trained Modular and
Monolithic systems, the former outperforms but not by ahuge gap.
Together, these two pie charts indicate that currentend-to-end trained modular systems do not achieve goodspecialization and are thus sub-optimal by a substantial margin.
We then look at the specific architectural choices (MLP, MHAand RNN cells for functions f and f m in Table 1) and analyse their performance and trends across increasing number ofrules.
Figure 4 shows that while there are concrete benefitsof a perfectly specializing system (GT-Modular) or even models that know what information to drive specialization from(Modular-op), typical end-to-end trained Modular systems arequite sub-optimal and not able to realize these benefits, especially with increasing number of rules which is where we seesubstantial benefits of good specialization (contrast Modular vs
GT-Modular and Modular-op)
Moreover, while such end-toend Modular systems do generally outperform the Monolithic ones, it is often only by a small margin.
We also see the training pattern of different models averaged over all other settings, with theaverage containing error for classification and loss for regression, in Figure 7.
We can see that goodspecialization not only leads to better performances but also faster training.
Collapse
We evaluate all the models on the two collapse metrics outlined in Section 5.
Figure5 shows the two collapse metrics, Collapse-Avg and Collapse-Worst , for different models against varying number of rules, averaged over the different model architectures (MLP, MHA and RNN),training settings (Classification and Regression), model capacities, tasks and seeds.
First, we noticethat a Random activation baseline and the GT-Modular system do not have any collapse, which isexpected.
Next, we notice that both Modular and Modular-op suffer from the problems of collapseand this problem becomes worse with increasing number of rules.
Figure 6 further shows similarinformation averaged over the number of rules too, highlighting that Modular-op has less collapsethan Modular in general.
However, we still see that the problem of collapse is significant wheneverback-propagation is tasked with finding the right activation patterns, especially in the regime oflarge number of rules.
This clearly indicates the need for investigation into different forms ofregularizations to alleviate some of the collapse problems.
Specialization.
Next, we evaluate through the proposed specialization metrics in Section 5 whetherthe end-to-end trained modular systems actually specialize according to the data-generating distribution.
Figure 5 shows the three specialization metrics, Alignment , Adaptation and Inverse Mutual
Information , for different models against varying number of rules, again averaged over differentmodel architectures, training settings, model capacities, tasks and seeds.
As expected, we see that
the
Random activation baseline has poor specialization (high metrics) while the GT-Modular system hasvery good specialization.
We further see that end-to-end trained Modular systems as well as Modularop suffer from sub-optimal specialization, as indicated by the high metrics.
As with collapse, we againsee that it becomes harder to reach optimal specialization with increasing number of rules.
Figure6 shows that while Modular-op has marginally better specialization than standard Modular systems,they are indeed quite sub-optimal when compared to a perfectly specializing system, i.e. GT-Modular.
We refer the readers to Appendix D, E and F for training details as well as additional experimentsregarding the effect of model sizes for MLP, MHA and RNN architectures respectively.
Conclusion and Discussion
We provide a benchmark suitable for the analysis of modular systems and provide metrics that notonly evaluate them on in-distribution and out-of-distribution performance, but also on collapse andspecialization.
Through our large-scale analysis, we uncover many intriguing properties of modularsystems and highlight potential issues that could lead to poor scaling properties of such systems.
Perfect Specialization.
We discover that perfect specialization indeed helps in boosting performanceboth in-distribution and out-of-distribution, especially in the regime of many rules.
On the contrary,monolithic systems often do comparatively or sometimes better when there are only a few rules,but do not rely on specialization to do so.
End-to-End Trained Modular systems .
While Modular systems outperform Monolithic ones, themargin of improvement is often small.
This is because when solely relying on back-propagationof the task-losses, these models do not discover perfect specialization.
In fact, the problem of poorspecialization and high collapse becomes worse with increasing number of rules.
This is slightlymitigated by allowing contextual information from the task to be used explicitly, as is the case for
Modular-op, but the problems still persist and get worse over large number of rules.
In summary, through systematic and extensive experiments, this work shows that modularity, whensupporting good and distributed specialization (i.e. little collapse), can outperform monolithic modelsboth in and out of distribution testing.
However, we also find that although perfectly specializedsolutions are attainable by modular networks, end-to-end training does not recover them, often evenwith explicit information about task context (as in Modular-op).
Since real-world data distributionsare often complex and unknown, we cannot get access to oracle networks like GT-Modular foranalysis.
An important conclusion is that additional inductive biases are required to learn adequatelyspecialized solutions.
These could include other architectural features to facilitate module routing,or regularization schemes (e.g. load-balancing Fedus et al.
or optimization strategies (e.g.learning rate scheduling) to promote module specialization.
We refer the reader to Appendices Aand B for further discussion on these exciting prospects and extensions to real-world domains.
Webelieve the framework proposed in this work is ideal to drive research into such inductive biasesand a necessary stepping stone for applications of these designs at scale.
Finally, we highlight that the use of network architectures that promote contextual specialization,such as the use of modules as studied here, could potentially promote unwanted biases when deployedin models use by the public due to collapse or ill-distributed specialization.
The framework proposedin this work could help mitigate this potentially problematic impact on society.
Checklist
For all authors...
(a) Do the main claims made in the abstract and introduction accurately reflect the paper’s
contributions and scope?
[Yes]
(b) Did you describe the limitations of your work?
[Yes] See Appendix A
(c) Did you discuss any potential negative societal impacts of your work?
[N/A]
(d) Have you read the ethics review guidelines and ensured that your paper conforms tothem?
[Yes]
If you are including theoretical results...
(a) Did you state the full set of assumptions of all theoretical results?
[Yes]
(b) Did you include complete proofs of all theoretical results?
[Yes]
If you ran experiments...
(a) Did you include the code, data, and instructions needed to reproduce the main experi-
mental results (either in the supplemental material or as a URL)?
[Yes]
(b) Did you specify all the training details (e.g., data splits, hyperparameters,
how they
were chosen)?
[Yes]
(c) Did you report error bars (e g., with respect to the random seed after running experiments multiple times)?
[Yes]
(d) Did you include the total amount of compute and the type of resources used (e
g., typeof GPUs, internal cluster, or cloud provider)?
[Yes]
If you are using existing assets (e g., code, data, models) or curating/releasing new assets...
(a) If your work uses existing assets, did you cite the creators?
[N/A]
(b) Did you mention the license of the assets?
[N/A]
(c) Did you include any new assets either in the supplemental material or as a URL?
[N/A]
(d) Did you discuss whether and how consent was obtained from people whose data you’re
using/curating?
[N/A]
(e) Did you discuss whether the data you are using/curating contains personally identifiable
information or offensive content?
[N/A
If you used crowdsourcing or conducted research with human subjects...
(a) Did you include the full text of instructions given to participants and screenshots, if
applicable?
[N/A]
(b) Did you describe any potential participant risks, with links to Institutional Review
Board (IRB) approvals, if applicable?
[N/A]
(c) Did you include the estimated hourly wage paid to participants and the total amount
spent on participant compensation?
[N/A]