Breaking Neural Network Scaling Laws with Modularity
Authors: Akhilan Boopathy
Sunshine Jiang
William Yue
Jaedong Hwang
Abhiram Iyer
Ila R Fiete
Venue: ICLR 2025 Poster (2025)
Keywords: scaling laws
modularity
neural network
generalization
compositionality
combinatorial generalization
Primary area: other topics in machine learning (i.e., none of the above)
TL;DR: We show theoretically that modular neural networks trained on modular tasks can generalize to high-dimensional tasks with a fixed number of training points; we propose a learning rule to exploit this advantage empirically.

Abstract: Modular neural networks outperform nonmodular neural networks on tasks ranging from visual question answering to robotics. These performance improvements are thought to be due to modular networks' superior ability to model the compositional and combinatorial structure of real-world problems. However, a theoretical explanation of how modularity improves generalizability, and how to leverage task modularity while training networks remains elusive. Using recent theoretical progress in explaining neural network generalization, we investigate how the amount of training data required to generalize on a task varies with the intrinsic dimensionality of a task's input. We show theoretically that when applied to modularly structured tasks, while nonmodular networks require an exponential number of samples with task dimensionality, modular networks' sample complexity is independent of task dimensionality: modular networks can generalize in high dimensions. We then develop a novel learning rule for modular networks to exploit this advantage and empirically show the improved generalization of the rule, both in- and out-of-distribution, on high-dimensional, modular tasks.

Published as a conference paper at ICLR 2025
BREAKING NEURAL NETWORK SCALING LAWS WITH MODULARITY
Akhilan Boopathy, Sunshine Jiang, William Yue, Jaedong Hwang, Abhiram Iyer, Ila Fiete Massachusetts Institute of Technology (akhi lan, sunshl 6e, will lue,, johwang abiyer, fiete)@mm deduu.edu
ABSTRACT
Modular neural networks outperform nonmodular neural networks on tasks rang- ing from visual question answering to robotics. These performance improvements are thought be due to modular networks' superior ability to model the composi- tional and combinatorial structure of real-worl wollworld problems. However, theoretical explanation of how modularity improves generalizability, and how to leverage task modularity while training networks remains elusive. Using recent theoreti- cal progress in explaining neural network generalization, we investigate how the amount of training data required to generalize a task varies with the intrinsic dimensionality of a task's input. We show theoretically that when applied to mod- ularly structured tasks, while nonmodular networks require exponential num- ber of samples with task dimensionality, modular networks' sample complexity is independent of task dimensionality: modular networks can generalize in high di- mensions. We then develop a novel learning rule for modular networks to exploit this advantage and empirically show the improved generalization of the rule, both in- and out-of-distribution, on high-dimensional, modular tasks.
INTRODUCTION
Modular neural network (NN) architectures have achieved impressive results in a variety of domains
ranging from visual question answering (VQA) (Andreas et al.) [2016a] b) |Hu al.] 2017]] Johnson
et al.
2017 Yi et al. 2018 Kim et al.
2019). reinforcement learning (Goyal et al. 2021 Madan
jet al. 2021). robotics (Alet al. 2018h 8h Pathak et al. 2019, Yang et al. 2020) and natural language
processing for which modular architectures based on attention (Bahdanau et al., 2015) are standard.
Modular NNs are thought to be better generalized by facilitating combinatorial generalization, phenomenon where learning system recombines previously learned components in novel ways to
generalize to unseen inputs (Alet et al. 2018b D Amario et al.) 2021 Mittal et al., 2022a by Jarvis al. 2023). Yet, we lack fundamental understanding of why modularity benefits generalization.
In parallel, the generalization properties monolithic (nonmodular) NNs have been increasingly well understood both theoretically and empirically. In particular, current theory explain the double descent phenomenon where NN generalization error decreases with increasingly large ca- pacity (Belkin al., 2019; Spigler al.) 2019 Neal et al.] 2019 Rocks Mehta (002)) NN learning in a certain regime can also be understood as kernel regression (Jacot et al. 2018). Exten-
sive empirical studies have also measured scaling laws of NN generalization error Kaplan et al.
2020), and moreover, these scaling laws can be explained theoretically (Bahri et al. 2021 Hutter 2021 Hastie al.) 2022). However, these laws indicate that the sample complexity required to gen- eralize on a task scales expomentially with the intrinsic dimensionality of the task's input (McRae al. 2020 Sharma Kaplan (022). This raises the question: how can we hope to generalize on high-dimensional problems with limited training data?
In this work, we investigate how modular NNs can circumvent this exponential number of samples. We first synthesize existing generalization results in simple theoretical model of NN generalization error and empirically validate on tasks with varying intrinsic dimensionality. We then use our
model to show theoretically that appropriately structured modular NNs avoid using an exponential number of samples on modular tasks. However, recent work shows that architectural modularity is in practice not sufficient its solve modular tasks efficiently (Csordás et al. 2021 Mittal

---

Published as a conference paper at ICLR 2025
solution to align NN modules to task's modularity is lacking. We propose propeepiiica learning rule that aligns NN modules to approach the underlying modules of a task and empirically demonstrate its improved generalization.
et al. 2022a);
novel
We summarize our contributions as follows:
We propose a simple, theoretical model of NN learning that synthesizes existing general- ization results. Our model predicts the generalization error of NNs under varying number of model parameters, number of training samples, and dimensions of variation in a task input. We empirically validate our theoretical model on novel parametrically controllable sine wave regression task and show that sample complexity varies exponentially with task dimension.
We apply the theoretical model to compute explicit, non-asymptotic expressions for gen- eralization error in modular architectures; to our knowledge, we are the first to do so. Our result demonstrates that sample complexity is independent of task dimension for modular NNs applied to modular tasks specific form.
Based on our theory, we develop a learning rule align NN modules to the modules under- lying high-dimensional modular tasks with the goal of promoting generalization on these tasks.
We empirically validate the improved generalizability (both in- and out-of-distributionn) of our modular learning approach parametrically controllable, high-dimensional tasks: sine-wave regression and Compositional CIFAR-10.
RELATED WORK
2.1 MODULAR NEURAL NETWORKS
Recent efforts to model cogmitive processes show that functional modules and compositional repre-
sentations emerge after training on task (Yang al. 2019 Yamashita & Tani 2008 [yer et al.,
2022). Partly inspired by this, recent works in Al propose using modular networks: networks com-
posed of sparsely connected, reusable modules (Alet et al.) 2018b a Chang al. 2019 Chaudhry
et al.
2020 Shazeer et al. 2017 Ashok et al.) 2022] Yang et al. 2022, Sax et al.] 2020 Preiffer
ct al. 2023). Empirically, modularity improves -utto--distribution generalization (Bengio et a
2020
Madan et al.] 2021 Mittal et al., 2020) Jarvis et al.) 2023), modular generative models are
effective unsupervised learners (Parascandolo et al., 2018:
Locatello et al. 2019) and modular ar- chitectures can be more interpretable (Agarwal al. 2022 In addition, meta- meta- veeaa-laannin algorithms
can discover and learn the modules without prespecifying them (Chen et al. 2020 Sikka et al. 2020
Chitnis et al. 2219)).
Recent empirical studies have investigated how modularity influences network performance and generalization. The degree of modularity increases systematic generalization performance in VQA tasks (D' Amario et al.) 2021) and sequence-based tasks (Mittal et al.] 2020). [Rosenbamm et al.
(2019) Cui & Jacch (2020) study routing networks (Rosenbaum et al. [0088). type of modular architecture, and identified several difficulties with training these architectures including training instability and module collapse. Csordás et al. (202T) and Mittal al.) (2022a) extend this type of analysis to more general networks to show that NN modules often may not be optimally used to promote task performance despite having the potential to do so. These analyses are primarily empirical; in contrast, in our work, we aim to provide a theoretical basis for how modularity may improve generalization. Moreover, given that architectural modularity may not sufficient to en- sure generalization, propose a learning rule designed to align NN modules to the modularity of the task.
2.2 NEURAL NETWORK SCALING LAWS
Many works present frameworks to quantify scaling laws that map NN's parameter count train- ing dataset size to an estimated testing loss. Empirically and theoretically, these works find that testing loss scales trained NNs (Bahri et al. [202 Rosenfeld et al. 2020). including transformer-bast based language mod- els (Sharma & Kaplan
power-law with respect to the dataset size and parameter count on well-
2022 Clark et al. 2022
Tay al. 2022)
2

---

