Breaking Neural Network Scaling Laws with Modularity

We really appreciate your positive feedback on our submission and your suggested directions for improvement.

Fixed compute budget

This is an important question of great practical relevance. We have not tried any experiments in which the compute budget for the randomly initialized vs. pre-initialized networks are fixed. In fact, in this work, we do not at all optimize our kernel-based modular learning rule for computational efficiency.

However, we note that approximating kernel methods for computational efficiency is a well-established field of research which can reduce the computational cost of kernel methods down to linear time in the size of a dataset (such as using random Fourier features). We expect that combining these methods with our proposal can help make our method significantly more competitive with random initialization under a fixed compute budget.

Mismatch between architectural module count and task module count

This is an insightful point.

In our experiments, in general, the number of architectural modules is always greater than the number of modules in the task: thus, these do not need to match. For example, in the case of Compositional CIFAR-10, we fix the number of architectural modules to 32 and vary the number of task modules from 1 to 8. Figure 10 shows an additional experiment demonstrating the effect of choosing different numbers of architectural modules on the sine wave regression task. We find that adding more architectural modules helps performance even beyond the number of task modules.

Relevance to modern architectures

We appreciate this point as well. In section 5, we have expanded our discussion of how our results apply to modern architectures.

In terms of practical relevance, we believe our results show that when training modular architectures, more emphasis should be placed on the optimization methods used to train the modules. Naively applying gradient descent may not be sufficient for effective training of these architectures. Our results also suggest a potential theoretical basis for why modularity is so effective in practice: namely, modularity breaks up a large high-dimensional task into easier to learn low-dimensional tasks.

Thank you for your detailed review and for highlighting areas where our work can be strengthened. We address your concerns point by point below.

Novelty of theoretical results

We understand that projecting high-dimensional data into lower-dimensional subspaces is a well-known technique. However, our contribution lies in providing a quantitative analysis of how modular architectures affect sample complexity in neural networks.

Our work is the first to derive explicit, non-asymptotic expressions quantifying how modular architectures can circumvent the exponential sample complexity associated with high-dimensional inputs. This provides a theoretical foundation for designing modular networks that are efficient in practice.

While our analysis assumes linearity and specific modular structures, these simplifications are essential for analytical tractability. Future work can build upon our framework to explore more general, nonlinear architectures and investigate how modularity affects their generalization properties.

Scalability of the proposed algorithm

You raise a valid concern regarding scalability.

We note that approximating kernel methods for computational efficiency is a well-established field of research which can reduce the computational cost of kernel methods down to linear time in the size of a dataset (such as using random Fourier features). We do not explore these methods in this work as it is orthogonal to our contributions; however, we believe combining these well-known methods with our proposal is a fruitful direction.

Potential ensemble effects

This is an interesting point.

In our compositional CIFAR-10 experiments, however, the modules all have the same weights (they are weight tied). Thus, the improved performance of the modular architecture cannot be attributed to an ensemble effect in this case.

Overclaiming and empirical comparisons

We apologize if any claims appeared overstated.

While our theoretical model captures the general trends observed empirically, certain deviations occur due to factors not accounted for in the simplified model, such as optimization dynamics. We discuss these factors in Section 3.3.

We are happy to soften any specific claims in our submission to match our empirical results.

Comparison to benchmark results

We emphasize that Jarvis et al. tests on Compositional MNIST while we test on Compositional CIFAR-10, a significantly harder task. Thus, their performance metrics cannot be compared with ours.

Confusion about test error increasing with larger datasets

We understand the confusion.

In Figure 1, the increase in test error with more data reflects the double descent phenomenon (Belkin et al.). This is a now well-established phenomenon in which test error increases with the number of training points until the interpolation threshold is reached (training points = number of parameters). Beyond this point, the test error decreases with the number of training points; this is also clearly illustrated in Figure 1. We are happy to provide further references on this phenomenon if helpful.

Independence of $y_j$

In equation 3, we wish to rescale the summation such that the left hand side has constant magnitude as the number of modules varies. As a heuristic, if we treat each of the terms in the summation as independent under the assumption that $x$ is a random variable, we then find that the variance of the summation scales with the number of modules $K$. Thus, it is natural to divide the summation by $\sqrt{K}$ to normalize.

Equation 4 clarification

That's correct: as mentioned in the text before and after equation 4, we make a linearizing assumption on the parameters of the model. We assume that the model output is a linear function of the model parameters (which now include both the module parameters and the module input projection). This is, in fact, the same assumption made in Section 3.

