Paper: Breaking Neural Network Scaling Laws with Modularity (ICLR-2024)

Reviewer comment (Reproducibility):
I am a bit worried about the reproducibility of the empirical part of the paper since no code was provided as part of the submission. I also encourage the authors to specify the exact number of seeds per experiment in Figure 3b as "up to five seeds" as stated in the caption could technically mean only one seed is reported.

Author response:
**Code**

Thank you for the suggestion. In our latest revision, we have included code to reproduce the empirical validation of Section 3.3 and the empirical results of Section 4.3.

**Number of seeds**

Due to experimental limitations at the time of submission, we were unable to run 5 seeds for each of the experimental settings for Compositional CIFAR-10. We have now updated our results with all 5 random seeds for each setting. Note that the error bars in Figure 3 indicate standard errors individually for each point.

Reviewer comment (Theory):
The modular learning rule minimizing the norm of the $\theta_i$ is applied as a pretraining step assuming that the $\varphi(X;\hat{U}_i$ are sufficiently expressive. Since this is before training, can you elaborate why this assumption might be justified and to what extent the algorithm is robust to a violation of it?

Author response:
**Sufficient expressivity assumption**

Thank you for highlighting this important point: indeed, we assume the features of Equation 18 are sufficiently expressive in the sense that $pK > dn$. Effectively, this assumption states that for any choice of module projections $\hat U_i$, some set of module parameters can perfectly fit the training data. This assumption is natural when modules are neural networks since they are typically capable of fitting a range of training datasets. When this assumption does not hold, the pseudoinverse solution of Equation 19 finds a solution minimizing the mean squared error between the predicted and actual training targets $y(X)$. In this case, the training set prediction error is generally not $0$ (except when $\hat U_i$ are at their correct values). Importantly, minimizing the norm of $\theta$ with respect to $\hat U_i$ *may not necessarily* yield a lower prediction error. Thus, we do not necessarily expect our algorithm to be effective when modules are not sufficiently expressive.

We have added a discussion of this point in our revised Section 4.2.

Reviewer comment (Theory):
There are discrepancies between the theory and toy model in Figure 5 as the paper points out in App A.2. Can you elaborate why this is not a matter of concern for the theory, i.e. what exactly causes the mismatch?

Author response:
**Discrepancy in Figure 5**

We note two key discrepancies between our theory and empirical results: first, the loss is empirically larger than predicted for small amounts of training data, and second, the error spike at the interpolation threshold is 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-convex, and the global optimum may not be found especially for small datasets (indeed, we 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 error in low data regimes. 

We hypothesize that the second discrepancy is also partly due to imperfect optimization. This is because the interpolation threshold spike can be viewed as highly adverse fitting to spurious training set patterns. This imperfect optimization is more pronounced at smaller $m$. Despite these discrepancies, we nevertheless find that our theory precisely captures the key trends of empirical test error.

We have added a more thorough discussion of this point in our revision.

Reviewer comment (Presentation):
Figure 5 is missing labels and the caption is a bit sparse. Could you specify how exactly the four plots differ? Maybe adding a colour bar to indicate the values of the light lines could be helpful? How do the theoretical predictions look like for individual (light) lines?

Author response:
**Comments on Figure 5**

Thank you for the suggestion. In response, we have edited the caption of Figure 5 and modified the Figure to indicate the meaning of the light vs dark lines as well as indicate the difference between the light lines. We do not include the theoretically predicted curves for the light lines for visual ease of understanding; however, for larger $m$, the predicted test loss is higher while the location of the interpolation threshold is retained, and for larger $p$ the predicted test loss shifts vertically up and down.

The four plots indicate trends of training and test loss along four different dimensions: $k$ (number of modules), $m$ (input dimensionality), $p$ (model size) and $n$ (training set size).

Reviewer comment (Presentation):
Figure 12 is missing a legend for what the colours encode. Could you please clarify?

Author response:
**Figure 12 clarification**

Note that the Compositional CIFAR-10 inputs have shape 32x32kx3. In our modular architecture, the module input dimensionality is 512 (see Appendix E.3). Therefore, each learned input projection $\hat{U}$ each has shape 32x32kx3x512. In Figure 12, we plot a representative slice of $\hat{U}$ with shape 32x32kx3; note that since it has the same shape as the original image, we may plot it as an image with channel dimensions encoding RGB values.

We would be happy to clarify this in our revision if necessary.

Reviewer comment (Evaluation):
What causes the large deviation of the test loss between actual and predicted in Figure 5?

Author response:
Thank you for your valuable comments.

**Concerns on the discrepancy between actual and predicted test loss**

First, we'd like to note that the prediction of NN test loss in Figure 5 is relatively strong considering that we are using a *linear* model to approximate the generalization trends of a highly *nonlinear* neural network. Importantly, we are able to predict the location of the interpolation threshold, and we believe the results are overall competitive with prior work modeling NN generalization.

We note two key discrepancies between our theory and empirical results: first, the error spike at the interpolation threshold is smaller than predicted by the theory, and second, the loss is empirically larger than predicted for small amounts of training data. We believe the discrepancies are due to imperfect optimization, especially in regimes of small training data and small model size. Note that the linearized analysis assumes that the linear model solution finds the exact global optimum. However, in the actual optimization landscape for modular architectures is highly non-convex, and the global optimum may not be found especially for small models and datasets (indeed, in Figure 5, note the discrepancy between predicted and actual training loss values for small model size p and small data size n). We believe this causes the large discrepancy between predicted and actual test error in low data regimes. Moreover, we believe this imperfect optimization leads to a smaller-than-expected test error at the interpolation threshold given that the interpolation threshold spike can be viewed as highly adverse fitting to spurious training set patterns.

We have added a more thorough discussion of this point in our revised Section 3.3.

Reviewer comment (Evaluation):
In figure 4 (also figure 3b), the total range of the similarity score is quite small, it is therefore difficult to say whether the result is a significant improvement from baseline.

Author response:
**Significance of improvement over baseline**

Regarding Figure 3b, note that results are averaged over 5 seeds and the standard error margins plotted indicate that the improvement over the baseline is statistically significant. Note that for the highest number of images tested, the improvement in accuracy is roughly 5% which we believe is a significant improvement.

Regarding Figure 4, we believe that while the difference in raw similarity score between our method and the baseline appears small, this corresponds to a large difference in the qualitative behavior of the two networks. To illustrate this, in Figure 12 we have plotted an illustration of a learned module direction under our method relative to the baseline. We find that the module direction learned by our method is visibly sensitive to only one component image while the baseline is not; thus, our method learns the underlying module structure of the task.

Reviewer comment (Experiments):
Including synthetic experiments in the linear model to demonstrate how the sample complexity changes for modular and non-modular networks in a specific setting.

Author response:
Thank you for your valuable comments. We are glad to hear of your positive impression of our work.

**Comments on synthetic experiments on linear model**

As the reviewer notes, our theoretical analysis hinges on modeling *nonlinear* neural networks as linear models. However, our experiments are conducted on nonlinear networks and tasks. This is because our theoretical analysis is aimed at capturing generalization trends in actual, nonlinear neural networks, and can do so accurately as demonstrated in Appendix A.2.

Moreover, our construction of modular tasks and models (in Equations 3 and 6) is *fundamentally nonlinear*: the input is linearly projected, fed through nonlinear functions, and then summed. The linearization analysis of Section 4.1 separately linearizes the network with respect to the projections and the parameters of the nonlinear functions; this separation of parameters ultimately yields a superior sample complexity scaling for modular networks relative to monolithic networks. However, for linear models, there is no corresponding notion of notion of modularity: fundamentally, there is no division of parameters into "projection" parameters and "module" parameters. We believe our linearization analysis is relevant only for nonlinear models.

Thus, while we understand the spirit of the reviewer's suggestion, we are unclear on what exact experiment would be most valuable to run. However, we would be happy to include any specific experiments suggested by the reviewer.