distributions under Published as a conference paper at ICLR 2025
Many previous works also conclude that generalizations of power-law nonpower-law-base dis- tributions can also model neural scaling laws well. in many cases better than vanilla power-law frameworks (Mahmood et al. 2022 Alabdulmohsin et al., 2022). For instance, instance|HHutter (202I) shows that countably infinite parameter models closely follow non- pon--ower-law--eee
unbounded data complexity regimes. In another case, [Sorscher al. (2022) show that exponential scaling works better than power-law scaling if the testing loss is associated with a pruned dataset size, given pruning metric that discards easy or hard examples under abundant or scarce data guarantees, respectively.
Some works approach this problem by modeling NN learning as manifold or kernel regression. For example, McRae et al. (2020) considers regression on manifolds and concludes that sample complexity scales based on the intrinsic manifold dimension of the data. In another case, Canatar et al. (2021) draws correlations between the study of kernel regression to how infinite-width deep networks can generalize based on the size of the training dataset and the suitability of particular kernel for task. Along these lines, several works use random matrix theory to derive scaling laws
for kernel regression (Hastie et al. 2022 Cui al. 2021 2022
Wei et al. 2022 Jin et al.) (02I).
Among other observations, this body of work shows that in the absence of strong inductive biases, high-dimensional tasks have sample complexity growing roughly exponentially with the intrinsic dimensionality of the data manifold. In this work, we borrow the theoretical techniques from this line of work to investigate if learning the modular structure of modular tasks will reduce the sample complexity of training.
3 MODELING NEURAL NETWORK GENERALIZATION
In this section, we present a toy model of NN learning that treats NNs as linear functions of their parameters: this is along with the lines of prior work such as Bahri et al. (2021); Canatar et al. (202T). Although this common theoretical assumption does not directly ap- ply to practical, non-linear architectures, the assumption provides analytical tractability and, morcover, can be shown to predict gen- eralization even in nonlinear networks (e.g. in the Neural Tangent Kernel literature (Jacot et al.] 2018)). Our specific analytical ap- proach follows that of Hastie al. 2022). Under this setting, we find exact closed-form expressions for expected training and test loss, representing a simplified version of the results in Hastie al.. (2022). We find that our toy model captures key features of NN gen- eralization applied to a sine wave regression task. We summarize our notation in Tab
Table 1: Table of symbols.
Symbol
Meuning
lask imput
((x)
Feature matrix
M(2)
Desired task output
W
Weights of target function
((2)
Output of model
0
Weights of model
Covariance
of
A
Element of # of training samples of model parameters e of total features
P
P
Task octpot dimensionality
m
Task intrinsic dimensionality
Module projection vector
U.
Module projection matrix
3.1 MODEL SETUP
We defer the full details of our theoretical model setup to App A Here we present an overview: we consider a regression task with input E R. and feature matrix œxx) E R such that over the data distribution, the features are distributed i.i.d. from a unit Gaussian: (fz),,) ~ NNO 1); no other assumptions
made about y. We consider the limit when P Suppose that our regression target function y Fm Rd (mapping from to a d dimensional output) is constructed linearly from spx)):
y(z) = W,
where W E Px1, To accommodate multidimensional outputs, note that we shape input features (xx) as matrix (with one row for each output dimension) and parameters W as vector. This choice is more general than parameterizing each output dimension independently (which can be cap- tured as a special case of our approach) and, moreover, aligns with prior theoretical literature (Jacot et al. 2018). Assume [[] = 0 and E EWWTT] = A, where A is diagonal and Tr(A) is finite. Sup-
pose we aim to approximate y(x) using a model (x:) constructed as follows: ((x) = (p()) [ 0,
where 0 E R are model parameters and is the number of parameters. This corresponds to the model only being able to control p of the P true underlying parameters in the construction of
3

---

Published as a conference paper at ICLR 2025
3012 101224161122222222200
1012 81012211111222222222222
2.00
2.00
Loss
1.50
1.50
1.25
1.25
threshold
Interpolation
1.00
14-
Interpolation
11000
18.
20
20
22
0.50
22
0.50
24
26
0.25
26
24
0.25
78
28
Figure 1: Expected training (left) and test (right) set error in a toy model of NN generalization as a function of the number of samples n and the number of model parameters p. The output dimensionality is set as d 1.
y. We decompose A blockwise as: A A1
0
where A E Rexp and A2 E R(P-P)x(P-P)
0
A2
To capture the dependence of the target function on the input dimensionality m, we parameterize as having individual elements A; i-9 (i + 1)-0- This corresponds the number of effective dimensions of variation of W scaling exponentially with m, which is consistent with prior work (McRae et al., (020). also produce versions of our theoretical results without setting specific form for λ; Finally, we consider learning the minimum norm interpolating solution.
3.2 THEORETICAL PROPERTIES
Next, we theoretically analyze the expected training and test set eтror of the above model. Theorem 1. Given target function and model ý estimated as described above, in the limit that P + oc, the expected test loss when averaging over x and is:
lim E [lv(x) ==
P+00
=dTr(A2)F(dnn d min(dn,,) p)
(r(AA)+dTTT
p
ddF, + +1) - min((n, (c c(p + 11-00" + de (2)
p
The expected training loss is:
lim I [[l(() - 0( x)
P-200
n
Tr(AA) - min(dn, p))
dn min (dn, p) clp + 1) -2 11-2
(3)
n
with F(n,p) defined as F(n,p) = [| where R E nxx has elements drawn i.i.d. from 10.
Please see App for a proof and App D for more details on F(n,p). p). Under general A,, training and error grow with Tr(A2): the specific rate at which they grow or shrink with parameters m and p depends on how rapidly Tr(A2) grows with m and shrinks with p. Under the specific parameterization for X; described above, Fig D the value of the training and test set error for varying and number of parameters, holding d 1. Observe that there is clear interpolation threshold at == n where the training loss becomes zero: for p 2 n the model has sufficient capacity to perfectly interpolate the training set. The training loss is positive and increases with in the underparameterized regime (p < n) since the model lacks the capacity to fit increasing amounts of data n. At the p threshold, the test loss dramatically increases, then decreases p increases beyond n. This is consistent with empirically observed behavior of overparameterizze NNs (Belkin) al. 2019). Similarly, test loss decreases increases beyond p.
4

---

Published as a conference paper at ICLR 2025
Verying
Varying
Averaged
Averaged
1-
m
Train, Actusl
Averaged
Test, Actual
Predicted
Incressing m
3-
2
Averaged
increasing m
107 107 10
10
10a 10a
10f 10a
Figure 2: Empirical trends of training (blue) and test (orange) loss over four parametric variations for a NN trained on a sine wave regression task. The parameters varied are: & (number modules), (input dimensionality), p (model size) and (training set size). In the first two plots, each line indicates a different model architecture, and in the last two plots, each line indicates a different choice of between and with n/p fixed at 1000/1 1100/1153 respectively (left/right). The light lines are averaged over all other parameters, and bold lines show averages over the light lines. Dashed lines show theoretical predictions.
3.3 EMPIRICAL VALIDATION
We empirically validate our theoretical model on a parametrically- variable modular sine regression task with targets constructed y(x)) = v Σ Ei az sin(weju of x IR are outputs,
102.
Φig), where x E R" are inputs, y(r) asj. Wij. diy E RR U, E R are parameters chosen randomly for each target function and l and T are fixed (see App for further task details). We train fully connected ReLU-activvted NNs of varying depth and width on the task. The task allows us to quantify how NN generalization depends on a number of factors such the dimensionality model parameters p. the number samples n and the number of modules in the construction of the target function. In Fig 2 we find that our theoretical model matches empirical trends of neural network training and test error (see App A for full details). Nevertheless, we note two key discrepancies between empirical and predicted trends: first, the test loss is empirically larger than predicted under low training data. We hypothesize that this may be because of difficulty optimizing for small indeed, we find that the training loss is larger than expected for small n (in the overparameterized regime (n < p), we expect training loss of 0). Second, the error spike at the interpolation threshold is smaller than theoretically predicted. This again may be due to incomplete optimization, given that the interpolation threshold spike can be viewed as highly adverse fitting to spurious training set patterns.
10y
of the task input, the number of
Figure Theoretically predicted trend of m vs. n to achieve loss of 1.2 on a sine wave re- gression task. Each line indicates a different fully connected NN with different width and depth. increases approximately ex- ponentially with
Our theoretical model predicts that the number of training samples required for generalization to fixed error rate scales exponentially with task dimensionality B. This raises practical challenge for dighedimensional problems, which may require massive amounts training data. Next, aim eircumvent this exponential scaling.
5

---

Published as a conference paper at ICLR 2025
4 USING MODULARITY TO GENERALIZE IN HIGH DIMENSIONS
So far, we have shown that theoretical model treating NN learning as linear regression can closely model the generalization trends of actual NNs. In our model, the number of training samples re- quired to generalize on a task with m dimensional input scales exponentially with m. Now, we demonstrate that modular NNs (in contrast to the monolithic NNs studied so far) can avoid this exponential dependence on m for tasks with an underlying modular structure. We will consider modular networks in which model parameters are divided into separate modules, each of which pro- cesses a projection of the input; monolithic (or nonmodular) networks in this context will correspond to networks without an explicit separation of parameters into modules. We first demonstrate the the- oretical advantages of modular NNs under a specific form of modularity, then develop a modular NN learning rule to learn the underlying modular structure of task. We then empirically validate our approach and demonstrate that our approach can learn the true modules underlying the task.
4.1 SAMPLE COMPLEXITY OF MODULAR LEARNING
Recall in Sec 3 modeling a NN as a linear function of its parameters successfully captured its generalization properties. We aim to use this model to demonstrate the improved generalization of modular learning. For analytical traceability, we restrict our analysis to a specific modular setting that captures crucial aspects of many real-worke world modular learning scenarios. In practical settings, modules often handle low-dimensional inputs, such as attention maps in Andreas et al. (2016b). As such, assume modules receive projected versions of the task input & IR Our theoretical analysis will assume linear projections, but our method and experiments are also applied to non- linear projections. Moreover, we assume that the module outputs are summed produce a final output, feature of architectures such as Mixture of Experts (Jacobs et al.) 1991] Shazeer et al., 2017). Specifically, consider modular NN constructed as linear combination of general NNs (each constituting a module) of low-dimensional projections of the input:
((x) = VK o)(0f)),
(4)
where Üg E Rmmb is a linear projection, úJ R + Rd is NN. We will assume that b is small, creating a bottleneck to each module's 's input. We normalize by VK to make the scale of j(x) invariant to K (treating each term 8,(OTx)) x) as independent, the sum of the terms has variance O(K), so dividing my K makes their variance constant). Note that each module JJ is itself a monolithic NN with arbitrary architecture; we do not restrict the form of the modules themselves. Assuming that the model is a linear function of its parameters, we may model this as:
(0) + (x)) 90.0-
(5)
where 02 E IXP are the parameters of 0j. (x)) E Rddmb and to (x) e RdxP are feature matrices, F(-) denotes flattening matrix into vector, and consider the limit when P co. Observe that this is derived simply by assuming the model output is linear in U, and f and defining the coefficients multiplying them as features (x))(x) and (x)). As before, we assume the fea- tures are distributed i.i.d from a unit Gaussian: (z)) ~ NO0, (z), ~ ~N( 1). We may then write 9 as a linear model:
11
ú(x) (0))
Σ F(U)) sp (x)
(z)
Σイ;
,--
(6)
Next, we assume that the regression target y has the same modular structure with k modules:
y(x)
Σ 3 x)
Vk
6

---

Published as a conference paper at ICLR 2025
where Ug E Rm are the true projection directions and y, RR underlying the target. We apply the same linearity assumption to find:
IR are the true modules
k
u(x) - ( (x)
Eu,+ x (w) (z)
Σ
W
(8)
VI
==1
where W Rº are the parameters of V/j- As in the case of monolithic networks, assume B[W,] = 0 and B[W, A where A is diagonal. In this case, observe that each W; parameter- izes a function with a b-dimensional input; thus is appropriate to assume that W is distributed as if m b:
A [ i- -- (i 1) 1)-a²
To preserve spherical symmetry in the distribution of Uy. we assume that B[U,] = 0 and ((F( = I. To complete our model definition, we define (p(x) E IRd (mb++) P) as the concatenation of (((0(z))
(x) and (x): plz) [000) (x)) We may then write:
u(z) ppx)) I 0
!1
(10)
0
where 0 E Rmb+p is defined as: 0 VR Σラ= F(U,), K 03]. Similarly, we may write:
((x) (( where W Rmb+ P is defined as: W vR _ F(U,)(U)))))))))))) E/ W.] Note that this model is nearly identical to that of monolithic NNs in Sec 3 the key difference is the different distribution of W. W has covariance A
1 where A is not dependent on m; A
0 A
is parameterized analogously to Section 3. Assuming is trained to minimize squared loss training set, we may adapt Theorem I to this setting to compute the expected training and test loss of modular networks:
Theorem 2. Given a target function y and model ý estimated as described above, the limit that P > co, the expected test loss when averaging over and W is:
lim E [1l((x) - -0(x))((1111)]1) = d'Tr(A22F((dn,)) - min(dn. p) Tr(Ai) + +dTr((A)
P+200
=dF(dn, p)c () 1 min (dd, (mb+ -c c- c0 1)-0--) dmb + de (11)
p
The expected training loss is:
dn min(dn,y) Tr(A2)
lim
E [lx(X) xx)
P n
n
dn min(dn.
c(p + 1- 1)- (12)
n
with F(n.p) p) defined as:
F(n.p) E ||
(13)
where R E Bnxp has elements drawn i.idd from N N(00 1).
The proof simply applies Theorem I with a different covariance matrix for W; see App for the full proof. When Tr(Az) is independent of m, we see that unlike the monolithic network, the training loss does not depend on m, and the dependence of the test loss on m is linear. This is because the module inputs have effective dimensionality 6 instead of due to the bottleneck caused by the module projections Ü, Furthermore, in the underparameterrized regime (dn > p), the test loss becomes (dF(ddn,)) p) + d) Tr(Az) which does not depend on m, implying that the n required to reach a specific loss can be bounded by function of only (assuming that some value of p can achieve the desired loss). Thus, the sample complexity of modular NNs is independent of the task dimensionality. This dimension independence holds regardless of the parameterization of A; the only condition is that the modules must have a dimension-inde iewpnsiiooniddeedeeenn input bottleneck (i.e. the covariance A of module parameters must be independent of m)) This result suggests that, unlike monolithic NNs, modular NNs can scale to digh-dimensional, modular problems without requiring intractable amounts of data.
7