Equation 15 clarification

Equation 15 simply solves equation 14 for $\theta$. Note that equation 14 is a linear equation in $\theta$. It is well known that the minimum norm solution for a variable in a linear equation (if a solution exists) can be computed by the pseudoinverse (denoted by $\dagger$ in our notation): the solution to $Y=X \theta$ minimizing $||\theta||$ is $\theta = X^\dagger Y$. We are happy to provide further references if helpful.

Mixture of Experts paper reference

Thank you for pointing this out. We have fixed this in our latest revision.

We really appreciate your thoughtful and positive feedback on our submission. We address the points you raised below:

How to extend to setting with hierarchical modules

We think this is an excellent point: indeed, our paper only considers a relatively simple form of modularity in which we use a single layer of modules and the module outputs are linearly combined to form the model output. As we note in the discussion section, practical modular architectures are often more complex, involving some level of hierarchy. We believe our results are a step in the direction of more theoretical analysis of how modularity can benefit generalization more broadly.

We do note however, that our experimental results include the case of nonlinear module projections (in which the module inputs are nonlinear projections of the task inputs). This is a step towards generalizing our results to more practical settings.

Concerns on CIFAR-10 experiment

This is a nice suggestion. As we note in Appendix E though, given the very large number of possible class permutations ($10^k$ where $k$ is the number of images), we expect that for large $k$, all class permutations in the test data will be unseen. However, in Table 3 we explicitly test accuracies in the case where test inputs have distinct class combinations and find similar results.

For this task, we opt to use accuracy (fixing training samples) instead of sample complexity (fixing accuracy) as our performance metric since it is more natural to treat the Compositional CIFAR-10 dataset as unlimited (given the combinatorial number of samples that can be drawn).

Comments on tuning

Our experiments involve a number of hyperparameters. Unless otherwise mentioned, we set the hyperparameters to the most natural choice without tuning them. For certain hyperparameters, we do sweep over different values, with the sweep range indicated in Appendix E. We are happy to clarify any specific hyperparameter choices if needed.

Experiment with no input projection

We appreciate this interesting suggestion. We note however, that in our Compositional CIFAR-10 experiment, all modules have the same weights (they are weight-tied). Thus, if they all had the same inputs, their outputs would also be the same. The final model output is the concatenation of the module outputs which in this case would perform very poorly.

Comment on validation set

For both the Compositional CIFAR-10 and sine wave regression tasks, we use a separate validation set. Note that for the sine wave regression task, each input is drawn completely independently from an infinite sized dataset; thus, any points not trained on can be treated as validation data.

Number of architectural modules

We treat the number of architectural modules as a (potentially tunable) architectural characteristic. It is fixed (at 32) for the Compositional CIFAR-10 task and tuned over for the sine wave regression task. Figure 10 illustrates performance for different choices of number of architectural modules for the sine wave regression ask.

Number of epochs for Compositional CIFAR-10 task

In this task, the number of possible training inputs can be treated as virtually infinite (for large $k$). Thus, we believe it is more practically reasonable to fix number of epochs as one rather than repeat training on the same input multiple times.

Comment on equation references

Thank you for pointing this out. We have fixed this in our latest revision.

Thank you for your thoughtful review and for highlighting areas where our paper can be improved. We appreciate your insights and address your concerns point by point below.

Confusion with notation and explanation of the feature matrix

We apologize for the confusion caused by the notation in our theoretical model, especially regarding the feature matrix. In our model, the feature matrix arises from applying a feature mapping to the input data, transforming each input $x$ lying in $m$ dimensions into a higher-dimensional feature matrix $\phi(x)$ lying in $d \times P$ dimensions. Here, $d$ denotes the number of output dimensions of the model and $P$ denotes the number of features per output dimension. We do not make any particular assumptions about exactly how $\phi$ is constructed (it can be arbitrarily nonlinear and complex)- the main property it must satisfy is that the features $\phi(x)$ are distributed as a Gaussian.

The reason why we use a feature matrix instead of a flattened feature vector is that it allows us to easily account for non-scalar output sizes ($d > 1$) with a fixed number of parameters $P$. Another option to account for multi-dimensional outputs is to have $d \times P$ parameters (one set of parameters per each output dimension) while the number of features is fixed at $P$. This is also valid; however, our choice is more consistent with some prior literature (e.g. Jacot et al.).

We are happy to revise any particular points of confusion in our revision.

Clarity of expected training and test loss formulas

