Breaking Neural Network Scaling Laws with Modularity

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.

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).

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.

Empirical result on Figure 2

Thank you for this suggestion. We are unable to use our current empirical results to generate an empirical sample complexity curve since our current experiments in this setting compute the test set error at a given pair of training set size n and input dimensionality m. However, we are currently running binary search over training set size n to compute the minimum n required to achieve the desired test set error of 1.2. We will update the figure with empirical results once available.

Page 7 typo

Thank you for pointing this out; we have corrected this in our revision.

Thank you for your valuable and constructive comments and suggestions. We are glad to hear of your positive impression of our work.

Practical relevance of scaling law in nonlinear setting

We first would like to clarify that the empirical validation of Section 3.3 is performed on nonlinear neural networks; thus, the generalization error of linear models can predict generalization in nonlinear networks.

We also note that although we primarily consider nonlinear model architectures containing linear module projections (as in Eqn 3), in Section 4.3 Table 2, we also conduct experiments on nonlinear module projections. As we discuss in Section 5, we believe that further studying nonlinear module projections is an important direction for future work, particularly in practical settings where nonlinear module projections may be more realistic.

Discrepancy with theory for sample complexity trend

Indeed, as the reviewer notes, while the theory predicts that in modular architectures, sample complexity scales as a constant with input dimensionality, our empirical results indicate that sample complexity grows with input dimensionality. We believe the discrepancy with the theory arises due to the optimization challenges associated with modular architectures (as previously found in Csordas et al. and Mittal et al.). Specifically, with higher dimensional inputs, the gap between the performance of the optimal modular model and of the modular model learned by gradient descent grows. Our modular learning rule introduced in Section 4.2 is aimed at easing these optimization difficulties and is effective as illustrated in Figure 3. However, it still does not find the optimal modular architecture parameters, thus maintaining an input dimensionality dependence for sample complexity. We regard further improving the optimization of modular architectures as an important future direction.

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

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.

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.

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.

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.

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.