---

Published as a conference paper at ICLR 2025
0.350
0.525
Our method
Baseline modular
Sample Complexity
0.500
60*
Baseline monolithic
0.45
Baselline monolithic
modular
Our method
0.450
0.425
0.400
2
Dimensions
Images
(a) Sine Wave Regression Task
(b) Compositional CIFAR-10
Figure 4: Comparison of our method with baselines of modular and monolithic architectures trained from random initialization the sine regression task (a) and Compositional CIFAR-10 (b). (a): Required training sample size to achieve a desired test error vs. # of input dimensions. Each light line indicates a different model architecture specified in App E averaged over five random seeds. The bold lines show averages over the light lines. (b): Accuracy vs. # of component images with fixed number training samples. Margins indicate standard errors over five random seeds.
4.2 MODULAR LEARNING RULE
Inspired by the improved theoretical generalizability modular NNs, and the finding that modular architectures trained with gradient descent on task often cannot exploit these efficiencies (Csordás) jet al., 2021; Mittal et al.] 2022a), we develop modular learning rule that practically exhibits this advantage. importantly, now relax the assumption that module input projections are linear.
We consider modular regression tasks with targets constructed as follows:
v( Σ (3;
(14)
where y functions that depend potentially nonlinear projection of the input represented by My depending on both module projections U, and inputs z. Observe that this generalizes the linear projections considered before (in Eqn 77) Suppose approximating the U, with U, and the yj with Dj (parameterized
aim model the target function by neural network). We propose kernel-based rule to learn the initializations of Ü, from the training data; this allows us to efficiently learn the modules Új- Assume we are provided a set training data y(X) E Rdax1, Given modular structure, we first aim to approximate the data as:
K
x(X)
o (()
(15)
where E Raxm and to is an arbitrary nonlinearity applied elementwise to the input data such that xpX;; U.) E Rdaxp, 0g IR and K the number of expected modules.
We expect that if Ú; = U; and & is sufficiently expressive, then yXX) can be well approximated. Assuming pk > dn, observe that the minimum norm solution for 8 can be computed as:
01
0
[+(X;ú)) v(X:t,)
p(x;0x)) v(x)
(16)
:
0x]
In general, such solution exists for any choice of Ú. However, we hypothesize that the Ü, is far from U, then the norm of the 0; will be large: intuitively, interpolating the data along the "incorrect" projection directions Üg will be more difficult. Thus, we optimize Ü; to minimize the squared norm
8

---

Published as a conference paper at ICLR 2025
of 0, Specifically, minimize:
Σ -
(17)
where K E Rduxdn is kernel matrix applicd to the data corresponding to the following ker- nel: x2)
= - pa (,)p(a";;,)"" s(x1, z2; U.), where x(11, z2; U,) is a module- conditional kernel between V2 and Z2. Experimentally, we tailor & to the modular structure of the problem we consider. Note that the above analysis only applies when e sufficiently expressive (i.e. pk > dn), which is a natural assumption for typically overparameterized models like neu- ral networks. When models not satisfy this assumption, (generally nonzero) error between the predicted and true y (X). Importantly, in this case, minimizing the norm of 0 with respect U, may not necessarily yield lower prediction error.
[ a solution minimizing the
App El describes the specific choice of k. Alg I shows the full procedure to find a single module projection Ü; each step of the algorithm simply applies gradient on Eqn 17 with respect to Ü. We repeat this procedure K times with different random initializations to find the initial values of all K module projections in our architecture. Then, we train all module parameters (including the Ú,) via gradient descent on the task loss. We stress that our approach is applicable to fairly general set of modular architectures of the form E, U2 (1; U;): it does not restrict modules to receive only linear projections of inputs, and, moreover, does not restrict the form of the modules.
4.3 EXPERIMENTAL RESULTS
We evaluate the generalizability of our method on modular NN vs. baselines of a randomly initialized monolithic and modular NN trained on 1) sine wave regression tasks of varying dimen- sionality (fixing k = m.. 2) nonlinear variant of the sine wave regression task where the task has a nonlinear module struc- ture, and 3) Compositional CIFAR-10 (based Compositional MNIST (Jarvis et al. 2023)), ), a modular task in which each input consists of multiple CIFAR-10 images and the goal is simul- lancously predict the classes of all images; see App E Fig for an illustration. In Compositional CIFAR-IO, each input is con- structed as a concatenation of k flattened CIFAR-10 images (re- sulting in a 3072k dimensional vector) and target outputs are hot encoded 10k dimensional vectors encoding the class of each component image. The modular architectures are constructed as ((x) = Σラ (0; U3) where yj are fully connected, ReLU- activated networks, and the projection operation Uj parameter- izes the first layer. Monolithic architectures normal fully connected See App for further details on the datasets and the full experimental setup.
Baseline
0.82
Smitarty Score
8.40
0.38
0.58
0.54
The Tranting Herattor
Figure 5: Average cosine simi- larity between learned and target module directions over training for a modular NN initialized with method random initial- ization (baseline).
networks.
Modular NNs empirically generalize better in and nutrof-distributtin As shown in Fig our modular method generalizes better compared with both the monolithic baseline method and the modular baseline method evaluated by sample complexity for the sine wave regression task and accuracy for Compositional CIFAR-10. On both tasks, our method's advantage persists even on higher-dimensional inputs. Interestingly, on the sine-wave task, the monolithic baseline outper- forms the modular bascline, highlighting the difficulty of optimizing modular architectures. We also conduct experiments on additional variants of Compositional CIFAR-10 that test -iitt--iistiittti generalization: 1) our method learns to classify unseen class combinations, thus generalizing com- binatorially FFppFTabb 3) and 2) our method robust to small amounts of Gaussian noise added to training inputs (Appl F Tab 4J, thus generalizing to small distribution shifts.
Our learning rule finds the true task modules Fig 5 computes a similarity score between our learned module projections (Ú) and the target module projections (U) on the sine wave regression task; our method indeed aligns with the target modules. App F Fig I1 plots a low-dimensional
representation of the target (U) and learned (Ü) module projections: our learned NN initializations

---

Published as a conference paper at ICLR 2025
closely cluster around the target modules without any task training. On Compositional CIFAR-IO, the learned module projections can be directly visualized as in App F
Fig 12 here, our module initializations make each module sensitive to single component images without any task training, suggesting that our approach promotes generalization by correctly learning the modular structure of task. We also conduct ablation studies in App F] Fig which show that 1) our method performs nearly as well as using the ground-truth module directions and 2) allowing our learned module directions to adapt directly to the task loss improves performance.
Table 2: Comparison of our method with baselines on a nonlinear variant of the sine wave regression task where k = m = 5. Test loss val- ues are evaluated with stan- dard errors over 5 trials.
Our learning rule extends to nonlinear module projections So far, we have considered modules (yj in Eqn T4p for which the input is a linear function of both U, and x. Next, we consider nonlin- car variant of the sine wave regression task in which modules are a function of l|, z1|2. where uy are the vector module projection directions, which is nonlinear in both uj and x. The modular archi- tecture is constructed non-linearly as: J Σÿ= C - =11l)): see App E for further details. Tab @illustrate that our method sig- nificantly outperforms the baselines, indicating that our method ex-
VK
Method
Test Loss
Baseline monolithic
1.302 + 0.223
Baseline modular
0.555 0.136
tends to nonlinear settings as well.
Our method
0.393 + 0.099
DISCUSSION
Existing NN scaling laws show that in order to generalize task, monolithic NNs require an exponential number of training samples with the task's dimensionality. this paper, we develop theory demonstrating that modular NN can break this scaling law: they only require only a constant number of samples to generalize in terms of task dimension when applied to modular tasks. To our knowledge, we the first to demonstrate such result using explicit expressions for generalization error in modular NNs. Based on this theoretical finding, we propose a novel learning rule for modular NNs and demonstrate its improved generalization, both in and out of distribution, sine wave regression task and Compositional CIFAR-10.
In pursuit of explicit, non-asymptotic expressions for generalization error in modular and monolithic NNs, we make strong theoretical assumptions consistent with previous literature, which we hope can be further relaxed in future work. Moreover, our results apply to a specific form of modularity that captures structures in common real-world modular architectures but is not fully general. Notably, our theory considers linear module projections (although our method is also applied to nonlinear projections), and both the theory and the experiments assume that the model output is the sum of module outputs. We expect that future analyses can demonstrate the benefits of modularity more widely. For example, routing mechanisms are popular type of modularity in which modules are flexibly composed or combined based on routing network. We expect analyses similar to ours to show that, to the extent that routing mechanisms allow process oow-dimensional inputs rather than the full task input, they also generalize better. Similar theory may also explain the generalization benefits of self-attention based architectures, which may learn modular substructures. Finally, we find that while the theory predicts a deasddiinnioonnddeeeeeeen sample complexity for modular NNs, empirically we do not eliminate this dependence due the difficulty of optimizing modular NNs in high dimensions. Nevertheless, our learning rule significantly eases this challenge. Further, our results suggest that more focus should be placed on optimization strategies for modern modular architectures beyond naively applying gradient descent: this could unlock further generalization benefits of commonly used modular motifs (such as routing mechanisms, attention, MoE etc.).
Practically, we expect that modularity provides the most benefit for modular tasks with high- dimensional inputs; this is because the relative sample complexity improvement between nonmodu- lar and modular tasks is greater when task dimensionality increases. Indeed, as discussed in See modularity empirically significantly improves generalization in domains ranging from reinforce- ment learning and robotics to visual question answering and language modeling, all of which can be highly compositional and can have digensinmnsi task inputs. We speculate that the modular structure of detfaattenion-baase architectures may explain their success in many of these domains. Our findings provide a step toward fundamentally understanding how modularity can be better ap- plied to solve high-dimensional generalization problems.
10

---