We clarify the forms in Theorem 1 as follows: the test loss consists of three terms. The first term is largest near the interpolation threshold ($dn \approx p$) and decreases away from it- this causes a spike near the interpolation threshold. Intuitively, this corresponds to overfitting to noise in the training data. The second term is negative and corresponds to the reduction in loss caused by capturing information about the true model. In the overparameterized regime, it linearly grows in magnitude with the number of training points, and in the underparameterized regime it is constant with amount of training data (as information in more training points can't be captured by the model). The last term is a constant offset corresponding to the loss when the number of parameters approaches infinity.

The training loss is a product of two terms: the first corresponds to the information about the underlying target function not captured by the model and the second corresponds to the amount of training data available that can't be captured by the model. If we have more parameters than training data, then the training loss is zero since all the information can be captured. Otherwise, there will be information loss proportional to the excess training data times the remaining information about the target function.

We are happy to revise any particular points of confusion in our revision.

Does modular data alone help without a modular architecture?

The reviewer raises an interesting question about whether it is possible to benefit from merely modular data without a modular architecture. Our experiments on Compositional CIFAR-10 (see Figure 4) suggest the answer is no: a non-modular architecture performs relatively worse on higher-dimensional modular tasks than a modular architecture. Of course, it is possible that the non-modular architecture has learned the modular structure of the task to some extent (just not as well as a modular architecture) and we think this is an interesting future direction to explore.

Definition of modularity used

We appreciate the need for a clear definition of modularity.

In this paper, we define modularity as the decomposition of a task or function into distinct, independently operating components or modules. Each module processes a subset of the input dimensions and contributes to the final output in a compositional manner. This structure allows for specialized processing within modules and recombination to solve complex tasks. Our modular neural network architecture mirrors this by having separate subnetworks (modules) that handle specific input projections.

We also highlight that in other literature, modularity may be used in other, potentially different ways.

Dependency on correct number of modules

This is an insightful question.

In our experiments, in general, the number of architectural modules is always greater than the number of modules in the task: thus, these do not need to match. For example, in the case of Compositional CIFAR-10, we fix the number of architectural modules to 32 and vary the number of task modules from 1 to 8. Figure 10 shows an additional experiment demonstrating the effect of choosing different numbers of architectural modules on the sine wave regression task. We find that adding more architectural modules helps performance even beyond the number of task modules.

Thank you for your continued engagement with our work and for providing additional feedback. We are grateful that you have increased your score and appreciate the opportunity to further clarify and improve our paper.

Modularity Definition and Attention Mechanisms

We apologize for the confusion regarding our inclusion of attention mechanisms as an example of modular architectures. Our intent was to illustrate that certain aspects of attention mechanisms exhibit modular properties, but we understand that this may not have been clear within the context of our specific definition.

We define modularity as the decomposition of a task or function into distinct, independently operating components (modules), each processing a subset of the input features and contributing to the final output in a compositional manner. In self-attention, each output token is computed as a linear combination of input tokens, with the coefficients of the linear combination being the computed self-attention "weights" (computed via softmax). In many NLP and CV tasks, the self-attention "weights" are sparse: each output token is primarily sensitive to only a small number of input tokens. In this sense, we may consider self-attention to be performing modular computation.

Clarity of Notation and Feature Matrix Explanation

We apologize for not incorporating the detailed explanations provided in our response into the paper itself. It is important to us that the paper is clear and accessible to all readers, regardless of their familiarity with prior work. We have revised the paper to include the clarifications regarding the feature matrix. Specifically, we have:

• Clearly defined how the feature matrix is constructed from the input data, including any assumptions made about its properties.
• Explicitly stated the shapes and roles of all matrices involved, including the feature matrix $\varphi(x)$ and the weight vector $W$, to eliminate ambiguity.
• Provided an explanation of why we use a feature matrix instead of a feature vector.

Regarding your question about the feature matrix already having an output dimension of $d$ and the role of the weight vector $W$, we realize that the notation may have been confusing. In our formulation, the feature matrix $\varphi(x)$ maps the input $x$ to a feature space of dimension $d P$. The weight vector $W$ then maps these features to the output space, and has dimensions $P \times 1$. The output is computed as $y = \varphi(x) W$. We have made sure to clearly state the dimensions of all quantities and explain how they interact to produce the output.

Once again, we appreciate your feedback and are committed to making our paper as clear and comprehensible as possible. We believe that the revisions we made will address your concerns and enhance the overall quality of our work.