Published as a conference paper at ICLR 2025
REFERENCES
R Agarwal, N Frosst, X Zhang, R Caruana, and GE Hinton. Neural additive models: interpretable machine learning with neural nets. NeurIPS, 2021.
Ibrahim Alabdulmohsin, Behnam Neyshabur, and Xiaohua Zhai. Revisiting neural scaling laws in language and vision. In NeurlPS, 2022.
Ferran Alet, Maria Bauza, Alberto Rodriguez, Tomas Lozano-Perez, and Leslie P. Kaelbling. Mod- ular meta-learning in abstract graph networks for combinatorial generalization. In NeurIPS meta- learning workshop, 2018a
Ferran Alet, Tomás Lozano- and Leslie P. Kaelbling. Modular meta-learning CoRL, 2018b.
Jacob Andreas, Marcus Rohrbach, Trevor Darrell, and Dan Klein. Learning to compose neural networks for question answering. In NAACI, 201
Jacob Andreas, Marcus Rohrbach, Trevor Darrell, and Dan Klein. Neural module networks. Pro- ceedings of the IEFE conference on computer vision and pattern recognition, pp. 11.39-45 48, 201
Arjun Ashok, Chaitanya Devaguptapu, and Vineeth N Balasubramanian. Learning modular struc- tures that generalize -isttr-ditt (student abstract). In AAAI, 2022.
Dzmitry Bahdanau, Kyunghyun Cho, and Yoshua Bengio. Neural machine translation by jointly learning align and translate. In ICLR, 2015.
Yasaman Bahri, Ethan Dyer, Jared Kaplan, Jaehoon Lee, and Utkarsh Sharma. Explaining neural scaling laws. arXiv preprint, 2021.
Mikhail Belkin, Daniel Hsu, Siyuan Ma, and Soumik Mandal. Reconciling modern machine- learning practice and the classical biai-variance trade-off. PNAS, 116(32):15849--5583 2019.
Yoshua Bengio, Tristan Deleu, Nasim Rahaman, Rosemary Ke, Sébastien Lachapelle, Olexa Bila- niuk, Anirudh Goyal, and Christopher Pal. A meta- meta-transfer objective for learning to disentangle causal mechanisms. ICLR, 2020.
Abdulkadir Canatar, Blake Bordelon, and Cengiz Pehlevan. Spectral bias and mode- alignment explain generalization in kernel regression and infinitely wide neural networks. Nature communi- cations, 12(1):1-12, 2021.
Michael B. Chang, Abhishek Gupta, Sergey Levine, and Thomas L. Griffiths. Automatically com-
posing representation transformations as means for generalization. In ICLR, 2019.
Arslan Chaudhry, Naeemullah Khan, Puneet K. Dokania, and Philip H. S. Tom. Continual learning in low-rank orthogonal subspaces. In NeurIPS, 2020.
Yutian Chen, Abram L. Friesen, Feryal Behbahani, Arnaud Doucet, David Budden, Matthew W. Hoffman, and Nando de Freitas. Modular meta-learning with shrinkage. In NeurIPS, 2020.
Rohan Chitnis, Leslie Pack Kaelbling, and Tomás Lozano-Pérez. Learning quickly plan quickly using modular meta-learning. In ICRA, 2019.
Aidan Clark, Diego de Las Casas, Aurelia Guy, Arthur Mensch, Michela Paganini, Jordan Hoff- mann, Bogdan Damoc, Blake Hechtman, Trevor Cai, Sebastian Borgeaud, et al. Unified scaling laws for routed language models. In International Conference on Machine Learning, pp. 4057- 4086. PMLR, 2022.
Róbert Csordás, Sjoerd van Steenkiste, and Jürgen Schmidhuber. Are neural nets modular? inspect- ing functional modularity through differentiable weight masks. In ICLR, 2021.
Hugo Cui, Bruno Loureiro, Florent Krzakala, and Lenka Zdeborová. Generalization error rates in kernel regression: The crossover from the noiseless to noisy regime. NeurIPS, 14:10131-10143, 2021.
11

---

Published as a conference paper at ICLR 2025
Hugo Cui, Bruno Loureiro, Florent Krzakala, and Lenka Zdeborová. Error rates for kernel classifi- cation under source and capacity conditions. arXiv preprint, 2022.
Limeng Cui and Aaron Jaech. Re-examining routing networks for multi-task learning. arXiv preprint, 2020.
Vanessa D' A'marioo Tomotake Sasaki, and Xavier Boix. How modular should neural module net- works be for systematic generalization? In NeurIPS, volume 34, pp. 23374-23385, 2021.
Anirudh Goyal, Alex Lamb, Jordan Hoffmann, Shagun Sodhani, Sergey Levine, Yoshua Bengio, and Bernhard Schölkopf. Recurrent independent mechanisms. ICLR, 2021.
Trevor Hastie, Andrea Montanari, Saharon Rosset, and Ryan Tibshirani. Surprises in high- dimensional ridgeless least squares interpolation. Annals of statistics, 50(2):949, 2022.
Ronghang Hu, Jacob Andreas, Marcus Rohrbach, Trevor Darrell, and Kate Saenko. Learning to reason: End-to-end module networks for visual question answering. In ICCV, pp.
2017.
Marcus Hutter. Learning curve theory. arXiv preprint, 2021.
Abhiram lyer, Karan Grewal, Akash Velu, Lucas Oliveira Souza, Jeremy Forest, and Subutai Ah- mad. Avoiding catastrophe: Active dendrites enable multi-task learning dynamic environments. Frontiers in neurorobotics, 16. 2022.
Robert A Jacobs, Michael I Jordan, Steven J Nowlan, and Geoffrey E Hinton. Adaptive mixtures of local experts. Neural computation, 3(1):79-88 87, 1991.
Arthur Jacot, Franck Gabriel, and Clément Hongler. Neural tangent kernel: Convergence and gen- eralization in neural networks. In NeurIPS, volume 31, 2018.
Devon Jarvis, Richard Klein, Benjamin Rosman, and Andrew Saxe. On the specialization of neural modules. In ICLR, 2023.
Hui Jin, Pradeep Kr Banerjee, and Guido Montúfar. Learning curves for gaussian process regression with power-law priors and targets. arXiv preprint, 2021.
Justin Johnson, Bharath Hariharan, Laurens der Maaten, Li Fei-Fei, C. Lawrence Zitnick, and Ross Girshick. Clevr: A diagnostic dataset for compositional language and elementary visual reasoning. CVPR, 2017.
Jared Kaplan, Sam McCandlish, Tom Henighan, Tom B Brown, Benjamin Chess, Rewon Child, Scott Gray, Alec Radford, Jeffrey Wu, and Dario Amodei. Scaling laws for neural language models. arXiv preprint, 2020.
Seung Wook Kim, Makarand Tapaswi, and Sanja Fidler. Visual reasoning by progressive module networks. In ICLR, 2019.
Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. In ICLR, 2015.
Dilip Krishnan, Terence Tay, and Rob Fergus. Blind deconvolution using normalized sparsity measure. In CVPR, pp. 233-240. IEEE, 2011.
Francesco Locatello, Damien Vincent, Ilya Tolstikhin, Gunnar Rätsch, Sylvain Gelly, and Bern- hard Schölkopf. Competitive training of mixtures of independent deep generative models. arXiv preprint, 2019.
Kanika Madan, Nan Rosemary Ke, Anirudh Goyal, Bernhard Schölkopf, and Yoshua Bengio. Fast and slow learning of recurrent independent mechanisms. In ICLR, 2021.
Rafid Mahmood, James Lucas, David Acuna, Daiqing Li, Jonah Philion, Jose M Alvarez, Zhiding Yu, Sanja Fidler, and Mare T Law. How much more data do i need? estimating requirements for downstream tasks. In CVPR, pp. 275-284, 2022.
Andrew McRac, Justin Romberg, and Mark Davenport. Sample complexity and effective dimension for regression manifolds. In NeurIPS, volume 33, pp. 1 13004, 2020.
12

---

Published as a conference paper at ICLR 2025
Sarthak Mittal, Alex Lamb, Anirudh Goyal, Vikram Voleti, Murray Shanahan, Guillaume Lajoie, Michael Mozer, and Yoshua Bengio. Learning to combine top-down and bottom-up signals in recurrent neural networks with attention over modules. In ICML, pp. 6972-6986. PMLR, 2020.
Sarthak Mittal, Yoshua Bengio, and Guillaume Lajoie. Is modular architecture enough? In NeurlPS, 2022a.
Sarthak Mittal, Sharath Chandra Raparthy, Irina Rish, Yoshua Bengio, and Guillaume Lajoie. Com- positional attention: Disentangling search and retrieval. ICLR, 2022b.
Brady Neal, Sarthak Mittal, Aristide Baratin, Vinayak Tantia, Matthew Scicluna, Simon Lacoste- Julien, and Ioannis Mitliagkas. A modern take on the bias-variance tradeoff in neural networks. ICML 2019 Workshop Identifying and Understanding Deep Learning Phenomena, 2019.
Giambattisa Parascandolo, Niki Kilbertus, Mateo independent causal mechanisms. In ICML, 2018.
and Bernhard Schölkopf. Learning
Deepak Pathak, Christopher Lu, Trevor Darrell, Phillip Isola, and Alexei A Efros. Learning to control self- assembbsnmllng morphologies: study of generalization via modularity. In NeurIPS, volume 32, 2019.
Jonas Pfeiffer, Sebastian Ruder, Ivan Vulić, and Edoardo Maria Ponti. Modular deep learning. arXiv preprint, 2023.
Jason W Rocks and Pankaj Mehta. Memorizing without overfitting: Bias, variance, and interpolation in overparameterized models. Physical Review Research, 4(1):013201, 2022.
Clemens Rosenbaum, Tim Klinger, and Matthew Riemer. Routing networks: Adaptive selection of non-linear functions for multi-task learning. In ICLR, 2018.
Clemens Rosenbaum, Ignacio Cases, Matthew Riemer, and Tim Klinger. Routing networks and the challenges of modular and compositional computation. arXiv preprint, 2019.
Jonathan S Rosenfeld, Amir Rosenfeld, Yonatan Belinkov, and Nir Shavit. A constructive prediction of the generalization error across scales. In ICLR, 2020.
Alexander Sax, Jeffrey Zhang, Amir Zamir, Silvio Savarese, and Jitendra Malik. Side-tuning: Net- work adaptation via additive side networks. ECCV, 2020.
Utkarsh Sharma and Jared Kaplan. Scaling laws from the data manifold dimension. JMLR, 23:9-1, 2022.
Noam Shazeer, Azalia Mirhoseini, Krzysztof Maziarz, Andy Davis, Quoc Le, Geoffrey Hinton, and Jeff Dean. Outrageously large neural networks: The sparsely-gated mixture-of-expprrs layer. In ICLR, 2017.
Harshvardhan Sikka, Atharva Tendle, and Amr Kayid. Multimodal modular meta-learning. OSF Preprints, Oct 2020. doi: 10.31219/osfio/66kkbbb URL 10/6ek2b
Ben Sorscher, Robert Geirhos, Shashank Shekhar, Surya Ganguli, and Ari Morcos. Beyond neural scaling laws: beating power law scaling via data pruning. In NeurlPS, 2022.
Stefano Spigler, Mario Geiger, Stéphane d' Ascoli, Levent Sagun, Giulio Biroli, and Matthieu Wyart. A jamming transition from under-to over seer-paramtrization affects generalization in deep learning. Journal of Physics A: Mathematicul and Theoretical, 52(47):474001, 2019.
Stanislaw J Szarek. Condition numbers of random matrices. Journal of Complexity, 7(2)1131-149 1991.
Yi Tay, Mostafa Dehghami, Jinfeng Rao, William Fedus, Samira Abnar, Hyung Won Chung, Sharan Narang, Dani Yogatama, Ashish Vaswani, and Donald Metzler. Scale efficiently: Insights from pre-training and fine-tuning transformers. In ICLR, 2022.
Dietrich Von Rosen. Moments for the inverted wishart distribution. Scandinavian Journal of Statis- fies, pp. 97-109, 1988.
13

---

Published as a conference paper at ICLR 2025
Alexander Wei, Wei Hu, and Jacob Steinhardt. More than toy: Random matrix models predict how real-world neural representations generalize. In ICML, pp.
PMLR, 2022.
Yuichi Yamashita and Jun Tani. Emergence of functional hierarchy in a multiple timescale neural network model: A humanoid robot experiment. PLOS Computational Biology, 4(11):1-18, 11 2008. doi: 10.13710 /fouuraalll..pb....... 1000200 URL https:///oi.org//// 122111111111 1371 1371/ pebi. 1000220
Guangyu Robert Yang. Madhura R. Joglekar, Francis Song, William T. Newsome, and Xiao-Jing Wang. Task representations in neural networks trained to perform many cognitive tasks. Nature Neuroscience, Jan 2019. URL https://www http://wwwaaarereeeeeeeeeeeea
nature. com/ com/articles/
101593-018-0310-1
Ruihan Yang. Huazhe Xu, Yi Wu, and Xiaolong Wang. Multi-task reinforcement learning with soft modularization. In NeurIPS, 2020.
Xingyi Yang, Jingwen Ye, and Xinchao Wang. Factorizing knowledge in neural networks. In ECCV, pp. 73-91. Springer, 2022.
Kexin Yi, Jiajun Wu, Chuang Gan, Antonio Torralba, Pushmeet Kohli, and Joshua B. Tenenbaum. Neural-symbolic vqa: Disentangling reasoning from vision and language understanding. In NeurIPS, 2018.
14

---

Published as a conference paper at ICLR 2025
A THEORETICAL MODEL OF NEURAL NETWORK GENERALIZATION
A.1 SETUP
We consider a regression task with input E R and a feature matrix (p()) E Rdx such that over the data distribution, the features are distributed i..d. from unit Gaussian: (p()); NO0 1). We consider the limit when P + oc. Suppose that our regression target function Rm + R is constructed linearly from pzz):
v(z) (
(18)
where W E Px1 To accommodate multidimensional outputs, note that we shape input features ((x) as matrix (with one for each output dimension) and parameters W as vector. This choice is more general than parameterizing each output dimension independently (this can be cap- tured as a special case of our approach) and moreover, aligns with prior theoretical literature (Jacot et al. 2018). Assume E[W] = and E[WW"] = A, where A is diagonal and Tr(A) is finite. Suppose we aim approximate y(x) using a model j(x)) constructed as follows:
[] 0,
g(z) = (ax)
(19)
where 0 E IRpx1 are model parameters and p is the number of parameters. This corresponds to the model only being able to control p of the P true underlying parameters in the construction of y. We
decompose A blockwise as: A A 0
where A₁ E RPxP and A2 E (R(P)))((-o))
A2
We make specific choice of parameterization for the individual elements A of A:
A = 2- (i Q
(20)
for some constants e and Ω. We justify this choice as follows: we define the effective dimensionality of A as (2,, A²)² (this measure approximates the to norm Krishnan et al. 2011). and thus can be used Σ, X
as a measure A's dimensionalityy. For large m, this can be approximated as: C A,2² X, Rs Ω we interpret this as y having 222m free parameters. This is consistent with the observation that regression an mn-dimensional input space has a function space that scales exponentially with m (McRae et al. 2020). Intuitively, this because the function must have enough free parameters
to express values at all points in volume of its input space, and volume scales exponentially with m (Sharma & Kaplan 2022).
given set of training data X with associated feature matrix transformation (xX) E Rdnx where is the number of data points. Suppose 0 is optimized to find the minimum norm interpolating solution to the data:
Next, suppose
min (1eex) -
x(X)w
+ 2 11012
(21)
where 7 + O. Recall that the solution can be found as:
= [x(x]) [ x(X).:
(22)
Thus, the model's prediction on a point .r is:
11 ((x)
[ ((X)W.
((x) p(x)
(23)
A.2 EMPIRICAL VALIDATION
Next, empirically validate our theoretical model on a parametrically modular sine wave regression task (see App [ for task details). The task allows us to quantify how NN generalization depends the number of dimensions of variation such as the dimensionality of the task input,
15

---

Published as a conference paper at ICLR 2025
the number of model parameters p. the number of samples and the number of modules & in the construction of the target function. Note that our theoretical model does not include the number of modules l since it does not explicitly construct the target modularly. Thus, directly applying our theory, we would expect the loss to be invariant to k. Intuitively, this because varying & increases the complexity of the task in way that is irrelevant for generalization.
Trends of generalization error of NNs Now, we train NNs of various architectures the task and observe error trends as a function of k, m. and n. We parameters e and f2 of our theoretical model our task. Furthermore, because each parameter in our NNs may not correspond parameter in our theoretical model, we use a linear scaling of the number of true NN parameters estimate the number of effective parameters in our theoretical model; specifically, we estimate p ap' where p' is the actual number of NN parameters. See App E for more details.
single
Fig shows that our theoretical model can capture many trends of the training and test loss function of k. m. and n. In particular, our model predicts the invariance of loss to k. the sub-linear increase in loss with m, and the double descent behavior of loss with p and n. Notably, our model predicts the empirical location of the interpolation threshold as seen in the last two plots of Fig 2
We note two key discrepancies between our theory and empirical results: first, the loss is empir- ically larger than predicted for small amounts of training data, and second, the error spike at the interpolation threshold smaller than predicted by the theory.
We believe the first discrepancy is due to imperfect optimization of neural networks, especially in low data regimes. Note that the linearized analysis assumes that the linear model solution finds the exact global optimum. However, the actual optimization landscape for modular architectures is highly non non-convex, and the global optimum may not be found especially for small datasets (indeed, find a significant discrepancy between predicted and actual training loss values for small data size n; in the overparameterized regime, the predicted training error is exactly 0). We believe this causes the discrepancy between predicted and actual test errors in low data regimes.
We hypothesize that the second discrepancy is also partly due to imperfect optimization. This is because the interpolation threshold spike be viewed highly adverse fitting to spurious training set patterns. This imperfect optimization is more pronounced at smaller m. Despite these discrep- ancies, we nevertheless find that our theory precisely captures the key trends of empirical test error.
Finally, we consider the trend between m and n implied by our model. As Fig 3 reveals, for vari- ous nonmodular architectures, the sample complexity grows approximately exponentially with the task dimensionality, consistent with prior theoretical observations (see Sec 2.2). This implies that generalizing on sige-dimensional problems can require a massive number of samples.
B PROOF OF THEOREM 1
Proof. Test set error We first compute the expected test set error. Note that the squared error can be written as:
l10(x) -(x)||== p(xx) 1) (xx)
[ x(X)W -
x(X)W-v(x)HH x(x)W
((x) - w 1
w/
= ((x) (x)
(24)
For notational convenience, define A as the first p columns of (xX) and B as the remaining P-p
columns such that (xX) [A, B). Also, define M = []] ((x) [
(x) = [ At[A, B| =
A
A
Then, using the cyclic property of trace, the squared error can be written
0
0
(M 2 = Tr
(M - xx))) = Tr ((M - IJWWT p0=) (())))(()))))))))))))))))))))())
(25)
16

---

Published as a conference paper at ICLR 2025
Next, take the expectation with to and W and use the fact that Elp(xx) = dI and B[WW7] = A to find that:
E [1lx(2) - -(x)] - de [Tr((M --MM- r)=
(26)
Finally, expanding:
de [T((M - -
- - -
[Tr(MP [T((MMA - 2dE [Tr(MA)) + dTTr(A) (27)
Next, we compute
[Tr(M MA]]: [T]]:
A A ATT
0
AAA
A'B
B [Tr(M" MA]] = E 'I
A
0
0
ATA AITAIATA
ATAATB
A)]
E
Tr
(28)
BT
ATA
B' A A
Next, we decompose A blockwise as:
A1 0
0
A =
(29)
A2
where A1 E Rpxp and A2 E (R(P-))((PP)) Then:
E [TT(M = = Tr (
A'AA B
A'A
1)
BT A AA A'BAA)
=E [TAA +E [Tr(Ate AABABB")
(30)
Next, consider, [[Tr(MA)]:
E [[T( E Th ([ATA A 0
=E
(31)
Combining this result with the earlier result, the expected squared error can be expressed as:
E (1) (:) - -(=)]l= dE [T>(AtTA A' }] dE [Tr(A' AAI)] dTr(AA)
(32)
Next, we evaluate E [Tr(A AT By linearity of trace, and the independence of A and B, we have:
E (Tr A' Tr (E[A [[] E[BA,B"])
(33)
Define the following quantities:
a= E
AtT AATA!
(34)
and
3 E A2
(35)
where ".y indicates the ith and jth column of the argument. Note that by symmetry the data points and output dimensions, both E [AtT AT] dn and E [B2BBT] dud must be proportional to the identity matrix. Thus, the expectation of their top left entry is the same as the expectation of any other entry:
E AA
A!]
al
(36)
and
IE [BA 31
(37)
Then,
E [T((A A BA,B") Tr(a131) a,3n To evaluate a, observe that since has elements distributed from (( 1): F(dn, = E [|A [ Tr (ELA)) (AI)
(38)
(39)
Using the definition of a:
F(dn, p) odn
(40)
17

---

Published as a conference paper at ICLR 2025
Therefore,
F(dn.p) dn
To evaluate B, note that B1: has elements distributed from (00 1). Thus,
is simply: 3 = A2.
Σ
BiB (2,, A2, Tr(A2)
(41)
1=1
Substituting a and | into the expression for E "Tr(ATAA' BA)
E [Tr( A'BAB") Tr(A2)
(42)
Next, we evaluate E ([r)) First, define the singular value decomposition of A as A UEVI, Then, expanding A and using the cyclic property of trace:
=
E (TT = [TVV"""""EVV"A = E [T<(EEEVTA))
(43)
Applying the linearity of trace:
E [T<(AAAA)]] = Tr (EIVE/EV\JA)
(44)
Now, we examine E VE E R First, note that Σ is diagonal matrix with entries and O: specifically, has min(dn, p) is and remaining entries (if any) O. Thus, we may write:
min(dn.p) p)
E [VEIEVT
Σ
V.V.
(45)
Next, note that the distribution of A is symmetric rotations of its p columns. Thus V., must also have rotationally symmetric distribution. Since ||V,ell2 1, V.. = 11. Thus:
p
min (in, p)
VEEENT
(46)
P
Substituting into the expression for E (T(AATA)))
E (TT(AAAA)) min(dn,p) p)
Tr(Ai)
(47)
p
Combining the results from earlier, the expected squared error can be written as:
min(dn.p)) +
E [11(()) v(>)] -' d'Tr(A2)
(48)
P
Next, using the definition of A --
(1 -6- 1) 0
observe that:
P
Tr(At) Σ
A/ =
Σ
0
- -
(i 1)
= c - c(p 1)
11- -0-
(49)
1=1
Tr(A2) Σ
A Σ
m
(i+1) 9m c(p + 1)
11-2 -2-
(50)
1=p+1
1=p+1
TT(A)
-
(51)
We finally use the expressions for A₁ and A to write the result in terms of c and Ω2:
mimmmnn (dn, p) (c- c(p + 1) 1)- dc (52)
EE [1lú(x) (((b(t))v(>)11]]] dF (d', 1 -
p
Training error Now, we compute the training set error. Writing out the summed training set error over all data points:
[[(()) ( p(x) [1] (x() [!]l x(X)W - ((X)) x(X)W
= =I[(AA' D)A, B|W
2
=Il(AAt - T)A, (AAT 1)B
(AAt 3 -T)BW/] (53)
18

---

Published as a conference paper at ICLR 2025
Expressing this squared norm trace and using the cyclic property of trace:
[[0, - Tr (w (AA) [0, I)B)W)
=Tr
(10. (IO,
I) (AA' - I)B]WW
(54)
Taking the expectation with respect to W:
E Tr ((0, (AAt [), (AA* - B) Tr 1 A2)
-
(55)
Note that (AAT 1)2 = I - AAt Again using the eyclic property of trace and the independence of and B, we find:
(E|B" -)) (ATTAA)))
(56)
From the calculations for test set error, we have:
E [BAyB"] Tr(A2)
(57)
Substituting into the earlier expression:
Tr (BI ) Tr (E|I- AA Tr((2) - AAT)]
(58)
To evaluate ([Tr(I-AA)) we use the singular value decomposition of A = UEV* and the cyclic
property of trace:
[Tr( - AAT) = E - II(I-UVVVVSiU E - -
|r ΣΣイ
(59)
Observe that A has full rank with probability 1. Thus, E also has full rank with probability 1, implying that ΣEt is with probability diagonal matrix with min(dn,p) p) 1s and remaining entries (if any) 0:
E [Ir(1 - yΣ' = dn min(dn, p)
(60)
Substituting into the earlier expression, we have a result for the total training set error over all training points:
E [ll((x) - x -r(A-)(dn-minddn. p))
-
(61)
To arrive at the final result for expected training set error simply divide by and express Tr(Ag)
in terms of e and Ω
B[))((X x(x)1 dn min (dn, p) Tr(A2) dn min (n. p)
c(p 1) 1)-
(62)
C PROOF OF THEOREM 2
Proof. Test set error Using the same techniques as in the proof of Theorem 1, may write the expected test set error in terms of A as:
min((d..)) p) Tr(An) + (TA
E [I(((>) - -000)(()11111 = -(=))
(63)
P
Using the definition >; = 1-2 (i -1+11-22 1 and the assumption > mb, observe that:
Tr(A,) = mb +
Σ
\ mb +
Σ
- - -ii+ - = = mb + - c(p +1
(64)
1=1
i=1
20
Tr(Ag) Σ
x= = Σ
c
i - (i 11+11= 1)- '1 =c(pp + 1
(65)
i=p+1
1=p+1
Tr(A)= Tr(Ai) Tr(A2) mb +
(66)
19

---

Published as a conference paper at ICLR 2025
Substituting these expressions, the expected test set error is:
E [1]((z) - = ()((p)))) -d min(dn, p)
(mb c(p c 10-++11--00 +dmmb= +dmb+dc
p
(67)
Training set error Again, using the techniques in the proof of Theorem we write the expected training set error in terms of A2
E [Il((X) - -(x))) dn min(dn, p) Tr(A2)
(68)
Using the expression for Tr(A2):
dn min (n., p)
E [l0(X) - v(x)
c(p 1) 1)-2
(69)
n
D PROPERTIES OF F(n,p)
In this section, we summarize some known properties about the function F(n,p), which appears in Theorem Recall that
F(n,p) E [ 11:) where R RRX has elements drawn i.i.d. from (0, 1). In the regime In pl 2 2, an exact closed form is given by
min p)
F(n,p) n pl
if In pl 2
Computing the square of the Frobenius norm of R is equivalent to finding Tr(RR = Tr(RRT R.).
When n 2- R R x matrix with Inverse-Wishart distribution of identity covariance, which has mean ----T/ (Von Rosen 1988). Therefore, the expected value of its trace n-p-I' Analogously, when p- 2 2, 18 R is a n x matrix with Inverse-Wishart distribution of identity covariance, which has mean p-n so the expected value of its trace p-a
In the case where for the cases where Ip n S 1. Fnn,p) has no known explicit form, so it was computed by averaging over 100 random trials.
bound on the Frobenius norm of R are known (Szarek 1991). However,
E EXPERIMENTAL DETAILS
E.
DATASET DETAILS
Sine Wave Regression Task We construct a regression problem where the regression target is constructed sum of k functions of -ieddimensional linear projections of the input 1 Specifically, the regression target function y Rm - IR is constructed as follows:
k
v(x)
oi sin (anwwyu + 2xy))
(70)
where dij, Wij, 03 E RR, u, E R.m is the number of sine waves that comprise each function of one-dimensional linear projection uf x. We sample these parameters of the target function indepen- dently from the following distributions: a, (O, 1), Wij (O, 4n²), Φij Ud(-x, T) where A denotes Gaussian distribution and 2/ denotes uniform distribution. w; is drawn uniformly from a unit sphere. Note that the full function is made of & separate functions of one-dimensional projec- tions of each of which has 5 sine components. We normalize by VE that v(z) does not scale
20

---

Published as a conference paper at ICLR 2025
with k. Note that the task is modular: although the task takes dimensional input, the target constructed as a combination of k modules operating on dimensional projections of the input. Given the projections, parameterizing functions of the projections are sufficient to parameterize the full task.
A training dataset is generated by first drawing training samples from mean-zero Gaussian: x ~ NO, l). Then, for each I, the regression target y(z) is computed. The test dataset con- structed analogously. See Fig for illustrated example target function in dimension.
Recall that any square integrable function can be approximated on a finite interval to arbitrary pre- cision by sufficiently large Fourier series. Thus, may expect that 7 approaches infinity, the functional form in Equation can can express any function constructed as sum of square integrable functions of the uff x: Ei mi(uf z) for any square integrable 1/:
Note that v controls the dimensionality of æ; thus, we may make the task test generalization over arbitrarily high dimensions by simply increasing that the number of samples required to generalize to a fixed precision on a regression problem scales exponentially with the intrinsic dimensionality of the task input (McRae et al. 20200 [Sharma & Kaplan 2022). Thus, even with a relatively simple task construction, we may expect to produce tasks with arbitrary difficulty measured by sample complexity.
This is significant because prior work shows
1.0
Target function
0.5
Output
0.5
components
-1.0
Training points
-2
0
x
Figure 6: An illustration of single one-dimensional projection in sine wave regression task with sampled training points. The target function in this projection direction is made three sine wave components summed together.
Nonlinear Sine Wave Regression Task our sine wave regression task. Recall that in the original sine wave regression task, outputs are constructed as:
We also test our approach on a non-linear variation of
v(x)
oig sin (Wiju x + dii))
(71)
where u; module projection directions. Note that this task has linear module input projections (the projections uf x are linear functions of x and u:). In our nonlinear variant, we consider the following outputs constructed with non-linear module input projections:
v(x) =
Oij sin(wi; l|u -:/l2 + dir))
(72)
where uTa x replaced with |/u - x 2/|2- which non-linear both u, and æ. Remaining task parameters set the way as in the original sine regression task.
21

---

Published as a conference paper at ICLR 2025
Figure 7: An example input from the Compositional CIFAR-IO task with 4 component images. The goal of the task is to predict the classes of all component images in the input. In this case, with component images, the output target would be a 4-hot encoded 40-dimensional vector representing the true class of each of the 10 possible classes for each component image.
Compositional CIFAR-10 dataset in- spired by the Compositional MNIST dataset of (Jarvis et al., 2023). In the task, combinations of CIFAR-10 images are concatenated together and the model is asked predict the class of all com- ponent images simultaneously. Fig 7 illustrates an example input. Inputs are flattened to remove all spatial structure. Outputs k-hot encoded vectors constructed by concatenating the 1-hot encoded labels for each component image: thus, targets are 10k dimensional. For this task, accuracies are reported on average over component images (for instance, if a model correctly guesses the class of two out of four images, the accuracy would be 50%).
Compositional CIFAR-10 We conduct experiments
Training and test sets for this task are constructed respectively as follows: each input the training (or test) set is produced by randomly selecting & images without replacement from the original CIFAR-10 10 training (or test) set and concatenating them in a random permutation. With images, there are 10 possible class permutations for each input. We use a fixed training set size of 10;; thus, the probability of a test set point having the same class permutation as training set point 10 at most Tor. For large k; we expect each test set point to test class permutation unobserved in the training set.
Note that this task fits the modular structure of Equation 14 (z; Uy) is 1-hot encoded label of dimensionality 10k indicating both the label of the j-th component image and which of the k images is being predicted by the module. The full output is constructed as a sum of the k labels yy (æ; Uy)- As with the sine wave regression task, by increasing the number of component images, we may test generalization in arbitrarily high dimensions.
Class combination experiments: Compositional CIFAR-10 We consider Compositional CIFAR-10 variant which the training inputs are constructed to have a distinct set of class la- bel combinations compared to the test inputs (e.g. with h = 3, if any training set input has the class combination cat, a sirples ane, ship, then this class combination is not permitted on any test set input). This is done by partitioning the full set of class combinations into ing inputs and another disjoint set allocated for the test inputs. Thus, this tests out-of out-of-distribe distribution generalization. All other dataset parameters are set the same way as in the original Compositional CIFAR-10 task.
allocated for the train-
Noisy inputs experiment: Compositional CIFAR-10 We consider a Compositional CIFAR-10 variant in which the training inputs have added Gaussian noise drawn from (0, of varying magnitude this is done after the concatenation of k images together. Test inputs do not have any added noise added. All other dataset settings are identical to the original Compositional CIFAR-10 task.
E.2 EXPERIMENTS ON MONOLITHIC NETWORKS
Architecture and hyperparameter settings: sine wave regression In our experiments, all neu- ral networks are fully connected and use ReLU activations except at the final layer. We do not use additional operations in the network such as batch normalization. Networks are trained using
Adam (Kingma & Ba 2015) to minimize mean squared error loss. We perform sweep
22

---

Published as a conference paper at ICLR 2025
learning rates in {0.001, 0.01, 0.1} and find that the learning rate of 0.01 performs best in general over all experiments, as justified by Fig in Appendix E. Our results are reported this setting. All networks are trained for 10000 iterations which we find to be generally sufficient for convergence of the training loss.
The network architectures are varied follows: the width of the hidden layers is selected from {8, 32, }28), and the number of layers is selected from {3, 5, 7). This yields total architectures. In order consistently measure the number of parameters for architecture ality varies, when we count the number of parameters we treat the input dimensionality m as fixed m 1. Note that this slightly underestimates the true number of parameters each NN. The values of l and m range from 2 to 9. The value of ranges from 100 to 100000.
the input dimension-
All experiments are run over 5 random seeds and results are averaged. Experiments are run on computing cluster with GPUs ranging in memory size from 11 GB to 80 GB.
Fitting our theoretical model: sine wave regression task Our theoretical model of generaliza- tion has three free parameters 12 and We select these parameters to best match the empirically observed trends of training set performance on the sine wave regression task; we find that c = 1.15, Ω = 1.57, and 0.85.
Architecture and hyperparameter settings: Compositional CIFAR-10 In our experiments, all neural networks are fully connected and use ReLU activations except at the final layer. Note that the inputs are flattened; our networks do not use the spatial structure of the input. Batch normalization applied before each ReLU. Images are normalized with standard normalization.
Networks are trained using Adam (Kingma & Ba 2015) to minimize softmax cross-entropy loss using a learning rate 0.0001 and batch size of 128. Note that the loss averaged over all com- ponent images for each input. All networks are trained for single epoch on 1000000 random training points; note that with a moderately large number of images in each input, the total number of possible training inputs can be much larger. We use test set of size 10000.
The network architecture consists of fully connected layers size 512, 512, 512, 256 and 128 before final fully connected layer to predict the output label.
Experiments are run over five random seeds for each hyperparameter configuration. Experiments computing cluster with GPUs ranging in memory size from 11 GB to 80 GB.
E.3 EXPERIMENTS ON MODULAR NETWORKS
Architecture and hyperparameter settings: sine wave regression task Each module (x: Ú,) constructed follows: Ü, consists of two vector components úz and ig. The module output is constructed as:
(73)
where fj represents neural network with scalar input and output.
We set the kernel x as follows:
k(x1, (ū,
2
(74)
of of
0-22+ $
ana
21-
+ e
where a is hyperparameter. Intuitively, E R.m corresponds to a direction along which projection directions à of other modules are not sensitive. This choice of kernel is motivated by the observation
új = x² û; and (x removes all the information in relevant module u, while retaining information relevant to all other modules.
that 6T ûn for j * i, then (x 0 3t
T ú, = 0: x Bra.
xT
Due to the computational cost of computing sample complexity via binary search (Algorithm we fix several hyperparameters by small-scale experiments before the final experiment to control the total runtime. We first sweep through some parameters of in module initialization method
23

---

Published as a conference paper at ICLR 2025
(Algorithm 1. Specifically, we perform sweep over module batch size in {32, 128, 512), module learning rate in {0.001, 0.01, 0.13, and module iteration number in {100, 1000}. The combination of module batch size 128, module learning rate 0.1, and 1000 module iterations has the lowest test emror in our experiment.
Then, sweep through number of architectural modules in {1, 2, x k. x k). learning rate in (0.001, 0.01, 0.13 and a in {0.3, 3.0). We find that the combination of module number k. learning rate 0.01 and main experiments on modular NNs, all networks for dimension (k = m) smaller than 7 are trained for 1000 iterations, while dimension 7 and networks are trained for 700 iterations. dimension networks are trained for 500 iterations and dimension 10 networks are trained for 200 iterations. The high-dimension networks are stopped early since smaller number of iterations was sufficient to converge on the training set; these numbers are determined based on small-scale experimental observations.
1.0 achieves best test performance. In addition, for our
In binary search, we stop the search when the higher bound and the lower bound are close enough (r 1r-1- < 0.3) to shorten our runtime. Also, due to GPU memory limitations, we can only support sample size up to 106, so we stop experiments when the current sample size reaches 2 (c = 22). We set the maximum search iteration (B) to be 18.
We test network in 9 network architectures: the width of the hidden layers is selected from {8, 32, 128]. and the number of layers is selected from {2, 4, 6}. We also have three different values (f0.5, 1.0, 1.5]) for the desired test error in binary search
pinpoint the most suitable value for further application of our method; we select a desired error of 1.5 for our results. We keep = m in all experiments and the values range from 9. All experiments run over 5 random seeds and on the same computing cluster described in the previous section.
Algorithm Finding single module projection Ú
Require:
Supervised training set (X, v(X)), iterations iters, number of training points n, learning rate 17, batch size b, kernel function
Randomly initialize Ú
for iler = 1, ... ilers do Initialize b b kernel matrix K
Randomly subsample b training points from X and store in X
for = I...., do
for 1....,b do
Kii ] n(ā,, By; U); index points from X
end for
end for
L = v(> x 1 x
Compute 96 V o C: can be found with automatic differentiation
0 Ú 790
end for Return
24

---

Published as a conference paper at ICLR 2025
Algorithm 2 Binary search for sample complexity
Require: A training algorithm 7(n) which outputs test loss model when trained on n samples; the desired enror the number of binary search iterations B.
Initialize = 0 and r = oo: our current guess for the number of required samples lies in [2',2] Initialize 12: our current guess for the number of required samples is 2".
for b = 1..... . B do
Find test loss if > e then
((2))
Increase number of samples
[ = c
if 30 then c c c 2
else
c= r)/2 end if
else
Decrease number of samples
c= (1 + c)/2
end if end for Return 2
Disentanglement experiments: sine wave regression task For our disentanglement experiments evaluating whether modular NNs can find the true modules underlying the task, we use the following 1000. We use modular NN with 20 modules. Each
hyperparameter settings: m = k 10, module uses a fully connected architecture with layers and a hidden layer width of 32. We train the modular NN with learning rate of 0.001 for 1000 iterations.
For learning our module initialization, we use 100 iteration steps with a learning rate of 0.01 and a batch size of 128. is set 1.0.
For constructing a t-SNE embedding, we use a perplexity 5. For computing similarity scores, we first compute the absolute value of the cosine similarity between each pair (u,, ûy) of learned and target module directions. For each learned module direction ú; we then find the target module with the largest absolute cosine similarity. Finally, average the maximum absolute cosine similarities
across all modules to produce similarity score: Σ max
uy
Ablation experiments: sine wave regression task For our ablation experiments, we train on 10000 points. & = m is varied from 2 to and the number of architectural modules is set to 5 x k. Each module uses fully connected architecture with layers and hidden layer width of 32, We train the modular NN with a learning rate of 0.001 for 1000 iterations.
For learning our module initialization, we use 100 iteration steps with learning rate of 0.01 and batch size of 128. a is set to 1.0.
Architecture and hyperparameter settings: nonlinear sine wave regression task the nonlinear sine wave regression function, we consider modular architectures of the form: Σ 9, (1 - æ112) where Új is a neural network and ú, are learned parameters. We apply our method to learn an initialization for modules ú, using the following kernel:
To learn
x[ 222 (1)) = - (llm - u. |||1|==--6111
(75)
We consider modules constructed as fully connected networks with 6 layers and width 128. We set = m and set n = 1000. All other hyperparameter settings are consistent with our original sine wave regression experiments.
25

---

Published as a conference paper at ICLR 2025
Architecture and hyperparameter settings: Compositional CIFAR-10 Note that given in- put composed of l images, the flattened input dimensionality is 3072k. Each module ūز (1; U,) is constructed as follows: Uj= IR 1233022x512 and the module output is constructed as:
(76)
where fJ represents a neural network with a 512 dimensional input and an output of dimension 10k. We set the kernel follows:
H 1(21, T2; () e- - -=:
(77)
where a is hyperparameter; we set o = 20.0 to match the scale the distances between projected inputs. To make kernel optimization more efficient, we stochastically optimize only the components of ( K "X) corresponding to a single class at time.
set to be consistent with the monolithic network. Additional hyperparameters are set as module batch size 128 and module learning rate 0.01. Module optimization is performed over single pass the training data. We fix the number of architec- tural modules as 32. Each module a ReLU-activated neural network with hidden layer sizes 256, 128 and 64. Each ReLU preceded by batch normalization. The module outputs all concate- nated and fed into final linear layer to produce the 10k dimensional output. Critically, all the modules have the same weights. This is done match the properties the task: each modular component of the task is the same, namely to predict the class of single CIFAR-10 image. This unlike the sine wave regression task, where each modular component corresponds to different function.
Unless otherwise specified, hyperparameters
We vary the number of images from 1 to 8. All experiments are run on 5 random seeds and on the same computing cluster described in the previous section.
26

---

Published as a conference paper at ICLR 2025
F ADDITIONAL EXPERIMENTS
Complexity
Sample Complesity
Semple Complexity Layers
10y
107
Conelestty
Samplo Conelisity
Compandyy
10
30
30*
1.0
102
10
10
2x100
3x 4x 10
Leurning
Layers
Figure 8: Performance monolithic architecture on sine wave regression task across learning rates of {0.001, 0.10, (left), widths of (8, 32, 128} (center), and number of layers (3,5, 7} (right). Each line represents different architecture and desired test error, and the average performance is shown in bold. Note in the left plot that learning rate of 0.01 tends to perform the best, all our experiments are reported in this setting.
Table 3: Comparison of our method with baselines Compositional CIFAR-IO 10 variant in which the training inputs are constructed to have a distinct set of class label combinations compared to the test inputs (e.g., if any training input has the class combination cat, airplane, ship, then this class combination is not permitted on any test set input). Thus, this tests combinatorial out-of-distribution generalization. We find test set accuracies for inputs with 6 component images. Standard errors over 5 trials are reported.
Method
Test Accuracy
Baseline monolithic
42.56% + 0.07%
Baseline modular
45.26% + 0.06%
Our method
49.90% + 0.15%
Table 4: Comparison of our method with baselines Compositional CIFAR-10 variant in which the training inputs have added Gaussian noise drawn from (O, af) of varying magnitude a. We find test set accuracies for inputs with component images (note that the test points do not have added noise). This tests out-o--distriuutio generalization to small distribution shifts. Standard errors over 4 trials reported.
Noise Level
Baseline Monolithic
Baseline Modular
Our Method
42.92 0.05%
45.66 0.14% 0.14%
50.49 + 0.09%
0.3
42.57 +: 0.05%
45.42 1./2+000 0.10%
50.27 : 0.08%
1
39.47 0.09%
43.07 0.08%
46.95 + 0.11%
3
31.75 + 0.03%
34.71 = 0.05%
34.67 = 0.13%
21.76 0.10%
24.31 + 0.34%
24.41 + 0.13%
30
11.09 + 0.33%
12.50 + 0.55%
12.29 + 0.40%
27

---

Published as a conference paper at ICLR 2025
modules
Ours
6
Fixed Initralizationn
5
LO55
4-
3
4
5
8
Figure 9: Test loss of various training methods of training a modular NN on a sine wave regression task as the dimensionality of the task &s m is varied. Margins indicate standard deviations over random seeds. Ours indicates that NN module directions are initialized using our method and then trained on the task. Fixed initialization indicates that module directions are learned with our method and are fixed during task training. True modules indicates that ù is set to the underlying module directions of the task u (which are generally unknown) and fixed during task training. Observe that when A m is large, using the true, ground-truth module directions slightly outperforms our method, although interestingly our method performs better for low-dimensional tasks. Fixing the initialization found by our method results in significantly worse performance relative to allowing the module directions to vary.
m_mod_les-10%
200
400
600
800 1000
herwion
Figure 10: Test loss of modular NN initialized with our method with different numbers of modules on a sine wave regression task. The experiment conducted with all hyperparameters fixed at the optimal value detailed in App E except the learning rate and a. The experiment is under random seeds. The lines are averaged over all runs. Model performance increases as the number of modules increases. Intuitively, the number of modules increases, the model can contain more information, thus achieving lower error.
28

---

Published as a conference paper at ICLR 2025
Target
150
-50
-50
-40
-150
-100
-50
Figure 11: t-SNE embedding of target (u) and learned module projections (0) learned by a randomly initialized modular NN without training (left), randomly initialized modular NN trained with gradi- ent descent (center), and our initialized modules (right) on a sine wave regression task. From left to right, learned modules cluster more closely around target modules.
(a) Baseline Modular, Initialization
(b) Baseline Modular, After Training
(c) Our Method, Initialization
(d) Our Method, After Training
Figure 12: Illustration of the weights ò before and after training using the baseline modular method and our method on the Compositional CIFAR- 10 task with 6 component images. A single column of the corresponding to a single module is plotted. The weights are organized into 6 groups corresponding to which of the component images each weight is sensitive to. Intuitively, this plots the sensitivity of a single module to each of the component images. (a) At initialization, the baseline modular method has randomly initialized Ú. (b) After training, it learns to be mostly sensitive to a particular image the input. (c) Our method learns be sensitive to only a single image (the
rightmost image) in the input before any training on the task, thus correctly learning the underlying modular structure of the task. (d) The module our method retains its sensitivity to the original rightmost image over the course of training.
29


Reviews:
[{"review_id": "HqILmnbtZd", "reviewer": "Reviewer_MzHa", "summary": "This paper (1) constructs a simplified theoretical model of generalization (focused on the case of linear regression from what I believe to be a set of not-strictly-defined features), (2) provides empirical demonstrations that, at least in broad strokes, a sine wave regression task follows the predictions made by that model, (3) argues that in cases where the simple problem under consideration is modular (here meaning that it has k modules of size b which interact, rather than a full set of P features), better sample complexity can be obtained by using an explicitly modular parameter structure, and (4) provides some empirical support for this generalization behavior on another set of (this time modular) sine wave regression tasks.", "questions": "- What is the explicit definition of modularity being used? This concept was referenced without ever being really explicitly defined in a general-but-still-technical sense (and attention was given as an example). I ended up being confused about whether the focus was on independently functioning parts of a network, or shared weights in a more general sense .\n- As mentioned in \"Weaknesses\": what assumptions are made in general about the structure of the feature matrix? It is indicated in the modular version of the model that arbitrary nonlinear transforms of the input are considered as valid feature matrices, but this isn't clarified for the first treatment of the model \n- Do the benefits cited require being correct about the number (k) of modules in the underlying data? How much do positive results depend on being correct in your choice of k relative to what is present in the underlying data?", "limitations": "", "rating": 5, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "- The paper focuses on an empirical task that lets them validate their ideas, but is realistic in terms of scope \n- The paper engages with an interesting problem of the structure of weights that allow for better generalization", "weaknesses": "- The paper's notation, particularly in the initial presentation of the theoretical model, was confusing and felt under-explained. In particular I was confused by feature matrix (how did it get constructed from the inputs? What assumptions are being made about it? Why is it a matrix to begin with rather than a feature vector), and this confusion made it hard to understand future claims made in the paper (especially since the central claim was about the effect on generalization of _input_ dimension, which is mediated by the function implied in this feature matrix)\n- The forms of the expected training and test loss could have been broken down in a clearer and more intuitive way, rather than simply being presented as not-very-comprehensible formulas \n- This paper assumes that the only way to benefit from the generalization behavior of modularity is to have explicitly modular structure; it would have been interesting if it had also engaged with whether modular data gives you generalization benefits without a parallel parameter structure (since in practice modern models seem to generalize well without the benefit of this)", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "lidSjWCsNc", "reviewer": "Reviewer_vQzj", "summary": "The paper shows that the sample complexity associated with training modular neural networks is independent (under certain conditions) of the input dimensionality and does not follow the same exponential increase with input dimension as in the case of monolithic or traditional neural networks (NNs). \n\nFirst, the authors present a derivation of training and test error in monolithic NNs when the task is non-modular. The task and the NN are modeled linearly based on features generated using the input. The authors then, empirically validate that NNs with different architectures (loosely) match the theoretical trends when varying the input dimension, number of samples and the number of parameters. Note that the task considered for this experiment is modular. \n\nThe authors then theoretically compute the training and generalization errors of modular NNs given that the underlying task is also modular with the same structure. Each module is associated with a small NN (modeled linearly) and an input projection mechanism that reduces the module input dimensionality. The task is also modeled in the same way where the parameters are randomly initialized. \n\nThe resulting closed form solution shows that the training error is independent of the input dimension and the test error under the condition of under-parametrization is also independent of the input dimension. This result hinges on the input projection associated with each module, where the dimensionality associated with the overall input is reduced. The authors then propose a method to initialize (or learn) the parameters associated with the input projections. Once initialized the modular NNs are trained end-to-end. \n\nEmpirical results show that modular NNs that are learned using the proposed initialization mechanism achieve significant improvements over monolithic NNs and modular NNs conventionally trained.", "questions": "Continuing with the previous weakness, the algorithm to learn or initialize the input projection parameters may not work in such a case as it is dependent on the initial module NN weights. \n\nThe generalization performance for the compositional CIFAR-10 experiment can be divided into input class permutations present in the training data vs. not present in the training data to further dissect the difference between monolithic NNs and modular NNs. The sample complexity experiments with compositional CIFAR-10 tasks are not present and should be added to further strengthen the claims. \n\nFor individual tasks considered, there appears to be a large amount of tuning of methodology to train the modular NNs and the module input projections. (Referring to appendix)\n\nHow would the solution to training and test errors change if the input projections from each module were removed, and the modules considered the input x in its entirety? This is consistent with the current mixture-of-exert (MoE) models. \n \nIs there a validation set used for the experiments or is the generalization performance reported from the last training iteration?\n\nDo the modular NNs treat the number of modules as a hyper-parameter and tune it to improve the performance. Or is it an architectural characteristic such as the width or depth in monolithic NNs that is fixed ? \n\nThe CIFAR-10 experiments are run only for a single epoch, will increasing the number of epochs result in better performance for networks? \n\n\nMinor: Equations are referred to the appendix when they are also present in the main part of the paper.", "limitations": "", "rating": 8, "confidence": 2, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "The theoretical derivations are sound and very well done, and the paper is well written. The authors did a good job showcasing how modular networks can be related to monolithic networks (based on certain linear assumptions).\n\nIntuitively, modeling of the data, and learning of parameters related to module input bottlenecks make sense. This is similar to learning the connectivity associated with module input or limiting the number of module inputs to avoid module collapse. \n\nThe experiments show a clear trend that the input projection mechanism results in better performance and sample complexity, as compared to monolithic NNs and end-to-end trained modular NNs.", "weaknesses": "The major weakness of the paper is the consideration of a single layer of modules and data generating system. In such a system the output is a linear composition of the module outputs. This may not be true for many real world systems where multiple such modular layers can exist in a hierarchy.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "458889POda", "reviewer": "Reviewer_uK9v", "summary": "The authors present both theoretical and algorithmic results regarding the generalization ability of a particular form of modular architectures. They first highlight theoretical results on generalization results showing exponentially large sample complexity as a function of the input dimension. Then they present a particular class of modular architectures as a sum of experts and assume that the training data was generated by this same architecture. They show that thanks to each module making a low-dimensional projection before processing the data, the scaling behavior is better behaved. Then they present a kernel-style algorithm to initialize such an architecture, to be then fine-tuned by usual supervised learning and SGD. Finally, they show results on a toy 1-dim sine-wave regression task and on the recently introduced compositional MNIST task.", "questions": "(1) I was confused by the results in figure 1, whereby test error INCREASES with larger datasets. This seems incompatible with empirical observartions and traditional statistical analyses of generalization.\n\n(2) I did not understand in what sense the y_j could be considered independent (and what is the random variable), after eq 3.\n\n(3) eqn 4 seems wrong: on the LHS y is a function of the linear projection U x, whereas on the RHS U and x only interact via the presumably non-linear function phi.\n\n(4) why should we expect eqn 15 to give the minimum norm solution? (i.e. why is it a solution and why is it minimum norm, from what class)\n\n(5) You should add the citation of the original mixture of expert papers (e.g. Jacobs et al 1991).", "limitations": "", "rating": 3, "confidence": 4, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "Much remains to be understood about the generalization behavior of neural nets, especially the types that have a modular architecture, so advances on the theory (in special cases) that are presented do seem useful.", "weaknesses": "(1) the theoretical results are not suprising: projecting the m-dim input to several b-dim low-dimensional representations unsurprisingly reduces the exponential badness from m to b. The theory is also of fairly limited scope, with lots of unreasonable assumptions (e.g., of linearity wrt parameters) that may not tell us as much as we would like for more general forms of modular architectures.\n\n(2) the proposed algorithm is unlikely to scale well in terms of computational efficiency beyond small-size problems and into frontier AI, given the use of kernel methods in the novel part of the method\n\n(3) the fact that all modules are initialized independently and using the same (randomized) procedure suggests that a significant part of the advantage could come from an ensemble effect (which always helps generalization)\n\n(4) the paper seems to overclaim in multiple places, e.g., suggesting that their results tracks empirical behavior of modern neural nets (even the empirical comparisons don't match the theory, e.g., fig 2 bottom right).\n\n(5) I did not find numerical comparisons against benchmark results from other papers, and when I look at Jarvis et al 2023, their figures show much lower errors. Hence the experimental results may not be that good after all.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "LX1tOmzlod", "reviewer": "Reviewer_2nhB", "summary": "In this paper, the authors study modular neural networks and their ability to learn functions that are themselves modular insofar as they have either a compositional nature or a modular composition in their statistics. The authors first provide a theorem outlining expected scaling laws for sample complaxitiy and trainign accuracy in a curated linear task setup. They then corroborate the predictions from this theorem with numerical experiments. Following this, they present a task setup with explicit ground-truth modular structure and prove a second theorem outlining similar sclaing propoerties for this novel tasks setup along with a particual modular NN architecture. The result is reduced scaling complexity for tusch tasks when NNs are appropriately modular. The authors then propose an initialization scheme for NNs which first learn module initialization in a self-supervised fashion using task statistics. Finally, the authors show that this approach behaves as expected on non-trivial tasks such as compositional CIFAR, and that the proposed method works well on other modular architectures beyond the one used for their theoretical results.", "questions": "See above", "limitations": "", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "This paper presents an excellent set of theoretical results outlining expected scaling laws for modular networks when task have a modular structure. It also presents a practical initialization scheme for potential models that is based on self-supervised alignment with task statistics. Lastly, it validates predictions with non-trivial and relevant experiments with architectures that go beyond the ones used for theory, outlining the potential generality of the result.", "weaknesses": "The paper is sound. There are potntials for improvements in two key areas:\n\n1. In experiments, it is unclear if the generalization advanatage of the modular networks remain if one factors in the pre-training (i.e. learning module initialization). In other words, for the same total compute, would a monolithic model do as well as the modular one for which some of the compute budget went toward initialization? My apologies if I missed it, but this is a key point that would factor into real-world scaling laws.\n\n2. What happens if there is a missmatch between the task modularity and the NNm odule count? In experiments, the models could accurately recover tasks modularity when the number of modules is known. In contrast, in other tasks, this is not known a priori. How would the models behave in this mismatched environment? To be fair, the authors acknowledge this issue, but I wonder if some rapid experiments could outline expected behaviors in this case. Once again, apologies if this has been explored and I missed it.\n\n3. While this is no doubt very relevant for the field, it's overall impact with respect to modern architectures remains unclear. Some discussion about the use of such methods in modern settings such as in the prsence of attention could greatly enhance the scope of the result. This is not necessary however, as this is a good paper in the current scope.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)